Embedding Layers

A token ID needs to lead to the vector associated with that token. An embedding layer provides the lookup table, and training adjusts its rows so the retrieved vectors become useful for the task.

13 sections · 10 questions · 3 coding challenges

1 / 13

The arrow keys move between sections too.

Could Related Words Share What the Model Learns?

BOTIE

A token ID tells us which item we have, but its numerical size says nothing about the item's meaning. If every word is treated as wholly separate, learning something from one sentence offers little help with a similar sentence containing a different word.

Bengio and colleagues' 2003 neural language model learned a vector for each word alongside the prediction model. Their motivation was the enormous number of possible word sequences. Related learned representations could allow evidence from one sequence to help with another that had never appeared in training.

An embedding layer stores those vectors in a table. A token ID selects a row, and training adjusts the row through the task's loss. Similarity is something the learning process may produce, not a property of adjacent IDs. The table gives the rest of the network useful numerical inputs without requiring us to hand-write a list of linguistic attributes.

Sources for this history

Bengio and colleagues (2003), A Neural Probabilistic Language Model

An ID Tells Us Which Word, but What Can We Do with It?

A token ID identifies a vocabulary entry. Two neighbouring IDs do not tell us that their words have similar meanings. We need another representation for the network to use.

An embedding layer gives each ID a row of numerical values that training can adjust.

Step 1 of 4

Treat the ID as an address

The word cat might receive ID 103. That number tells us which entry to select, much like a numbered drawer. Its size is not a measurement of the word.

1
Text piece
cat
2
Vocabulary address
103

An illustrative ID assignment.

Step 2 of 4

Open the corresponding row

The layer stores a table with one row per vocabulary entry. The ID selects its row. Lookup does not calculate the word's meaning from the digits in its ID.

1
Table entries
- Row 102
- Highlighted: Row 103
- Row 104
2
Lookup result
The values stored in row 103

A row is selected by identity.

Step 3 of 4

Pass the row's values into the network

The retrieved values form a vector. Later layers can compare and transform it. A sequence of token IDs becomes a sequence of vectors, with one position for each token occurrence.

Repeated IDs select the same table row.
Token position	IDs	Vectors
Position 1	103	Cat row
Position 2	104	Dog row
Position 3	103	Cat row

Position 1
IDs
103
Vectors
Cat row
Position 2
IDs
104
Vectors
Dog row
Position 3
IDs
103
Vectors
Cat row

Repeated IDs select the same table row.

Step 4 of 4

Let training adjust selected rows

The task's loss sends gradients back to the retrieved entries. Repeated occurrences contribute to the same row's update. Useful relationships emerge from training and data, not from the original numbering scheme.

1
Selected vector
Used by later layers
2
Loss gradient
Returns to the selected row

The task determines how the table changes.

An embedding is a learned table lookup

We can now inspect the table, output shape, and accumulation of gradients for repeated IDs.

A stored vector need not have a human meaning attached to each coordinate. Its usefulness depends on what the larger model learns to do with it.

Look up the row belonging to each token

Keep cat in both the first and third positions. Compare the two lookups, then examine how their gradients combine in the cat row. Replace one occurrence with mat to see which rows receive an update.

For this exercise, we assign a target vector to each word and take one training step from the same starting table. Real tasks usually supply prediction targets, from which embedding gradients are calculated.

Word 1 Word 2 Word 3

Learning rate: 0.1

Calculating the example…

Part 1. A Lookup Before the Network

1. Separate an identifier from a measurement

A tokenizer hands a model whole numbers. In this lesson’s small vocabulary cat is 0, dog is 1 and mat is 2. Those numbers say which word is meant and nothing else. Hand them to a neuron as they are and the arithmetic treats them as amounts. Mat would count for twice as much as dog, and cat, at zero, would add nothing to any weighted sum, although the three words could have been numbered in any other order.

We need numerical values that training can adjust independently of the arbitrary IDs. An embedding layer stores a vector for each ID in a table. Reading it is a lookup: ID zero selects row zero, and ID two selects row two. The selected values then enter the model's calculations, where a training signal can tell us how to change them.

For a small example, use three vocabulary entries: cat, dog, and mat. Assign each a two-number vector. The entries begin as chosen values here; another model might initialize them randomly or start from pretrained vectors.

