Hybrid Gated Attention

# 混合门控注意力
来源：https://arxiv.org/html/2608.11805
###### 摘要

门控注意力是缓解注意力汇点并增强注意力表征能力的有效方法。为了进一步拓展其效果-效率帕累托前沿，我们提出了一种混合门控注意力（HyGA）框架，包含三种类型的门控策略。具体来说，这些门控利用注意力多个阶段的不同信息，从多个视角协同构建逐元素/逐头门控，捕捉头内和跨头的信息交互。通过我们的混合门控组件，HyGA 能够提供多源调制信号，实现对信息流更全面的控制，并提升注意力的表征能力。我们还引入了低秩矩阵分解和可学习的注意力汇点，以进一步提升训练效率和稳定性。在实验中，我们基于不同骨干网络在广泛使用的基准上评估了 HyGA。实验结果表明，与门控注意力相比，HyGA 在训练损失和各项下游性能上均有全面提升。HyGA 还被验证在不同计算成本下均能达到最佳性能，并提供了全面的模型分析以加深理解。所提出的 HyGA 为一种更有效、更高效且更稳定的注意力机制提供了新思路。

## 1 引言

多头注意力（[37](https://arxiv.org/html/2608.11805#bib.bib1)）已成为大语言模型中用于上下文信息交互的核心组件。因此，注意力的信息处理能力直接影响大语言模型的上下文理解和信息整合能力，从而在决定模型表征能力方面起着关键作用。

近年来，越来越多的研究工作致力于改进注意力机制。一个主要研究方向聚焦于提升注意力效率并降低 KV 缓存开销，由此产生了线性注意力（[42](https://arxiv.org/html/2608.11805#bib.bib16)；[41](https://arxiv.org/html/2608.11805#bib.bib17)）、SSM（[10](https://arxiv.org/html/2608.11805#bib.bib39)）、稀疏注意力（[27](https://arxiv.org/html/2608.11805#bib.bib42)；[39](https://arxiv.org/html/2608.11805#bib.bib43)；[43](https://arxiv.org/html/2608.11805#bib.bib44)）以及 KV 共享或压缩机制（如 GQA（[2](https://arxiv.org/html/2608.11805#bib.bib2)）和 MLA（[26](https://arxiv.org/html/2608.11805#bib.bib3)））等代表性方法。

另一个重要方向旨在解决与注意力相关的性能和稳定性问题，其中注意力汇点现象已得到广泛研究。注意力汇点（[38](https://arxiv.org/html/2608.11805#bib.bib5)；[15](https://arxiv.org/html/2608.11805#bib.bib35)；[4](https://arxiv.org/html/2608.11805#bib.bib36)；[35](https://arxiv.org/html/2608.11805#bib.bib34)；[36](https://arxiv.org/html/2608.11805#bib.bib45)）指的是注意力模块将过多的注意力权重分配给少数语义相关性有限的标记（如 BOS 标记）的现象。这些标记通过吸收多余的注意力来充当注意力汇点，实际上使某些注意力头近似于无操作。然而，过度集中的注意力可能会降低上下文信息的有效利用，这一点值得关注。

图 1：我们提出的 HyGA 包含三个门控，即 X-gate、H-gate 和 C-gate，通过更具信息量的输入和多样化的控制来增强注意力的门控策略。它还进一步引入了低秩矩阵分解和可学习的注意力汇点，以实现更高效、更稳定的训练。

门控注意力是缓解注意力汇点和提升注意力能力的代表性方法（[31](https://arxiv.org/html/2608.11805#bib.bib4)）。它对缩放点积注意力（SDPA）输出应用逐元素门控，引入非线性和稀疏性，以缓解注意力汇点并增强模型表达能力。尽管如此，仍存在几个局限：（a）原始门控注意力依赖原始输入 X 进行门控，忽略了其他可作为当前门控策略良好补充的信息。（b）逐元素门控引入了相当多的额外参数和计算开销，在提升门控效率方面仍有很大空间。（c）现有的门控注意力仍表现出不可忽视的 BOS 标记汇点比率，这促使我们进一步缓解注意力汇点现象以实现稳定训练。

基于上述观察，我们大幅扩展了门控注意力，并提出了我们的**混合门控注意力（HyGA）**。具体来说，我们的 HyGA 包含三个协调良好的注意力门控，从不同方面智能地调整 SDPA 的输出幅度。除了由输入 X 控制的经典注意力门控（*X-gate*）之外，我们首先引入了一个以 SDPA 输出 H_i 作为特定头门控控制器的逐元素注意力门控，称为 *H-gate*，它在注意力之后包含更丰富的上下文信息，从而提供额外的信息增益。由于这两个门控以来自注意力不同阶段的表示为条件，它们可以提供互补的调制信号，实现对信息流更全面的控制，从而提升注意力的表征能力。为了实现良好的效果-效率帕累托前沿，我们对两个门控的投影应用了低秩矩阵分解。在这种情况下，HyGA 能够以比经典门控注意力少得多的注意力参数达到可比的训练损失和下游性能。

除了专注于头内交互以生成逐元素门控分数的 H-gate 之外，我们还设计了跨头门控（*C-gate*）来捕捉头间相关性。它联合采用所有头的 SDPA 输出来计算逐头门控分数，作为上述细粒度逐元素门控的良好补充。这三个门控相互配合良好，通过更全面地考虑多样特征来丰富注意力的门控策略。为防止三个门控叠加相乘可能导致训练不稳定和输出消失问题，我们采用了一种门控融合策略。我们还发现，配备可学习的注意力汇点（[1](https://arxiv.org/html/2608.11805#bib.bib8)）后，HyGA 可以更加稳定，这似乎与门控注意力存在潜在的功能重叠。

在实验中，我们在广泛使用的基准上基于不同模型评估了 HyGA，结果表明 HyGA 持续优于原始门控注意力。通过适当的低秩压缩设置，HyGA 仅使用门控注意力 26% 的门控计算成本即可获得略微更好的性能。贡献总结如下：

- • 我们提出了 HyGA，它联合采用三种混合门控策略，从不同因素计算逐头和逐元素门控分数。配备可学习的汇点后，HyGA 实现了更稳定的训练。
- • 我们探索了混合门控中不同的低秩矩阵分解设置，以拓展门控注意力方法的效果-效率帕累托前沿。
- • 总体而言，与不同骨干网络和模型设置的基线相比，HyGA 取得了显著改进，为实践中更有效、更高效、更稳定的门控注意力模块提供了新思路。

## 2 预备知识

在本节中，我们首先简要介绍和阐述本工作中使用的经典注意力机制。

**分组查询注意力（GQA）。** GQA 在流行的大语言模型中被广泛使用。传统的多头注意力机制采用查询（Q）、键（K）和值（V）来计算上下文相关的标记表示，其中每个查询头独立地关注其对应的键值头。GQA 允许多个注意力头的查询共享一组键值投影。设 H_i 表示第 i 个头的注意力输出，公式如下：

H_i = Softmax(Q_i K_{g(i)}^T / √d_h) V_{g(i)}，(1)

其中 g(i) 将第 i 个查询头映射到其对应的键值组。同一组内的查询头共享相同的键和值表示，从而大幅减少 KV 缓存的内存占用，同时保留模型的大部分推理能力。

**多潜在注意力（MLA）。** MLA 通过下投影将 K 和 V 压缩到低维潜在空间，从而改进 GQA。KV 缓存中仅缓存低维潜在表示，在注意力计算时再将完整的 K 和 V 矩阵上投影。具体公式如下：c_t^KV = W^{DKV} h_t，k_t^C = W^{UK} c_t^KV，v_t^C = W^{UV} c_t^KV，其中 W^{DKV} 是下投影矩阵，c_t^KV 是键和值共享的压缩潜在向量。与 GQA 相比，MLA 实现了更强的 KV 缓存压缩，在长上下文推理中可以节省更多内存。

**门控注意力。** 门控注意力（[31](https://arxiv.org/html/2608.11805#bib.bib4)）是指在原始逐头注意力之后对输出矩阵应用逐元素门控操作。典型的门控注意力函数在 SDPA 输出之后采用：

H'_i = H_i ⊙ σ(X W_i)，O = H'W_O，W_i ∈ R^{d_model × d}。(2)

其中 d 是头维度，d_model 是模型维度，X 是输入矩阵，H' 是在 W_O 之前的门控输出。W_i 表示第 i 个头的投影矩阵，为 d 个维度提供逐元素门控。它同时提供了非线性表达能力，并抑制了注意力汇点和大规模激活，因此可以进一步提升 GQA 和 MLA 的性能。

## 3 方法

在本工作中，我们提出了混合门控注意力，试图进一步优化门控注意力的有效性、效率和训练稳定性。

图 2：所提出的 HyGA 概述，包含三个门控：X-gate（G_X）、H-gate（G_H）和 C-gate（G_C），并配备低秩矩阵分解（在 H-gate 和 X-gate 中）和可学习的注意力汇点（SDPA 中的 LS）。

### 3.1 总体框架

如图 2 所示，HyGA 主要通过三种门控机制扩展了原始门控注意力：（a）以注意力层之前的原始输入 X 作为输入的原始门控（X-gate），（b）以 SDPA 之后的输出 H 作为输入的所提出的门控（H-gate），以及（c）所提出的跨头门控（C-gate），它捕捉头间连接并提供头级重新加权。我们发现这三个注意力门控能够良好协作，共同提升性能。接下来，我们采用门控融合策略来获得最终的混合门控注意力输出，避免不希望出现的门控主导，确保训练平稳。此外，为了提升训练稳定性，我们使用可学习的注意力汇点来实现 HyGA，为训练稳定性提供双重保障。基于上述机制，我们的 HyGA 有效提升了模型能力和训练稳定性。

### 3.2 混合门控策略

#### 混合门控输入与函数

原始门控注意力采用原始输入 X 来控制逐元素门控，如公式 (2) 所述。通过对原始门控注意力的分析，我们观察到输出 H 在注意力之后包含更丰富的标记交互信息，可能除了 X 之外还能提供额外的信息增益。基于这一观察，我们提出使用 SDPA 之后的输出 H 作为控制注意力门控的输入，作为 X 的补充。我们有：

H'_i = H_i ⊙ σ(X W_i + SiLU(H_i W_i^d) W_i^u)。(3)

对于新添加的 H-gate，我们采用带有 SiLU（[12](https://arxiv.org/html/2608.11805#bib.bib6)）的两层 MLP 形式来增强非线性调制。这是因为上述形式可以通过调整中间维度更方便地进行低秩参数压缩，为下文讨论的效率改进做准备。这里，W_i^d ∈ R^{d × d_int} 和 W_i^u ∈ R^{d_int × d} 分别是第 i 个头的下投影和上投影门控矩阵，其中 d_int 表示控制下投影-上投影瓶颈的中间维度。

从上述公式可以看出，基于 X 的原始 X-gate 和基于 H 的 H-gate 表现出很强的功能互补性。前者主要捕捉注意力之前输入标记的内在特征，而后者主要捕捉通过注意力进行上下文交互之后的特征。因此，两个门控建模的信息在一定程度上是不同的，并且它们在不同位置调制信息。它们的联合效果因此可以更全面地捕捉和建模注意力输出中的信息，使模型能够通过逐元素门控更好地控制信息流。

#### 融合或分离的 X-/H- 门控。

与公式 (3) 中的融合形式不同，混合 X-/H- 门控也可以是分离的：

H'_i = H_i ⊙ σ(X W_i) ⊙ σ(SiLU(H_i W_i^d) W_i^u)，(4)

然而，我们选择公式 (3) 中的加性耦合形式，而不是公式 (4) 中的形式。这是因为公式 (4) 中的乘性耦合形式将两个门控相乘。由于 sigmoid 激活函数的取值范围为 (0,1)，每个门控通常作为信息流的抑制性调制器。因此，相乘过多的门控可能导致过度抑制，这可能削弱梯度传播并限制有效的表示学习（尤其是结合第 3.3 节的跨头门控时）。相比之下，公式 (3) 中的融合门控首先将激活前的门控 logits 相加，然后应用激活函数，从而允许不同因素共同控制信息流，而不是独立调制。这可以带来更稳定的信息流和更强的表达能力。