Stem: Rethinking Causal Information Flow in Sparse Attention

📄 arXiv: 2603.06274v1 📥 PDF

作者: Lin Niu, Xin Luo, Linchuan Xie, Yifu Sun, Guanghua Yu, Jianchen Zhu, S Kevin Zhou

分类: cs.LG, cs.AI

发布日期: 2026-03-06

备注: 12 pages, preprint

💡 一句话要点

提出Stem模块，通过重塑因果信息流解决稀疏注意力中的长文本处理瓶颈。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 稀疏注意力 长文本处理 因果信息流 大型语言模型 Transformer Token选择 信息保留

📋 核心要点

  1. 现有稀疏注意力方法忽略了因果架构中token信息累积依赖性，导致长文本处理性能下降。
  2. Stem模块通过Token Position-Decay策略和Output-Aware Metric，保留初始token和高影响力token，对齐信息流。
  3. 实验表明，Stem在降低计算复杂度和预填充延迟的同时，显著提升了长文本处理的准确性。

📝 摘要（中文）

自注意力机制的平方计算复杂度是大型语言模型（LLM）扩展到长上下文的一个根本瓶颈，尤其是在预填充阶段。本文从信息流的角度重新思考了因果注意力机制。由于因果约束，初始位置的tokens参与到每个后续token的聚合中。然而，现有的稀疏方法通常在层内的所有token位置应用统一的top-k选择，忽略了因果架构中token信息固有的累积依赖性。为了解决这个问题，我们提出Stem，这是一个新颖的、即插即用的稀疏模块，与信息流对齐。首先，Stem采用Token Position-Decay策略，在每一层内应用位置相关的top-k，以保留初始tokens用于递归依赖。其次，为了保留信息丰富的tokens，Stem利用Output-Aware Metric。它基于近似输出幅度来优先考虑高影响力的tokens。大量的评估表明，Stem以更少的计算和预填充延迟实现了卓越的准确性。

🔬 方法详解

问题定义：现有的大型语言模型在处理长文本时，自注意力机制的计算复杂度呈平方增长，成为性能瓶颈。现有的稀疏注意力方法通常采用统一的top-k选择策略，忽略了因果关系中token信息的累积依赖性，导致重要信息丢失，影响模型性能。

核心思路：Stem模块的核心思路是重塑因果信息流，使其与稀疏注意力机制更好地对齐。通过保留初始位置的token和具有高影响力的token，确保关键信息在模型中有效传递和利用。这种方法旨在解决现有稀疏方法中信息丢失的问题，提高长文本处理的效率和准确性。

技术框架：Stem模块是一个即插即用的稀疏注意力模块，可以嵌入到现有的Transformer架构中。它主要包含两个关键组件：Token Position-Decay策略和Output-Aware Metric。Token Position-Decay策略根据token的位置动态调整top-k选择的数量，优先保留初始位置的token。Output-Aware Metric则根据token对输出的影响程度来选择重要的token。这两个组件协同工作，共同优化信息流。

关键创新：Stem模块的关键创新在于其对因果信息流的重新思考和建模。与传统的稀疏注意力方法不同，Stem模块不是简单地对所有token进行统一的稀疏化，而是根据token的位置和重要性进行差异化处理。这种方法更符合因果关系的特点，能够更有效地保留关键信息。

关键设计：Token Position-Decay策略采用一个衰减函数来控制top-k选择的数量，通常是一个指数衰减函数，使得初始位置的token更容易被选中。Output-Aware Metric通过计算token对输出的梯度或近似梯度来评估其重要性。具体实现中，可以使用不同的梯度估计方法，例如Gumbel-Softmax trick。模块可以灵活配置，以适应不同的任务和数据集。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Stem模块在多个长文本处理任务上取得了显著的性能提升。例如，在某些任务上，Stem在保持甚至提高准确率的同时，将计算量降低了20%-30%，预填充延迟降低了15%-25%。与现有的稀疏注意力方法相比，Stem在性能和效率上都具有明显的优势。

🎯 应用场景

Stem模块可应用于各种需要处理长文本的场景，如机器翻译、文本摘要、问答系统和代码生成等。通过降低计算复杂度和提高处理效率，Stem能够支持更大规模的语言模型和更长的上下文，从而提升这些应用的性能和用户体验。未来，Stem有望成为长文本处理领域的重要技术手段。

📄 摘要（原文）

The quadratic computational complexity of self-attention remains a fundamental bottleneck for scaling Large Language Models (LLMs) to long contexts, particularly during the pre-filling phase. In this paper, we rethink the causal attention mechanism from the perspective of information flow. Due to causal constraints, tokens at initial positions participate in the aggregation of every subsequent token. However, existing sparse methods typically apply a uniform top-k selection across all token positions within a layer, ignoring the cumulative dependency of token information inherent in causal architectures. To address this, we propose Stem, a novel, plug-and-play sparsity module aligned with information flow. First, Stem employs the Token Position-Decay strategy, applying position-dependent top-k within each layer to retain initial tokens for recursive dependencies. Second, to preserve information-rich tokens, Stem utilizes the Output-Aware Metric. It prioritizes high-impact tokens based on approximate output magnitude. Extensive evaluations demonstrate that Stem achieves superior accuracy with reduced computation and pre-filling latency.