Multi-head attention
Q/K/V, masks, and self/cross-attention derivatives.
01Shape contract#
MultiHeadAttention(embed,num_heads,dropout=0,causal=false). num_heads must be positive and divide embed; division precedes validation, so never pass zero. Q/K/V use rank-3 [B,T,E]. All four projections are Linear layers with bias.
02Scaled dot-product#
Heads split into [B·H,T,D], D=E/H. Scaled masked scores pass through softmax, optional dropout, and multiplication by V. Heads merge and pass through the output projection.
03Masks#
forward(query,key,value,mask) takes nullptr or an additive float [Tq,Tk] mask. Zero permits, large negative suppresses; it is not boolean. causal adds −1e9 to future positions and requires Tq=Tk. Leave a valid key in each row.
MultiHeadAttention attn(64,4,0.f,true);auto x=Tensor::zeros({4,32,64});auto y=attn.forward(x);auto dx=attn.backward(Tensor::full({4,32,64},1.f));04Self / cross gradients#
forward(x) returns dq+dk+dv. The multi-argument overload detects self by object addresses; separate Tensor headers sharing storage are not self. Cross backward returns dq. dk/dv are in grad_key_input()/grad_value_input(); sum them if they refer to the same memory.
05Quadratic score memory#
Scores store B·H·Tq·Tk values. This is not FlashAttention. Analytic Linear projections work alongside a local MicroTape softmax/matmul chain.
06Attention derivative chain#
With dropout disabled, let C=PV, P=softmax(S), S=QKᵀ/√d. dP=dC Vᵀ, dV=Pᵀ dC; obtain dS through the softmax vector-Jacobian product. Then compute dQ=dS K/√d and dK=dSᵀ Q/√d. MicroTape implements this chain; Linear projections compute parameter derivatives.