SSA: Sparse Sparse Attention by Aligning Full and Sparse Attention Outputs in Feature Space

• Full self-attention has quadratic complexity in context length, making training and inference prohibitive for extended contexts.Sparse attention restricts the tokens each query attends to in pursuit of full-attention-level performance with sub-quadratic complexity.Source: §1 · p. 1
• Full-only training can incur a distribution-mismatch penalty at sparse inference, while sparse-only training has weaker learning signals and can underperform full attention.The paper links sparse-only training to weaker signals because it excludes some tokens from attention.Source: §1 · pp. 1–2
• SSA achieves the highest attention sparsity and best performance under both inference modes, according to the introduction’s figure and contribution summary.The paper also reports state-of-the-art performance across multiple benchmarks.Source: Figure 1 · p. 2, §1 · p. 2
• SSA randomly alternates between sparse and full attention training and aligns their outputs in both directions.Sparse training directly optimizes sparse inference; full attention supplies complete gradient flow, while alignment encourages sparser distributions and limits drift.Source: §1 · p. 2

• Training-free sparse methods use fixed or query-dependent token selection, but the training–inference mismatch can cause error compounding on long-context tasks [Hu et al., 2026].Examples include sliding windows, attention sinks, and block-sparse top-k selection [Child et al., 2019, Beltagy et al., 2020, Brown et al., 2020] [Xiao et al., 2024, Jiang et al., 2024, Xu et al., 2025, Tang et al., 2024, Zhang et al., 2025].Source: §2 · p. 3
• Trainable sparse methods include block-sparse attention, multi-pattern gating, hierarchical block selection, and differentiable routing [Lu et al., 2025] [Yuan et al., 2025] [Zhao et al., 2025] [Huang et al., 2026].The passage notes that NSA’s multi-component structure makes sparsity adjustment and reversion to full attention difficult, and attributes its reported gains to added gating rather than sparse design [Qiu et al., 2025].Source: §2 · p. 3
• Unlike full-attention distillation and hybrid architectures that treat full attention as a fixed teacher, SSA aligns both directions and explicitly drives full attention to become sparser.The paper says this design supports performance under either inference mode.Source: §2 · p. 3

• Full attention uses all preceding tokens, but its computation grows quadratically with context length.The paper describes this cost as prohibitive for large contexts.Source: §3 · p. 3
• SSA’s framework alternates with equal probability between sparse and full attention streams, aligning each stream’s output with the other’s counterpart.The figure caption notes that skip connections, normalization, and dropout are omitted.Source: Figure 2 · p. 4
• Sparse attention partitions tokens into contiguous blocks, scores preceding blocks using mean-pooled keys, and selects the top-k blocks for each query.The selected blocks’ tokens form reduced key and value matrices.Source: §3 · p. 4
• Selecting k of n blocks gives a receptive field of ks tokens and reduces attention computation in proportion to the sparsity ratio.The supplied complexity expression is split across passages; the prose states the proportional reduction from full attention’s quadratic cost.Source: §3 · p. 4
• Despite sub-quadratic scaling in long-context settings, sparse attention has remained a modeling-quality bottleneck in inference-only and natively trained use.The passage attributes this concern to quality degradation reported in prior work [Lu et al., 2025, Yuan et al., 2025].Source: §3 · p. 4

4 Theoretical Framework of Sparse Attention

• The paper’s framework identifies gradient deficiencies under exclusive sparse training and distribution mismatch under exclusive full training; approximation error scales linearly with dropped attention mass.These findings motivate training inherently sparser attention distributions.Source: §4 · p. 4

4.1 Problem Setup

• The problem setup indexes hidden representations by the training attention mechanism and the inference mechanism, and states a target closeness condition for sparse attention.The passage identifies T as sequence length and gives the condition for some small epsilon greater than zero.Source: §4.1 · p. 5

4.2 Theoretical Framework

