Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SPLA: Block Sparse Plus Linear Attention for Long Context Modeling

Bailin Wang, Dan Friedman, Tao Lei, Chong Wang

TL;DR

SPLA addresses the long-context decoding bottleneck by combining principled block-level selection with a residual linear attention path that compresses the unselected tail instead of discarding it. Block selection is derived from a second-order Taylor expansion around block statistics, enabling accurate retrieval of relevant blocks, while residual linear attention preserves global context with minimal IO overhead. The framework partitions the context into exact (selected) and approximate (unselected) components, computing the unselected contribution via a recurrent, IO-efficient mechanism that avoids explicit memory access to unselected blocks. Empirically, SPLA closes the gap to dense attention on long-context benchmarks (e.g., up to 256k tokens) and adapts pretrained dense models with few extra parameters, offering a scalable path for efficient long-context modeling in continual pretraining. Overall, SPLA delivers robust performance across general knowledge, long-context, and reasoning tasks with pronounced speedups at large context lengths while maintaining compatibility with existing attention architectures.

Abstract

Block-wise sparse attention offers significant efficiency gains for long-context modeling, yet existing methods often suffer from low selection fidelity and cumulative contextual loss by completely discarding unselected blocks. To address these limitations, we introduce Sparse Plus Linear Attention (SPLA), a framework that utilizes a selection metric derived from second-order Taylor expansions to accurately identify relevant blocks for exact attention. Instead of discarding the remaining "long tail," SPLA compresses unselected blocks into a compact recurrent state via a residual linear attention (RLA) module. Crucially, to avoid IO overhead, we derive an optimized subtraction-based formulation for RLA -- calculating the residual as the difference between global and selected linear attention -- ensuring that unselected blocks are never explicitly accessed during inference. Our experiments demonstrate that SPLA closes the performance gap in continual pretraining, surpassing dense attention models on long-context benchmarks like RULER while maintaining competitive general knowledge and reasoning capabilities.

SPLA: Block Sparse Plus Linear Attention for Long Context Modeling

TL;DR

SPLA addresses the long-context decoding bottleneck by combining principled block-level selection with a residual linear attention path that compresses the unselected tail instead of discarding it. Block selection is derived from a second-order Taylor expansion around block statistics, enabling accurate retrieval of relevant blocks, while residual linear attention preserves global context with minimal IO overhead. The framework partitions the context into exact (selected) and approximate (unselected) components, computing the unselected contribution via a recurrent, IO-efficient mechanism that avoids explicit memory access to unselected blocks. Empirically, SPLA closes the gap to dense attention on long-context benchmarks (e.g., up to 256k tokens) and adapts pretrained dense models with few extra parameters, offering a scalable path for efficient long-context modeling in continual pretraining. Overall, SPLA delivers robust performance across general knowledge, long-context, and reasoning tasks with pronounced speedups at large context lengths while maintaining compatibility with existing attention architectures.

Abstract

Block-wise sparse attention offers significant efficiency gains for long-context modeling, yet existing methods often suffer from low selection fidelity and cumulative contextual loss by completely discarding unselected blocks. To address these limitations, we introduce Sparse Plus Linear Attention (SPLA), a framework that utilizes a selection metric derived from second-order Taylor expansions to accurately identify relevant blocks for exact attention. Instead of discarding the remaining "long tail," SPLA compresses unselected blocks into a compact recurrent state via a residual linear attention (RLA) module. Crucially, to avoid IO overhead, we derive an optimized subtraction-based formulation for RLA -- calculating the residual as the difference between global and selected linear attention -- ensuring that unselected blocks are never explicitly accessed during inference. Our experiments demonstrate that SPLA closes the performance gap in continual pretraining, surpassing dense attention models on long-context benchmarks like RULER while maintaining competitive general knowledge and reasoning capabilities.
Paper Structure (39 sections, 14 equations, 2 figures, 4 tables)

This paper contains 39 sections, 14 equations, 2 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Background and Motivation
  3. Notation.
  4. Page Attention and Selection
  5. Grouped Sparsity
  6. Open Challenges
  7. 1. Block Selection Fidelity.
  8. 2. Long Tail Divergence.
  9. SPLA: Block Sparse Attention with Residual Linear Attention
  10. Principled Block Selection via Taylor Expansion
  11. Strided Block Selection
  12. Residual Linear Attention (RLA)
  13. Recurrence
  14. IO-Efficient Implementation
  15. Parameterization
  16. ...and 24 more sections

Figures (2)

  • Figure 1: (a) Block-level selection illustrates the initial identification of relevant blocks. (b)SPLA performs exact attention on selected blocks (teal) while processing unselected blocks via residual linear attention (purple). (c) Existing sparse attention mechanisms (e.g., InfLLM-v2, NSA) typically either discard unselected blocks entirely or process all blocks redundantly.
  • Figure 2: Decoding step time speedup (%) of SPLA models compared to dense-attention models at 13B scale. As the batch size increases, the theoretical speedup of SPLA increases more quickly, since the KV cache size increasingly dominates the memory cost relative to the model parameter size.