Exclusive Self Attention — Detailed Summary

Shuangfei Zhai | Apple | Technical report (arXiv:2603.09078v1 [cs.LG]) | 10 Mar 2026

Per-section summary organized by paper headings. Each section includes paragraph-level bullet points and exact quantitative results where the paper provides them.

Abstract

• Introduces Exclusive Self Attention (XSA), a minimal modification of standard self attention (SA) for Transformers, intended to improve sequence-modeling performance.
• Core idea: constrain attention to capture only information that is orthogonal to the token's own value vector, thereby excluding information about the self position from the attention output.
• The exclusion of self-value information is intended to encourage better context modeling — separating SA's role (contextual aggregation) from FFN's role (point-wise feature transformation).
• Evaluated on standard language modeling at three model sizes up to 2.7 B parameters, XSA consistently outperforms SA, with the margin growing as sequence length grows.

1. Introduction

Background — the canonical Transformer block:

• Transformers (Vaswani et al., 2017) interleave Self Attention (SA) layers with Feed-Forward (FFN) layers; SA aggregates information across context while FFN performs position-wise feature updates.
• This SA/FFN division of labor has been the de-facto building block for modern Transformer variants and has remained stable for nearly a decade.

Hypothesis and the attention similarity bias:

• The paper hypothesizes that the SA/FFN division of labor can be sharpened — specifically, that SA is implicitly leaking into FFN's role.
• The authors identify a phenomenon they call the attention similarity bias: in trained Transformers, the output of an attention layer y_i tends to have a high cosine similarity with the token's own value vector v_i.

Why the bias is harmful:

• Because v_i already has a direct residual path to the next FFN layer, re-routing v_i-aligned information through SA is redundant — the FFN can see it for free via the residual stream.
• Worse, it is harmful: SA spends capacity producing direction-along-v_i signal, which competes with point-wise transformation that FFN should be doing on the residual stream. This creates a mis-allocation of capacity between SA (contextual) and FFN (point-wise).
• This reasoning motivates the proposed fix: explicitly subtract the v_i-aligned component from SA's output, so that SA outputs only context-orthogonal information.

Summary of empirical contributions:

• The paper claims that XSA (1) introduces minimal computational overhead; (2) achieves lower training and validation loss across three model sizes; (3) achieves better downstream evaluation results; (4) maintains consistent gains across different learning rates; (5) shows larger gains as sequence length increases; and (6) is robust with respect to the use of attention sinks.

2. Motivation

Formal definition of standard SA (causal):

• The paper defines a single-head causal SA as y = f(x) with per-token computation given by Equation (1):

  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.

• Here W_q, W_k, W_v are query/key/value projections; the sum is over the causal window j <= i.

Empirical demonstration of the attention similarity bias:

• The authors take a trained 1.3 B-parameter language model with sequence length 2048 (see Sect. 4 for details) and analyze each attention layer.
• For each layer, three quantities are computed and averaged over all attention heads and 1024 random training sequences (Figure 1):
  1. Avg <v_i, v_j> for i < j within a sequence — average cosine similarity of value vectors within a sequence.
  2. Avg a_{i,i} — average diagonal entry of the attention matrix (the weight a token assigns to itself).
  3. Avg <y_i, v_i> — average cosine similarity of attention output y_i with the corresponding self value vector v_i.
• Observation 1: value vectors within a sequence are positively correlated on average (left panel of Figure 1, mean cosine ~0.05-0.15 across layers, rising with depth).
• Observation 2: attention scores to the current position a_{i,i} are relatively high (middle panel; ~0.025-0.125 across layers, rising sharply at deep layers).
• Consequence: the average cosine <y_i, v_i> is correspondingly large and increases with layer depth (right panel; from ~0.2 at shallow layers to ~0.6 at the deepest layer).
• Interpretation: standard SA tends to aggregate value vectors similar to what the self value v_i already encodes. This implicitly overlaps with the role of FFN and diminishes SA's goal of contextual modeling — the attention output is partially re-encoding "what is at this position" instead of "what surrounds this position".

