Token Embedding Lookup

The embedding matrix is a learned lookup table

Start with the object itself. A model with a vocabulary of V distinct tokens and a hidden width of d_model stores one matrix E with shape [V, d_model]. Row i of that matrix, written E[i], is a vector of length d_model that is the model’s representation of token i. There is nothing else to a token embedding: token id 7 means ‘go read row 7.’

The word learned is the important one. E is initialized randomly — typically small Gaussian noise — and then trained by gradient descent exactly like every other weight in the network. Over training, rows drift so that tokens used in similar ways end up with similar vectors; the geometry of E becomes a map of the vocabulary. So the lookup table is not a fixed dictionary you hand-write, and it is not a frozen algorithm like one-hot or hashing. It is a block of V × d_model trainable numbers that the model shapes to whatever internal coordinate system makes the downstream layers’ job easiest. Everything else in this article is a consequence of that one fact: an embedding is a big, trained, indexed-by-integer table.

Why a lookup is exactly one-hot times E

It looks like indexing an array should have nothing to do with linear algebra, but the two are identical, and seeing why unlocks the rest of the math. Represent token i as a one-hot vector e_i: a row of length V that is 1 in position i and 0 everywhere else. Now multiply it by E:

e_i : [1, V]   (one-hot, 1 at position i)
E   : [V, d_model]
e_i · E = Σ_v e_i[v] * E[v]  =  E[i]   : [1, d_model]

Because every entry of e_i is zero except the one at position i, the sum collapses to a single term: row E[i]. The matrix product selects exactly one row. So the embedding lookup is mathematically a matrix multiply by a one-hot vector — and that equality is not a curiosity. It tells you the layer is linear, it explains the shapes on both sides, it makes the gradient fall out cleanly (later), and it justifies why frameworks implement it as a cheap gather instead of an actual V-wide multiply: the multiply would be almost entirely multiplying by zero.

Shapes: from token ids to [N, d_model]

Real inputs are sequences, so trace the shapes end to end. A tokenizer turns text into a list of N integer ids, a tensor of shape [N] (or [B, N] with a batch dimension). The embedding layer maps each id to a row, producing [N, d_model] — the sequence of vectors the transformer blocks actually operate on.

ids   : [N]                 e.g. [8, 8291, 262, 1049]
E     : [V, d_model]
X = E[ids]  : [N, d_model]  (row-gather, one row per id)

one-hot view:  OneHot(ids) : [N, V]
               OneHot(ids) · E : [N, V] × [V, d_model] = [N, d_model]

Both views land on the same [N, d_model] tensor. Note what the lookup does not do: it does not mix positions and it does not look at neighbours — each id is embedded independently, in parallel, in a single gather. It also carries no order information; the vector for the token ‘dog’ is the same whether it appears first or fiftieth. That is why a separate positional signal is added right after this step. The token embedding answers ‘what is this token?’ and leaves ‘where is it?’ to positional encoding.

Counting the parameters

The parameter count of the embedding table is the easiest count in the whole model: it is just V × d_model, because that is how many numbers the matrix holds. There is no bias and no nonlinearity. Plug in realistic numbers for a small language model — a 32,000-token vocabulary and a hidden size of 512:

params(E) = V * d_model
          = 32000 * 512
          = 16,384,000  ≈ 16.4 M parameters

Over sixteen million parameters, and not one of them is in an attention head or a feed-forward layer — they are all in the front-door lookup. Widen the model to d_model = 768 and the same vocabulary costs 32000 × 768 ≈ 24.6 M. The count scales linearly in both factors, but the vocabulary term is the one people forget: doubling the vocabulary doubles the embedding parameters outright, which is a real reason small-model designers keep vocabularies modest and lean on subword tokenization rather than a huge word-level vocabulary. Knowing this count in your head is useful — it is the first line of any honest memory estimate for an SLM.

Why embeddings dominate small models

Here is the asymmetry that surprises people. The embedding count V × d_model grows linearly in width, but the transformer stack — attention plus feed-forward, repeated over L layers — grows roughly with L × d_model^2, quadratically in width. So as models get bigger and wider, the stack outgrows the embedding table and embeddings shrink to a rounding error. But run that logic backwards: in a small model, d_model and L are small, the quadratic stack is cheap, and the embedding table — anchored to a vocabulary that does not shrink — stays stubbornly large.

The canonical illustration is GPT-2 small: 124 M parameters total, of which the token embedding 50257 × 768 ≈ 38.6 M is nearly a third. For a genuinely tiny SLM the fraction can be worse — a few-layer, narrow model can spend the majority of its parameters just turning ids into vectors. The practical consequences are direct: on a CPU-hosted SLM the embedding matrix is a big slab of RAM, it is a prime target for quantization or factorization, and weight tying (below) is not a minor optimization but a way to delete a third of your parameters. You cannot reason about a small model’s footprint without reasoning about its embeddings first.

The gradient is sparse: only used rows update

Now the subtle part. In a training step you process a batch containing some set of token ids — a few thousand distinct ones, perhaps, out of a 32,000-token vocabulary. What is the gradient of the loss with respect to E? For any row whose token never appeared in the batch, it is exactly zero. That token was never read, so it could not have influenced the loss, so its embedding gets no update this step.

This falls straight out of the one-hot identity. Since X[t] = E[id_t], the only row that X[t] depends on is row id_t. By the chain rule the gradient ∂L/∂E[i] is a sum over exactly the sequence positions whose token equals i, and rows never used contribute nothing:

∂L/∂E[i] = Σ_{t : id_t = i}  ∂L/∂X[t]

if token i appears 0 times in the batch  →  ∂L/∂E[i] = 0

