Multi-Head Attention Residuals
Cheng Luo ⋅ Zefan Cai ⋅ Junjie Hu
Abstract
Transformers propagate information across depth through a single additive residual stream: every sublayer reads only the most recent state. \emph{Attention residuals} relax this by letting each sublayer attend, through a learned softmax, over the outputs of \emph{all} previous layers---an attention over depth. But that read uses a \emph{single} query shared across the entire width, so every feature subspace must read the depth history through one distribution. The resulting \emph{forced compromise} costs more the more subspaces disagree about which layers to read---and disagreement grows with model width. We introduce \textbf{Multi-Head Attention Residuals (MHAR)}: reshape the routing query into $H$ per-subspace heads, each with its own softmax over the depth history. The read becomes block-diagonal, the reshape adds \emph{zero} parameters and negligible compute, and $H{=}1$ recovers attention residuals exactly. Trained from scratch on FineWeb-Edu, MHAR improves validation loss over a standard Transformer at 100M, 350M, and 1B ($-0.049$, $-0.080$, $-0.063$). What grows with scale is the \emph{necessity} of multiple heads: a single routing query helps at 100M but by 1B is $0.105$ \emph{worse} than the plain additive residual---even at its own optimal learning rate---so MHAR's advantage over single-head routing widens from $0.010$ to $0.168$. Setting the number of routing heads equal to the number of KV heads is a hyperparameter-free, near-optimal default, and a direct probe of the trained queries supports the mechanism: learned subspace disagreement grows monotonically with scale, peaks at the regression scale, and is isolated to width by a fixed-depth widening control.
Chat is not available.
Successful Page Load