Exclusive Self Attention
Shuangfei Zhai | Apple | Technical report (arXiv:2603.09078v1 [cs.LG]) | 10 Mar 2026
Problem
Standard self attention (SA) in Transformers exhibits an empirically observable pathology: the attention output y_i has a high cosine similarity with the token's own value vector v_i, growing from roughly 0.2 at shallow layers to roughly 0.6 at deep layers in a trained 1.3 B-parameter language model. The authors term this the attention similarity bias. Since v_i already reaches the next FFN block via the residual stream, having SA re-emit a v_i-aligned component is redundant; worse, it competes with FFN's point-wise transformation role, mis-allocating capacity between contextual modeling (SA's job) and point-wise feature updates (FFN's job). The question is whether a minimal fix to SA that removes this bias improves language-modeling performance without disturbing the rest of the Transformer block.
Core Insight
Constrain attention to output only the component orthogonal to v_i. A two-line projection-removal step appended to SA fully eliminates <y_i, v_i> by construction; in the presence of residual connections and FFN, expressiveness is preserved (v_i still flows downstream via the residual) while the attention layer is forced to exclusively encode contextual information.
Method
Standard causal SA is defined as
q_i = W_q x_i, k_j = W_k x_j, v_j = W_v x_j,
a_{i,j} = exp(q_i^T k_j) / sum_{j'=1..i} exp(q_i^T k_{j'}),
y_i = sum_{j=1..i} a_{i,j} v_j.Exclusive Self Attention (XSA) appends a single projection-removal step:
z_i = y_i - (y_i^T v_i) * v_i / ||v_i||_2^2.Implementation is two extra lines on top of multi-head causal SA: L2- normalize V, then subtract (Y . V_n) projected along V_n. Algorithm 1 in the paper gives the PyTorch-style pseudocode; the rest of the Transformer block (residual stream, FFN, RoPE position embeddings, LayerNorm) is unchanged.
Experimental Setup
| Component | Value |
|---|---|
| Codebase | NanoGPT (Karpathy), with RoPE position embeddings and an extra post-embedding LayerNorm |
| Optimizer | AdamW, peak LR per Table 1, 2 K-step linear warm-up, cosine decay to LR/10 |
| Context length (default) | 2048 |
| Global batch size | 256 sequences = 0.5 M tokens |
| Iterations | 200 K (~100 B tokens; ~1 epoch over FineWeb-100BT) |
| Dataset | FineWeb-100BT (Penedo et al., 2024), GPT-2 tokenizer, 0.05% held-out validation |
| Model sizes | 0.7 B (24 layers, d=1536, 6 heads x 256), 1.4 B (24 layers, d=2048, 24 heads x 128), 2.7 B (32 layers, d=2560, 24 heads x 128) |
| Speed/memory benchmark | NVIDIA B200 GPU, batch 32, bf16 |
| Downstream evaluation | LM Evaluation Harness on ARC-Easy, BoolQ, HellaSwag, LAMBADA, OpenBookQA, PIQA, SocialIQA, WinoGrande |
| LR robustness sweep | 1.3 B model, peak LR in {1e-4, 2e-4, 4e-4, 6e-4} |
| Sequence-length sweep | 1.3 B model, lengths in {512, 1024, 2048, 4096, 8192, 16384}, tokens-per-batch held constant at 0.5 M |
| Attention sink sweep | 1.3 B model, {0, 1, 4} explicit learned sinks |
Headline Quantitative Results
Attention similarity bias in baseline 1.3 B model (Figure 1):
- Avg cosine <v_i, v_j> within a sequence: ~0.05-0.15, rising with depth.
- Avg diagonal attention a_{i,i}: ~0.025-0.125, rising sharply at deep layers.
- Avg cosine <y_i, v_i>: ~0.2 (shallow) -> ~0.6 (deep).
Computational overhead (Figure 2):
- Forward+backward time and peak memory curves for XSA closely track the baseline at every tested sequence length (512 to 16384) and at every tested d_model (with seq=2048). Overhead is minimal.
End-of-training downstream accuracy (Table 2):
| Model size | Variant | ARC-E | BoolQ | HSwag | LAMBADA | OBQA | PIQA | SocIQA | WinoGr | Avg | DAvg |
|---|---|---|---|---|---|---|---|---|---|---|---|
| 0.7 B | Baseline | 51.26 | 61.07 | 55.68 | 52.82 | 35.00 | 74.05 | 40.02 | 55.88 | 53.22 | |
| 0.7 B | XSA | 52.69 | 61.19 | 56.29 | 54.07 | 32.20 | 73.78 | 41.45 | 56.20 | 53.48 | +0.26 |
| 1.3 B | Baseline | 56.19 | 65.47 | 60.69 | 56.24 | 34.60 | 75.90 | 41.40 | 58.80 | 56.16 | |
| 1.3 B | XSA | 58.84 | 62.29 | 62.41 | 58.57 | 36.00 | 76.61 | 42.84 | 59.98 | 57.19 | +1.03 |
| 2.7 B | Baseline | 58.59 | 60.98 | 66.20 | 60.18 | 37.00 | 76.61 | 42.94 | 60.80 | 58.06 | |
| 2.7 B | XSA | 60.65 | 64.86 | 67.40 | 62.04 | 38.40 | 77.80 | 41.45 | 62.75 | 59.42 | +1.36 |
- XSA wins on average at every model size, and the gain grows with model size (+0.26 -> +1.03 -> +1.36).
- Loss curves (Figure 3) show XSA below baseline at every checkpoint for all three sizes.
Robustness sweeps:
- LR (Figure 4, 1.3 B): XSA is below baseline at all four LRs by an approximately constant 0.02-0.03 margin in both training and validation loss.
- Sequence length (Figure 5, 1.3 B): XSA's margin grows as length scales from 512 to 16384.
- Attention sinks (Figure 6, 1.3 B): XSA's ~0.02 loss margin persists for {0, 1, 4} explicit learned sinks; sinks neither close the gap nor destabilize XSA.
Limitations
- The paper provides only empirical justification; a theoretical argument for why removing the v_i projection preserves expressiveness while improving efficiency is explicitly deferred.
- Scale ceiling: 2.7 B parameters, 100 B tokens, ~1 epoch.
- Single modality (autoregressive language modeling) and single optimizer family (AdamW).
- Single codebase / tokenizer / dataset (NanoGPT + GPT-2 BPE + FineWeb-100BT).
- Only full projection removal is tested; no ablation on a partial / fractional removal of the v_i direction.
- Per-task losses to baseline exist and are not analyzed: OBQA at 0.7 B (32.20 vs 35.00), BoolQ at 1.3 B (62.29 vs 65.47), SocIQA at 2.7 B (41.45 vs 42.94).
Open Problems
- Scale behavior. Does the XSA advantage persist or widen at substantially larger model and data scale than 2.7 B / 100 B tokens?
- Optimizer compatibility. Does XSA compose with non-AdamW optimizers, in particular Muon (Jordan et al., 2024)?
- Modality and task generality. Does XSA help beyond autoregressive language modeling — e.g., vision Transformers, speech, multi-modal?