In §1.1 we decided to give every item a learned vector, stored as row $i$ of a matrix $\mathbf{E} \in \mathbb{R}^{V \times d}$, where $V$ is the vocabulary size and $d$ the embedding dimension. This section answers the engineering questions. How does a network use that matrix? How do gradients flow into it? Why is the PyTorch layer nn.Embedding an indexing operation rather than a matrix multiply? How big should the numbers be, and how big should $d$ be?

The punchline fits in one line, and it's worth burning into memory. With $\mathbf{e}_i \in \mathbb{R}^V$ the one-hot vector for item $i$,

In plain words, multiplying a one-hot row vector by the embedding matrix just picks out row $i$.

One-hot times a matrix is a lookup

Let's prove it. The $j$-th entry of the product $\mathbf{e}_i^\top\mathbf{E}$ is a sum over the $V$ rows,

Every term vanishes except $k = i$. So the "embedding layer" is exactly a linear layer with no bias applied to a one-hot input. We just never build the one-hot vector, because indexing gives the same answer while doing $d$ copies instead of $Vd$ multiply-adds per token.

This matters conceptually. An embedding is not a special kind of object. It's an ordinary weight matrix, trained by ordinary backpropagation, with whatever loss sits downstream.

Let's confirm the equality and measure the speed difference.

The exact timings depend on your machine, but the gather is typically one to two orders of magnitude faster here, and the gap grows linearly with $V$.

Batches and sequences

In practice, ids has a shape like (B, T), meaning a batch of $B$ sequences, each $T$ tokens long. Indexing broadcasts over any leading shape and appends the embedding dimension,

In plain words, every integer in the batch is replaced by its row, so a (B, T) grid of IDs becomes a (B, T, d) block of vectors. That one line, H = E[X], is the first layer of essentially every language model. In GPT-2 it's called wte ("word token embedding"), and in Llama it's embed_tokens. You'll write it yourself when you build a GPT.

The backward pass is sparse, and it adds up

Now the part that trips people up. Suppose the rest of the network produces a loss $\mathcal{L}$, and backprop has already computed the gradient with respect to the embedding outputs, $\mathbf{G} = \partial\mathcal{L}/\partial\mathbf{H}$, one $d$-vector per token. What is $\partial\mathcal{L}/\partial\mathbf{E}$?

Flatten the batch into a list of $N$ token positions with IDs $x_1, \dots, x_N$. Each output row is $\mathbf{h}_n = \mathbf{E}_{x_n,:}$, so row $i$ of $\mathbf{E}$ influences exactly those positions where $x_n = i$, and it influences each of them with Jacobian $\partial\mathbf{h}_n/\partial\mathbf{E}_{i,:} = \mathbf{I}$ (it is copied verbatim). By the chain rule, summing over every path from $\mathbf{E}_{i,:}$ to the loss,

where $\mathbf{g}_n$ is row $n$ of $\mathbf{G}$. In plain words, the gradient for word $i$'s vector is the sum of the upstream gradients at every position where word $i$ appears. Words that don't appear in the batch get exactly zero.

A worked example you can check by hand

Take a tiny table with $V = 4$, $d = 3$, a batch of IDs $[2, 0, 2]$, and upstream gradients

Row 2 appears at positions 1 and 3, so its gradient is $[1, 0, -1] + [2, 1, 0] = [3, 1, -1]$. Row 0 appears once and gets $[0.5, 0.5, 0.5]$. Rows 1 and 3 never appear and get $[0, 0, 0]$. The full gradient is

Two consequences follow.

1. The gradient is sparse. In a batch of 4,096 tokens from a 128k vocabulary, at most 4,096 rows (usually far fewer) are nonzero. PyTorch can exploit this with nn.Embedding(..., sparse=True).
2. Repeated IDs must accumulate. "the" may appear 300 times in a batch, and its row must receive the sum of all 300 gradients. This is where a very common bug lives.

The fancy-indexing trap

In NumPy (and PyTorch), G_E[ids] += grads does not accumulate duplicates. With repeated indices, each write overwrites the previous one, and only one survives. You need np.add.at (in PyTorch, index_add_ or scatter_add_). Here is the worked example above, done both ways.

Row 2 should be $[3, 1, -1]$. The buggy version keeps only the last write, $[2, 1, 0]$. A bug like this doesn't crash anything. Training just quietly gets worse, with frequent words learning too slowly, and it is exactly the kind of bug an interviewer loves.

import numpy as np

