Embeddings

An embedding is a table whose rows are indexed by a finite vocabulary. The primary definition in this file uses Fin vocab, so an invalid token id cannot be supplied to the mathematical model. The one-hot presentation is retained as a differentiable linear map and is useful when comparing an indexed lookup with matrix multiplication.

References / analogies:

• In most ML frameworks, an embedding table is a matrix weight : (vocab x embedDim) and an index-based lookup returns weight[token_id]. One-hot embeddings are the equivalent linear map oneHot @ weight.
• Bengio et al., "A Neural Probabilistic Language Model" (2003) for the classic embedding-table framing in neural language models.
• Mikolov et al., "Efficient Estimation of Word Representations in Vector Space" (2013) for the modern word-embedding perspective.
• PyTorch API reference:
  • torch.nn.Embedding: https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html
  • torch.nn.functional.one_hot: https://pytorch.org/docs/stable/generated/torch.nn.functional.one_hot.html

structure Spec.Embedding (vocab embedDim : ℕ) (α : Type) [TorchLean.Storage α] : Type

A trainable table with vocab rows of width embedDim.

• weight : TorchLean.Tensor α [vocab, embedDim]

  Embedding table, stored row-major by executable backends.

def Spec.Embedding.lookup {α : Type} [TorchLean.Storage α] {vocab embedDim : ℕ} (embedding : Embedding vocab embedDim α) {shape : Shape} : TorchLean.Tensor (Fin vocab) shape → TorchLean.Tensor α (shape.appendDim embedDim)

Look up every finite index in a tensor and append the embedding dimension to its shape.

The index type is Fin vocab, rather than Nat, so this specification has no out-of-range case. File and tokenizer boundaries may begin with natural-number identifiers, but executable APIs must validate them before interpreting them as indices into this table.

def Spec.Embedding.oneHot {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {vocab embedDim seqLen : ℕ} (embedding : Embedding vocab embedDim α) (oneHot : TorchLean.Tensor α [seqLen, vocab]) : TorchLean.Tensor α [seqLen, embedDim]

Embed a batch/sequence of one-hot vectors:

oneHot : (seqLen × vocab) and W : (vocab × embedDim) gives (seqLen × embedDim).

Gradients

Embedding.oneHot is matrix multiplication:

Y = oneHot @ weight.

So the reverse-mode derivatives are the standard ones:

• dOneHot = dY @ Wᵀ
• dWeight = oneHotᵀ @ dY

Even though "true" one-hot tensors are often treated as non-differentiable in practice, having a named VJP is useful for:

• treating embeddings as a pure linear map in proofs,
• debugging equivalences (one-hot vs index-based embeddings),
• and keeping this layer consistent with the rest of the spec library.

def Spec.Embedding.oneHotVjp {α : Type} [TorchLean.Storage α] [Add α] [Mul α] [Zero α] {vocab embedDim seqLen : ℕ} (embedding : Embedding vocab embedDim α) (oneHot : TorchLean.Tensor α [seqLen, vocab]) (dY : TorchLean.Tensor α [seqLen, embedDim]) : TorchLean.Tensor α [seqLen, vocab] × TorchLean.Tensor α [vocab, embedDim]

VJP for Embedding.oneHot, returning gradients for the input and table.