• Sparse-only training gives dropped keys zero gradient and removes them from competition during attention normalization, weakening learning and global suppression of irrelevant tokens.The paper predicts that sparse-trained models may underperform full-trained models and fail to benefit from full-attention inference because of high-entropy attention patterns.Source: §4.2 · p. 5
• Switching a full-attention-trained model to sparse inference shifts its output distribution by a KL term, while sparse-trained models avoid that particular train–inference shift.A sufficiently large mismatch can make the full-trained model worse than a sparse-trained model under sparse inference, despite its stronger full-inference performance.Source: §4.2 · p. 6
• The framework defines attention sparsity through the attention mass retained by selected blocks and states that the approximation-error bound scales linearly with dropped mass.Greater retained mass suggests a closer approximation to full attention; reducing dropped mass is the stated motivation for learning inherently sparser distributions.Source: §4.2 · p. 6

5 Sparse Sparse Attention (SSA)

• SSA alternates full- and sparse-attention streams with equal probability, combining complete gradient flow from full attention with direct optimization for sparse inference.Alternating rather than jointly computing both streams reduces computational cost by half while keeping the number of processed tokens equal to baselines.Source: §5 · p. 7
• At each layer, counterpart attention outputs are computed for alignment, with full-attention outputs trained toward sparse-attention outputs to encourage sparser distributions.Because top-k selection constrains sparse outputs, this direction is intended to reduce dropped attention mass and tighten the approximation bound.Source: §5 · pp. 6–7
• A commitment loss keeps sparse-attention outputs close to full-attention outputs, complementing the sparsity loss in the bidirectional alignment objective.The paper relates this role to commitment loss in VQ-VAE and KL regularization in RLHF.Source: §5 · p. 7
• The final objective combines cross-entropy under the sampled attention mode with the alignment loss, weighted by α.SSA uses SmoothL1 for both alignment components; inference uses the same sparse operations as MoBA, with 2× speedup over full attention at 128k context length and about 17% additional training cost for 8k-context pretraining.Source: §5 · pp. 7–8
• SSA consistently outperforms other methods on average across sparsity levels, and its performance extrapolates well as more tokens become visible.MoBA exhibits poor extrapolation as the receptive field grows.Source: Table 1 · p. 8, Figure 3 · p. 8

6 Experimental Setup

• The model follows Llama-3.2-1B architecture and is pretrained for 100B tokens at 8k context, then continued for 10B tokens at 32k context.Continued training uses RoPE scaling factor 4 and DCLM-sampled data.Source: §6 · pp. 8–9
• SSA outperforms all baselines on LongBench and NIAH across the reported long-context settings, except NIAH with full-attention inference after continued training.Table 2 separates 32K RoPE-scaled extrapolation from models continually trained to 32K.Source: Table 2 · p. 9
• Baselines are FullAttn, MoBA, and NSA, evaluated at receptive fields of 1024 and 256, with SSA's 1024 results extrapolated from its 256 model.NSA uses a window size of 128 in both configurations.Source: §6 · p. 9
• Evaluation covers commonsense reasoning, WikiText perplexity, LongBench, NIAH retrieval, and PG-19 long-context perplexity under sparse and full inference where applicable.Commonsense benchmarks average five seeds; RoPE scaling factor 4 is used for 32k extrapolation when models lack continued 32k training.Source: §6 · p. 9
• On PG-19, SSA has the lowest perplexity under sparse inference and the strongest RoPE-scaling robustness under full attention.Figure 4 compares perplexity across context lengths and inference modes.Source: Figure 4 · p. 10

7.1 Language Modeling

• SSA achieves the lowest perplexity at RF=256, though it slightly trails NSA at RF=1024, and has a smaller sparse/full perplexity gap than FullAttn.The authors link the smaller gap to higher attention sparsity and reduced dropped attention mass.Source: §7.1 · p. 10

7.2 Commonsense Reasoning

