HardNLPnlp-050

Scaled Dot-Product Attention

Problem

Compute single-head scaled dot-product attention. query is a vector, and keys and values are lists of equal-length vectors.

Scores are dot(query, keys[i]) divided by the square root of len(query), passed through a softmax, and the output is the weighted sum of the value vectors. Return an empty list when keys is empty.

Hint

The softmax weights are per key; use them to average the rows of values column by column.

Test cases

case 1

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

expected: [0.6697615493266569,0.3302384506733431]

case 2

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

expected: [3,4]

case 3

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

expected: [2,4,6]

case 4

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

expected: []

Starter codePython

import math
def scaled_dot_product_attention(query, keys, values):
    # Your code here
    pass

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