Embedding Layers as Learned Lookup Tables

An embedding layer is usually introduced as a way to turn a token id into a vector, as if the layer performed some learned transformation of the id. It does not. The id is an index, the vector is a row of a matrix, and the only thing learned is the contents of that matrix. Nearly every property that matters later, including memory cost, update frequency, and the behaviour of rare inputs, follows from that single fact.

The matrix behind the API

An embedding module owns one parameter tensor of shape vocab × dim. The forward pass takes a tensor of integer ids and returns the corresponding rows. This is exactly equivalent to multiplying a one-hot vector by the matrix, but the equivalence is only conceptual: the implementation reads dim contiguous floats per id instead of performing vocab × dim multiply-adds. For a vocabulary of 50,000 and a width of 768, the one-hot formulation would touch 38.4 million weights per token while the lookup touches 768.

import torch
from torch import nn

table = nn.Embedding(1000, 64)
ids = torch.tensor([3, 17, 3])
rows = table(ids)                      # gathers rows 3, 17, 3

one_hot = torch.zeros(3, 1000)
one_hot[torch.arange(3), ids] = 1.0
same = one_hot @ table.weight          # identical values, 3 * 1000 * 64 products

table.weight.grad = None
rows.sum().backward()
touched = (table.weight.grad.abs().sum(dim=1) > 0).sum()
print(touched)                         # 2: only rows 3 and 17 moved

The ids themselves carry no arithmetic meaning. Id 17 is not “between” 16 and 18, and no part of the layer assumes otherwise. Positional or structural information has to come from somewhere else, which is why sequence models add a separate positional encoding and graph models add structural features.

The width dim is a capacity choice, not a property of the data. A wider table can represent finer distinctions but costs memory linearly and gives each row more parameters to fit from the same number of observations. Common designs use a few hundred dimensions for token tables in transformer stacks whose residual width is larger, and the table is often the single largest parameter tensor in a small model.

Gradients arrive one row at a time

Because the forward pass reads only the rows that appear in the batch, the backward pass writes only those rows. The gradient of the loss with respect to the table is a matrix that is zero everywhere except for the rows whose ids occurred. With a batch of 32 sequences of 512 tokens, at most 16,384 rows are touched out of a vocabulary that may hold millions, and in practice far fewer distinct ids occur.

This sparsity has consequences beyond speed. An optimizer that maintains per-parameter state, such as Adam with its first and second moment estimates, allocates state for the entire table whether or not a row is updated. A row that appears once in a training run still occupies its share of optimizer memory, and its moment estimates remain at their initial values until that single update. Framework support for sparse gradients exists precisely for this case: the optimizer applies the update only to indices present in the gradient, which changes both the arithmetic and, for some optimizers, the semantics of bias correction and weight decay.

Token frequency is heavy-tailed. Word frequencies in natural text follow a distribution close to Zipf’s law, where frequency is roughly proportional to the inverse of rank, and subword vocabularies inherit a similar skew. The most frequent few hundred tokens may account for a large fraction of all occurrences, while tens of thousands of tail tokens appear rarely. A tail row therefore receives a handful of gradient updates over an entire training run, and its final vector is dominated by initialization and by those few contexts. This is the mechanism behind the observation that rare-word embeddings are poorly estimated, and it is also why the geometry of the table at initialization matters more for the tail than for the head.

Padding is a special case worth handling explicitly. A padding id that is allowed to receive gradients will drift into a vector that encodes “absence”, which is harmless until the same row is reused as a real token or until a downstream component compares padding positions against content positions. Frameworks provide a padding index that pins that row to a fixed vector and excludes it from the gradient.

Memory is vocabulary times dimension

The table’s size in bytes is vocab × dim × 4 for 32-bit floats. A 50,000-token vocabulary at width 768 is about 154 MB. Adding Adam’s two moment buffers triples that to roughly 460 MB, and mixed-precision training that keeps a 32-bit master copy pushes it higher still. None of these figures depend on batch size, sequence length, or how often a token is used, which is what makes the table a fixed tax on the model rather than a variable one.