ID  Word  Starting vector
0   cat   [ 0.2,  0.0]
1   dog   [-0.2,  0.0]
2   mat   [ 0.0, -0.2]

Input IDs: [0, 1, 0]
Lookups:   [[0.2, 0.0], [-0.2, 0.0], [0.2, 0.0]]

No arithmetic happens in the lookup. Row 0 is copied out wherever ID 0 appears, so three IDs become three vectors of two numbers each. The layers that follow work with those vectors and not with the IDs.

The two cat occurrences initially receive exactly the same vector. This layer has not read their neighbors. Later layers, such as attention, can use the context to produce different representations for those occurrences.

2. Give the table a reason to change

A lookup alone cannot discover useful relationships. The network must make a prediction, compare it with a target, and calculate how its error depends on the embedding values. Those derivatives tell us how to update the table.

To make that update visible without introducing a second network, the playground supplies a target vector for each word. These targets are chosen for the exercise. Their coordinates are not measured properties of cats or dogs.

The targets are [1, 1] for cat, [−1, 1] for dog and [0, −1] for mat. The loss is half the sum of the squared differences between every looked-up coordinate and its target, taken over the whole sequence as one observation. For the default sequence of cat, dog, cat the starting table is a long way from those targets.

cat   [0.2, 0.0] − [1, 1]     = [−0.8, −1.0]
dog   [−0.2, 0.0] − [−1, 1]   = [0.8, −1.0]
cat   [0.2, 0.0] − [1, 1]     = [−0.8, −1.0]

Loss = ½ × (0.64 + 1 + 0.64 + 1 + 0.64 + 1)
     = ½ × 4.92 = 2.46

That 2.46 is the loss the playground reports before its step. The table now has a reason to change, since moving a row toward its word’s target lowers the loss, and Part 2 follows how far each row moves.

Part 2. Training the Table

Two appearances contribute to one update

The target for cat is a vector with both entries equal to one. With half the sum of squared errors, the derivative at each looked-up coordinate is its current value minus its target. Both cat occurrences contribute to the same row.

One cat occurrence: [0.2, 0.0] − [1.0, 1.0] = [−0.8, −1.0]
Two occurrences:    [−0.8, −1.0] + [−0.8, −1.0] = [−1.6, −2.0]

Learning rate = 0.1
New cat row = [0.2, 0.0] − 0.1 × [−1.6, −2.0]
            = [0.36, 0.2]

The mat row receives no gradient in the default sequence because it was never read. The example measures one sequence as one training observation; it sums the coordinate errors within that observation.

Dog appears once, so its row receives one contribution, and the whole step can be read off the playground’s second table.

Word	Row before	Gradient	Row after
cat	[0.2, 0.0]	[−1.6, −2.0]	[0.36, 0.2]
dog	[−0.2, 0.0]	[0.8, −1.0]	[−0.28, 0.1]
mat	[0.0, −0.2]	[0, 0]	[0.0, −0.2]

One step at a learning rate of 0.1 on the sequence cat, dog, cat.

Looking the three words up again in the updated table gives a loss of 1.7138, down from 2.46. Replace the third word with mat and the cat row has one occurrence to learn from. Its gradient halves to [−0.8, −1.0] and the row lands on [0.28, 0.1]. The mat row, now read once, moves from [0.0, −0.2] to [0.0, −0.28].

3. Accumulate by ID

In the general case, each occurrence can have a different arriving gradient because its context and prediction differ. We still add all contributions for a particular ID. In NumPy, repeated advanced indices require explicit accumulation.

table_slopes = np.zeros_like(table)
np.add.at(table_slopes, token_ids, arriving_gradients)
updated_table = table - learning_rate * table_slopes

Ordinary indexed addition can lose contributions at repeated indices. On the default sequence it would leave the cat row’s slope at [−0.8, −1.0], one occurrence’s worth, and the row would land on [0.28, 0.1] where the correct step lands on [0.36, 0.2]. Nothing would fail and the loss would still fall, only by less. This is why the implementation is checked both with repeated tokens and by numerical differentiation of the table entries.

4. Stop the derivative at the discrete ID

