✦星尘知识库

GatedAttention

NLP📅 2026-06-30👁 4 阅读♥ 0 点赞⏱ 约 12 分钟

论文：Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

这篇论文《Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free》的核心做法是对标准的Transformer注意力机制进行了一个简单但极具破坏力（正面意义）的修改：在缩放点积注意力（SDPA）的输出之后，加入了一个门控机制（Gating Mechanism）。

以下是对其核心做法的详细原理阐述，分为具体实现方式、生效机制（原理）以及带来的核心优势三个维度。

1. 核心做法：SDPA 输出门控 (SDPA Output Gating)

作者在测试了30多种变体后，确定了效果最好的配置，并在论文中被称为 G_1** 位置的 Head-Specific Sigmoid Gating**。

**具体公式**

在标准的多头注意力（MHA）中，注意力输出 Y 通常直接进入输出投影层 W_O。这篇论文引入了一个门控操作：

Y' = Y \odot \sigma(X W_\theta)

其中：

  • Y：是缩放点积注意力（SDPA）的输出结果。Y = \text{Softmax}(\frac{QK^T}{\sqrt{d_k}})V。
  • X：是该注意力层的输入（即Query对应的Hidden State，通常是Pre-Norm之后的输入）。
  • W_\theta：是门控机制学习到的权重矩阵。
  • \sigma：是 Sigmoid 激活函数（将值映射到 [0, 1]）。
  • \odot：是逐元素乘法（Hadamard product）。

**关键配置细节**

  1. 位置（Position G_1）：门控被放置在 SDPA计算之后，但在 最终输出投影层（W_O）之前。这是论文对比后发现的最佳位置。
  2. 激活函数：必须使用 Sigmoid，而非 SiLU 或 ReLU。
  3. 独立性（Head-Specific）：每个注意力头（Head）拥有自己独立的门控参数，而不是所有头共享。这意味着每个头可以根据输入独立决定保留或丢弃多少信息。
  4. 粒度（Granularity）：
    • Element-wise（最佳）：门控分数为向量，维度与 Y 相同，对每个维度进行精细调控。
    • Head-wise（高效）：每个头只生成一个标量分数，对整个头的输出进行缩放。效果略差但参数更少。

2. 生效原理：非线性与稀疏性

论文深入分析了为什么仅仅加入这一个门，就能带来显著的性能提升和训练稳定性。主要归因于两个因素：

**原理一：引入非线性，打破低秩瓶颈 (Non-linearity & Low-Rank Mapping)**

  • 问题：在标准Transformer中，Value投影矩阵 (W_V) 和输出投影矩阵 (W_O) 是两个连续的线性层。根据线性代数性质，两个连续线性变换 W_O \cdot W_V 等价于单个线性变换。由于注意力头的维度 d_k 通常远小于模型维度 d_{model}，这实际上形成了一个低秩瓶颈（Low-Rank Bottleneck），限制了模型的表达能力。
  • 解决：在 W_V 的输出（经过SDPA加权后）和 W_O 之间插入一个非线性的门控操作，使得这两个线性层无法被合并。这增加了变换的非线性表达能力，类似于在两层MLP之间加激活函数。
  • 证据：作者发现，即使只是在这里加一个 RMSNorm（也是非线性的），也能提升性能，证明了非线性的重要性。

**原理二：输入依赖的稀疏性 (Input-Dependent Sparsity)**

  • 机制：Sigmoid函数会将门控值压缩在 [0, 1] 之间。实验发现，训练后的模型倾向于让大部分门控分数接近于 0（非常稀疏）。
  • Query依赖：门控分数是基于输入 X（即当前的 Query 信息）计算的。这意味着模型学会了根据当前 token 的需求，动态地过滤掉那些由 Softmax 聚合进来的、但不相关的信息。
  • 对比：相比于仅仅修改 Value (G_2 位置)，在 SDPA 输出处 (G_1 位置) 加门控效果更好，因为这里的门控是由当前查询 (Query) 决定的，而不是由历史的 Key/Value 决定的。

3. 核心优势：消除“注意力汇” (Attention-Sink-Free)

这是该论文最引人注目的发现之一。这个简单的门控机制顺带解决了大模型中常见的**Attention Sink（注意力汇）**问题。

**什么是 Attention Sink？**

在标准 Softmax Attention 中，\sum \text{Softmax} = 1。如果当前的 Token 发现上下文中没有任何相关信息，它依然被迫要把这 "1" 的概率分配出去。模型通常学会把这部分无意义的概率分配给序列的第一个 Token（起始符）或特定的标点符号。这导致首个 Token 拥有极高的注意力分数和巨大的激活值。