Multiple tables multiply the tax. A language model with a separate output projection has two large matrices over the vocabulary. A multilingual model may hold one table per script or per language family. A recommendation model may hold one table per categorical feature, and the combined vocabulary of user ids, item ids, and context features can dwarf the dense layers. Memory planning for such models is mostly a matter of counting rows.

Several techniques reduce the count. Factoring the table into two smaller matrices whose product approximates it trades expressiveness for memory, and is most useful when the vocabulary is large and the effective rank of the learned table is low. Hashing several ids into one row, with a per-id sign or projection to disambiguate, bounds memory regardless of vocabulary size at the cost of collisions. Sharding splits rows across devices so that each device stores a slice, which requires an all-to-all gather when a batch’s ids are spread unevenly. All three accept some loss of fidelity in exchange for a bound on the largest tensor in the model.

Tying the input and output matrices

A language model projects the final hidden state back to vocabulary size with a dim × vocab matrix. That matrix has the same shape, transposed, as the embedding table, and tying the two so they share storage halves the vocabulary parameters. The published motivation for tying is that it improves perplexity, particularly for rare words, because the same vector serves as both the input representation of a token and the direction that the output layer scores for that token. The gradient of a token’s output logit and the gradient of its input row then accumulate into one parameter rather than two.

class TiedHead(nn.Module):
    def __init__(self, vocab, dim):
        super().__init__()
        self.embed = nn.Embedding(vocab, dim)
        self.proj = nn.Linear(dim, vocab, bias=False)
        self.proj.weight = self.embed.weight   # shared storage, not a copy

    def forward(self, ids):
        h = self.embed(ids)
        return self.proj(h)                    # logits = h @ embed.weight.T

Tying constrains the geometry: the output score for token t becomes the dot product between the hidden state and the input vector for t. That is a reasonable assumption for language modelling, where the same lexical identity appears on both sides, and a poor one when the input and output alphabets differ, as in translation between scripts or in a model that emits a different tokenization than it consumes. It also requires the two widths to match, which forces the residual width and the embedding width to be the same number.

Initialization sets the early geometry

With the usual small normal initialization, all rows start near the origin and pairwise dot products are near zero. Cosine similarity is then dominated by noise: two random rows have an angle close to 90 degrees but with substantial variance, and the ranking of similarities carries no information. Early in training the table behaves less like a metric space and more like a set of nearly independent codes that gradually acquire structure as co-occurring tokens are pulled together by the shared contexts in which they appear.

Normalization interacts with this. Applying layer normalization to the looked-up vector fixes its scale, so downstream layers see a consistent magnitude regardless of how large a row has grown. Without normalization, frequently updated rows tend to grow in norm, and a linear layer reading them sees an implicit frequency signal. That signal can be useful, but it is not what the layer was meant to encode, and it makes the table’s geometry depend on corpus statistics in a way that is hard to disentangle later.

Where the abstraction leaks

The lookup abstraction assumes a closed vocabulary fixed before training. Any token outside it has no row, so inference must map unknown input to a reserved id whose vector was trained on almost nothing, or fall back to a character-level or hashed representation. Subword tokenization reduces the frequency of this case by composing rare words from frequent fragments, but it raises the vocabulary size and shifts the long tail onto the fragments rather than removing it.

Vocabulary drift is the slower version of the same problem. Adding tokens after training requires new rows, which start at initialization while the surrounding rows are trained, and the new rows are effectively untrained for as long as the old ones saw data. Remapping ids between a tokenizer and a checkpoint is a silent failure mode: the table loads without error, every id is valid, and the model’s behaviour degrades in ways that look like a modelling problem rather than an indexing one.

The representations in the table are also not a general-purpose semantic space. They are whatever minimizes the training objective under the contexts that occurred. Two rows that are close in cosine similarity are close in the sense the objective found useful, which for a language model is largely substitutability in context. Reading more than that from the geometry requires the kind of evaluation discussed in Linear Probing and Fine-Tuning as Different Measurements, and the distance behaviour of such spaces is not Euclidean intuition, as Distance Concentration and Hubness in Embedding Spaces sets out.