PyTorch nn.Embedding: shapes, padding, and gradients

PyTorch nn.Embedding maps integer IDs to rows in a weight matrix. With num_embeddings=4 and embedding_dim=3, the matrix contains four vectors with three values each. Looking up an ID selects its vector; a batch of IDs produces a batch of vectors.

This guide tests the layer itself: lookup values, output shapes, repeated-index gradients, padding, frozen weights, and max_norm. We ran 41 checks using Python 3.12.14 and PyTorch 2.15.0+cpu. The six examples use fixed weights and include captured output. They do not measure GPU performance or the quality of learned representations.

For the role of embeddings within a larger model, see our neural network architecture guide. Here, the goal is to know exactly which rows the layer reads and what can change those rows.

Read the weight shape and output shape

We supply a small matrix through from_pretrained so every lookup can be checked by eye. The name of that constructor does not imply that these particular numbers were learned; they are an original test fixture.

import torch

weights = torch.tensor([[0., 0., 0.], [1., 2., 3.],
                        [4., 5., 6.], [7., 8., 9.]])
layer = torch.nn.Embedding.from_pretrained(weights.clone(), freeze=False)
ids = torch.tensor([[2, 1], [0, 2]], dtype=torch.long)
output = layer(ids)
print('Weight shape:', tuple(layer.weight.shape))
print('Output shape:', tuple(output.shape))
print('Output:', output.tolist())

Weight shape: (4, 3)
Output shape: (2, 2, 3)
Output: [[[4.0, 5.0, 6.0], [1.0, 2.0, 3.0]], [[0.0, 0.0, 0.0], [4.0, 5.0, 6.0]]]

The input has shape (2, 2). The output keeps those two axes and appends a vector width of three, giving (2, 2, 3). ID 2 selects the third row, [4, 5, 6], both times it occurs. ID 0 selects the first row.

For token batches shaped (batch, sequence_length), the same rule gives (batch, sequence_length, embedding_dim). The Embedding API reference documents this shape contract. Our checks also covered a scalar ID, a one-dimensional list, and an empty ID axis. An independent one-hot matrix multiplication matched every value in this plain lookup example.

The layer needs a consistent ID-to-row mapping. If your vocabulary changes order, the same integer will select a different token’s vector. Keep that mapping with the model configuration and checkpoint. Inspect the ID tensors produced by your PyTorch DataLoader before they reach the layer.

Check index types and bounds first

These examples use torch.long. Both int64 and int32 IDs worked in our CPU checks; floating-point IDs were rejected. The valid IDs for the four-row matrix are 0 through 3. Both -1 and 4 raised errors.

An unknown token needs an explicitly assigned valid ID. The layer does not automatically create an extra row for an unseen value or translate strings into IDs. The integer ID tensor is also not the parameter being learned: gradients flow to the floating-point weight matrix.

A negative padding_idx has a separate meaning at construction. Our padding_idx=-1 check selected the last row as padding, but an input tensor containing -1 still failed. Do not use Python-style negative indexing in the input tensor.

See how repeated IDs affect gradients

import torch

weights = torch.tensor([[0., 0.], [1., 2.], [3., 4.]])
ids = torch.tensor([1, 1, 2], dtype=torch.long)
plain = torch.nn.Embedding.from_pretrained(weights.clone(), freeze=False)
scaled = torch.nn.Embedding.from_pretrained(
    weights.clone(), freeze=False, scale_grad_by_freq=True)
plain(ids).sum().backward()
scaled(ids).sum().backward()
print('Default gradient:', plain.weight.grad.tolist())
print('Frequency-scaled gradient:', scaled.weight.grad.tolist())

Default gradient: [[0.0, 0.0], [2.0, 2.0], [1.0, 1.0]]
Frequency-scaled gradient: [[0.0, 0.0], [1.0, 1.0], [1.0, 1.0]]

The loss is the sum of all retrieved values, so each occurrence contributes a gradient of one to every component of its row. ID 1 occurs twice, giving [2, 2]; ID 2 occurs once, giving [1, 1]. Row 0 is unused and receives zeros.

With scale_grad_by_freq=True, the repeated row’s gradient is divided by its frequency in this input. Both used rows then have [1, 1]. Our checks confirmed that the forward values stayed identical. This option changes the backward calculation; it does not change which vector is retrieved.

Each example creates fresh layers. In a training loop, clear accumulated gradients at the intended step boundary. The counts above describe one backward call with this sum loss, not every possible loss or accumulation scheme.

Test padding with an explicit optimizer setup

import torch

layer = torch.nn.Embedding(4, 2, padding_idx=0)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([[0., 0.], [1., 2.],
                                     [3., 4.], [5., 6.]]))
ids = torch.tensor([0, 1, 1, 2], dtype=torch.long)
optimizer = torch.optim.SGD(layer.parameters(), lr=0.1,
                            momentum=0, weight_decay=0)
optimizer.zero_grad(set_to_none=True)
output = layer(ids)
output.sum().backward()
gradient = layer.weight.grad.clone()
before = layer.weight.detach().clone()
optimizer.step()
print('Padding output:', output[0].tolist())
print('Gradient:', gradient.tolist())
print('Padding row unchanged:', torch.equal(before[0], layer.weight[0]))
print('Unused row unchanged:', torch.equal(before[3], layer.weight[3]))
print('Used row changed:', not torch.equal(before[1], layer.weight[1]))

Padding output: [0.0, 0.0]
Gradient: [[0.0, 0.0], [2.0, 2.0], [1.0, 1.0], [0.0, 0.0]]
Padding row unchanged: True
Unused row unchanged: True
Used row changed: True

