Embedding Layer Forward & Backward

The problem

A language model receives token ids, integers such as [2, 0, 2], but it computes with vectors. The embedding layer is its first layer: a learned table $E$ with $V$ rows and $D$ columns, one vector of $D$ numbers for each of the $V$ token ids. The vector of token $i$ is simply row $i$ of $E$ .

Forward: look up the rows

Example. $V = 3$ , $D = 2$ , $E$ = [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]].

ids = [2, 0, 2] gives rows 2, 0, 2: [[0.5, 0.6], [0.1, 0.2], [0.5, 0.6]], shape (3, 2).

ids can have any shape. For a batch of 2 texts of 3 tokens, ids has shape (2, 3) and the result has shape (2, 3, 2): every id is replaced by its vector.

Backward: send each gradient to its row

The backward pass receives grad_out, the gradient of the loss with respect to each output vector, so it has the same shape as the output. Each output vector was a copy of one row of $E$ , so its gradient goes back to that row:

rows of $E$ that were not used get a gradient of 0;
a row used once gets the gradient of that output;
a row used several times gets the sum of the gradients of all its outputs (gradients add up, see section 1).

Example. ids = [2, 0, 2] and grad_out = [[1, 1], [2, 2], [3, 3]], $V = 3$ .

Row 0 was used by output 1: gradient [2, 2].

Row 1 was not used: gradient [0, 0].

Row 2 was used by outputs 0 and 2: gradient [1, 1] + [3, 3] = [4, 4].

The result is [[2, 2], [0, 0], [4, 4]], shape $(V, D)$ .

What to implement

embedding_forward(E, ids): ids is an integer array of any shape. Return the vectors, shape ids.shape + (D,).
embedding_backward(ids, grad_out, vocab_size): grad_out has shape ids.shape + (D,). Return the gradient of $E$ , shape (vocab_size, D).

Requirements

Rows of $E$ that were not used get a zero gradient.
A token that appears several times gets the sum of all its gradients. This is the whole point of the exercise.
No Python loops over tokens.

Hint

dE[ids] += grad_out looks right but is wrong when ids has repeats. NumPy first computes every dE[i] + grad from the old value of dE, then writes the results back one after the other, so for a repeated id only the last write survives. In the example, row 2 would get [0, 0] + [3, 3] = [3, 3] instead of [4, 4]. np.add.at(dE, ids, grad_out) really accumulates.