Look up token embeddings

Learn it

Words to know

token

One unit of text as the model sees it: a character, a piece of a word or a whole word.

token ID

The whole number that stands for a token.

embedding

A list of numbers that stands for a token. The model learns these numbers during training.

How it works

1. The table has one row per token ID. Row 0 belongs to token 0, row 1 to token 1, and so on.
2. For each ID in the sequence, take the row at that index.
3. A token that appears twice gets the same row twice. A sequence of 3 IDs becomes 3 rows.

Worked example

Input
ids = [1, 0, 1]
table = [[1, 2], [3, 4]]

Output
[[3, 4], [1, 2], [3, 4]]

Token 1 appears at positions 0 and 2 and gets the same row both times. The lookup knows nothing about position yet.

Why it matters

The first layer of a GPT is this lookup, done with torch.nn.Embedding. The lesson on positions adds the missing sense of order.

Your turn

Replace each token ID with its row from an embedding table.

ids

a list of token IDs; each is a valid row index, and the list may be empty

table

the embedding table: one row of numbers per token

Returns

A new list with a copy of the table row for each ID, in the same order.

Input
ids = [1, 0, 1]
table = [[1, 2], [3, 4]]

Output
[[3, 4], [1, 2], [3, 4]]

Input
ids = []
table = [[1]]

Output
[]

Input
ids = [2, 2]
table = [[0, 0, 0], [1, 1, 1], [5, 6, 7]]

Output
[[5, 6, 7], [5, 6, 7]]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   This is a lookup, not arithmetic. Each ID is an index into table.

2. Hint 2: The key step

   For each ID, take table[i]. Wrap it in list(...) so the answer holds copies, not the table's own rows.

3. Hint 3: The outline

   Return a list with list(table[i]) for each i in ids.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. The ID is a row number: pick that row. Nothing is computed from the ID's value. Token 1 is not smaller than token 2 in any useful sense.
2. Copy it, so changing the answer cannot change the table. Without the copy, two positions with the same token would share one list, and editing one would edit both.

The same numbered steps are the comments in the solution.

def solve(ids, table):
    rows = []
    for token_id in ids:
        # 1. The ID is a row number: pick that row
        row = table[token_id]
        # 2. Copy it, so changing the answer cannot change the table
        rows.append(list(row))
    return rows

Watch out

• Do not multiply the row by the ID or use the ID as a number in any formula. It only says which row to take.
• Keep duplicates and keep the order. The output has one row per position in the sequence.
• The checks compare values only, so they cannot see whether you copied the rows. Copy anyway: an answer that shares its lists with the table is a bug that shows up much later.

Check yourself