Attention Score Gradient

Problem

In scaled dot-product attention the scores are s_j = (q dot k_j) / sqrt(d) and the attention weights are a = softmax(s).

For attention weight i, the gradient with respect to the query is

d a_i / d q = a_i * (k_i - sum_j a_j k_j) / sqrt(d)

Given q, the list of key vectors, and the index i, return that gradient vector.

Hint

This is the softmax Jacobian contracted with the key gradients k_j / sqrt(d).

case 1

in: [[1,0],[[1,0],[0,1]],0]

expected: [0.1563985965451104,-0.1563985965451104]

case 2

in: [[1,1],[[1,0],[0,1]],1]

expected: [-0.17677669529663687,0.17677669529663687]

case 3

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

expected: [-0.22733165093293675,0.05703056094213042]

Starter codePython

def attention_score_grad(q, keys, i):
    # Your code here
    pass

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