So the embedding gradient is sparse: most of its V rows are zero on any given step. This is qualitatively unlike a dense weight matrix in the stack, every element of which gets a gradient every step. It is why embedding layers are implemented with sparse-gather backward passes, and why the update touches only the handful of rows the batch actually referenced.

What sparse updates mean in practice

The sparsity has real training consequences. First, frequency drives learning: a common token appears in nearly every batch and its row is nudged constantly, so it converges to a rich, well-placed vector quickly. A rare token might appear once in thousands of steps; its row is updated only on those rare visits and stays close to its random initialization for a long time. The embedding table therefore learns at wildly different rates across rows, governed by the long-tailed frequency distribution of language.

Second, it interacts with the optimizer in ways that bite. Momentum and Adam-style optimizers keep running statistics per parameter. A naive dense implementation would apply weight decay and momentum to every row every step — including rows with zero gradient — slowly decaying the embeddings of tokens that simply have not appeared yet, which is usually not what you want. Correct sparse handling updates the optimizer state only for the rows that were touched. Third, the sparsity is a genuine efficiency win: the backward pass for the embedding scatters gradients into a handful of rows rather than writing the full V × d_model matrix, which for a large vocabulary is a big saving on both compute and memory traffic.

Scaling embeddings by sqrt(d_model)

Read the original Transformer paper closely and you find a line that trips people up: ‘in the embedding layers, we multiply those weights by √d_model.’ So the vector actually handed to the first layer is not E[i] but √d_model · E[i]. Why scale up by a factor of, say, √512 ≈ 22.6?

The reason is a magnitude match. Embeddings are commonly initialized with a small variance — on the order of 1/d_model, so each component has standard deviation around 1/√d_model. Multiplying by √d_model rescales the components to roughly unit standard deviation, which puts the token embedding on the same scale as the sinusoidal positional encodings that are about to be added to it (those live in [-1, 1]). Without the boost, a tiny token vector would be swamped by the positional signal. There is a second motive when the embedding is tied to the output layer: the same matrix must serve as an input table (wants small values) and as an output projection producing logits (wants a controlled scale), and the √d_model factor helps reconcile the two. It is a small constant, but drop it and early training can behave very differently. Not every implementation uses it — it depends on the initialization and whether weights are tied — so always check what a given codebase does.

Weight tying: the output layer is the input table

At the other end of the model, after the final transformer block, the hidden state h of shape [N, d_model] has to be turned back into a distribution over the vocabulary. That needs a projection of shape [d_model, V] producing logits of shape [N, V], one score per vocabulary token. Notice the shape: [d_model, V] is just E transposed. The input embedding and the output projection want to be the same size — and, it turns out, the same numbers.

input:   X      = E[ids]        E  : [V, d_model]
output:  logits = h · E^T     E^T: [d_model, V]
                                logits : [N, V]

Weight tying (Press & Wolf; Inan et al.) makes the output projection literally reuse E instead of learning a separate matrix: the logit for token v is the dot product of the final hidden state with that token’s embedding row. This is intuitive — ‘how much does the model’s output point at token v?’ becomes ‘how aligned is h with v’s embedding?’ — and empirically it tends to improve perplexity, because the same word-geometry is learned once from both directions instead of twice from half the signal.

Tying as a pointer, not a copy

The word ‘tying’ is worth taking literally. In implementation, weight tying is not copying E into a second tensor and keeping them in sync — it is making the output layer’s weight be the same tensor object as the input embedding. The output layer holds a pointer to E, so there is exactly one array of numbers in memory, referenced from two places in the computation graph. Reading it forwards gathers rows; reading it backwards, transposed, produces logits.

Two consequences follow. The obvious one is memory: you delete an entire V × d_model matrix. On GPT-2 small that is the ~38 M output-projection parameters gone — on a tiny CPU SLM where embeddings are the dominant cost, tying can shrink the whole model by a third with no architectural loss. The subtler one is gradients: because both uses share one tensor, the backward pass accumulates gradient into E from both ends — the sparse input-side gather and the dense output-side projection. That output-side term is not sparse, so a tied embedding gets a gradient contribution for every row every step from the logits, on top of the sparse input contribution. Tying changes not just the parameter count but the gradient dynamics of the table.

A worked example, by hand

Make it concrete with a toy model: vocabulary V = 6, width d_model = 4. The embedding matrix is 6 × 4:

       d0     d1     d2     d3
E[0]  0.10  -0.20   0.30   0.05   <- token 0
E[1]  0.40   0.10  -0.10   0.20   <- token 1
E[2] -0.30   0.25   0.15  -0.05   <- token 2
E[3]  0.05   0.60  -0.40   0.10   <- token 3
E[4]  0.20  -0.10   0.00   0.35   <- token 4
E[5] -0.15   0.05   0.45  -0.20   <- token 5

Embed the id sequence [3, 1, 3] (N = 3). The lookup gathers rows 3, 1, 3:

X = E[[3,1,3]] =
  [ 0.05   0.60  -0.40   0.10 ]   (row 3)
  [ 0.40   0.10  -0.10   0.20 ]   (row 1)
  [ 0.05   0.60  -0.40   0.10 ]   (row 3)     X : [3, 4]

Check the one-hot view for the first position: e_3 = [0,0,0,1,0,0], and e_3 · E sums 0·E[0] + … + 1·E[3] + … = E[3] = [0.05, 0.60, -0.40, 0.10] — the same row. Positions 0 and 2 are identical because the token is the same; the lookup carries no notion of order. And the sparse gradient is visible here too: this batch used only rows 1 and 3, so ∂L/∂E is nonzero only on those two rows; rows 0, 2, 4, 5 get exactly zero this step.