arXiv:2507.16577cs.LGcs.CL2025-07被引 8

提出稀疏状态扩展,让线性注意力更好处理长文本。

Scaling Linear Attention with Sparse State Expansion

  • 用稀疏更新机制分类信息,扩大感受野。
  • 状态容量可扩展,2B模型在AIME上达64.5分。
  • 适合需要长上下文的推理与检索任务。

Transformer架构在长文本场景下因二次计算和线性内存增长而受限。现有线性注意力通过固定大小状态压缩上下文,但常导致检索与推理性能下降。本文提出两项创新:首先,将状态更新视为信息分类,采用softmax-top-k硬分类实现行稀疏更新,提升感受野并减少类间干扰;其次,在稀疏框架内引入稀疏状态扩展(SSE),将上下文状态划分为多个分区,使参数量与状态容量解耦,同时保持稀疏分类范式。结合高效并行实现,该设计实现高判别性状态表示。我们在语言建模、上下文检索与数学推理任务上验证SSE,结果表明其在纯线性与混合架构(SSE-H)中均表现优异。经强化学习训练后,2B规模的SSE-H模型在AIME24上达64.5分,AIME25上达50.2分,显著优于同规模开源Transformer,证明SSE是长上下文建模的有效架构。

原文摘要 · Abstract (English)

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. Supported by efficient parallelized implementations, our design achieves effective classification and highly 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.5 on AIME24 and 50.2 on AIME25, significantly outperforming similarly sized open-source Transformers. These results highlight SSE as a promising and efficient architecture for long-context modeling.

线性注意力长文本状态扩展数学推理

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。