• SSA generally outperforms baselines at matched sparsity budgets and surpasses FullAttn with a receptive field of 256 on commonsense benchmarks.Most benchmarks use contexts shorter than 1024 tokens, making RF=1024 effectively near full attention.Source: §7.2 · p. 10

7.3 Extrapolation between Different Levels of Sparsity

• Across four tasks, SSA improves largely monotonically as more tokens enter sparse attention, whereas MoBA fails to extrapolate on some tasks and budgets.MoBA also degrades at RF=4096 on Commonsense and NIAH; the authors attribute this to insufficient suppression of irrelevant tokens.Source: §7.3 · p. 11

7.4 Long Context Evaluation

• On PG-19, SSA outperforms baselines under sparse attention; without continued training it also exceeds FullAttn under full attention, but after continued training FullAttn surpasses SSA in that mode.The authors attribute the initial full-attention result to greater robustness of SSA's sparser patterns to RoPE scaling degradation; SSA retains its advantage under sparse inference.Source: §7.4 · p. 11
• Within 8k context, SSA is the strongest sparse NIAH method across receptive fields and reaches 100% accuracy with full-attention inference, but its performance declines beyond 8k.MoBA retrieves better at longer contexts, while SSA averages the best performance in five of six settings.Source: §7.4 · p. 11
• SSA achieves the best LongBench performance across inference modes and context-extension settings.LongBench evaluates diverse long-document reasoning tasks, beyond average prediction quality or retrieval of a single fact.Source: §7.4 · p. 11

7.5 Ablation Studies

• A larger receptive field does not improve SSA, while shrinking it too aggressively also fails to help, indicating a receptive-field sweet spot.The authors hypothesize that smaller fields impose stronger structural constraints that regularize sparse attention learning.Source: §7.5 · p. 12
• Moderate sparse-stream inclusion gives near-optimal perplexity, more full-stream weight generally improves downstream results, and removing either stream degrades performance.Removing alignment substantially harms performance, one-way alignment destabilizes training, and the authors conclude bidirectional alignment is needed; α requires tuning, while adjusted L2 alignment slightly trails default SSA but beats MoBA.Source: §7.5 · pp. 12–13, Table 3 · p. 12

8 Conclusion

• SSA jointly aligns sparse and full attention bidirectionally and achieves the highest attention sparsity among the methods discussed.The conclusion reports superior perplexity and downstream-task performance in both inference modes, plus robustness across sparsity levels.Source: §8 · p. 13
• Sparse-attention-trained models show stronger long-context extrapolation than full-attention models, and SSA achieves the best long-context understanding results.The conclusion states that SSA's high attention sparsity benefits both sparse- and full-attention inference.Source: §8 · p. 13
• The paper identifies insufficient attention sparsity in native sparse-attention methods, attributing it to gradient-update deficiency on excluded key–value pairs.The conclusion describes this as a paradox in sparse-attention training.Source: §8 · p. 13

Impact Statement

• SSA is presented as a feasible framework addressing both the capability gap and train-inference mismatch while reducing inference computation and memory costs.The impact statement notes potential relevance to resource-limited settings and future attention-architecture design.Source: Impact Statement · p. 13

A Proof of Proposition 4.2

• The Proposition 4.2 proof uses the full-attention next-token distribution as the reference so the cross-entropy gap becomes a KL divergence.Both distributions are evaluated at the same query position and context.Source: A Proof of Proposition 4.2 · p. 18

B Proof of Theorem 1

• The Theorem 1 proof decomposes full attention into selected and dropped-token contributions, with δ(t) denoting the dropped full-attention mass.For selected tokens, the proof relates full and sparse attention weights through their normalization over all versus selected keys.Source: B Proof of Theorem 1 · p. 19

C Implementation Details

• The implementation trains with a randomly routed main attention mode and the opposite-mode auxiliary stream, then adds bidirectional alignment losses to cross-entropy.The implementation notes reducing key-value heads to two and applying Gated Attention to all methods for fair comparison.Source: B Proof of Theorem 1 · p. 20, C Implementation Details · p. 20

