Scaled dot-product
- Machine Learning course
- Expert
- Lesson 4/44
- 5 min
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, AThis 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.
Practice 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).
Starter code
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.]])
# buraya yazSample 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))Expected output
4 2.0 1.0 1024 32.0 1.0 [[5. 5. ] [1.956 8.044]]
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 By what are the scores divided when d_k = 64?
- 4096
- 2
- 8 Correct answer
- 64
Explanation: ��64 = 8.
2 Scale the scores.
A = softmax(Q @ K.T / np.____(d_k), eksen=-1)A = softmax(Q @ K.T / np.sqrt(d_k), eksen=-1)Explanation: The divisor is √d_k.
3 Match the values.
| Term | Meaning |
|---|---|
d_k = 16 |
√d = 4 |
d_k = 64 |
√d = 8 |
large unscaled scores |
One-hot softmax |
scaled score std |
≈ 1 |