Machine Learning
Scaled dot-product
Scaled dot-product attention divides the query-key product by the square root of dimension.
Scaled dot product
As the dimension d grows, the variance of the dot products grows too (a sum of d independent terms). Large scores pile softmax onto a single element and the gradient vanishes. The fix is dividing the scores by √d_k:
def attention(Q, K, V):
d_k = Q.shape[-1]
A = softmax(Q @ K.T / np.sqrt(d_k), eksen=-1)
return A @ V, A
This is the transformer's core formula: Attention(Q, K, V) = softmax(QKᵀ/√d_k) V.
The effect of scaling
For random vectors the standard deviation of the dot product is about √d. After scaling, the scores keep a standard deviation of ~1 regardless of d:
rng = np.random.default_rng(0)
for d in [4, 64, 1024]:
q, k = rng.normal(size=(2, 10000, d))
s = (q * k).sum(axis=1)
print(d, round(s.std(), 1), round((s / np.sqrt(d)).std(), 2))
With large unscaled scores softmax becomes almost one-hot; this can lock up training at the start.
Exercise The attention function
Write the scaled dot-product attention(Q, K, V) function (returning the output and the weights). For d = 4 and 1024, print the dot-product standard deviation for random (10000, d) Q and K before and after scaling with round(..., 1) on one line. Finally print the output for the given small Q, K, V with np.round(..., 3).
Show example solution
import numpy as np
def softmax(z, eksen=-1):
e = np.exp(z - z.max(axis=eksen, keepdims=True))
return e / e.sum(axis=eksen, keepdims=True)
rng = np.random.default_rng(0)
Qk = np.array([[1., 1.], [0., 2.]])
Kk = np.array([[1., 0.], [0., 1.]])
Vk = np.array([[10., 0.], [0., 10.]])
def attention(Q, K, V):
d_k = Q.shape[-1]
A = softmax(Q @ K.T / np.sqrt(d_k), eksen=-1)
return A @ V, A
for d in [4, 1024]:
q, k = rng.normal(size=(2, 10000, d))
s = (q * k).sum(axis=1)
print(d, round(s.std(), 1), round((s / np.sqrt(d)).std(), 1))
print(np.round(attention(Qk, Kk, Vk)[0], 3))
Test yourself
-
By what are the scores divided when d_k = 64?
Show answer
8 — ��64 = 8.
-
Scale the scores.
A = softmax(Q @ K.T / np.____(d_k), eksen=-1)
Show answer
A = softmax(Q @ K.T / np.sqrt(d_k), eksen=-1)
-
Match the values.
- d_k = 16
- √d = 4
- d_k = 64
- √d = 8
- large unscaled scores
- One-hot softmax
- scaled score std
- ≈ 1
Next lesson: Multi-head
