PyTorch Scaled Dot-Product Attention
A score is the dot product of one query with one key, and high means the key matches the query. Scaled dot-product attention divides every score by the square root of the head size, then softmaxes each row into weights.
Big scores put the softmax to sleep
A softmax over scores of 10, 20, and 30 hands nearly all the weight to the largest score. Dividing the same scores by 10 first spreads the weight across all three entries.
import torch
from torch.nn import functional as F
scores = torch.tensor([10.0, 20.0, 30.0])
print("unscaled:", [round(v, 4) for v in F.softmax(scores, dim=-1).tolist()])
print("scaled by 10:", [round(v, 4) for v in F.softmax(scores / 10.0, dim=-1).tolist()])Queries meet keys in one product
Transposing the keys lets one matrix product score every pair of positions at once. The result holds one row of four scores per reading position.
import torch
from torch import nn
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
head_size = n_embd // n_head
torch.manual_seed(42)
x = torch.randn(1, 4, n_embd)
query = nn.Linear(n_embd, head_size, bias=False)
key = nn.Linear(n_embd, head_size, bias=False)
q = query(x)
k = key(x)
scores = q @ k.transpose(-2, -1)
print("scores shape:", tuple(scores.shape))
print("row 3:", [round(v, 2) for v in scores[0, 3].tolist()])Example
import torch
from torch import nn
from torch.nn import functional as F
batch_size = 32
block_size = 64
n_embd = 128
n_head = 4
n_layer = 4
head_size = n_embd // n_head
torch.manual_seed(42)
x = torch.randn(1, 4, n_embd)
query = nn.Linear(n_embd, head_size, bias=False)
key = nn.Linear(n_embd, head_size, bias=False)
q = query(x)
k = key(x)
raw = q @ k.transpose(-2, -1)
scaled = raw * head_size ** -0.5
hot = F.softmax(raw, dim=-1)[0]
calm = F.softmax(scaled, dim=-1)[0]
print("scale factor:", round(head_size ** -0.5, 4))
print("unscaled weights:")
for row in hot.tolist():
print(" ".join(f"{v:.2f}" for v in row))
print("scaled weights:")
for row in calm.tolist():
print(" ".join(f"{v:.2f}" for v in row))
print("scaled row sums:", [round(v, 2) for v in calm.sum(dim=-1).tolist()])Key points
- Dot product: a sum of matching products that grows as two vectors agree
- Softmax: a function that turns a row of scores into weights that sum to one
- Long vectors give large dot products, even when the match is weak
- Dividing by the square root of the head size pulls the scores back to a calm range
- An unscaled softmax can collapse onto one position and ignore the rest
- Each row of the weight grid belongs to one reading position
In this example, one seeded window produces queries and keys with a head size of 32. The raw scores feed a softmax directly. The same scores, multiplied by the scale factor of about 0.18, feed a second softmax. The unscaled grid is peakier, with top weights near one half in every row. The scaled grid stays close to even, with every weight between 0.18 and 0.31. Each position keeps a real share of every mix. The separate snippet with scores of 10, 20, and 30 shows the extreme case. There the unscaled softmax collapses to a single 1.0. The last line confirms that every scaled row sums to 1.0 after rounding.
Try it Yourself
Edit the code, then press Run.
Running is simulated · Python Playground
Exercise
Why does attention divide the scores by the square root of the head size?