**门控如何解决？**

  • 二次过滤：虽然 Softmax 依然归一化为 1，但后续的 Sigmoid Gate 可以输出 0。
  • 逻辑：即使 Softmax 错误地把 90% 的注意力给了第一个 Token（因为没别的地方给），Gate 可以通过乘以一个极小的系数（例如 0.001），将这部分无效的信息流直接截断。
  • 结果：
    1. 无需 Sink Token：模型不再需要专门维护一个“垃圾回收站”式的 Sink Token。
    2. 消除数值尖峰：模型内部不再出现异常巨大的激活值（Massive Activations），这大大提升了训练稳定性，允许使用更大的学习率和Batch Size。
    3. 长文本外推：由于消除了对首个 Token 的过度依赖，模型在推理长度超过训练长度时（Length Extrapolation），表现出更强的鲁棒性。

总结

这篇论文的核心做法可以概括为：在 Transformer 的 Attention 模块内部，利用“Query依赖的 Sigmoid 门”对“Softmax后的结果”进行二次筛选。

这一做法本质上是将 Attention 从单纯的“加权求和（必须为1）”变成了“加权求和 + 动态门控（可以为0）”，从而同时提升了模型的非线性表达能力和对无关信息的过滤能力。

补充说明：$ W_O^k $的定义与作用

在论文第4.1节（Non-linearity Improves the Expressiveness of Attention Mechanisms）及公式(6)中，符号W_O^k具有明确的数学定义和关键理论意义。以下从数学定义、低秩瓶颈分析及非线性门控作用三方面系统阐述：

1. 数学定义

在标准多头注意力（Multi-Head Attention）机制中，所有注意力头的输出先拼接后通过全局投影矩阵W_O进行转换：

\text{Output} = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) \cdot W_O,

其中W_O \in \mathbb{R}^{(h \cdot d_k) \times d_{\text{model}}}，h为头数，d_k为单个头的维度，d_{\text{model}}为模型维度。

通过矩阵分解，该操作等价于每个头独立投影后求和：

\text{Output} = \sum_{k=1}^{h} \left( \text{head}_k \cdot W_O^k \right),

其中W_O^k \in \mathbb{R}^{d_k \times d_{\text{model}}}是W_O中对应第k个注意力头的子矩阵。每个\text{head}_k的维度为d_k，因此W_O^k的维度需匹配该输入维度（即d_k行、d_{\text{model}}列）。

2. 低秩瓶颈的数学分析

论文通过W_O^k揭示了传统多头注意力的线性结构限制。以第k个头为例：

  • Value投影：输入X \in \mathbb{R}^{n \times d_{\text{model}}}经W_V^k \in \mathbb{R}^{d_{\text{model}} \times d_k}投影为Value（维度降至d_k）；
  • 注意力加权：通过注意力权重S_{ij}（标量）对Value进行加权求和；
  • 输出投影：经W_O^k投影回d_{\text{model}}维度。

若无非线性门控，整个变换可表示为：

\sum_{j} S_{ij} X_j \cdot \left( W_V^k W_O^k \right),

其中W_V^k W_O^k \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}的秩受限于：

\text{rank}(W_V^k W_O^k) \leq \min\left( \text{rank}(W_V^k), \text{rank}(W_O^k) \right) \leq d_k.

由于d_k \ll d_{\text{model}}（例如d_k=128，d_{\text{model}}=4096），该乘积矩阵必然为低秩矩阵，导致线性变换能力严重受限。

示例说明：
若d_k=128，d_{\text{model}}=4096，则W_V^k W_O^k的最大秩仅为128，无法表示秩为4096的任意线性变换。这相当于将高维空间的复杂关系压缩到低维子空间，严重削弱模型表达能力。

3. 非线性门控的突破性作用

论文指出，通过在W_V^k与W_O^k之间插入非线性门控（如公式(6)中的G_1），可彻底打破低秩瓶颈：

  • 门控机制：例如，若G_1为ReLU或GLU（Gated Linear Unit），则变换变为：

\sum_{j} S_{ij} X_j \cdot W_V^k \cdot G_1(\cdot) \cdot W_O^k,

  • 其中G_1引入非线性激活，使中间表示突破秩的约束；
  • 数学本质：非线性函数G_1将低秩线性变换转化为高秩非线性变换。例如，ReLU的分段线性特性允许模型在不同输入区域动态切换子空间，显著提升表达能力。

关键对比：

  • 无门控：W_V^k W_O^k的秩上限为d_k，仅能表示低维流形上的线性关系；
  • 有门控：G_1的非线性特性使整体变换的秩可接近d_{\text{model}}，充分捕捉高维数据的复杂结构。

4. 结论

W_O^k的引入为分析多头注意力的表达能力提供了关键视角。其与W_V^k的乘积清晰揭示了线性结构的固有局限，而通过非线性门控的介入，模型得以突破低秩瓶颈，显著提升表达能力。这一发现直接支持了论文中关于非线性增强注意力机制有效性的核心论点，为后续改进注意力架构（如引入门控、激活函数等）提供了理论依据。

评论（0）

评论加载中…