LinearAttention系列解读(三)：线性注意力中的Gating机制

引言

标准attention在自回归推理时需要缓存全部历史的key，value (即kv-cache)，kv-cache随序列线性增加，在长程生成任务上有较大的显存压力。为此，linear attention提出用可分解的核函数来代替标准attention的softmax kernel

从而把随序列递增的kv-cache，转为固定大小key-value association的累加和，其注意力公式为：

为了方便阐述，将记为identity function，并忽略normalization term，此时

写成递归形式

其中：

但正如前文所言，以累积外积和作为memory是有局限的。随序列长度的增加，新旧association不断叠加，使得不同信息之间的crosstalk不断增大，最终导致memory collision。

为了提升有限memory的利用效率，在deltanet中引入delta rule机制。每次memory更新前，先从旧状态检索出与当前相关的，然后在更新阶段只做增量更新

本文将介绍另一种提升memory利用率且与delta-rule互补的方法，gating机制。

Linear Attention的Gating机制

对于Gating机制，在rnn-era就有广泛应用，如代表方法GRU。它的核心思想是：让模型自己决定，在当下输入下信息应当通过多少。

将gating机制(forgetting gate)放到linear attention的memory更新机制中，可以抽象成

其中。它决定旧状态的保留程度。越接近1保留的越多，越接近0遗忘的越多。设初始状态 \(S_0=0\)。表示hadamard乘法，逐元素乘积。

有很多设计思路。大体上分类两大类：

Data independent gate

Data dependent gate

下文具体介绍。

Data independent gate

最简单的方法，采用与输入无关的衰减率

它本质上引入一种recency bias：越远的association的权重越低，越近的越高。并且由于是data-independent的，这个规则应用于任意输入。

采用data-independent gate的代表方法如retnet。

Data dependent gate

这个方法假定gate 规则与当前的输入有关，即：

如何设计是核心差异。根据gate的粒度，大致可以分为4类方法：

(一) Full Matrix Gate

最直接的方法，直接用一个linear layer预测整个gating matrix

其中：, 负责将值域限制在(0,1)。

其状态更新为：

这个方式的优点是模型可以提供feature-wise 的forgetting，但参数量和计算量都很大。

(二) Scalar Gate

这个方式类似前文data-independent gate的方法，差异在于scalar由输入构造

其中。状态更新：

与固定不同，动态标量门可以根据当前token调整整个memory的遗忘速度。Mamba2、mLSTM、Gated Retention和早期Gated RFA（它还额外有个写入强度的gate）都可以归入这一类。

其优点是参数少、计算简单，并且容易转换为矩阵乘法；缺点是整个状态只能统一缩放。如果模型只想删除一小部分过时信息，其他记忆也会一起衰减。

(三) Two-sided Separable Matrix Gate

这个方法是上述两者的折中。分别给key, value两个维度生成门控向量。

再用向量外积的方式构造gating matrix

其中。通过上述处理，将matrix-level所需的的参数量降到

其状态更新为：

还可以用low-rank的思路进一步降低计算量，如构造

其中。通过low-rank处理，进一步将参数量降低至

(四) One-sided Separable Matrix Gate

不少linear attention的相关论文发现value维度的gate向量收益有限，因此不少工作舍弃了value维度的gate，即让。此时，每个 key channel 对应一个衰减率，所有value channel共享同一个衰减率。

其状态更新为：

GateLoop、HGRN-2、RWKV-6和GLA都采用了这种one-sided vector gate。与 scalar gate相比，它可以按memory channel控制遗忘；与完整矩阵门相比，它又保留了较低的参数量和较好的计算结构。

Gating Linear Attention的局限性

引入gating机制只能提升对有限memory的利用率，但并没有消除linear attention的基本限制：

短序列上，缓存和计算可能不如标准attention

key-value association累加造成的memory collision，导致长序列表现不佳。实际使用中，通常会用softmax based attention和linear attention的混合结构。

LinearAttention系列解读(一)：Transformers Are RNNs

LinearAttention系列解读(二)：DeltaNet技术剖析

另外，forget gate通常作用于状态坐标或memory channel，而不是独立的历史association，无法像delta rule那样沿给定key的方向进行定点擦除和重写。后面Gated-DeltaNet将gating机制引入deltanet中，有效弥补了这个不足。

在性能上，data-dependent gating机制的引入会额外增加训练复杂度，需要重新设计并行化算法。GLA 这篇工作对gating linear attention的并行化做了深度优化。

小结

本文从输入依赖性和门控粒度两个维度，梳理了linear attention中的forgetting gate。Data-independent decay提供固定的时间衰减；data-dependent gate则根据当前输入动态管理记忆，并进一步发展出 scalar、one-sided vector、two-sided separable matrix和full matrix等不同参数化。

这些设计本质上都是表达能力与计算效率之间trade-off：更细粒度的gate能提供更灵活的状态控制，但也会带来更多参数、更复杂的数值问题和更高的硬件实现成本。

如有疏漏之处，欢迎指出～