Scaling Linear Attention with Sparse State Expansion

• URL: http://arxiv.org/abs/2507.16577v1
• Date: Tue, 22 Jul 2025 13:27:31 GMT
• Title: Scaling Linear Attention with Sparse State Expansion
• Authors: Yuqi Pan, Yongqi An, Zheng Li, Yuhong Chou, Ruijie Zhu, Xiaohui Wang, Mingxuan Wang, Jinqiao Wang, Guoqi Li,
• Abstract summary: Transformer architecture struggles with long-context scenarios due to quadratic computation and linear memory growth.<n>We introduce a row-sparse update formulation for linear attention by conceptualizing state updating as information classification.<n>Second, we present Sparse State Expansion (SSE) within the sparse framework, which expands the contextual state into multiple partitions.
• Score: 58.161410995744596
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: The Transformer architecture, despite its widespread success, struggles with long-context scenarios due to quadratic computation and linear memory growth. While various linear attention variants mitigate these efficiency constraints by compressing context into fixed-size states, they often degrade performance in tasks such as in-context retrieval and reasoning. To address this limitation and achieve more effective context compression, we propose two key innovations. First, we introduce a row-sparse update formulation for linear attention by conceptualizing state updating as information classification. This enables sparse state updates via softmax-based top-$k$ hard classification, thereby extending receptive fields and reducing inter-class interference. Second, we present Sparse State Expansion (SSE) within the sparse framework, which expands the contextual state into multiple partitions, effectively decoupling parameter size from state capacity while maintaining the sparse classification paradigm. Our design, supported by efficient parallelized implementations, yields effective classification and discriminative state representations. We extensively validate SSE in both pure linear and hybrid (SSE-H) architectures across language modeling, in-context retrieval, and mathematical reasoning benchmarks. SSE demonstrates strong retrieval performance and scales favorably with state size. Moreover, after reinforcement learning (RL) training, our 2B SSE-H model achieves state-of-the-art mathematical reasoning performance among small reasoning models, scoring 64.7 on AIME24 and 51.3 on AIME25, significantly outperforming similarly sized open-source Transformers. These results highlight SSE as a promising and efficient architecture for long-context modeling.

Related papers

• Pimba: A Processing-in-Memory Acceleration for Post-Transformer Large Language Model Serving [1.9508863993381267]
  Transformers are the driving force behind today's Large Language Models (LLMs), serving as the foundation for their performance and versatility.<n>In response, the algorithm community is exploring alternative architectures, such as state space models (SSMs), linear attention, and recurrent neural networks (RNNs)<n>This paper presents a serving system that efficiently supports both transformer and post-transformer LLMs within a unified framework.<n>Our evaluation shows that, compared to LLM-optimized GPU and GPU+PIM systems, Pimba achieves up to 3.2x and 2.1x higher token generation throughput.
  arXiv  Detail & Related papers  (2025-07-14T11:40:17Z)
• Decoder-Hybrid-Decoder Architecture for Efficient Reasoning with Long Generation [129.45368843861917]
  We introduce the Gated Memory Unit (GMU), a simple yet effective mechanism for efficient memory sharing across layers.<n>We apply it to create SambaY, a decoder-hybrid-decoder architecture that incorporates GMUs to share memory readout states from a Samba-based self-decoder.
  arXiv  Detail & Related papers  (2025-07-09T07:27:00Z)
• Compress, Gather, and Recompute: REFORMing Long-Context Processing in Transformers [58.98923344096319]
  REFORM is a novel inference framework that efficiently handles long contexts through a two-phase approach.<n>It achieves over 50% and 27% performance gains on RULER and BABILong respectively at 1M context length.<n>It also outperforms baselines on Infinite-Bench and MM-NIAH, demonstrating flexibility across diverse tasks and domains.
  arXiv  Detail & Related papers  (2025-06-01T23:49:14Z)
• Sliding Window Attention Training for Efficient Large Language Models [55.56483740523027]
  We introduce SWAT, which enables efficient long-context handling via Sliding Window Attention Training.<n>This paper first attributes the inefficiency of Transformers to the attention sink phenomenon.<n>We replace softmax with the sigmoid function and utilize a balanced ALiBi and Rotary Position Embedding for efficient information compression and retention.
  arXiv  Detail & Related papers  (2025-02-26T05:31:44Z)
• Pivoting Factorization: A Compact Meta Low-Rank Representation of Sparsity for Efficient Inference in Large Language Models [1.6385815610837167]
  Pivoting Factorization (PIFA) is a novel low-rank representation that unsupervisedly learns a compact form of any low-rank representation.<n>PIFA achieves 24.2% additional memory savings and 24.6% faster inference over low-rank layers at rank = 50% of dimension.<n>MPIFA, combining M and PIFA into an end-to-end framework, significantly outperforms existing low-rank pruning methods.
  arXiv  Detail & Related papers  (2025-01-31T12:36:31Z)
• Linear Attention via Orthogonal Memory [46.14195464583495]
  We propose textbfLinear textbfAttention textbfVia textbfOrthogonal memory(shortname) to address these limitations. shortname achieves strong performance while maintaining linear complexity.
  arXiv  Detail & Related papers  (2023-12-18T12:26:27Z)
• Towards A Unified View of Sparse Feed-Forward Network in Pretraining Large Language Model [58.9100867327305]
  Large and sparse feed-forward layers (S-FFN) have proven effective in scaling up Transformers model size for textitpretraining large language models. We analyzed two major design choices of S-FFN: the memory block (a.k.a. expert) size and the memory block selection method. We found a simpler selection method -- textbftextttAvg-K that selects blocks through their mean aggregated hidden states, achieving lower perplexity in language model pretraining.
  arXiv  Detail & Related papers  (2023-05-23T12:28:37Z)
• Efficient Long Sequence Modeling via State Space Augmented Transformer [92.74707853711374]
  We propose SPADE, short for $underlinetextbfS$tate sunderlinetextbfP$ace. We augment a SSM into the bottom layer of SPADE, and we employ efficient local attention methods for the other layers. Experimental results on the Long Range Arena benchmark and language modeling tasks demonstrate the effectiveness of the proposed method.
  arXiv  Detail & Related papers  (2022-12-15T20:51:27Z)
• DCT-Former: Efficient Self-Attention with Discrete Cosine Transform [4.622165486890318]
  An intrinsic limitation of the Trasformer architectures arises from the computation of the dot-product attention. Our idea takes inspiration from the world of lossy data compression (such as the JPEG algorithm) to derive an approximation of the attention module. An extensive section of experiments shows that our method takes up less memory for the same performance, while also drastically reducing inference time.
  arXiv  Detail & Related papers  (2022-03-02T15:25:27Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.