Architecture & Measurement-Design Analysis

Exclusive Self Attention (XSA)

Source: Zhai, S. Exclusive Self Attention. Technical report, Apple. arXiv:2603.09078v1 [cs.LG], 10 Mar 2026. Author: Shuangfei Zhai (Apple, [email protected]). Codebase used: NanoGPT (https://github.com/karpathy/nanoGPT). Dataset used: FineWeb-100BT (Penedo et al., 2024), GPT-2 tokenizer. Reader: Direct PDF read (Read tool, full 7 pages in one request). Analyst: Vishwakarma Date: 2026-05-14


Table of Contents

  1. Evaluation Harness Architecture (the "instrument")
  2. System-Under-Test Architecture (the "specimen") — XSA vs SA block
  3. Design-Space Diagram (axes swept, axes held fixed)
  4. Algorithm / Control Flow Diagrams (key procedures)
  5. Quantitative Results — Empirical Findings by Regime
  6. Configuration-Regime Trade-off Tables
  7. Bottlenecks & Insights Surfaced by the Measurements
  8. Limitations of the Methodology
  9. Analogy

1. Evaluation Harness Architecture (the "instrument")

The harness is a paired A/B comparator that re-runs a fixed NanoGPT training recipe with exactly one architectural change — the self-attention block — toggled between standard SA (baseline) and exclusive SA (XSA). Every observed difference in loss, downstream accuracy, or runtime is attributable to that single two-line code change (Algorithm 1 in the paper). Unlike systems papers that sweep collective primitives or transports, this paper sweeps a mathematical operator inside the model and measures whether that operator change persists across model scale, learning rate, sequence length, and the presence of attention sinks. The harness is therefore not a microbenchmark of a kernel; it is a training-curve comparator that asks "does this loss-margin survive the four most common ways a language-modeling result fails to replicate?"

+------------------------------------------------------------------+
|                  Paired-A/B Training Harness                     |
|                                                                  |
|  +---------------------+      +-------------------------------+  |
|  | Training Driver     |----->| Model Builder (NanoGPT)       |  |
|  | (NanoGPT loop;      |      |  - n_layers, d_model, n_heads |  |
|  |  AdamW; 200K iters; |      |  - RoPE position embeddings   |  |
|  |  global batch 256;  |      |  - extra LayerNorm after embed|  |
|  |  ctx 2048; 100B tok)|      |  - Attention block: SA or XSA |  |
|  +---------------------+      +---------------+---------------+  |
|              |                                |                  |
|              v                                v                  |
|  +------------------------------------------------------------+  |
|  |          Attention-Block Switch (two-line diff)            |  |
|  |                                                            |  |
|  |   Baseline path (SA):                                      |  |
|  |     y_i = sum_j a_{i,j} v_j                                |  |
|  |                                                            |  |
|  |   XSA path:                                                |  |
|  |     y_i = sum_j a_{i,j} v_j                                |  |
|  |     z_i = y_i - (y_i^T v_i) * v_i / ||v_i||_2^2            |  |
|  |     (project out the self-value direction)                 |  |
|  +------------------------------------------------------------+  |
|              |                                                   |
|              v                                                   |
|  +------------------------------------------------------------+  |
|  |  Common Software Stack                                     |  |
|  |    PyTorch + torch.nn.functional.scaled_dot_product_attn   |  |
|  |    bfloat16 numerics, B200 GPU (efficiency benchmark only) |  |
|  +------------------------------------------------------------+  |
|              |                                                   |
|              v                                                   |
|  +------------------------------------------------------------+  |
|  |  Measurement Protocol                                      |  |
|  |    - Train 200K iters (~100B tokens, ~1 epoch)             |  |
|  |    - Record training loss and validation loss every step  |  |
|  |    - Save final checkpoint; eval on 8 downstream tasks    |  |
|  |    - Efficiency micro-benchmark: fwd+bwd time, peak mem    |  |
|  |      on B200 with bs=32, varying seq len and d_model       |  |
|  +------------------------------------------------------------+  |
|              |                                                   |
|              v                                                   |
|  +------------------------------------------------------------+  |
|  |  Result Aggregator                                         |  |
|  |    - Loss curves (Fig. 3, 4, 5, 6)                         |  |
|  |    - Downstream accuracy table (Table 2)                   |  |
|  |    - Efficiency curves (Fig. 2)                            |  |
|  +------------------------------------------------------------+  |
+------------------------------------------------------------------+
^ Fig 1: Measurement harness for XSA. The single switch in the
  middle is the entire experiment: every other component — driver,
  builder, data, hardware — is shared between the baseline and XSA
  arms. Every reported delta is therefore attributable to the
  attention-block change alone.

The harness is unusually clean for a method paper. Three properties matter. First, the NanoGPT lineage means the recipe is shared with a large community of replication attempts, so a reader can re-run exactly the same baseline. Second, the batch-size discipline — global batch 256, 0.5M tokens/iter, 200K iters — is held identical across SA and XSA arms so that loss curves are directly comparable without re-normalization. Third, the harness explicitly varies four axes orthogonal to the architectural change: model size, learning rate, sequence length, and number of attention sinks. Each of these is a known failure mode for "improved attention" papers (many proposed mechanisms collapse outside the original hyperparameter recipe), and the paper budgets one ablation figure per failure mode.

Methodology specifics extracted verbatim:

Knob Value
Codebase NanoGPT (Karpathy)
Position embedding RoPE (Su et al., 2023)
Extra normalization LayerNorm after token embeddings
Tokenizer GPT-2 (Radford et al., 2019)
Training data FineWeb-100BT (~100B tokens)
Validation split 0.05% of tokens
Optimizer AdamW (Loshchilov & Hutter, 2017)
LR schedule Linear warmup (2K steps) + cosine to 1/10 max
Context length (default) 2048
Global batch (default) 256 sequences (= 0.5M tokens / step)
Training duration 200K iterations (~100B tokens; ~1 epoch)
Numerics bfloat16 (efficiency bench), training default
Efficiency-bench hardware Single B200 GPU
Efficiency-bench batch 32
Downstream eval harness LM Evaluation Harness (Gao et al., 2024)
Variance / seeds reported None — single run per (size, method)
Open-source release Not announced in the paper

The absence of multi-seed runs is the harness's most important limitation; it is partly compensated by the four orthogonal sweeps (LR, seq-len, sinks, scale), each of which acts as a soft replication of the headline 1.3B / 200K result. Each sweep is an opportunity for the reported margin to disappear, and the paper's contribution is precisely that it does not.


2. System-Under-Test Architecture (the "specimen") — XSA vs SA block

The system under test is one attention layer, not a cluster. The architectural delta is two lines of PyTorch wrapped around the standard scaled-dot-product attention call (Algorithm 1 in the paper).

+------------------------------------------------------------------+
|             Standard Self Attention Block (baseline)             |
|                                                                  |
|   x  ----+                                                       |
|         |                                                        |
|         |   Q = x W_q                                            |
|         +---> K = x W_k                                          |
|         |   V = x W_v                                            |
|         |                                                        |
|         v                                                        |
|   +---------------------------+                                  |
|   | scaled_dot_product_attn  |  a_{i,j} = softmax(q_i^T k_j)    |
|   |  (causal mask, multi-hd) |  y_i = sum_{j<=i} a_{i,j} v_j    |
|   +-----------+---------------+                                  |
|               |                                                  |
|               v                                                  |
|        y  -->  W_o  -->  out                                     |
|                                                                  |
+------------------------------------------------------------------+

                              versus

+------------------------------------------------------------------+
|             Exclusive Self Attention (XSA) Block                 |
|                                                                  |
|   x  ----+                                                       |
|         |    Q, K, V  (same as SA)                               |
|         v                                                        |
|   +---------------------------+                                  |
|   | scaled_dot_product_attn  |  y_i = sum_{j<=i} a_{i,j} v_j    |
|   +-----------+---------------+                                  |
|               |                                                  |
|               |  +---------------------------------+             |
|               |  |     Self-value projector        |             |
|               +->|   V_n = V / ||V||_2  (unit dir) |             |
|               |  |   s_i = (y_i^T v_{n,i}) v_{n,i} |             |
|               |  +---------------------------------+             |
|               v                                                  |
|        z_i = y_i - s_i  (remove component along v_i)             |
|               |                                                  |
|               v                                                  |
|        z  -->  W_o  -->  out                                     |
|                                                                  |
+------------------------------------------------------------------+
^ Fig 2: SA vs XSA. The XSA block is identical to SA up to the
  attention output y, then subtracts the projection of y onto the
  unit self-value direction v_n. This zeroes out the component of
  the attention output that is parallel to the token's own value
  vector — what the paper names the "attention similarity bias."

The geometric intuition is direct. In the baseline, the diagonal attention weight a_{i,i} and the empirical correlation between value vectors in a sequence together produce a strong systematic component along v_i in the attention output y_i. The paper's Figure 1 demonstrates this empirically on a trained 1.3B model: the average cosine similarity <y_i, v_i> climbs from ~0.2 at the first layer to ~0.6 at the last layer. XSA removes precisely that component in one Gram-Schmidt-style step.

The reason this is interesting as a system is the division-of- labor argument: the residual path already carries the token's positional/value content forward to the FFN, so the attention block sending another copy of v_i is redundant at best and competes with contextual modeling at worst. XSA enforces a structural orthogonality between what attention contributes (context) and what the residual carries (self), without requiring a new parameter, a new loss term, or a new training schedule.

The PyTorch implementation reproduced verbatim from Algorithm 1:

# x: (B, T, D); Wq, Wk, Wv, Wo: (D, D); H = number of heads

def exclusive_self_attention(x, Wq, Wk, Wv, Wo, H):
    B, T, D = x.shape
    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: project out the self-value direction
    Vn = torch.nn.functional.normalize(V, dim=-1)
    Z  = Y - (Y * Vn).sum(dim=-1, keepdim=True) * Vn

    out = Z.transpose(1, 2).reshape(B, T, D) @ Wo
    return out

Notice what is not in the diff: no extra parameters (Wo is the same), no extra collective communication, no extra loss term, no extra activation memory beyond what Vn and Z consume (and Z overwrites Y in practice). The complexity is O(B*T*H*d_head) for the projection — sub-leading to the O(B*T^2*d_head) attention itself.


3. Design-Space Diagram (axes swept, axes held fixed)

The independent variables form a small but well-targeted sweep. Every figure in the paper fixes everything except one axis at a time.

                  DESIGN SPACE (5 axes + held-fixed)
  +---------------------------------------------------------------+
  |                                                               |
  |  Axis 1: ATTENTION MECHANISM (2 levels)                       |
  |    [Standard SA]   <- baseline                                |
  |    [XSA]           <- proposed                                |
  |                                                               |
  |  Axis 2: MODEL SIZE (3 levels) — Table 1                      |
  |    [0.7B]  n_layers=24, d_model=1536, n_heads= 6, d_head=256  |
  |    [1.4B]  n_layers=24, d_model=2048, n_heads=24, d_head=128  |
  |    [2.7B]  n_layers=32, d_model=2560, n_heads=24, d_head=128  |
  |                                                               |
  |  Axis 3: LEARNING RATE (4 levels) — Fig. 4                    |
  |    [1e-4]  [2e-4]  [4e-4]  [6e-4]   (at 1.3B model)           |
  |                                                               |
  |  Axis 4: SEQUENCE LENGTH (6 levels) — Fig. 5                  |
  |    [512] [1024] [2048] [4096] [8192] [16384]                  |
  |    (batch size adjusted so tokens/step = 0.5M)                |
  |                                                               |
  |  Axis 5: NUMBER OF ATTENTION SINKS (3 levels) — Fig. 6        |
  |    [0]  [1]  [4]   (at 1.3B model)                            |
  |                                                               |
  |  Held FIXED (no sweep):                                       |
  |    - Optimizer: AdamW (Muon flagged as future work)           |
  |    - LR schedule: cosine with 2K warmup, decay to 1/10 max    |
  |    - Tokenizer: GPT-2                                         |
  |    - Dataset: FineWeb-100BT                                   |
  |    - Training duration: 200K iters (~100B tokens)             |
  |    - Global batch tokens: 0.5M (256 seq * 2048 ctx)           |
  |    - Position embedding: RoPE                                 |
  |    - Modality: language modeling only                         |
  |    - Random seed: single seed per cell                        |
  |                                                               |
  +---------------------------------------------------------------+
^ Fig 3: 5-axis design space. The headline experiment is the
  (SA, XSA) x (0.7B, 1.4B, 2.7B) cross. The three remaining axes
  (LR, seq-len, sinks) are robustness ablations holding model
  size fixed at 1.3B.

The held-fixed list is the more revealing column. Three absences shape what XSA's evidence can claim. Optimizer fixed at AdamW: the discussion section names Muon as an explicit open question — XSA might interact with second-order-flavored optimizers differently. Single modality (language): the paper makes no vision or multimodal claim. Single seed per cell: the ablations are doing the heavy lifting on variance estimation rather than direct replication. The sweep is therefore a "robustness ladder" rather than a "confidence interval" — the question being answered is "does the margin survive perturbations along axes that typically erase attention-paper margins?" not "what is the standard error of the margin?"


4. Algorithm / Control Flow Diagrams (key procedures)

4.1 Per-step control flow — one attention layer, one forward pass

  START (one attention block, one input tensor x : (B,T,D))
       |
       v
  (1) Compute Q, K, V via three linear projections; reshape to
      (B, H, T, d_head)
       |
       v
  (2) Run causal scaled-dot-product attention:
        a_{i,j} = softmax(q_i^T k_j) over j <= i
        Y_i    = sum_{j<=i} a_{i,j} V_j
       |
       v
  (3) Mode switch
       |
       +-- if SA (baseline) ---> Z = Y                  --> step (5)
       |
       +-- if XSA           ---> step (4)
       |
       v
  (4) XSA correction:
        Vn = normalize(V, dim=-1)              # unit value vector
        s  = (Y * Vn).sum(dim=-1, keepdim=True) * Vn
        Z  = Y - s                              # remove self component
       |
       v
  (5) Reshape Z to (B, T, D); project through W_o
       |
       v
  END  -> return out
^ Fig 4: Per-step control flow inside one attention block. The
  XSA branch adds three element-wise ops (normalize, dot, sub)
  after the attention kernel; numerics-wise it is a small
  rank-one Gram-Schmidt step.

The branch at step (3) is what makes XSA cheap. Standard fused attention kernels (FlashAttention, scaled_dot_product_attention) already compute Y; the XSA branch is a pure post-processing pass that touches Y once and V once, without re-reading K or the attention scores. The kernel-fusion implication: an XSA-aware attention kernel could fuse the self-value projection into the attention epilogue, eliminating the extra global-memory pass; the paper does not do this but the cost numbers in Fig. 2 are without that fusion and already show negligible overhead.

4.2 Training-time experiment control flow

  START (one (size, method) cell, e.g. 1.4B + XSA)
       |
       v
  (1) Build model from Table 1 spec; insert RoPE + extra LayerNorm
      after token embeddings; select SA or XSA at each layer
       |
       v
  (2) Initialize AdamW; LR schedule = linear warmup 2K steps,
      cosine decay to (1/10) * max_LR
       |
       v
  (3) Stream FineWeb-100BT tokens; pack into seq=2048; batch=256
      -> 0.5M tokens per iteration
       |
       v
  (4) FOR iter = 1 to 200000:
         forward pass    (compute, single device or DDP)
         backward pass
         AdamW step
         log (train_loss, val_loss)  every step / interval
       |
       v
  (5) Save final checkpoint
       |
       v
  (6) Run LM Evaluation Harness on 8 downstream tasks:
        ARC-E, BoolQ, HellaSwag, LAMBADA,
        OBQA, PIQA, SocialIQA, WinoGrande
       |
       v
  END  -> one row in Table 2 + one curve in Fig. 3
^ Fig 5: Training pipeline for a single (size, method) cell. The
  pipeline is identical for SA and XSA; only step (1)'s attention
  factory function differs.

The cleanliness of this pipeline matters because it bounds the kinds of confounding the paper can have. There is no architecture-aware learning-rate tuning ("we did a grid search per baseline and reused it for XSA" — Section 4.1) and no architecture-aware data filtering. This is a stronger evidentiary stance than a paper that re-tunes for each method: re-tuning hides margin contributions that come from architecture-LR interaction rather than the architecture itself.

4.3 The "attention similarity bias" diagnostic flow

The paper's empirical motivation comes from Figure 1, which is the output of a small diagnostic harness run on a pre-existing trained 1.3B-parameter baseline. The flow:

  START (trained 1.3B SA baseline, ctx=2048)
       |
       v
  (1) Sample 1024 random training sequences
       |
       v
  (2) For each sequence, for each layer, for each head:
         compute v_j = W_v x_j  for j = 1..T
         compute a_{i,j} attention scores
         compute y_i = sum_j a_{i,j} v_j
       |
       v
  (3) Aggregate three statistics per layer:
        - mean_{i<j} cos(v_i, v_j)        (left panel of Fig. 1)
        - mean a_{i,i}  over i > 1        (middle panel)
        - mean cos(y_i, v_i)              (right panel)
       |
       v
  (4) Plot all three across the 24 layers
       |
       v
  END  -> Fig. 1: the three correlated rising curves
^ Fig 6: How the "attention similarity bias" was measured. Two
  separately-rising trends (value-vector correlation and diagonal
  attention strength) combine into a third rising trend (output
  alignment with self-value), motivating the XSA correction.

This is the evidence-collection harness, distinct from the training harness. It is what justifies the existence of the XSA correction before any training experiment is run — without Figure 1, the XSA equation would look like an arbitrary regularizer.


5. Quantitative Results — Empirical Findings by Regime

5.1 Headline downstream-evaluation table (Table 2 verbatim)

Model Method ARC-E BoolQ HSwag LAMBADA OBQA PIQA SocIQA WinoGr Avg dAvg
0.7B Baseline 51.26 61.07 55.68 52.82 35.00 74.05 40.02 55.88 53.22
0.7B XSA 52.69 61.19 56.29 54.07 32.20 73.78 41.45 56.20 53.48 +0.26
1.3B Baseline 56.19 65.47 60.69 56.24 34.60 75.90 41.40 58.80 56.16
1.3B XSA 58.84 62.29 62.41 58.57 36.00 76.61 42.84 59.98 57.19 +1.03
2.7B Baseline 58.59 60.98 66.20 60.18 37.00 76.61 42.94 61.96 58.06
2.7B XSA 60.65 64.86 67.40 62.04 38.40 77.80 41.45 62.75 59.42 +1.36

The margin grows monotonically with model size: +0.26 -> +1.03 -> +1.36 average accuracy. This is the paper's strongest empirical claim — XSA's benefit is not a tuning artifact at small scale that washes out at larger scale; it grows with scale across the tested range.

5.2 Computational overhead (Fig. 2 prose)

"XSA introduces minimal overhead in terms of both speed and memory."

Measured on B200 with batch=32, bfloat16, the fwd+bwd time curves for SA and XSA are visually indistinguishable across seq_len in {512, 1024, 2048, 4096, 8192, 16384} and across d_model. Peak memory likewise. The largest measured point — seq=16384, d_model=2048 — is where XSA's three extra elementwise ops would have most absolute impact, and the gap is on the order of 1-3% (read from Fig. 2; not tabulated in the paper).

5.3 Loss curves across model sizes (Fig. 3 prose)

"XSA maintains a clear margin over the baseline in all three model sizes, across both training and validation."

The three loss curves for {0.7B, 1.3B, 2.7B} show XSA below baseline from approximately iteration 25K onward and staying there through 200K. No curve crossing is reported.

5.4 Learning-rate robustness (Fig. 4 prose)

"There is a near constant margin across all learning rates, demonstrating the robustness of the XSA architecture."

At 1.3B, with LR in {1e-4, 2e-4, 4e-4, 6e-4}, the XSA curve sits roughly 0.02-0.03 below the baseline curve at every LR setting, in both training and validation loss. This is the most important robustness result — many proposed attention variants are best at one LR and worse at others, and a single point can be cherry-picked. A constant offset across four LRs cannot be cherry-picked.

5.5 Sequence-length scaling (Fig. 5 prose)

"XSA claims larger gains as sequence length increases. We suspect that this is due to the increasing tension on context modeling for longer sequences."

At 1.3B with token-count-equalized batches (so global tokens/step stays at 0.5M), the loss gap widens monotonically from seq=512 to seq=16384. This is consistent with the bias-magnitude story in Figure 1: longer sequences give more opportunity for value-vector correlation to accumulate, so removing it has more leverage.

5.6 Attention-sink compatibility (Fig. 6 prose)

"XSA maintains the loss margin in the existence of attention sinks."

With 0, 1, or 4 learned sink tokens, XSA stays below baseline by approximately the same constant margin. The paper interprets XSA itself as an "implicit attention sink" because the projection effectively allocates any unneeded attention to the diagonal (which is then projected out), but the empirical claim is that XSA does not conflict with explicit sinks either.

5.7 The bias measurement (Fig. 1 prose)

The three panels of Figure 1 show, for a trained 1.3B baseline:

The right panel is the consequence the paper names the attention similarity bias. It is the joint product of the two left panels and is large enough that XSA's correction is not numerically trivial.


6. Configuration-Regime Trade-off Tables

6.1 SA vs XSA across regimes

Dimension Standard SA (baseline) Exclusive SA (XSA) Winner
Parameter count Same Same (no new params) tie
FLOPs per layer O(BT^2d) O(BT^2d) + O(BTd) SA (marginal)
Activation memory Same Same (overwrite Y -> Z) tie
Implementation complexity Standard 2-line diff tie
Train + val loss (0.7B-2.7B) Higher Lower (constant offset) XSA
Downstream avg accuracy Lower +0.26 / +1.03 / +1.36 XSA
Robustness to LR n/a baseline Margin constant 1e-4..6e-4 XSA
Robustness to seq-len n/a baseline Margin grows with seq-len XSA
Compat with attention sinks n/a baseline Margin preserved XSA
Compat with Muon optimizer n/a Untested (open question) unknown
Modalities other than language n/a Untested (open question) unknown

For language-model training with AdamW + RoPE + NanoGPT recipe, prefer XSA. The margin is positive across all four tested robustness axes and the overhead is below the measurement floor in Fig. 2.

6.2 Model-size regime sensitivity (extracted from Table 2)

Dimension 0.7B (small) 1.3B (medium) 2.7B (large) Trend
Avg accuracy lift (dAvg) +0.26 +1.03 +1.36 Monotonic up
Tasks with XSA better 6 of 8 7 of 8 7 of 8 Stable
Tasks with XSA worse OBQA, PIQA BoolQ SocIQA Different each
Margin direction Mostly positive Mostly positive Mostly positive Stable

Implication: No model size in the tested range reverses the overall sign of the margin, and the magnitude grows with scale. The worst-performing tasks for XSA shuffle between sizes (OBQA at 0.7B, BoolQ at 1.3B, SocIQA at 2.7B), suggesting the per-task variation is seed/noise rather than a systematic weakness.

6.3 Sequence-length regime sensitivity (Fig. 5)

Dimension Short (512-2048) Medium (4096) Long (8192-16384) Winner
Bias accumulation Modest Larger Largest --
XSA loss-margin magnitude Small Medium Large XSA strongest
Memory cost ratio (XSA/SA) ~1.00 ~1.00 ~1.00 tie
Latency cost ratio ~1.01 ~1.02 ~1.03 SA (tiny)

Implication: XSA's value proposition strengthens in the long- context regime — exactly the regime industry is currently scaling into. This is the most forward-looking quantitative claim in the paper.

6.4 Learning-rate regime (Fig. 4)

LR setting SA val loss (approx) XSA val loss (approx) Margin (approx)
1e-4 ~2.605 ~2.575 ~0.030
2e-4 ~2.570 ~2.545 ~0.025
4e-4 ~2.565 ~2.540 ~0.025
6e-4 ~2.570 ~2.550 ~0.020

(values read off Fig. 4 right panel)

Implication: No LR makes XSA worse; no LR makes XSA an order of magnitude better. The margin is a structural property, not an LR-interaction artifact.


7. Bottlenecks & Insights Surfaced by the Measurements

7.1 Attention output is partially redundant with the residual path

The central insight: in a Transformer with a residual connection, the component of y_i that points along v_i is information already delivered to the FFN through the residual skip. Attention spending its expressive capacity to re-deliver that information is a waste of capacity that grows with depth. Figure 1 quantifies this waste — 60% cosine alignment at the last layer of a 1.3B model is a substantial fraction of attention's output magnitude pointing in a direction the residual already covers.

7.2 Diagonal attention is not "the bug" — it's the symptom

A naive reading of Figure 1's middle panel ("a_{i,i} rises with layer") might suggest fixing diagonal attention itself, e.g., by masking it or normalizing it away. The paper's framing is sharper: the direction of the leaked information (parallel to v_i) is the problem, not the score of a_{i,i} itself. XSA projects out the direction regardless of which off-diagonal positions also happen to have value vectors correlated with v_i. This is why XSA is more than "masking the diagonal": even with a_{i,i} = 0, the rising <v_i, v_j> correlation (Fig. 1 left panel) would still produce self-aligned output.

7.3 The longer the context, the larger the leverage

Fig. 5's monotonic widening of the XSA margin with sequence length is a non-obvious systems implication. Standard intuition would predict the gap shrinks with longer context because attention has more positions to attend to and the per-position leakage averages out. The opposite happens: longer sequences accumulate more value-vector correlation, the diagonal-attention component compounds, and the bias grows. This means XSA is most valuable exactly where the Transformer is currently bottlenecked — long-context modeling.

7.4 The two-line implementation matters for adoption

The XSA change is two lines in Algorithm 1. It does not touch the attention kernel, does not change W_o, does not require a custom loss term, and does not require a tuned regularizer coefficient. This is a deployment-cost insight: the bar for adopting an attention variant in a production codebase is not its peak benefit but its peak benefit per unit of risk-of-something-going-wrong. XSA's risk surface is small — kernel-fusable, parameter-count-preserving, LR-stable.

7.5 The implicit attention-sink interpretation

The paper notes (Sec. on Attention Sink) that XSA, by removing self-aligned content from the output, effectively lets attention spend weight on the diagonal without that weight contributing to the output. In effect, the diagonal becomes a tunable "sink" for excess softmax mass. This is the same role that external sink tokens (Xiao et al., 2023) were introduced to fill — but here it emerges from the math rather than from an extra token. The implication is that XSA may obviate the need for explicit sinks in streaming-attention deployments.

7.6 The bias measurement generalizes the diagnostic toolkit

Fig. 1's three-panel diagnostic — value-vector correlation, diagonal attention strength, output-self alignment — is reusable beyond this paper. Any future attention variant that claims to modify the role of the diagonal can be evaluated by re-running this exact diagnostic on its trained baseline. The paper hands future work a measurement instrument as well as a method.


8. Limitations of the Methodology

Limitation Implication
Single seed per (size, method) cell Margin error bars not quantified; ablation sweeps act as proxy
Three model sizes only (0.7B-2.7B) No data at 7B+ or smaller than 0.7B
Single modality (language) No vision / multimodal / code-domain validation
Single optimizer (AdamW) Muon and Lion family untested; flagged in Discussion
Single position-embedding scheme (RoPE) ALiBi, NoPE, and learned absolute embeddings untested
Single tokenizer (GPT-2) Tokenizer-attention interaction not isolated
Single dataset (FineWeb-100BT) Code / math / multilingual not validated
Training duration capped at 200K iters / ~1 ep Multi-epoch and over-training regime unmeasured
Downstream eval = 8 zero/few-shot tasks No instruction-tuned or RLHF-tuned downstream eval
No analysis of head-level bias variation Some heads may have larger/smaller bias; uniform XSA may be loose
No theoretical guarantee of expressivity Paper defers "theoretical groundings" to future work
No comparison to other diagonal-mod variants No SA-without-diagonal or learned-self-gate baselines
Efficiency-bench on B200 only A100 / H100 / TPU numbers not reported
No code release announced in the paper Replication depends on the NanoGPT recipe being correctly applied
No ablation on the projection target Only v_i direction tested; q_i or k_i directions untried

The strongest limitation is the lack of variance reporting. The reported deltas at 0.7B (+0.26 average accuracy across 8 tasks) are within plausible single-seed noise for tasks like OBQA where the per-task swing is more than 2 points between SA and XSA in opposite directions. The 1.3B and 2.7B deltas are large enough to survive seed noise, but the 0.7B claim is the most fragile.

The second-most-important limitation is the single-modality scope. The bias mechanism described in Section 2 (rising <v_i, v_j> correlation across layers) is plausibly modality- dependent: vision transformers operating on patch embeddings may exhibit very different value-vector correlation structure, and the XSA correction may have different magnitude or sign in that regime.


9. Analogy

The paper is a carpentry instruction that changes one joint in a load-bearing frame — a 90-degree mortise replaced by a chamfered mortise that prevents the post and the beam from sharing a single plane of stress. Imagine a wooden frame where each vertical post already carries its own weight to the floor (the residual connection), while each horizontal beam is supposed to redistribute load across the frame (attention's job). In the standard joint, the beam ends up partly leaning on the very post it is attached to — the load it transfers to the post is partly load the post was already carrying, redundantly. The frame stands, but the beam is doing two jobs instead of one. The carpenter, who has been measuring joints in built frames for years, finally notices the redundancy (Fig. 1), and proposes a small chamfer (the XSA projection): cut away the part of the beam-end that shares a plane with the post, so the beam is mechanically forced to transmit load to other posts. The new joint takes two extra cuts (Algorithm 1's two lines) and uses no extra wood (no extra parameters). Once the chamfered frames are built and load-tested across three sizes of structure, four adhesives (learning rates), and six beam lengths (sequence lengths), the chamfered frames carry more load — and the longer the beam, the larger the improvement, because longer beams accumulate more of the shared-plane redundancy. The carpenter's claim is not that the chamfer is theoretically optimal; it is that across every load condition tested, the chamfered joint outperforms the square joint by a small, consistent, scale-growing margin. That margin, with no extra wood and almost no extra cutting time, is the entire paper.