D Efficiency Analysis

• 17% more time than MoBA; SSA’s dynamic attention computation introduces minimal training overhead.The comparison reports training time; FullAttn has the shortest time because its implementation is highly optimized.Source: D Efficiency Analysis · p. 20, Table A2 · p. 22
• Pre-training is shorter for RF=256 than RF=1024, but continual training reverses this trend under the two-stage sparse-attention computation.The stages compute block-level scores and then attend over selected tokens, with approximate complexity N^2 + M^2.Source: D Efficiency Analysis · p. 21, Table A4 · p. 22
• 2× speed-up at 128k length; SSA matches MoBA’s speed in sparse mode and FullAttn when full attention is triggered.The sparse-mode gain is modest at 8k and widens as context grows.Source: D Efficiency Analysis · p. 21, Table A3 · p. 22

E Evaluation and Datasets

• LongBench and NIAH use greedy decoding and single-run results; NIAH can vary slightly because its evaluation data is synthesized dynamically.The sparse-attention kernel does not currently support KV caching, making decoding slow; commonsense benchmarks use five fixed random seeds.Source: E Evaluation and Datasets · p. 22, Table A5 · p. 23

F NSA Is Not Suitable for Full Attention Evaluation

• NSA’s compression and sliding-window modules fail to extrapolate reliably to full attention; replacing all three modules with full attention yields perplexity 191.3.Only the selection module extrapolates well, so NSA is excluded from full-attention evaluation.Source: F NSA Is Not Suitable for Full… · p. 23, Table A6 · p. 23

G NSA with and without Sliding Window

• SSA consistently outperforms NSA without a sliding window on long-context perplexity, while removing that module degrades most downstream tasks.The reported drops include Hellaswag (-1.18) and ARC-E (-1.1); NSA’s perplexity advantage largely depends on the extra local tokens from its sliding window.Source: G NSA with and without Sliding Window · p. 23, Table A8 · p. 24, Table A7 · p. 24

H KL Divergence

• SSA achieves the smallest KL divergence among SSA, MoBA, and FullAttn, which the authors attribute to the alignment loss.Table A9 also reports attention sparsity, perplexity, and benchmark accuracy.Source: H KL Divergence · p. 24, Table A9 · p. 24

I Spasity Extrapolation Results; J Full Longbench Results

• The evaluation reports commonsense reasoning and LongBench scores across receptive fields, plus LongBench per-dataset scores for RoPE-scaled and continual-trained models.MoBA and SSA are trained with RF=256 and extrapolated to larger receptive fields; the commonsense reasoning results average five runs.Source: I Spasity Extrapolation Results · p. 24, J Full Longbench Results · p. 24, Table A10 · p. 25

K Sparse Attention is Robust to RoPE Scaling

• Sparse-attention perplexity degradation under RoPE scaling follows sparse 256 tokens < sparse 1024 tokens < full attention.The authors attribute sparse attention’s robustness to fewer attended tokens and a local-token skew, where RoPE scaling causes less distortion.Source: K Sparse Attention is Robust to RoPE… · p. 25

K.1 Aliased Frequency Bands in Length Extrapolation

• RoPE frequency aliasing makes distant positions appear local after extrapolation from 8k to 32k, contributing to perplexity explosion; scaling stretches low-frequency periods by 4 and eliminates the aliasing artifacts.The appendix illustrates the attention-score and QK-product patterns and reports frequency periods before and after scaling.Source: §K.1 · p. 25, Figure A1 · p. 28, Figure A2 · p. 29

K.2 Local Token Preference for Long Context

• FullAttn, MoBA, and SSA all prefer local tokens; under RoPE scaling, SSA consistently has higher sparsity and lower entropy.The appendix includes distance histograms for all three methods and a table of attention patterns across four context lengths.Source: §K.2 · p. 31, Table A15 · p. 31, Figure A4 · p. 32