def embedding_forward(E, ids):
    """E: (V, d), ids: int array of any shape -> ids.shape + (d,)"""
    raise NotImplementedError

def embedding_backward(ids, grad_out, V, padding_idx=None):
    """dL/dE: (V, d). grad_out has shape ids.shape + (d,)."""

Now in PyTorch

nn.Embedding is exactly the code you just wrote, plus a few production options that are worth knowing cold.

Bags of IDs, and their backward pass

Recommenders and bag-of-words models often embed a variable-length list of IDs per example (the items a user clicked, or the words in a review) and pool them into one vector. Storing ragged lists as a padded matrix wastes memory, so the standard format (used by nn.EmbeddingBag) is a flat ids array plus an offsets array saying where each bag starts. For example, ids = [4, 1, 1, 7, 2] with offsets = [0, 3] means bag 0 is [4, 1, 1] and bag 1 is [7, 2].

With mean pooling, bag $b$ with members $x_n$ for $n$ in a set $\mathcal{B}_b$ of size $\lvert\mathcal{B}_b\rvert$ produces $\mathbf{o}_b = \frac{1}{\lvert\mathcal{B}_b\rvert}\sum_{n\in\mathcal{B}_b}\mathbf{E}_{x_n,:}$. By the same chain-rule argument as before, each member receives the bag's upstream gradient divided by the bag size,

In plain words, every time item $i$ sits in a bag, it collects that bag's gradient, scaled by one over the bag size. In the example above, if the bag gradients are $\mathbf{q}_0$ and $\mathbf{q}_1$, row 1 (which appears twice in bag 0) gets $\tfrac{2}{3}\mathbf{q}_0$, and row 7 gets $\tfrac12\mathbf{q}_1$.

import numpy as np

def embedding_bag_mean(E, ids, offsets):
    """E: (V, d); ids: (N,) flat ids; offsets: (B,) start index of each bag -> (B, d)"""
    raise NotImplementedError

import numpy as np

def embedding_bag_backward(ids, offsets, grad_out, V, mode="mean"):
    """ids: (N,); offsets: (B,); grad_out: (B, d) -> dL/dE: (V, d)"""
    raise NotImplementedError

One table, two jobs, weight tying

A language model ends with an output layer that turns its final hidden vector $\mathbf{h} \in \mathbb{R}^d$ into one score per token, $\mathbf{z} = \mathbf{W}_{\text{out}}\mathbf{h}$ with $\mathbf{W}_{\text{out}} \in \mathbb{R}^{V\times d}$. That matrix has exactly the shape of the input table, and its row $j$ plays a similar role, a vector that "means" token $j$. Weight tying (Press & Wolf, 2017; Inan et al., 2017) sets $\mathbf{W}_{\text{out}} = \mathbf{E}$, so one table both embeds the input and scores the output. GPT-2 and the original Transformer tie, while many recent large models (Llama 3 8B and 70B among them) keep the two separate. We'll come back to it in §1.5, where word2vec's two tables are the same question in miniature.

Tying changes the backward pass in a way people often get wrong. The table now sits on two paths to the loss, and the chain rule says to add the contributions of all paths.

Take the smallest possible tied model, a bigram predictor. Position $n$ has input token $x_n$ and target token $y_n$. The model looks up $\mathbf{h}_n = \mathbf{E}_{x_n,:}$, computes logits $\mathbf{z}_n = \mathbf{E}\mathbf{h}_n \in \mathbb{R}^V$, and pays cross-entropy $\mathcal{L} = \frac1N\sum_n -\log\operatorname{softmax}(\mathbf{z}_n)_{y_n}$. With $\boldsymbol{\Delta} \in \mathbb{R}^{N\times V}$ holding the familiar "probabilities minus one-hot target" rows divided by $N$ (§0.2), the two paths give

where $\mathbf{H} \in \mathbb{R}^{N\times d}$ stacks the $\mathbf{h}_n$ and $\mathbf{S}$ is the one-hot selection matrix from the derivation above. In plain words, every row gets gradient from its role as a possible output (every row, every step, because the softmax touches the whole vocabulary) plus, for the rows that were looked up, gradient from its role as an input. That first term is why a tied table's gradient is dense even when the input lookups are sparse.

import numpy as np

def tied_bigram_loss_grad(E, x, y):
    """E: (V, d); x, y: (N,) -> (mean cross-entropy, dL/dE of shape (V, d))"""
    raise NotImplementedError

In PyTorch, tying is one assignment, and autograd adds the two paths for you.

Initialization, or how big the numbers should be