3. Method

Formal definition of XSA:

• The paper defines XSA as z = f(x) with the same query/key/value setup as SA, followed by an additional projection-removal step. Equation (2):

  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, z_i = y_i - (y_i^T v_i) * v_i / ||v_i||_2^2.

• The first line is exactly standard SA from Eq. (1).

• The second line projects out the component of y_i that lies along v_i, leaving z_i with no projection onto v_i.

What XSA removes:

• After the subtraction, z_i no longer contains v_i itself nor any component from context that is correlated with v_i.
• Therefore, XSA completely removes the attention similarity bias as measured in Section 2: <z_i, v_i> = 0 by construction.

Core hypothesis underlying XSA:

• In the presence of (a) residual connections that already pass v_i forward and (b) the FFN block that already does point-wise feature transformation, XSA (1) preserves the expressiveness of standard SA (the residual path ensures v_i is still available downstream), and (2) promotes modeling efficiency by forcing the attention layer to exclusively encode contextual information.
• The paper acknowledges that one could attempt a theoretical analysis but defers it; empirical evaluation is offered as the main justification.

Implementation cost:

• XSA is implemented in essentially two extra lines on top of standard SA. Algorithm 1 gives PyTorch-style pseudocode for multi-head causal XSA:

def exclusive_self_attention(x, Wq, Wk, Wv, Wo, H):
    B, T, D = x.shape
    # linear projections
    Q = (x @ Wq).reshape(B, T, H, D // H).transpose(1, 2)
    K = (x @ Wk).reshape(B, T, H, D // H).transpose(1, 2)
    V = (x @ Wv).reshape(B, T, H, D // H).transpose(1, 2)
    # standard multi-head attention
    Y = torch.nn.functional.scaled_dot_product_attention(Q, K, V, is_causal=True)
    # XSA mode
    Vn = torch.nn.functional.normalize(V, dim=-1)
    Z = Y - (Y * Vn).sum(dim=-1, keepdim=True) * Vn
    # output projection
    out = Z.transpose(1, 2).reshape(B, T, D) @ Wo
    return out

• The two XSA-specific lines are the L2-normalization of V and the subtraction of the (Y . V_n) projection along V_n.

4. Experiments

4.1 Setup

Codebase:

• All experiments use the NanoGPT codebase (https://github.com/karpathy/nanoGPT) for reproducibility.
• Two changes are made to NanoGPT's Transformer implementation:
  • Learned position embeddings are replaced with RoPE (Su et al., 2023), which is a common practice in modern LMs.
  • An additional LayerNorm (Ba et al., 2016) is inserted right after the token embeddings; this is found to improve training stability.
  • The number of attention heads and the head dimension are allowed to be configured independently, enabling more flexible model-size choices.

Architectures and learning rates (Table 1):

• All models are trained with a batch size of 0.5 M tokens for 200 K iterations.

Dataset:

• The training corpus is FineWeb-100BT (Penedo et al., 2024), which contains ~100 billion tokens.
• Preprocessing follows NanoGPT's protocol: tokenize with the GPT-2 tokenizer (Radford et al., 2019) and randomly hold out 0.05% of the tokens as the validation set.

Training details:

• Default context length is 2048; global batch size is 256; training duration is 200 K iterations, yielding 100 B training tokens, which is roughly one epoch over FineWeb-100BT.
• Optimizer is AdamW (Loshchilov and Hutter, 2017).
• Learning rate schedule: linear warm-up over 2 K steps to the peak LR, then cosine decay to 1/10 of the peak LR.
• A grid search of learning rate is performed for each baseline model configuration; the resulting LR is then reused for the corresponding XSA variant (shown in Table 1).
• Three model sizes are evaluated: 0.7 B, 1.4 B, and 2.7 B non-embedding parameters (note: the body text says 1.4 B while Section 2's analysis references a "1.3 B" model; the paper uses these labels for the same configuration in different places).

4.2 Results

Computational overhead (Figure 2):

• XSA is benchmarked against standard SA on an attention block (attention + FFN) across varied sequence lengths and model widths.
• Hardware: a single NVIDIA B200 GPU; batch size 32; numerical precision bfloat16.
• Sweeps are over sequence length in {512, 1024, 2048, 4096, 8192, 16384} with d_model = 2048 fixed, and over d_model with sequence length = 2048 fixed.
• Result: forward+backward time and peak memory of XSA closely track the baseline — XSA introduces minimal overhead in both speed and memory at all tested settings (qualitative; the curves nearly overlap).

Model size — training/validation loss (Figure 3):

• Training and validation loss curves are shown for the 0.7 B, 1.3 B, and 2.7 B models, plotted over 200 K training iterations.
• XSA maintains a clear margin over the baseline at every checkpoint for every model size, in both training and validation loss.
• The validation curves at the end of training are approximately 2.45-2.65 for XSA vs. 2.55-2.75 for baseline depending on model size (read from Figure 3).

Downstream evaluation (Table 2):

• Final checkpoints are evaluated on 8 downstream tasks covering language, knowledge, and reasoning: ARC-Easy (Clark et al., 2018), BoolQ (Clark et al., 2018), HellaSwag (Zellers et al., 2019), LAMBADA (Paperno et al., 2016), OpenBookQA (Mihaylov et al., 2018), PIQA (Ba et al., 2016), SocialIQA (Sap et al., 2019), and WinoGrande (Sakaguchi et al., 2021).
• Evaluation is run with the Language Model Evaluation Harness (Gao et al., 2024). Accuracy is reported for BoolQ, LAMBADA, SocialIQA, WinoGrande; length-normalized accuracy is reported for ARC-Easy, HellaSwag, OpenBookQA, PIQA.

• Bold entries indicate the better of the pair per column.
• Across all three model sizes, XSA improves average downstream accuracy: +0.26 (0.7 B), +1.03 (1.3 B), +1.36 (2.7 B).
• The DAvg gain grows with model size; the authors speculate that XSA will remain advantageous at even larger scale (both model size and data size).

Learning rate robustness (Figure 4):

• Using the 1.3 B model, XSA is compared with the baseline at four different peak learning rates: 1e-4, 2e-4, 4e-4, and 6e-4.
• For all four LRs, XSA's training and validation loss is uniformly below baseline by an approximately constant margin (roughly 0.02-0.03 nats).
• Conclusion: the XSA gain is not a side-effect of the chosen LR — the margin is robust across LR settings.

Sequence-length sweep (Figure 5):

• Using the 1.3 B model, XSA is trained at six sequence lengths: {512, 1024, 2048, 4096, 8192, 16384}.
• The same peak learning rate is used for all settings; the batch size is adjusted so that tokens per batch remains constant at 0.5 M.
• Result: XSA outperforms the baseline at every sequence length, and the margin grows with sequence length — both training and validation loss show a widening gap as length scales from 512 to 16384.
• The authors interpret this as evidence that longer sequences impose greater pressure on the attention layer to do context modeling well, and XSA's removal of the self-projection becomes more valuable in that regime.
• They argue XSA is a promising technique for long-context modeling, one of the critical problems of scaling Transformers.

Attention sinks (Figure 6):

• XSA can be viewed as an implicit form of attention sink: instead of prepending a set of learned sink tokens (Xiao et al., 2023), XSA allocates some of the attention mass to a_{i,i} that is then projected out of the output — effectively a way to dump unwanted attention.
• The authors run additional experiments where both the baseline and XSA models include 0, 1, or 4 explicit learned attention sinks, and report the result in Figure 6.
• Conclusion: XSA's loss margin over the baseline persists in the presence of explicit attention sinks. Both training and validation loss curves show that XSA stays roughly 0.02 below baseline regardless of the number of sinks; explicit sinks neither close the gap nor disrupt XSA.

5. Discussions

• The paper concludes that XSA shows promising performance on standard language modeling tasks.
• Many open questions remain that the paper does not address:
  • How does XSA behave at substantially larger model and data scale than 2.7 B / 100 B tokens?
  • Is XSA compatible with newer optimizers such as Muon (Jordan et al., 2024)?
  • Does XSA work for tasks or modalities beyond language modeling (e.g., vision, speech, multi-modal)?
• The authors explicitly invite future work to answer these.

6. Related Work and Differentiation (synthesized from the body)

• Vaswani et al., 2017 — the original Transformer; XSA is a minimal patch on top of the canonical SA defined here.
• RoPE (Su et al., 2023) — XSA uses RoPE as its position embedding scheme, following modern LM practice.
• LayerNorm (Ba et al., 2016) — an extra LN is inserted after token embeddings for stability.
• AdamW (Loshchilov and Hutter, 2017) — optimizer used; Muon (Jordan et al., 2024) is referenced as a future-work compatibility question.
• FineWeb (Penedo et al., 2024) — training data.
• NanoGPT (Karpathy) — the codebase upon which experiments are built.
• Attention Sink (Xiao et al., 2023) — XSA is explicitly compared and framed as related: XSA acts as an implicit sink by routing self-projected attention out of the output.
• The differentiating contribution is the attention similarity bias diagnosis — an empirical phenomenon and a single-equation fix that the paper claims is novel and previously unaddressed.

7. Limitations (synthesized; not explicitly enumerated)

• Empirical justification only. The paper explicitly defers any theoretical analysis of why excluding the v_i direction should preserve expressiveness and improve efficiency.
• Scale ceiling: experiments stop at 2.7 B parameters and 100 B tokens. Behavior at frontier-scale (10s of billions of parameters, trillion-token corpora) is conjectured but not measured.
• Single modality: only autoregressive language modeling is evaluated. The Discussions section explicitly notes that other modalities are open.
• Optimizer coverage: only AdamW is tested. Compatibility with second-order or geometry-aware optimizers (Muon) is unverified.
• One codebase / one tokenizer / one dataset: results are tied to NanoGPT + GPT-2 tokenizer + FineWeb-100BT; potential interaction with other tooling is unexplored.
• No ablation on how much of the v_i direction to subtract: XSA fully removes the v_i projection; a partial-removal variant (subtract a fraction) is not tested.
• The 0.7 B model shows a smaller gain (+0.26 average) than 1.3 B and 2.7 B (+1.03, +1.36), and on individual tasks XSA can lose to the baseline (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). The paper does not analyze these per-task losses.

8. Open Problems Explicitly Identified

The Discussions section names three concrete open questions:

1. Scale behavior. Does XSA's advantage persist or widen at substantially larger scale in both model size and training data?
2. Optimizer compatibility. Does XSA compose with optimizers other than AdamW, in particular Muon (Jordan et al., 2024)?
3. Modality/task generality. Does the XSA modification provide gains beyond language modeling — e.g., vision Transformers, speech, or multi-modal models?

9. Cross-Cutting Empirical Take-aways

10. Conclusion (paper's own)

• The paper presents XSA as a simple, drop-in replacement for SA in Transformer blocks: it explicitly excludes the self-value direction from the attention output, eliminating an empirically observed bias and improving language-modeling performance.
• The empirical case is built on three model sizes, eight downstream tasks, learning-rate sweeps, sequence-length sweeps, and an attention-sink comparison — all pointing in the same direction (XSA <= baseline in loss, XSA >= baseline in downstream accuracy).
• The authors leave theoretical analysis, larger-scale validation, optimizer studies, and other-modality studies as future work.