Embedding Backward Scatter-Add

Problem

Backward pass of an embedding lookup. Build a (vocab_size, dim) zero gradient table and scatter-add each upstream row dout[k] into the row at indices[k]. Repeated indices accumulate.

The signature is embedding_backward(indices, dout, vocab_size).

Hint

Use += so repeated token indices accumulate gradient rather than overwrite.

Test cases

case 1

in: [[0,2,0],[[1,2],[3,4],[5,6]],3]

expected: [[6,8],[0,0],[3,4]]

case 2

in: [[1],[[0.5]],2]

expected: [[0],[0.5]]

case 3

in: [[0,1,2,1],[[1,1],[1,1],[1,1],[2,2]],3]

expected: [[1,1],[3,3],[1,1]]

case 4

in: [[2,2],[[1],[2]],3]

expected: [[0],[0],[3]]

Starter codePython

def embedding_backward(indices, dout, vocab_size):
    # Your code here
    pass

4 test cases · run your code against them in the editor