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.

10 sections · 10 questions · 3 coding challenges

1 / 10

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.

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.

Text piece

Vocabulary address

An illustrative ID assignment.

1 of 4

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.

Explore the full example

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.

What is missing is a representation that arithmetic can use and that training can adjust. An embedding layer supplies one. It keeps a table with one row for each ID, and each row is a vector of ordinary numbers. Reading the layer is a lookup, in which ID 0 selects row 0 and ID 2 selects row 2.

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.

Keep in mind

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.