L4.4 · LESSON 4 OF 6

Softmax weights

Calculate query-key scores, scaled softmax weights and the weighted value output for one attention head.

THE IDEA

From scores to weights

The scaled scores of one query go through a softmax: exponentiate each, then divide by the sum. The weights are positive and add up to 1 across the keys of that query. Every query has its own row; rows are independent.

In code you subtract the row maximum before exponentiating. That changes no weight (it cancels in the division) but stops exp() from overflowing when scores are large.

score 10.71score 20score 3-1

key	score	score − max	exp(score − max)	weight
1	0.71	0	1	0.5972
2	0	-0.71	0.4931	0.2945
3	-1	-1.71	0.1814	0.1083

weights add up to 1 Subtracting the maximum changes no weight, but stops exp() from overflowing on large scores.

With several queries, each query has its own row of weights and each row sums to 1 separately, so the whole matrix does not sum to 1.

Three scores, their weights, and the maximum-subtraction trick.

Predict first. Must attention weights sum to one across every query in the whole batch?

Chapter context

Attention computes a data-dependent mixture of value vectors. A query scores candidate keys, a normalization turns scores into weights, and those weights combine values. Keys decide how to look; values supply what is gathered.

TRY IT YOURSELF

Part 1. Do it hands-on

Compute softmax weights by hand, then write a stable, row-wise softmax and see the naive one fail.

score 10.71score 20score 3-1

key	score	score − max	exp(score − max)	weight
1	0.71	0	1	0.5972
2	0	-0.71	0.4931	0.2945
3	-1	-1.71	0.1814	0.1083

weights add up to 1 Subtracting the maximum changes no weight, but stops exp() from overflowing on large scores.

With several queries, each query has its own row of weights and each row sums to 1 separately, so the whole matrix does not sum to 1.

first weight for scores [0.7071, 0] sum of all four weights (2 queries × 2 keys) each weight for scores [1000, 1000]

To do: Scaled scores [0.7071, 0]. Softmax weight of the first key (3 decimals)
To do: A 2-query, 2-key attention has two weight rows. Each row is normalised separately. What do ALL four weights add up to?
To do: Scores [1000, 1000]. A correct implementation subtracts the row maximum first, so exp(0) and exp(0). What is each weight?

Part 2. Do it in code

softmax(x, axis=-1): numerically stable softmax along the axis (subtract the maximum along that axis first). Entries equal to −inf get weight 0. Raise ValueError if a whole row is −inf.
naive_softmax(x): exp(x) / sum(exp(x)) with NO subtraction (to see why stability matters; allow numpy overflow warnings).
rows_normalised(W, tol=1e-9): True if every row of W sums to 1 within tol.
entropy(w): −sum(w * ln w) of one weight vector, treating 0·ln 0 as 0. Low entropy = focused attention.

import numpy as np

def softmax(x, axis=-1):
    pass

def naive_softmax(x):
    pass

def rows_normalised(W, tol=1e-9):
    pass

def entropy(w):
    pass

Ctrl/⌘ + Enter runs it

Stuck, or something not working? You can come back to it any time.

CHECK YOUR UNDERSTANDING

Must attention weights sum to one across every query in the whole batch?

AYesBNo

Chapter reference: math, worked example & code

This section belongs to the whole chapter, so it is the same on every lesson in it and may cover a different lesson's topic. The code for the experiment you just ran is inside the experiment itself.

Math and notation

Q\in\mathbb R^{n_q\times d_k},\ K\in\mathbb R^{n_k\times d_k},\ V\in\mathbb R^{n_k\times d_v}.

S=QK^T/\sqrt{d_k}+M,\quad A=\operatorname{softmax}_{rows}(S),\quad O=AV.

M is zero at allowed positions and conceptually negative infinity at excluded positions. Every query needs a defined policy if all its keys are masked.

Worked example

Use q=[1,0], K=[[1,0],[0,1]], V=[[2,0],[0,2]], and d_k=2. Dot scores [1,0] become [1/sqrt(2),0]. Exponentiate and normalize to approximately [0.669762,0.330238]. The output is [2×0.669762,2×0.330238]≈[1.339523,0.660477]. If only the first key is causally allowed, the output becomes [2,0].

Reference code

The following example uses the Python standard library and does not call a remote service.

from math import sqrt,exp,isclose
q=[1,0]; keys=[[1,0],[0,1]]; values=[[2,0],[0,2]]
scores=[sum(a*b for a,b in zip(q,k))/sqrt(2) for k in keys]
weights=[exp(s-max(scores)) for s in scores]
weights=[w/sum(weights) for w in weights]
out=[sum(weights[i]*values[i][j] for i in range(2)) for j in range(2)]
assert isclose(sum(weights),1)
assert isclose(weights[0],.6697615493266569)
assert isclose(out[0],1.3395230986533138)
assert isclose(out[1],.6604769013466862)

Python examples are provided for study and local execution; this page does not execute Python.

Common mistakes

Attention scores are not weights before normalization. Causal masking happens before softmax. Attention weights alone are not a full causal explanation of a model decision.