The table entries are continuous numbers, so we can differentiate with respect to them. A token ID is a discrete selection. Moving a little from one ID toward the next is not a meaningful operation for this lookup. The SDK therefore passes zeros back to the ID inputs while returning the gradient of the embedding table.

Part 3. What the Vectors Can Represent

Vector similarity becomes useful when the training objective makes it useful. For example, predicting neighboring words can encourage related contexts to produce related representations. The word2vec lesson studies one way to learn such a table. An embedding layer can also learn directly inside a classifier.

The rows go where the objective sends them

The starting rows for cat and dog are [0.2, 0.0] and [−0.2, 0.0], which point in exactly opposite directions. Nothing about cats and dogs put them there. The values were chosen for the exercise. The playground’s targets describe a different arrangement, in which cat at [1, 1] and dog at [−1, 1] agree in their second coordinate and differ in their first. Repeat the playground’s step on the sequence cat, dog, cat, each time starting from the table the last step left.

Steps taken	cat row	dog row	mat row	Loss
0	[0.2, 0.0]	[−0.2, 0.0]	[0.0, −0.2]	2.4600
1	[0.36, 0.2]	[−0.28, 0.1]	[0.0, −0.2]	1.7138
5	[0.7379, 0.6723]	[−0.5276, 0.4095]	[0.0, −0.2]	0.4620
20	[0.9908, 0.9885]	[−0.9027, 0.8784]	[0.0, −0.2]	0.0123

The same step repeated at a learning rate of 0.1. The playground itself takes the first of these steps.

Three things in that table hold beyond this example. The rows end up arranged the way the objective asked, and no longer the way they started. The cat row arrives sooner than the dog row, because cat is read twice in the sequence and its gradient is twice as large. And the mat row has not moved at all, although it has a target like the others, because the sequence never read it. A word the training data never contains keeps the vector it started with, and a word it seldom contains moves only a little.

The second coordinate came to hold something cat and dog share only because these targets were written that way. Turn every target by the same angle and the rows would settle on different numbers with the same distances between them, so no single coordinate can be said to mean anything by itself.

There is no fixed dictionary meaning for each coordinate. Changing the vocabulary order requires changing the corresponding table rows. Unknown tokens need an explicit vocabulary policy, and sequence position requires additional information. The layer itself only performs lookup and learning of the selected rows.

The SDK fixes the number of tokens per sequence when the layer is constructed. Its output has one vector per token, with a leading axis for the batch. For the default sequence that is one sequence of three vectors of two numbers each, an arrangement of 1 by 3 by 2. Padding and the treatment of padding tokens belong to the surrounding data and model design.

A Lookup Selects Rows, Not Numerical Magnitudes

A token ID identifies one row of the embedding table. Changing cat's ID from three to seven and moving its embedding row with it should leave the model's behaviour unchanged. The ID itself does not tell the model that cat is closer to one word than another.

The selected vector enters later computations, whose loss supplies a gradient for that row. Rows absent from the batch receive no lookup contribution from it. Their possible regularisation updates are a separate part of training.

Repeated Tokens Send Several Contributions to One Row

If the same token occurs twice, both positions read the same parameter row. During the backward pass, their gradient contributions must be added. Overwriting the row with the last contribution loses part of the learning signal.

Use an accumulation operation that handles repeated indices. This connects the embedding update directly to the repeated-destination case from NumPy indexing.

Accumulate two occurrences of the same token

Token one appears twice and token three once. Each occurrence has its own two-coordinate upstream gradient.

Python

import numpy as np
ids = np.array([1,3,1])
upstream = np.array([[1.,2.],[3.,4.],[5.,6.]])
gradient = np.zeros((4,2))
np.add.at(gradient, ids, upstream)
print(gradient.tolist())

Output

[[0.0, 0.0], [6.0, 8.0], [0.0, 0.0], [3.0, 4.0]]

The two contributions to row one add together. Their positions in the sentence remain separate even though their parameter row is shared.

Padding Is a Storage Token with a Special Role

A padding ID lets different-length sequences share a batch shape. Decide whether its embedding is fixed and ensure padded positions are excluded from the relevant loss and attention operations. A zero embedding alone does not automatically mask all later computations.

Keep unknown tokens, padding and end-of-sequence tokens distinct. An unknown token represents content outside the vocabulary; padding represents no content at that position; an end marker represents an actual sequence boundary.