🎓 A1Learner
🔍

PyTorch Scaled Dot-Product Attention

Scaled dot-product attention is a computation that mixes the vectors of a sequence into a new vector for each position. Each position makes a query, a key, and a value from its own vector.

One score from one pair

A query is what a position looks for, and a key is what a position advertises. A value is what a position delivers when its key is matched. The code dots q with k to get one score.

import torch

q = torch.tensor([1.0, 2.0])
k = torch.tensor([3.0, 1.0])
score = (q * k).sum()
print("Score:", score.item())

Why the scores are scaled

A dot product grows as the vectors get longer, and big scores push softmax toward one winner. Dividing by 8 here, the square root of 64, keeps the losing score in play.

raw = torch.tensor([8.0, 0.0])
print("Unscaled:", torch.softmax(raw, dim=0))
print("Scaled by 8:", torch.softmax(raw / 8, dim=0))

Weights mix the values

The weights say how much of each value row goes into the mix. The first value row counts three times as much as the second one here.

weights = torch.tensor([0.75, 0.25])
values = torch.tensor([[2.0, 0.0],
                       [0.0, 4.0]])
mixed = weights @ values
print("Mixed value:", mixed)

Example

import math

import torch

X = torch.tensor([[1.0, 0.0],
                  [0.0, 1.0],
                  [1.0, 1.0],
                  [0.0, 2.0]])
W_q = torch.tensor([[1.0, 1.0],
                    [0.0, 1.0]])
W_k = torch.tensor([[1.0, 0.0],
                    [1.0, 1.0]])
W_v = torch.tensor([[2.0, 0.0],
                    [1.0, 1.0]])

Q = X @ W_q
K = X @ W_k
V = X @ W_v

scores = Q @ K.T / math.sqrt(2)
weights = torch.softmax(scores, dim=1)
output = weights @ V

print("Scores:")
print(scores)
print("Weights:")
print(weights)
print("Output:")
print(output)

Key points

  • Dot product: multiply two vectors item by item and add the results up
  • Query: the vector a position sends out to state what it looks for
  • Key: the vector a position advertises so that queries can match it
  • Value: the vector a position delivers when its key is matched
  • Scores are divided by the square root of the vector length before softmax
  • Softmax turns each row of scores into weights that add up to 1
  • The output row is the value rows mixed by the weights and added up
  • In short, attention matches queries to keys and mixes values by the strength of the match

In this example, four token vectors sit in X as rows. A token is one item in a sequence, like one word in a line. W_q and W_k turn each row into a query and a key. W_v turns each row into a value. The scores table holds every query dotted with every key, divided by the square root of 2. Softmax turns each score row into weights that add up to 1. The output rows are the value rows mixed by those weights. Row 3 gives token 4 a weight of 0.7183, so the value of token 4 dominates that mix.

Note

The projections in this lesson are fixed small tables, while a trained model learns its own

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why are the scores divided by the square root of the vector length?