The padding row returns zeros and receives no lookup gradient. Row 1 receives two contributions, row 2 one contribution, and the unused row 3 receives none. With the specified SGD setup, both padding and unused rows stay unchanged while a used row changes. We also verified the entire update against before - 0.1 * gradient.

A zero padding vector does not automatically remove padded positions from downstream reductions. For the output above, an ordinary mean across all four positions was [1.25, 2.0]. Selecting only positions whose IDs are nonzero before averaging gave approximately [1.667, 2.667]. The first calculation divides by four; the second divides by three.

Choose masking or sequence handling appropriate to the next layer and the loss. The embedding’s padding rule alone does not prove that the whole model ignores padding.

Do not confuse zero lookup gradient with an immutable row

This example deliberately supplies a nonzero padding vector and enables SGD weight decay:

import torch

weights = torch.tensor([[5., 5.], [1., 2.]])
layer = torch.nn.Embedding.from_pretrained(
    weights.clone(), freeze=False, padding_idx=0)
optimizer = torch.optim.SGD(layer.parameters(), lr=0.1,
                            momentum=0, weight_decay=0.1)
optimizer.zero_grad(set_to_none=True)
layer(torch.tensor([0, 1], dtype=torch.long)).sum().backward()
padding_gradient = layer.weight.grad[0].clone()
optimizer.step()
print('Padding gradient:', padding_gradient.tolist())
print('Padding weights after decay:',
      [round(v, 3) for v in layer.weight[0].tolist()])

Padding gradient: [0.0, 0.0]
Padding weights after decay: [4.95, 4.95]

The lookup gradient for row 0 is zero, but the stored values move from 5 to 4.95. With these settings, each component changes by 0.1 * 0.1 * 5. This is an optimizer update, separate from the gradient contributed by the lookup.

The SGD reference describes its weight-decay term. This test uses dense gradients, zero momentum, and one update. It does not establish what every optimizer or optimizer state will do. If an exact padding value is required, verify it after the updates your application actually performs.

Supplying weights with from_pretrained preserves the supplied padding row in this example. A newly constructed nn.Embedding initializes its designated padding row to zero, which we tested separately. The functional embedding reference also distinguishes lookup behavior from initialization.

Make the freezing choice explicit

import torch

weights = torch.tensor([[1., 2.], [3., 4.]])
frozen = torch.nn.Embedding.from_pretrained(weights.clone())
trainable = torch.nn.Embedding.from_pretrained(weights.clone(), freeze=False)
ids = torch.tensor([1], dtype=torch.long)
frozen_output = frozen(ids)
trainable_output = trainable(ids)
trainable_output.sum().backward()
print('Frozen weight requires grad:', frozen.weight.requires_grad)
print('Frozen output requires grad:', frozen_output.requires_grad)
print('Trainable weight requires grad:', trainable.weight.requires_grad)
print('Trainable gradient:', trainable.weight.grad.tolist())

Frozen weight requires grad: False
Frozen output requires grad: False
Trainable weight requires grad: True
Trainable gradient: [[0.0, 0.0], [1.0, 1.0]]

from_pretrained freezes its weight matrix by default. In this isolated lookup, the frozen weight and output do not require gradients. Setting freeze=False enables gradients, and only the selected row receives them in the example.

We pass weights.clone() to each layer so they have separate storage from the fixture. Freezing controls gradient tracking for that parameter; it does not make the tensor read-only. The next example shows why that distinction matters.

When saving a model, retain the vocabulary mapping and embedding configuration alongside the weights. Our PyTorch save-and-load guide covers model state and training checkpoints. A correctly restored matrix still needs the same meaning assigned to each row.

Account for max_norm changing stored weights

import torch

weights = torch.tensor([[3., 4.], [0., 2.], [6., 8.]])
layer = torch.nn.Embedding.from_pretrained(
    weights.clone(), freeze=True, max_norm=1.0)
before = layer.weight.detach().clone()
output = layer(torch.tensor([0, 1], dtype=torch.long))
rounded = lambda rows: [[round(v, 3) for v in row] for row in rows]
print('Output:', rounded(output.tolist()))
print('Stored weights:', rounded(layer.weight.tolist()))
print('Frozen:', not layer.weight.requires_grad)
print('Stored weights changed:', not torch.equal(before, layer.weight))

Output: [[0.6, 0.8], [0.0, 1.0]]
Stored weights: [[0.6, 0.8], [0.0, 1.0], [6.0, 8.0]]
Frozen: True
Stored weights changed: True

The requested vectors initially have lengths five and two. The forward call reduces their norms to approximately one and writes the adjusted vectors back into the weight matrix. Row 2, whose vector has length ten, was not requested and stays [6, 8].

This happened with freeze=True. No backward call or optimizer step was involved. If you require supplied vectors to remain unchanged during lookup, do not enable this behavior without accounting for the mutation.

The API reference warns about in-place modification when max_norm is set, including interactions with differentiable operations that use the weight tensor before lookup. Follow that guidance if your model shares or transforms the weights. Our test verifies stored values in a frozen layer; it does not audit every shared-weight autograd graph.

Use small fixtures before training a larger model

Check one known ID, one repeated ID, a padding position, and an invalid index. Confirm the output shape and inspect the gradient of a used, unused, and padding row. Then verify the stored weights after your actual optimizer step and any max_norm lookup.

The examples here passed on CPU, including negative tests for index type and bounds. We did not test sparse-gradient optimizers, distributed execution, GPU speed, or semantic similarity. Those are separate questions from whether a layer reads and updates the intended rows.