arXiv:2603.22300cs.LGcs.AI2026-03中稿 · ICLR被引 3

通过特征稀疏性突破注意力计算瓶颈,实现超长文本高效建模。

Scaling Attention via Feature Sparsity

  • 用k稀疏编码表示查询与键,降低注意力计算复杂度
  • 在GPT-2和Qwen3上提速达2.5倍,减少近50%计算量
  • 适合需要处理超长序列且追求高效率的场景

将Transformer扩展至超长上下文面临自注意力 $O(n^2 d)$ 的计算瓶颈。现有方法通过局部窗口、核近似或标记级稀疏性降低序列维度开销,但均导致精度下降。本文提出正交思路:特征稀疏性。我们设计稀疏特征注意力(SFA),将查询与键表示为 $k$-稀疏码,在保持高维表达力的同时,将注意力复杂度从 $Θ(n^2 d)$ 降至 $Θ(n^2 k^2/d)$。为实现大规模高效运行,引入FlashSFA——一种面向输入输出的内核,可直接在稀疏重叠区域操作,无需生成稠密分数矩阵。在GPT-2与Qwen3预训练中,SFA性能媲美稠密基线,速度提升最高达2.5倍,计算量与KV缓存减少近50%。在合成与下游任务中,SFA在长上下文下保持检索准确率与鲁棒性,优于压缩特征多样性的短嵌入基线。结果表明,特征稀疏性是高效注意力的互补且未被充分探索方向,使Transformer在极长上下文上实现数量级扩展,仅带来微小质量损失。代码已开源。

原文摘要 · Abstract (English)

Scaling Transformers to ultra-long contexts is bottlenecked by the $O(n^2 d)$ 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 $k$-sparse codes that preserve high-dimensional expressivity while reducing the cost of attention from $Θ(n^2 d)$ to $Θ(n^2 k^2/d)$. 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 $2.5\times$ 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.

注意力机制稀疏性模型加速长序列建模

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