Attention Scores With Explicit Scale

Problem

Compute attention scores with an explicit multiplicative scale: out[i][j] = scale * (Q[i] . K[j]), where Q has shape (seq_q, d) and K has shape (seq_k, d).

The signature is attention_scores_scaled(Q, K, scale). Use this when the scale is set by hand instead of the usual 1 / sqrt(d).

Hint

Each score is the dot product of a query row and a key row, then multiplied by the shared scale.

Test cases

case 1

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

expected: [[0.5,0],[0,0.5]]

case 2

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

expected: [[5]]

case 3

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

expected: [[0.5,0,0.5],[0,1,1]]

case 4

in: [[[2]],[[3]],0.25]

expected: [[1.5]]

Starter codePython

def attention_scores_scaled(Q, K, scale):
    # Your code here
    pass

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