揭示门控线性注意力如何通过加权实现上下文学习。
Gating is Weighting: Understanding Gated Linear Attention through In-context Learning
- 门控机制生成数据相关权重，动态调节输入贡献。
- 多层门控线性注意力可实现加权预条件梯度下降算法。
- 理论证明其在特定条件下优于普通线性注意力。
线性注意力因其在循环解码中的高效性，成为软注意力的有力替代方案。近期研究通过引入门控机制提升标准线性注意力，同时保持计算优势，代表性模型包括 Mamba 与 RWKV。本文研究门控线性注意力（GLA）的上下文学习能力，发现多层 GLA 可实现一类数据依赖权重的加权预条件梯度下降（WPGD）算法。这些权重由门控机制与输入共同决定，使模型能动态控制各标记对预测的影响。为深入理解该加权机制，我们提出一种包含多任务提示的数据模型，并刻画了学习 WPGD 算法的优化景观。在温和条件下，我们证明全局最小值存在且唯一（至标量缩放），对应唯一最优 WPGD 解。最后，将这些发现应用于 GLA 的优化景观分析，阐明门控如何促进上下文感知学习，并在理论上证明其优于原始线性注意力的场景。
原文摘要 · Abstract (English)
Linear attention methods offer a compelling alternative to softmax attention due to their efficiency in recurrent decoding. Recent research has focused on enhancing standard linear attention by incorporating gating while retaining its computational benefits. Such Gated Linear Attention (GLA) architectures include competitive models such as Mamba and RWKV. In this work, we investigate the in-context learning capabilities of the GLA model and make the following contributions. We show that a multilayer GLA can implement a general class of Weighted Preconditioned Gradient Descent (WPGD) algorithms with data-dependent weights. These weights are induced by the gating mechanism and the input, enabling the model to control the contribution of individual tokens to prediction. To further understand the mechanics of this weighting, we introduce a novel data model with multitask prompts and characterize the optimization landscape of learning a WPGD algorithm. Under mild conditions, we establish the existence and uniqueness (up to scaling) of a global minimum, corresponding to a unique WPGD solution. Finally, we translate these findings to explore the optimization landscape of GLA and shed light on how gating facilitates context-aware learning and when it is provably better than vanilla linear attention.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。