cs.LGMar 10, 2026
SaveExclusive Self Attention
Organizations: Apple
Abstract
We introduce exclusive self attention (XSA), a simple modification of self attention (SA) that improves Transformer's sequence modeling performance. The key idea is to constrain attention to capture only information orthogonal to the token's own value vector (thus excluding information of self position), encouraging better context modeling. Evaluated on the standard language modeling task, XSA consistently outperforms SA across model sizes up to 2.7B parameters and shows increasingly larger gains as sequence length grows.
Figures & tables
Figure 1: Visualization of the attention similarity bias of a 1.3B parameter language model of sequence length 2048 trained for 100B tokens, aggregated on 1024 random training sequences. Left : the average cosine similarity of value vectors , within a sequence; middle : the average diagonal attention value ; right : the average cosine similarity of attention output and the self value vector . See Eq. 1 for notations.
Figure 2: Time and memory efficiency of XSA compared to standard attention. XSA introduces minimal computational overhead across various sequence lengths and model sizes .
| Model Size | Learning Rate | ||||
|---|---|---|---|---|---|
| 0.7B | 24 | 1536 | 6 | 256 | |
| 1.4B | 24 | 2048 | 24 | 128 | |
| 2.7B | 32 | 2560 | 24 | 128 |
Table 1: Architectures and learning rates. All models are trained with a batch size of tokens, for a total of 200K iterations.
Figure 3: Training and validation loss curves of XSA against the baseline Transformer for three model sizes.
| Model | ARC-E | BoolQ | HSwag | LAMBADA | OBQA | PIQA | SocIQA | WinoGr | Avg | Avg |
|---|---|---|---|---|---|---|---|---|---|---|
| 0.7B | ||||||||||
| Baseline | 51.26 | 61.07 | 55.68 | 52.82 | 35.00 | 74.05 | 40.02 | 55.88 | 53.22 | +0.26 |
| XSA | 52.69 | 61.19 | 56.29 | 54.07 | 32.20 | 73.78 | 41.45 | 56.20 | 53.48 | |
| 1.3B | ||||||||||
| Baseline | 56.19 | 65.47 | 60.69 | 56.24 | 34.60 | 75.90 | 41.40 | 58.80 | 56.16 | +1.03 |
| XSA | 58.84 | 62.29 | 62.41 | 58.57 | 36.00 | 76.61 | 42.84 | 59.98 | 57.19 | |
Table 2: Downstream evaluation results of XSA vs baseline Transformer.
Figure 4: Training and validation loss of XSA against the baseline Transformer for various learning rates evaluated with the 1.3B model.
Figure 5: Training and validation loss of XSA against the baseline Transformer for various sequence lengths evaluated with the 1.3B model.
Figure 6: Training and validation loss of XSA against the baseline Transformer for various number of Attention Sinks evaluated with the 1.3B model.
More in this topic
Self-Attention
Self-attention lets sequence elements exchange context through learned query-key matches, powering Transformers and their efficient variants.
See all 220 papersTraining causal transformers at extreme sequence lengths is bottlenecked by the quadratic time and memory of scaled dot-product attention (SDPA). In this work, we propose Lighthouse Attention, a training-only symmetrical selection-based hierarchical attention algorithm that wraps around ordinary SDPA and can be easily removed towards the end of the training. Our hierarchical selection is also gradient-free, which exempts us from dealing with a complicated and potentially inefficient backward pass kernel. Our contribution is three-fold: (i) A subquadratic hierarchical pre- and post-processing step that does adaptive compression and decompression of the sequence. (ii) A symmetrical compression strategy that pools queries, keys and values at the same time, while preserving left-to-right causality, which greatly improves parallelism. (iii) A two stage training approach which we pre-train for the majority of the time with Lighthouse Attention and recover a full attention model at the end with a short training. We run preliminary small scale LLM pre-training experiments that show the effectiveness of our method compared to full attention training with all other settings matched, where we achieve a faster total training time and lower final loss after the recovery phase. Full code is available at: https://github.com/ighoshsubho/lighthouse-attention
Scaling Attention via Feature Sparsity
Scaling Transformers to ultra-long contexts is bottlenecked by the cost of self-attention. Existing methods reduce this cost along the sequence axis through local windows, kernel approximations, or token-level sparsity, but these approaches consistently degrade accuracy. In this paper, we instead explore an orthogonal axis: feature sparsity. We propose Sparse Feature Attention (SFA), where queries and keys are represented as -sparse codes that preserve high-dimensional expressivity while reducing the cost of attention from to . To make this efficient at scale, we introduce FlashSFA, an IO-aware kernel that extends FlashAttention to operate directly on sparse overlaps without materializing dense score matrices. Across GPT-2 and Qwen3 pretraining, SFA matches dense baselines while improving speed by up to and reducing FLOPs and KV-cache by nearly 50%. On synthetic and downstream benchmarks, SFA preserves retrieval accuracy and robustness at long contexts, outperforming short-embedding baselines that collapse feature diversity. These results establish feature-level sparsity as a complementary and underexplored axis for efficient attention, enabling Transformers to scale to orders-of-magnitude longer contexts with minimal quality loss. Code is available at https://github.com/YannX1e/Sparse-Feature-Attention.
Sessa: Selective State Space Attention
Modern sequence modeling is dominated by two families: Transformers, whose self-attention can access arbitrary elements of the visible sequence, and structured state-space models, which propagate information through an explicit recurrent state. These mechanisms face different limitations on long contexts: when attention is diffuse, the influence of individual tokens is diluted across the effective support, while recurrent state propagation can lose long-range sensitivity unless information is actively preserved. As a result, both mechanisms face challenges in preserving and selectively retrieving information over long contexts. We propose Sessa, a decoder that places attention inside a recurrent feedback path. This creates many attention-based paths through which past tokens can influence future states, rather than relying on a single attention read or a single recurrent chain. We prove that, under explicit assumptions and matched regimes, Sessa admits power-law memory tails for , with slower decay than in the corresponding Transformer and Mamba-style baselines. We further give an explicit construction that achieves this power-law rate. Under the same assumptions, Sessa is the only model class among those considered that realizes flexible selective retrieval, including profiles whose influence does not decay with distance. Consistent with this theoretical advantage, across matched experiments, Sessa achieves the strongest performance on long-context benchmarks while remaining competitive with Transformer and Mamba-style baselines on short-context language modeling.