Causal Attention Scores

Problem

Compute scaled dot-product attention scores with a causal mask.

For query row i and key row j, score = dot(q_i, k_j) / sqrt(d_k) where d_k is the key dimension. Set every entry with j > i to -1000000000.0 so future tokens are masked out while every value stays finite. Return a len(q) x len(k) matrix.

Hint

Divide by sqrt(d_k), then block every key position past the current query position.

Test cases

case 1

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

expected: [[0.7071067811865475,-1000000000],[0,0.7071067811865475]]

case 2

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

expected: [[0.5773502691896258,-1000000000,-1000000000],[0,0.5773502691896258,-1000000000]]

case 3

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

expected: [[0]]

Starter codePython

import math
def causal_attention_scores(q, k):
    # Your code here
    pass

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