提出Stem模块，让长文本生成更高效
Stem: Rethinking Causal Information Flow in Sparse Attention
- 按位置衰减策略动态选择关键词，保留早期信息
- 用输出影响度筛选高价值token，减少计算量
- 适合需要长上下文处理的LLM部署场景
自注意力的二次计算复杂度是将大语言模型扩展到长文本上下文的主要瓶颈，尤其在预填充阶段。本文从信息流动角度重新审视因果注意力机制：由于因果约束，初始位置的标记会参与后续每个标记的聚合。然而，现有稀疏方法通常在每层中对所有标记位置采用统一的top-k选择，忽略了因果结构中固有的累积依赖性。为此，我们提出Stem——一种与信息流动对齐的新型即插即用稀疏模块。首先，Stem采用标记位置衰减策略，在每层内实施位置相关的top-k选择，以保留初始标记的递归依赖关系；其次，为保持信息丰富的标记，Stem引入输出感知度量，基于近似输出幅值优先选择高影响力标记。大量实验表明，Stem在减少计算量和预填充延迟的同时，实现了更高的精度。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。