Exclusive Self Attention

TL;DR

A two-line modification of self attention (XSA) that subtracts off the projection of each token’s attention output onto its own value vector, so attention is forced to carry only context-orthogonal information. The motivation is the empirical “attention similarity bias”: in standard SA the aggregated output $y_i$ is highly cosine-similar to the self value $v_i$ , meaning SA spends capacity replicating point-wise features that the FFN+residual path already provides. Across 0.7B/1.4B/2.7B NanoGPT LMs trained on 100B FineWeb tokens, XSA improves training and validation loss, lifts average downstream accuracy by +0.26 / +1.03 / +1.36 points, holds across four learning rates, and produces increasingly larger gains as sequence length grows up to long contexts.

Key claims

Standard self attention exhibits an “attention similarity bias”: across layers, the average diagonal of the attention matrix is high and the aggregated value $y_i$ has high cosine similarity with $v_i$ , with the cosine rising as a function of layer index [§2, Fig. 1].
The bias is wasted capacity: the FFN already updates each position pointwise via the residual path, so SA modeling self-correlated content overlaps with the FFN’s role and crowds out contextual modeling [§2].
XSA = z_i = y_i - (y_iᵀv_i) v_i / ‖v_i‖², applied per token per head; this completely removes the attention similarity bias by construction [§3, Eq. 2, Algorithm 1].
Compute overhead is minimal — speed and memory of (attention + FFN) are essentially unchanged on a B200 at bfloat16 across sequence lengths [§4.2 “Computational overhead”, Fig. 2].
XSA beats standard SA on both training and validation loss at 0.7B, 1.4B, and 2.7B non-embedding parameters trained for 100B FineWeb tokens; downstream-average gain grows with model size: +0.26 (0.7B), +1.03 (1.4B), +1.36 (2.7B) points across an 8-task harness (ARC-E, BoolQ, HellaSwag, LAMBADA, OpenBookQA, PIQA, SocialIQA, WinoGrande) [§4.2, Table 2, Fig. 3].
The gain is robust to learning rate: at 1.3B across four LRs the XSA-vs-SA loss margin is near-constant, ruling out “XSA just needs a different LR” as a confound [§4.2 “Learning rate”, Fig. 4].
The gain grows with sequence length: at 1.3B trained with constant 0.5M tokens/batch across six sequence lengths, XSA’s lead over SA increases monotonically — suggesting attention similarity bias becomes more harmful as context modeling gets harder [§4.2 “Sequence length”, Fig. 5].
XSA is compatible with — and behaves as an implicit form of — attention sinks: adding learned sink tokens to both baseline and XSA leaves XSA’s loss margin intact, so the two mechanisms are complementary rather than substitutes [§4.2 “Comparison to Attention Sink”, Fig. 6].

Method

The full method is the two-line change in Equation 2: compute standard $y_i = \mathrm{softmax}(QK^\top)V$ , then project out the component along the self value with $z_i = y_i - (y_i^\top v_i) v_i / \|v_i\|^2$ . The rest of the architecture is unchanged.

The motivating diagnostic in §2 averages three quantities per attention layer across 1024 random training sequences for a trained 1.3B / 2048-context LM: (i) cosine similarity of value vectors within a head, (ii) the diagonal of the softmax attention map, and (iii) cosine similarity of the aggregated $y_i$ with $v_i$ . All three are positive and tend to grow with depth, motivating the claim that SA is spending capacity on the point-wise direction that the FFN/residual already covers. XSA removes that component directly rather than hoping the model learns to avoid it.

Experiments are all run inside the NanoGPT codebase with three modifications — RoPE replaces learned positional embeddings, an extra LayerNorm sits right after token embeddings for stability, and the head count / head dim are decoupled. Training uses AdamW, 2K-step linear warmup, cosine decay to 10× lower, 2048 context, 256 global batch, 200K iterations ≈ 100B tokens (≈ one epoch of FineWeb-100BT tokenized with GPT-2 BPE). Three sizes (0.7B / 1.4B / 2.7B non-embedding params) are grid-searched for the baseline’s best LR and that LR is reused for XSA — so all XSA gains are at fixed baseline-optimal LR, not from re-tuning.

Results

2.7B, 100B FineWeb tokens: XSA 59.42 vs. baseline 58.06 average across the 8-task harness, +1.36 points [Table 2]. The gap widens with model size (0.7B +0.26 → 1.4B +1.03 → 2.7B +1.36).
Training/validation loss curves at all three sizes show a clear, persistent XSA margin from early in training that does not close [Fig. 3].
LR sweep at 1.3B across four learning rates: roughly constant XSA loss advantage at every LR [Fig. 4].
Sequence-length sweep at 1.3B across six lengths (with batch sized so tokens/step ≈ 0.5M): XSA’s margin grows monotonically with sequence length [Fig. 5].
Speed and memory of one transformer block (attention + FFN) on a B200 with batch 32, bfloat16: XSA on/off curves are visually indistinguishable [Fig. 2].
With learned attention sinks added to both baseline and XSA, the XSA loss margin persists — i.e. the two are stacking, not redundant [Fig. 6].

Why it’s interesting

XSA is the attention-side member of the same architectural-stability family already on the wiki and indexed in Training stability at scale. The recipe — “diagnose a load-bearing geometric artifact in modern transformers, then propose a minimal architectural change that removes the artifact while preserving (or improving) training dynamics” — is exactly the playbook used by A Unified View of Attention and Residual Sinks: Outlier-Driven Rescaling is Essential for Transformer Training (replace outlier-driven implicit rescaling with explicit Gated Attention / GatedNorm) and The Spike, the Sparse and the Sink: Anatomy of Massive Activations and Attention Sinks (trace SwiGLU → RMSNorm → attention-head pathway and decouple via normalization configuration). The interesting tension is that XSA’s diagnostic — high cosine( $y_i$ , $v_i$ ) caused partly by elevated diagonal attention scores — overlaps with the attention sink phenomenon those two papers analyze, but XSA’s fix attacks the consequence (the self-aligned output) rather than the source (outliers in residual stream or attention logits). The §4.2 sinks experiment shows XSA composes with explicit sink tokens, which is a strong hint that “XSA + GatedNorm” is a sensible thing to try on a Nano-Diffusion-style stack. For Luma the practical hook is the long-context result: the gap grows with sequence length, which is the regime Luma’s video DiTs spend their inference budget on. The fix is two lines on top of any FlashAttention call.