Transformer Decoder Masked Self-Attention

Problem

Single-head masked self-attention for a transformer decoder: compute Q = x Wq, K = x Wk, V = x Wv, scale scores by 1 / sqrt(d), replace entries above the diagonal with -1e9, apply a row-wise softmax and return the weighted sum of V.

The signature is masked_self_attention(x, Wq, Wk, Wv). All weight matrices are square d x d.

Hint

Each position may only attend to itself and earlier positions, which is what the causal mask enforces.

Test cases

case 1

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

expected: [[1,0],[0.3302384507,0.6697615493]]

case 2

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

expected: [[1,2],[2.999899605,3.999899605]]

case 3

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

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

case 4

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

expected: [[1,0,0],[0.3595425243,0.6404574757,0],[0.2644584615,0.2644584615,0.471083077]]

Starter codePython

import math
def masked_self_attention(x, Wq, Wk, Wv):
    # Your code here
    pass

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