A matrix multiply that nobody performs. The most common layer in modern models is a lookup wearing a linear algebra costume.
Overview
The problem
Neural networks consume numbers. A word, a user id or a product code is a symbol, and there's no real number to assign it — encoding "cat" as 7 and "dog" as 8 would tell the network they're adjacent. This Is a claim no one made.
One-hot encoding avoids that: a vector of zeros with a single 1 at the token's index. No false sorting, and no handy structure either. each pair of tokens is equally distant, and the vector is if the vocab.
An embedding layer maps each token to a dense learned vector instead. The vector has a few hundred dims rather than fifty thousand, and its contents are params, so training can put similar tokens near each other.
Embedding Layers
This module needs JavaScript: the
numbers are computed in the page rather than recorded.
Worth knowing
An embedding layer is a table with one row per token. Looking up a token is indexing that table.
Formally it is a one-hot vector times a weight matrix — and that product is precisely 'take one row'.
Nobody computes it that way. Doing so would touch every weight in the table to retrieve one row of it.
The rows are ordinary parameters and are learned by gradient descent. Only the rows that appear in a batch get a gradient.
Embedding Layers
The layer that turns discrete symbols into vectors. The optimisation that makes it hands-on.
The matmul that is a lookup
The formal meaning is a one-hot vector times'd by a weight mat. Write it out and the reason no one does it turns obvious: timesing a vector that is all zeros except one spot by a mat selects a single row and times all else by zero.
The readout does the math. At a 50,000-token vocab and 512 dims, the honest matmul touches 25.6 million numbers to pull 512 of them.
So each framework runs the layer as an indexing op. nn.Embedding in PyTorch and tf.keras.layers.Embedding are table lookups with a backward pass that scatters grads to the rows that were used.
That is also why only the rows appearing in a batch receive a grad. An embedding table is a very large param tensor that is almost fully idle on any given step, which has consequences for tuners keeping per-param state and for anyone writing spread training.
The table is the model, mostly
Drag the vocab and dim controls and watch the count. The table alone is vocab * dim params, and for large vocab lists that is a substantial fraction of the entire model. Often more than any other single layer.
Two standard responses:
Weight tying. In a language model the output layer also has one row per token. Sharing it with the input embedding halves the cost and often raises quality, since both are learning what tokens mean.
Subword tokenisation.Byte-pair encoding keeps the vocab in the tens of thousands rather than the millions a word-level vocab would need, and drops the out-of-vocab problem fully.
What the vectors learn
Nothing is imposed on the geometry. The rows start random and become what cuts the loss, and the handy structure is a consequence rather than a design.
Word2vec's famous math — king minus man plus woman — is that consequence showing. No one built the word puzzle in. It fell out of a training goal about predicting neighbouring words.
Which also means the structure reflects the training data, including its biases. Embeddings trained on a corpus inherit that corpus's links, and a great deal of fairness work in NLP is about that fact.
Beyond words
cat features in tabular models. A user id, a postcode, a product category — all embed, and this is why neural networks became competitive on tabular data with high-cardinality categories.
Spot-based embeddings in transformers. This Embed a spot rather than a symbol.
Recommendation. Users and items each get a table, and their dot product is a predicted bond. mat factorisation is just this.
Why the lookup replaces the matmul
The formal meaning of an embedding layer is a one-hot vector times a weight mat, and writing out the sizes shows why no one finds it that way. At a 50,000-token vocab and 512 dims:
operation
numbers touched
one-hot vector × weight matrix
50,000 × 512 = 25.6 million
the row it actually selects
512
The one-hot is all zeros but one. So The times reads 25.6 million values to return 512 and times all else by zero. each framework so runs the layer as an indexed table lookup, with a backward pass that scatters grads only to the rows a batch really used — which means the table is a huge param tensor that sits almost fully idle on any given step. The table itself is vocab × dim = 25.6 million params, often the largest single component of a model. Two economies attack that size. Weight tying reuses the input table as the output layer, halving the token params and often raising quality since both are learning what tokens mean. Subword tokenisation keeps the vocab in the tens of thousands rather than the millions a word-level scheme would need. The geometry the rows end up with — king − man + woman landing near queen — is never imposed. It falls out of the training goal. This Also means it inherits the training corpus's biases.
Because only the rows appearing in a batch receive a grad, an embedding table is a large param tensor that is almost fully idle on any given step, which complicates tuners that keep per-param state and anyone sharding a model across devices. The same lookup mechanism generalises well beyond words: a user id, a postcode or a product category embeds the same way. This Is what made neural networks competitive on tabular data with high-cardinality categories. A transformer embeds spots and tokens. And a recommender gives users and items each a table whose dot product is a predicted bond, which is mat factorisation wearing a different name. The nearness of two rows should be read as co-hit in the training data, not as meaning — a gap that matters wherever those links encode bias.
Where it goes wrong with Embedding Layers
Building the one-hot matmul truly. Correct, and enormously wasteful.
Choosing the dim by feel. A rule of thumb is the fourth root of the vocab size, and it's worth tuning.
Losing an out-of-vocab row. a thing has to handle unseen tokens; subword tokenisation avoids the question.
Reading likeness as meaning. Embedding nearness reflects co-hit in the training data, which isn't the same thing.
Questions people ask about Embedding Layers
What is an embedding layer? A trainable lookup table that maps each integer ID to a dense vector. The vectors are learned by backpropagation like any other weight.
How is it different from one-hot plus a linear layer? It is mathematically the same, but a lookup reads one row instead of multiplying by a mostly-zero matrix. It is far faster and uses less memory.
How large should the embedding dimension be? A rough rule is the fourth root of the vocabulary size, but it is tuned in practice. Common values are 64 to 1024 depending on the task.
What happens to unknown IDs? You must reserve an index for them, often 0 or a special unknown token. An out-of-range ID raises an error in most frameworks.
Do embeddings only work for words? No. Users, products, categories, genes and positions can all be embedded. Anything with a discrete ID is a candidate.
Should I use pretrained embeddings? They help when you have little data. With plenty of data, training from scratch or fine-tuning usually matches or beats them.
A lookup table that learns what things mean
An embedding layer is one-hot encoding times a weight mat, with the times skipped. This shows the equivalence, the cost saving, and what the learned vectors end up encoding.
example_01.pyNumPy
import numpy as np
rng = np.random.default_rng(0)
VOCAB, DIM = 8, 4
words = ["cat", "dog", "wolf", "car", "truck", "bus", "the", "a"]
E = rng.normal(0, 0.5, (VOCAB, DIM))
print("an embedding table: %d words x %d dimensions = %d parameters."
% (VOCAB, DIM, E.size))
print()
print("look up word 2 ('wolf') two ways:")
onehot = np.zeros(VOCAB); onehot[2] = 1
print(" one-hot @ E :", np.round(onehot @ E, 4))
print(" E[2] :", np.round(E[2], 4))
print(" identical:", np.allclose(onehot @ E, E[2]))
print()
print("that is the whole trick. multiplying by a one-hot vector selects one")
print("row, so you index instead. the maths is unchanged, the cost is not:")
for v, d in ((8, 4), (50_000, 256), (200_000, 1024)):
mults = v * d
print(" vocab %7d, dim %5d: one-hot matmul does %12d multiplications"
% (v, d, mults))
print(" %28s a lookup does %12d" % ("", d))
print()
print("the gradient is just as sparse. only the rows you used get updated:")
grad_out = np.zeros_like(E)
batch_ids = [2, 5, 2]
upstream = rng.normal(size=(3, DIM))
for i, wid in enumerate(batch_ids):
grad_out[wid] += upstream[i]
print(" a batch using word ids %s" % batch_ids)
print(" rows with a non-zero gradient:", np.where(np.abs(grad_out).sum(axis=1) > 0)[0])
print(" word 2 appeared twice, so its gradients ADD -- that is why a")
print(" frequent word's embedding moves further per batch than a rare one's,")
print(" and why rare words end up close to their initial random values.")
Output
The next block trains an embedding on a small vocabulary and shows how the learned vectors cluster by meaning.
example_02.pyNumPy
import numpy as np
rng = np.random.default_rng(0)
VOCAB, DIM = 8, 4
words = ["cat", "dog", "wolf", "car", "truck", "bus", "the", "a"]
groups = np.array([0, 0, 0, 1, 1, 1, 2, 2])
Emb = rng.normal(0, 0.3, (VOCAB, DIM))
W = rng.normal(0, 0.3, (DIM, 3))
for step in range(600):
logits = Emb @ W
p = np.exp(logits - logits.max(axis=1, keepdims=True))
p /= p.sum(axis=1, keepdims=True)
d = p.copy(); d[np.arange(VOCAB), groups] -= 1; d /= VOCAB
W -= 0.5 * (Emb.T @ d); Emb -= 0.5 * (d @ W.T)
print("predict group (animal, vehicle, article). accuracy: %.4f"
% ((Emb @ W).argmax(axis=1) == groups).mean())
print()
def cos(a, b):
return a @ b / (np.linalg.norm(a) * np.linalg.norm(b))
print("cosine similarity between the learned vectors:")
print("%8s %s" % ("", "".join("%8s" % w for w in words)))
for i, wi in enumerate(words):
print("%8s %s" % (wi, "".join("%8.3f" % cos(Emb[i], Emb[j])
for j in range(VOCAB))))
print()
print("nothing told the model that cat and dog are related. it only ever saw")
print("a group label. the similarity appeared because words that need the")
print("same output were pushed toward the same region of the space.")
print()
same = np.mean([cos(Emb[i], Emb[j]) for i in range(VOCAB) for j in range(VOCAB)
if i != j and groups[i] == groups[j]])
diff = np.mean([cos(Emb[i], Emb[j]) for i in range(VOCAB) for j in range(VOCAB)
if groups[i] != groups[j]])
print(" average similarity within a group : %+.4f" % same)
print(" average similarity across groups : %+.4f" % diff)
print()
print("that is what 'learned representation' means, and it is why embeddings")
print("transfer: the geometry is a by-product of the task, so a table trained")
print("on one task is often useful for another.")
Output
The next block shows how to choose the embedding dimension.
example_03.pyNumPy
import numpy as np
print("choosing the dimension is the only real decision. too small and")
print("distinct words are forced to share space; too large and rare words")
print("never see enough gradient to move. the usual rule of thumb is")
print("a common starting point is four times the fourth root of the vocabulary,")
print("then rounded to a power of two:")
for v in (1_000, 50_000, 500_000):
raw = 4 * v ** 0.25
print(" vocab %7d -> 4 * %5.1f = %5.1f, so 64 to 256 in practice"
% (v, v ** 0.25, raw))
print(" in real models the number is chosen by what fits in memory and what")
print(" the downstream layers expect, far more often than by any formula.")
Output
Check yourself
0 of 3
Answer without scrolling back up.
Why is an embedding layer implemented as a lookup rather than a matmul?
At 50,000 tokens and 512 dimensions that is 25.6 million numbers touched to get 512. The definitions are identical; only the implementation differs.
Which rows of an embedding table receive a gradient on a given step?
The backward pass scatters gradients to the rows that were used, which means a very large parameter tensor sits almost entirely idle on any step.
What does weight tying do?
Both layers have one row per token and both are learning what tokens mean, so sharing them is natural as well as cheaper.
Cheat sheet
Embedding Layers
Replace a one-hot vector with a learned embedding lookup, compare the memory and compute saved, and see why every NLP pipeline starts here.
DEEP LEARNING · vizlearn.in/deep_learning/embedding_layers.html
Ashish Jangra builds and maintains VizLearn. Every module here is written and the visualisation behind it hand-built, so the numbers in a readout come from the same code that draws the picture. Corrections are genuinely welcome and get priority over everything else — if a page states something wrong, or an animation misrepresents what the algorithm does, get in touch.