Abstract
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.
Abstract (translated)
线性注意力方法由于其在递归解码中的效率，为传统的softmax注意力提供了一种有吸引力的替代方案。最近的研究集中在通过引入门控机制来增强标准线性注意力，同时保持计算效益不变。这种门控线性注意（GLA）架构包括竞争性的模型如Mamba和RWKV。在这项工作中，我们探讨了GLA模型在情境学习中的能力，并作出了以下贡献： 1. 我们展示了多层的GLA可以实现一类具有数据依赖权重的加权预处理梯度下降(WPGD)算法。这些权重是由门控机制和输入共同诱导出来的，使模型能够控制各个令牌对预测的贡献。 2. 为进一步理解这种加权机制的工作原理，我们引入了一种新的数据模型，该模型使用多任务提示，并描述了学习WPGD算法时优化景观的特点。在适度条件下，我们证明了一个全局最小值的存在性和唯一性（直到缩放为止），这对应于一个唯一的WPGD解决方案。 3. 最后，我们将这些发现转化为探讨GLA的优化景观，并阐明门控机制如何促进上下文感知学习以及何时可以严格证明其优于传统的线性注意力。 这项研究提供了关于门控线性注意模型内部运作的新见解，并为进一步改进语言模型中的注意力机制指明了方向。
URL
https://arxiv.org/abs/2504.04308
