One-Hot Encoding & Gather

The problem

A language model reads a piece of text and, at every position, guesses the token that comes next. Its guess is not a single token: it is a probability for every token of the vocabulary. With $N$ positions and a vocabulary of $V$ tokens, all the guesses together form a matrix probs with $N$ rows and $V$ columns:

• row $i$ is the guess at position $i$;
• column $j$ is token number $j$ of the vocabulary (its id);
• probs[i, j] is how likely the model thinks token $j$ is at position $i$. Every row adds up to 1.

During training we know the right answer at every position: the token that really came next. Its id is stored in targets, one id per row. To score the model we need one number per row: the probability it gave to the right token. The rest of the row is not needed.

Example. A tiny vocabulary of $V = 3$ tokens: id 0 is "cat", id 1 is "dog", id 2 is "sun". The model makes guesses at $N = 2$ positions:

	cat (0)	dog (1)	sun (2)
position 0	0.2	0.7	0.1
position 1	0.5	0.1	0.4

The right answers are "dog" at position 0 and "sun" at position 1, so targets = [1, 2]. The numbers we want are 0.7 (row 0, column 1) and 0.4 (row 1, column 2): the result is [0.7, 0.4].

There are two ways to get those numbers. You will implement both.

Method 1: one-hot

The idea: build a mask with a 1 where the right token is and 0 everywhere else, multiply it with probs, and add up each row. The zeros erase every wrong token, so only the right one survives.

A one-hot vector is a row of $V$ zeros with a single 1 at the position of an id. With $V = 3$, id 1 becomes [0, 1, 0] and id 2 becomes [0, 0, 1].

Example. Same probs and targets = [1, 2] as above, step by step.

1. One-hot row for each target, stacked into a matrix with the same shape as probs: [[0, 1, 0], [0, 0, 1]].
2. Multiply it with probs, number by number: row 0 becomes [0.2·0, 0.7·1, 0.1·0] = [0, 0.7, 0], row 1 becomes [0.5·0, 0.1·0, 0.4·1] = [0, 0, 0.4].
3. Add up each row: $0 + 0.7 + 0 = 0.7$ and $0 + 0 + 0.4 = 0.4$. The result is [0.7, 0.4].

As a formula, the number picked from row $i$ is:

where:

• $\text{onehot}_{ij}$ is 1 if $j$ is the right token of row $i$, and 0 otherwise;
• $\text{probs}_{ij}$ is the probability of token $j$ at position $i$;
• the sum runs over all $V$ tokens, but every term is 0 except the right one.

The downside: the mask is as big as probs, $N \times V$ numbers, and almost all of them are 0. A real vocabulary has 50,000 tokens or more, so for 1,000 positions the mask holds 50 million numbers, only to pick out 1,000 of them.

Method 2: gather

The idea: build nothing and read the right cell directly. In row $i$, the right token is targets[i], so the number we want is probs[i, targets[i]].

Example. Same probs and targets = [1, 2]:

• row 0: targets[0] is 1, so read probs[0, 1], which is 0.7;
• row 1: targets[1] is 2, so read probs[1, 2], which is 0.4.

The result is [0.7, 0.4], the same as method 1, but it reads only 2 numbers instead of building a matrix of 6.

This is what real loss functions do.

What to implement

• one_hot(targets, num_classes): return the one-hot vectors of the ids in targets, as floats (1.0, not 1), so they can be multiplied with probabilities. num_classes is the vocabulary size $V$. targets can have any shape: every id is replaced by its one-hot vector, so the result has one more axis, of length num_classes. For a list of 2 ids the result has shape (2, num_classes). For a $2 \times 2$ table of ids it has shape (2, 2, num_classes).
• gather(probs, targets): probs has $N$ rows and $V$ columns, targets holds $N$ ids. Return the $N$ numbers probs[i, targets[i]], one per row.

Requirements

• No Python loops: use NumPy operations on whole arrays.

Hint

• For one_hot: np.eye(V) is the $V \times V$ identity matrix, and its row $j$ is exactly the one-hot vector of id $j$. Indexing it with an array of ids picks those rows.
• For gather: indexing with two arrays of the same length pairs them up position by position: probs[[0, 1], [1, 2]] means [probs[0, 1], probs[1, 2]]. Watch out: probs[:, targets] is different. It takes columns 1 and 2 from every row, and returns a $2 \times 2$ table, [[0.7, 0.1], [0.1, 0.4]].