重新思考高效注意力在混合架构中的作用

# 重新思考高效注意力在混合架构中的作用

来源：https://arxiv.org/html/2606.15378

Ziqing Qiao\(^1\)，Yinuo Xu\(^{1\*}\)，Chaojun Xiao\(^1\)，Zhou Su\(^2\)，Zihan Zhou\(^2\)，Yingfa Chen\(^1\)，Xiaoyue Xu\(^2\)，Xu Han\(^{2\dagger}\)，Zhiyuan Liu\(^{2\dagger}\)

\(^1\)清华大学 \(^2\)OpenBMB

qzq24@mails\.tsinghua\.edu\.cn，\{xcj,han\-xu,liuzy\}@tsinghua\.edu\.cn

###### 摘要

现代语言模型越来越多地采用混合架构，将全注意力与高效注意力模块（如滑动窗口注意力 (SWA) 和循环序列混合器）相结合。然而，这些高效模块如何塑造模型能力仍未被充分理解。为填补这一空白，我们从三个角度对混合架构进行了系统分析：缩放行为、机制分析和架构设计。首先，从缩放角度来看，我们发现高效注意力设计主要影响长上下文能力出现的速度，而不同的混合架构在充分训练后最终会收敛到相当的长上下文性能。其次，从机制上，我们表明长距离检索主要由全注意力承载，而高效注意力则塑造其优化轨迹。这解释了一个我们称之为**大窗口惰性**的反直觉现象：更大的 SWA 窗口会延迟全注意力层中检索头的形成。第三，在此机制的指导下，我们表明，对于小窗口 SWA 混合模型，仅在全注意力层上应用 NoPE，就能在几乎不影响短上下文性能的情况下，显著提升长上下文性能。\(^1\) 我们已在 rethinking-hybrid-attattention (https://github.com/thunlp/rethinking-hybrid-attention) 上发布代码。

## 1 引言

随着大型语言模型越来越多地用于长文档理解和智能体工作流，处理扩展上下文已成为近期模型发布的核心需求 (DeepSeek-AI, 2026；Singh et al., 2025)。然而，标准的 softmax 注意力（我们称之为全注意力）在长序列长度下成本高昂 (Vaswani et al., 2017)。这促使了一系列混合注意力架构的出现，这些架构将全注意力与高效注意力（如滑动窗口注意力 (SWA) (Beltagy et al., 2020) 和循环序列混合器 (Gu and Dao, 2023; Yang et al., 2024a)）相结合，这一设计现已被近期语言模型广泛采用 (Agarwal et al., 2025; Gemma Team, 2025; Cao et al., 2026)。尽管这些架构很普遍，但高效注意力在混合架构中的作用仍不清楚。现有研究缺乏对不同高效注意力设计如何塑造混合架构的能力和训练动态，特别是其长上下文性能的统一机制分析 (Xiao et al., 2026; Li et al., 2025; Wang et al., 2025; Bae et al., 2025)。为填补这一空白，我们研究三个研究问题：

*   **RQ1 - 缩放行为：混合架构在短上下文和长上下文性能上如何缩放？**
*   **RQ2 - 机制分析：高效注意力设计如何影响长上下文性能？**
*   **RQ3 - 架构设计：哪些设计原则能产生更有效的混合架构？**

#### 短上下文和长上下文能力的缩放规律

我们通过**缩放定律**的视角，研究混合架构在短上下文和长上下文性能上随多个模型规模和训练预算缩放的情况 (Kaplan et al., 2020; Hoffmann et al., 2022)。考虑到下游基准分数的离散性和不稳定性 (Liang et al., 2026)，我们使用验证损失 \(\mathrm{Loss}\) 和 \(\log(\mathrm{LongPPL})\) (Fang et al., 2025) 作为两个连续的拟合目标。前者捕获一般的短上下文建模质量，后者则为长上下文能力提供平滑的代理指标。拟合的缩放曲线清楚地表明，高效注意力设计对验证损失 \(\mathrm{Loss}\) 几乎没有影响，但会导致 \(\log(\mathrm{LongPPL})\) 出现更显著的差异。具体而言，在有限的训练预算下，不同的混合架构表现出显著的差距，其中大窗口 SWA 混合模型表现明显更差。然而，随着训练变得足够充分，这些差距显著缩小，并最终接近相似水平。

#### 作为优化先验的高效注意力

上述缩放模式给我们留下了两个看似矛盾的问题。第一，为什么具有不同高效注意力的混合架构最终会收敛到相似的长上下文水平？第二，为什么它们的收敛速度差异如此之大，尤其是在具有不同窗口大小的 SWA 变体之间？我们的机制分析表明，这两个问题有一个共同的解释：高效注意力并不直接决定长上下文能力；相反，它充当了一个**优化先验**，塑造了全注意力的训练方式。

**为什么混合架构会收敛？** 通过感受野约束和逐层探测实验，我们发现长距离信息主要由全注意力承载，而不是由高效注意力模块承载，即使是那些原则上具有无界感受野的循环序列混合器也是如此。由于共享相同的全注意力组件，不同的混合架构无论其高效注意力设计如何，都会收敛到相似的长上下文水平。

**为什么收敛速度不同？** 虽然全注意力决定了最终收敛水平，但高效注意力通过影响全注意力在训练期间发展其长距离检索行为的速度来影响长上下文能力。作为一个具体例子，通过追踪检索头 (Wu et al., 2025)，我们发现检索头在配备更大 SWA 窗口的混合模型中明显形成得更晚：一旦局部窗口已经为下一 token 预测提供了足够的上下文，推动全注意力学习长距离检索的梯度信号就会减弱。我们将这种现象称为**大窗口惰性**。

#### 超越高效注意力的混合架构设计

这些发现表明，混合架构设计应较少关注增加高效注意力的内在能力，而更多关注帮助全注意力更有效地学习长距离检索。从这个角度，我们重新审视了高效注意力模块之外的几个设计选择。作为一个简单而有效的实例，我们将 NoPE (Kazemnejad et al., 2023) 应用于小窗口 SWA 混合模型的全注意力层。这一简单的修改带来了明确的长上下文能力提升，而对短上下文性能的影响可以忽略不计，这在下游基准评估中一致地反映出来。图 1 总结了我们的主要发现及其设计启示。总之，我们的结果重新定义了高效注意力在混合架构中的作用。长上下文能力的实际瓶颈并不简单在于高效注意力模块有多强大，而在于它如何影响全注意力中长距离检索的出现。这一观点解释了跨混合架构的缩放模式，并将全注意力指向为改进长上下文混合模型的关键目标。

**见图说明**
图 1：概览。**缩放**：不同的高效注意力设计产生不同的 \(\log(\mathrm{LongPPL})\) 曲线，在充分训练后收敛到相似水平。**机制**：长距离检索主要由全注意力承载，而高效注意力充当**优化先验**，其中大窗口 SWA 滞后最多。**设计**：加强全注意力本身（例如，在全注意力中将 RoPE → NoPE）进一步改善了长上下文性能。

## 2 相关工作

#### 混合注意力架构

现有的混合架构主要遵循两条路线。一种使用 SWA (Beltagy et al., 2020) 作为高效注意力，近期设计已倾向于在整体性能有限退化的情况下使用更小的窗口和更稀疏的全注意力比例 (Agarwal et al., 2025; Huang et al., 2026)。另一种使用循环序列混合器，将过去历史压缩成紧凑的循环状态，如 Lightning Attention (Qin et al., 2024)、Mamba-2 (Dao and Gu, 2024) 和 Gated DeltaNet (Yang et al., 2025b)，这些正越来越多地被近期模型采用 (Li et al., 2025; Blakeman et al., 2025; Cao et al., 2026; Team et al., 2026)。除了高效注意力模块的选择，近期工作还探索了头部混合 (Dong et al., 2025; Xiao et al., 2025b) 和全注意力层的位置编码 (Yang et al., 2025a; Puvvada et al., 2025; Chen et al., 2026)。然而，这些研究大多数仅呈现最终结果或特定系统内的有限消融实验 (Gemma Team, 2025; Xiao et al., 2026)，缺乏跨高效注意力架构的受控比较。一些研究已开始更系统地审视混合架构中的结构选择。Wang et al. (2025) 比较了多种线性注意力变体和混合比例，而 Waleffe et al. (2024) 和 Bae et al. (2025) 分析了 Mamba-Transformer 混合模型中的层组成和放置。然而，这些研究仍局限于基于循环混合器的混合模型，并且缺乏机制解释。我们通过在受控的缩放规律设置下比较不同的高效注意力设计，并分析它们如何塑造混合架构的长上下文能力，来弥合这一差距。

#### 缩放规律与长上下文评估

缩放定律描述了预训练性能如何随模型和数据规模变化 (Kaplan et al., 2020; Hoffmann et al., 2022)，随后扩展到迁移学习 (Hernandez et al., 2021) 和下游能力预测 (Chen et al., 2024)。然而，针对长上下文能力的缩放定律仍未得到充分探索。现有的长上下文评估通常依赖于离散基准，如 RULER 和 LongBench (Hsieh et al., 2024; Bai et al., 2024)，这些基准衡量最终性能，但不适合跟踪预训练动态。一条互补的机制研究表明，检索头是长上下文事实检索的基础 (Wu et al., 2025; Xiao et al., 2025a)，并通过追踪检索头的形成来观察预训练期间长上下文能力的发展 (Liang et al., 2026)，但此类信号描述的是机制而非量化能力。相比之下，LongPPL (Fang et al., 2025) 提供了一种连续的困惑度风格指标，与长上下文基准强相关，并已被近期长上下文研究所采用 (Song et al., 2026; Willette et al., 2025)。我们进一步利用这一指标来拟合长上下文性能的缩放定律，从而能够更全面地比较长上下文能力如何在混合架构中出现。

## 3 预备知识

### 3.1 混合架构

我们涵盖两种常见的高效注意力形式：**滑动窗口注意力 (SWA)**，其中每个 token 仅关注有限的局部窗口；以及**循环序列混合器**，包括 **Lightning Attention**、**Mamba-2** 和 **Gated DeltaNet (GDN)**，它们通过不同的衰减策略和更新规则将过去 token 压缩到循环状态中。我们使用 \(q_t, k_t, v_t \in \mathbb{R}^{d_h}\) 表示位置 \(t\) 处的每头查询、键和值向量（为符号简化假设 \(d_k = d_v = d_h\)），并让 \(\mathrm{softmax}_s\) 表示在索引 \(s\) 上归一化的 softmax。下面的公式给出了这些机制的规范形式，用于匹配不同混合模型大小的实现级参数选择见附录 B。

#### 全注意力

对于每个位置 \(t\)，输出 \(O_t\) 在所有先前位置 \(s \le t\) 上计算：
\[
O_t = \sum_{s \le t} \mathrm{softmax}_s \left( q_t^\top k_s / \sqrt{d_h} \right) v_s
\tag{1}
\]

#### 滑动窗口注意力

SWA 将求和范围限制在大小为 \(w\) 的窗口内：
\[
O_t = \sum_{s \in [t-w+1, t]} \mathrm{softmax}_s \left( q_t^\top k_s / \sqrt{d_h} \right) v_s
\tag{2}
\]

下面三种循环混合器均共享形式 \(O_t = S_t q_t\)，其中循环状态 \(S_t \in \mathbb{R}^{d_h \times d_h}\)；它们主要区别在于 \(S_t\) 的更新方式。

#### Lightning Attention

Lightning 是一种线性注意力，具有固定的每头衰减 \(\gamma \in (0,1)\)：
\[
S_t = \gamma S_{t-1} + v_t k_t^\top.
\tag{3}
\]

#### Mamba-2

遵循结构化状态空间对偶 (SSD) 形式，Mamba-2 可以写为：
\[
S_t = \gamma_t S_{t-1} + v_t k_t^\top.
\tag{4}
\]
数据相关的 \(\gamma_t\) 允许对过去状态保留多少进行每 token 控制。

#### Gated DeltaNet

GDN 进一步增加了受控的... (翻译待继续，因输入截断)

(注意：原始输入在“GDN further adds controlled f”处截断，可能不完整。按现有内容翻译至此处。)