L4.3 · LESSON 3 OF 6

Scaling

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

THE IDEA

Why divide by √d_k

If query and key entries are random with unit variance, a dot product of two size-d vectors has variance d, so its typical size is √d. In a model with d_k = 64 raw scores of size 8 are routine. The softmax of such large scores is almost one-hot, and an almost one-hot softmax has almost no gradient. Scaled dot-product attention divides the scores by √d_k, which keeps their typical size near 1 whatever the dimension.

key size d_k64

dot products of random vectors of size 64 have a typical size of √d_k = 8 divisor √d_k = 8

	scores	weight on the best key	weights of 4 keys
unscaled	[8, 0, 0, 0]	0.999
scaled	[1, 0, 0, 0]	0.4754

As d_k grows the unscaled scores grow like √d_k and the softmax collapses onto one key (weight near 1, gradients near 0). Dividing by √d_k keeps the typical score near 1 at any size. This relies on common assumptions about the vector entries, so treat it as the reason, not a theorem about every trained model.

Change d_k. Compare the weight on the best key with and without the division.

Predict first. For d_k = 4, what divisor is used?

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 divisors and a scaled weight by hand, then measure how dot products grow with dimension by simulation.

key size d_k64

dot products of random vectors of size 64 have a typical size of √d_k = 8 divisor √d_k = 8

	scores	weight on the best key	weights of 4 keys
unscaled	[8, 0, 0, 0]	0.999
scaled	[1, 0, 0, 0]	0.4754

divisor for d_k = 4 divisor for d_k = 64 first weight for scores [2, 0], d_k = 4

To do: Scaled attention divides the scores by √d_k. For d_k = 4, what is the divisor?
To do: And for d_k = 64?
To do: d_k = 4 and raw scores [2, 0]. After dividing by 2 the scores are [1, 0]. Type the softmax weight of the first key (3 decimals)

Part 2. Do it in code

scale_factor(d_k): sqrt(d_k). d_k must be at least 1 (ValueError).
scaled_scores(Q, K): (Q @ K.T) divided by sqrt(d_k), where d_k is the number of columns of K.
score_std(d, trials, seed): draw `trials` pairs of random vectors with standard-normal entries of length d (np.random.default_rng(seed).standard_normal((trials, d)) for q and a second draw for k), compute each pair's dot product and return the standard deviation of those dot products.
scaled_score_std(d, trials, seed): the same, but divide each dot product by sqrt(d) first.
A softmax(z) helper is provided; max_weight(scores) returns the largest softmax weight.

import numpy as np

def softmax(z):
    z = np.asarray(z, dtype=float)
    e = np.exp(z - z.max())
    return e / e.sum()

def scale_factor(d_k):
    pass

def scaled_scores(Q, K):
    pass

def score_std(d, trials, seed):
    pass

def scaled_score_std(d, trials, seed):
    pass

def max_weight(scores):
    pass

Ctrl/⌘ + Enter runs it

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

CHECK YOUR UNDERSTANDING

For d_k=4, what divisor is used?

ATwo, because the divisor is √d_k = √4, not d_k itself.BFour, because the divisor is the key dimension.CSixteen, because the divisor is the key dimension squared.

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.