nn.Embedding initializes every entry from $\mathcal{N}(0, 1)$. Big models usually use something smaller. GPT-2 initializes token embeddings with standard deviation $0.02$. The original Transformer (Vaswani et al., 2017) multiplies the embedding output by $\sqrt{d}$ before adding positional encodings. Why does scale matter? Because the next layers see these numbers. If an embedding row $\mathbf{h}$ has $d$ independent entries of mean 0 and standard deviation $\sigma$, its expected squared norm is

so its length is about $\sqrt d\,\sigma$. In plain words, the vector gets longer with the square root of the dimension. With $\sigma = 1$ and $d = 4096$, every vector has length about 64, which can saturate whatever comes next. With $\sigma=0.02$ the length is about 1.3, and for GPT-2 small ($d = 768$) about 0.55.

The $\sqrt{d}$ multiplier goes the other way and scales the vectors up. The paper does not explain why. A common reading is that the shared input and output table is initialized with standard deviation about $d^{-1/2}$, which keeps output logits at a sensible scale, and multiplying by $\sqrt{d}$ on the way in brings the input entries back to roughly unit size, comparable to the sinusoidal position encodings (whose entries lie in $[-1, 1]$) they are added to. Treat that as a plausible rationale rather than established fact. We'll come back to scale in §5.2 (normalization) and §8.8 (initialization at scale).

Choosing the embedding dimension $d$

There is no formula, but there are good mental models.

• Capacity. $d$ is how many soft features each item can have. Too small and distinct items are forced to overlap. Too large and you waste parameters and memory, and with little data you overfit. Yin and Shen (2018) make this a bias and variance trade-off and show that, for count-based embeddings, the error is minimized at a finite $d$ that depends on the noise in the data.
• Geometry says you need fewer dimensions than you'd think. The Johnson–Lindenstrauss lemma (§1.3) says $n$ points can be placed in $O(\log n / \varepsilon^2)$ dimensions while keeping all pairwise distances within a factor $1 \pm \varepsilon$. That's logarithmic in the number of items, not linear.
• Rules of thumb, which disagree. Google's TensorFlow feature-column guide suggests $d \approx V^{1/4}$ for categorical features. The fastai library's default is $\min(600, \operatorname{round}(1.6\,V^{0.56}))$. For $V = 10{,}000$ categories the first gives 10 and the second 278. A factor of 28 between two popular heuristics tells you what they are, which is starting points for a sweep, not laws.
• What real models use. Word2vec and GloVe used $d$ from 50 to 300. BERT-base uses 768. Llama-3-8B uses 4096, the same as the width of the whole network.

The Llama row is striking. In an 8B-parameter model, more than a billion parameters live in the input and output token tables. GPT-3's table is larger than Llama's in absolute terms, yet it is only 0.35% of GPT-3's 175 billion parameters, because the rest of a model grows with its depth and with the square of its width while the table grows only linearly with the width. For small models the share is even larger, which is one reason people tie the input and output tables, or factorize them. ALBERT (Lan et al., 2020) replaces the $V\times d$ table with a product of a $V\times m$ table and an $m\times d$ projection. With $V = 30{,}000$, $m = 128$, and $d = 4{,}096$, that is $3.84\text{M} + 0.52\text{M} \approx 4.4$ M parameters instead of 123 M.

Memory during training is larger still. With Adam in float32, every parameter carries its value, its gradient, and two moment estimates, so 16 bytes. Llama-3-8B's input table alone then needs $525\text{M} \times 16 \approx 8.4$ GB. Where the rest of a model's training memory goes is the subject of §10.1.

def table_cost(V, d, tied, factor_m=None):
    """-> (number of parameters, bytes for weights + grads + two Adam moments in float32)"""
    raise NotImplementedError

Read more

• PyTorch documentation. torch.nn.Embedding. Initialization, padding_idx, max_norm, and which optimizers support sparse gradients.
• Raschka (2024). Build a Large Language Model (From Scratch), Chapter 2. Token embeddings, positions, and a notebook showing that an embedding equals one-hot times a matrix.
• Press & Wolf (2017). Using the Output Embedding to Improve Language Models. Weight tying, and why the output table is the better one to keep.
• Lan et al. (2020). ALBERT. Factorized embedding parameterization, separating the vocabulary table from the hidden width.
• Yin & Shen (2018). On the Dimensionality of Word Embedding. A principled bias and variance view of choosing $d$.
• Karpathy. Neural Networks, Zero to Hero. The makemore lectures build an embedding table live.

Watch next

▶