Attention and Multi-Head Attention: write the shapes first
Posted 26 October 2026. Original preparation material, not a claim about any company’s interview requirements.
What you should be able to do
Explain the objective, trace a small example on paper, implement the core without library algorithm shortcuts, and defend boundary behavior. Prerequisites: Python, arrays, basic linear algebra and derivatives.
Intuition
A query asks what information to retrieve. Keys define matching directions; values carry the information to mix. The score matrix connects query positions to key positions. Multi-head attention applies this operation in several learned projected subspaces; it is not multiple identical copies averaged together.
Formula and derivation
With independent zero-mean unit-variance query and key components, a dot product of d_k terms has variance d_k. Dividing by √d_k keeps that variance near one under these assumptions. This is motivation, not a proof that real learned components are independent. Project B×T×D into B×T×D, reshape B×T×h×d, then transpose B×h×T×d. Scores have shape B×h×T×T. Merge heads only after transposing back.
Paper-and-pencil example
Let
Implementation contract and common errors
Causal masks block future keys before softmax. For square full-sequence self-attention the mask is lower triangular. Cached decoding has different query/key lengths and needs position-aware masking. Our implementation deliberately rejects that case. Attention costs O(BT²D), score storage O(BhT²), and projections cost O(BTD²). No dropout, padding mask, gradients or KV cache are implemented here.
Complete runnable reference
The shared file contains all five algorithms, including multi-head attention. Requires Python 3.11+ and NumPy; validated with Python 3.12.14 and NumPy 2.3.5. Inputs are deterministic; no random seed is needed. No paid API or model download.
Download Python implementation and tests
python -m pip install numpy==2.3.5
python core.py
# PASS: 10 algorithm invariants (synthetic fixtures)Show full source
import numpy as np
def softmax(x):
x = np.asarray(x, dtype=float)
if x.size == 0 or not np.isfinite(x).all():
raise ValueError('nonempty finite logits required')
z = np.exp(x - x.max(axis=-1, keepdims=True))
return z / z.sum(axis=-1, keepdims=True)
def attention(q, k, v, causal=False):
if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
raise ValueError('incompatible shapes')
scores = q @ k.swapaxes(-1, -2) / np.sqrt(q.shape[-1])
if causal:
if q.shape[-2] != k.shape[-2]:
raise ValueError('this teaching mask requires equal sequence lengths')
mask = np.tril(np.ones(scores.shape[-2:], dtype=bool))
scores = np.where(mask, scores, -np.inf)
z = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = z / z.sum(axis=-1, keepdims=True)
return weights @ v
def multi_head(x, wq, wk, wv, wo, heads, causal=True):
b, t, d = x.shape
if heads < 1 or d % heads: raise ValueError('heads must divide model dimension')
def split(w):
return (x @ w).reshape(b, t, heads, d // heads).transpose(0, 2, 1, 3)
y = attention(split(wq), split(wk), split(wv), causal)
return y.transpose(0, 2, 1, 3).reshape(b, t, d) @ wo
def kmeans(x, centers, max_iter=100, tol=1e-8):
x, c = np.asarray(x, float), np.array(centers, float, copy=True)
if x.ndim!=2 or c.ndim!=2 or not len(x) or not len(c) or x.shape[1]!=c.shape[1] or max_iter<1 or not np.isfinite(x).all() or not np.isfinite(c).all():
raise ValueError('invalid data or centers')
for _ in range(max_iter):
labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
new = np.array([x[labels==j].mean(0) if np.any(labels==j) else c[j] for j in range(len(c))])
shift = np.max(np.linalg.norm(new-c,axis=1))
c = new
if shift <= tol: break
labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
return c, labels, ((x-c[labels])**2).sum()
class Adam:
def __init__(self, shape, lr=.001, b1=.9, b2=.999, eps=1e-8):
self.m=np.zeros(shape); self.v=np.zeros(shape); self.t=0
self.lr,self.b1,self.b2,self.eps=lr,b1,b2,eps
def step(self, p, g):
if p.shape!=self.m.shape or g.shape!=p.shape or not np.isfinite(g).all(): raise ValueError('invalid gradient')
self.t+=1
self.m=self.b1*self.m+(1-self.b1)*g
self.v=self.b2*self.v+(1-self.b2)*g*g
mh=self.m/(1-self.b1**self.t); vh=self.v/(1-self.b2**self.t)
return p-self.lr*mh/(np.sqrt(vh)+self.eps)
def beam_search(next_log_probs, width=2, max_steps=4, eos='EOS'):
if width<1 or max_steps<1: raise ValueError('positive limits required')
beams=[((),0.,False)]
for _ in range(max_steps):
candidates=[]
for seq,score,done in beams:
if done: candidates.append((seq,score,True)); continue
for token,lp in next_log_probs(seq).items():
if not np.isfinite(lp) or lp>0: raise ValueError('finite log probabilities required')
candidates.append((seq+(token,),score+lp,token==eos))
if not candidates: raise ValueError('no candidates')
beams=sorted(candidates,key=lambda x:(-x[1],x[0]))[:width]
if all(x[2] for x in beams): break
return beams
if __name__=='__main__':
assert np.allclose(softmax([1000,1000]),[.5,.5])
assert np.allclose(softmax([1,2]),softmax([101,102]))
x=np.arange(24,dtype=float).reshape(1,3,8)/24; eye=np.eye(8)
y=multi_head(x,eye,eye,eye,eye,2); changed=x.copy(); changed[:,2]+=10
assert np.allclose(y[:,:2],multi_head(changed,eye,eye,eye,eye,2)[:,:2])
assert np.allclose(multi_head(x,eye,eye,eye,eye,1),attention(x,x,x,True))
c,l,s=kmeans([[0],[2],[8],[10]],[[0],[10]])
assert np.allclose(c,[[1],[9]]) and s==4
c,l,s=kmeans([[0],[0]],[[0],[8]]); assert np.isfinite(c).all()
opt=Adam((1,)); assert np.allclose(opt.step(np.array([1.]),np.array([2.])),[.999])
assert opt.t==1
def probs(seq):
p={'A':.6,'B':.4} if not seq else ({'EOS':.5,'X':.5} if seq[-1]=='A' else {'EOS':.9,'X':.1})
return {k:np.log(v) for k,v in p.items()}
assert beam_search(probs,max_steps=2)[0][0]==('B','EOS')
assert beam_search(probs,width=1,max_steps=2)[0][0]==('A','EOS')
print('PASS: 10 algorithm invariants (synthetic fixtures)')
Interview follow-up
Why must future-token changes leave past outputs unchanged?
Hint: Trace the allowed keys.
Answer and reasoning
A causal row can only use keys and values up to its position. Test this invariant by perturbing future inputs; checking only output shape will miss a reversed mask.
Practice without looking
- Spend 3 minutes writing shapes, assumptions and an example.
- Spend 12 minutes writing the core implementation on paper.
- Spend 5 minutes tracing an ordinary input and one failure case.
- Type your handwritten version and run the reference invariants. Explain every mismatch before consulting the solution.
Transfer exercise: change one assumption in the implementation contract above. Write a new test that fails before your modification and passes afterward. A convincing explanation of a limitation is better than silently inventing unsupported behavior.
Sources and next lesson
Primary paper or official documentation · Paper PDF
Softmax: stable probabilities and gradients
Attention and Multi-Head Attention: write the shapes first
K-Means: assignment, centroid updates and empty clusters