Compute scaled dot-product attention

Learn it

Words to know

attention

The step in which each token gathers information from other tokens, taking more from those that match what it is looking for.

query, key, value

The three roles in attention: a query is what a token looks for, a key is what it can be matched on, a value is what it passes on.

causal mask

The rule that a token may attend only to itself and earlier tokens, never to later ones.

softmax

The function that turns a list of scores into probabilities: all positive, adding up to 1.

How it works

1. Score: the dot product of one query with each key, divided by sqrt(d), where d is the width of a query: how many numbers it holds.
2. Weights: the softmax of those scores. They are positive and add up to 1.
3. Output for that query: each value row times its weight, added together.
4. Causal: for query number i, use only keys and values 0 to i. Drop the later ones before the softmax, so the remaining weights still add up to 1.

Worked example

Input
Q = [[1, 0], [0, 1]]
K = [[1, 0], [0, 1]]
V = [[10, 0], [0, 10]]
causal = False

Output
[[6.69762, 3.30238], [3.30238, 6.69762]]

Each query matches its own key best and takes about two thirds of that value. With the causal mask, the first query could see only key 0 and its output would be [10, 0].

Why it matters

This one calculation is the centre of every transformer. The next two lessons run several copies of it side by side and wrap it into a block.

Your turn

Compute scaled dot-product attention, with an optional causal mask.

Q

the queries: one row per token

K

the keys: one row per token, the same width as the queries

V

the values: one row per token

causal

when true, query i may use only keys and values 0 to i

Returns

One output row per query: the value rows mixed with weights softmax(scores), where each score is dot(query, key) / sqrt(d) and d is the width of a query.

Input
Q = [[1, 0], [0, 1]]
K = [[1, 0], [0, 1]]
V = [[10, 0], [0, 10]]
causal = False

Output
[[6.69762, 3.30238], [3.30238, 6.69762]]

Input
Q = [[1, 0], [0, 1]]
K = [[1, 0], [0, 1]]
V = [[10, 0], [0, 10]]
causal = True

Output
[[10, 0], [3.30238, 6.69762]]

Input
Q = [[0], [0]]
K = [[0], [0]]
V = [[2], [6]]
causal = False

Output
[[4], [4]]

Input
Q = [[0], [0]]
K = [[0], [0]]
V = [[2], [6]]
causal = True

Output
[[2], [4]]

Input
Q = [[2, 0, 0, 0], [0, 0, 0, 0]]
K = [[1, 0, 0, 0], [0, 0, 0, 0]]
V = [[8], [0]]
causal = False

Output
[[5.84847], [4]]

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Where to start

   Handle one query at a time. For each one you need a list of scores, then a list of weights, then one output row.

2. Hint 2: The key step

   For query i, decide how many keys it may use: i + 1 when causal is true, all of them otherwise. Compute scores for only those keys, then apply softmax as in chapter 1.

3. Hint 3: The outline

   For each query: take count keys; score each as dot(q, k) / math.sqrt(len(q)); subtract the largest score and raise e to each; divide by the total to get weights. Each output column is the sum of weight * V[j][column] over the allowed rows.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

1. A causal query sees only keys 0 to i. Leaving the later keys out before the softmax means the weights of the remaining keys still add up to 1.
2. Score each visible key against the query. A large dot product means the query and the key point the same way. Dividing by sqrt(d) keeps scores from growing as the vectors get wider.
3. Softmax turns the scores into weights that add up to 1. The same stable softmax as in chapter 1: subtract the largest score, raise e to each, divide by the total.
4. Mix the value rows with those weights. Each output number is a weighted average of one column of the values, over the rows this query could see.

The same numbered steps are the comments in the solution.

import math


def solve(Q, K, V, causal):
    output = []
    for i, q in enumerate(Q):
        # 1. A causal query sees only keys 0 to i
        count = i + 1 if causal else len(K)
        # 2. Score each visible key against the query
        scores = []
        for k in K[:count]:
            dot = sum(a * b for a, b in zip(q, k))
            scores.append(dot / math.sqrt(len(q)))
        # 3. Softmax turns the scores into weights that add up to 1
        top = max(scores)
        exps = [math.exp(score - top) for score in scores]
        total = sum(exps)
        weights = [e / total for e in exps]
        # 4. Mix the value rows with those weights
        row = []
        for column in range(len(V[0])):
            row.append(sum(w * V[j][column] for j, w in enumerate(weights)))
        output.append(row)
    return output

Watch out

• Apply the mask before the softmax. Setting weights to 0 afterwards leaves the rest adding up to less than 1.
• Divide by the square root of the query width, len(q), not by the number of tokens.

Check yourself