循环状态中的学习：使用线性循环网络进行梯度下降

arXiv:2410.11687v4 公告类型：替换 摘要：上下文学习使序列模型能够根据输入中的示例适应新任务。一个重要的研究方向表明，可以构造自注意力机制，使其在前向传播过程中对根据上下文示例拟合的线性预测器执行梯度下降。状态空间模型（SSM）和其他线性循环网络（LRNN）以线性时间成本对序列进行建模，但其循环更新如何执行同样的上下文梯度下降尚不明确。我们提出 Gradient-based Recurrent In-context Learner（GRIL），这是一种对角线性循环网络，它将有监督梯度步骤分解为短窗口交叉乘积写入和对下一个查询的乘法读出。对于线性回归，该构造将上下文梯度累积在矩阵状态中，并在一次前向传播中应用它，同时仅需 $O(f^2)$ 个可学习自由度。相同的设计还可扩展到多步更新和交叉熵分类，并通过基于 MLP 的有限扩展应用于非线性回归。我们通过实验证明，在合成 ICL 任务上训练的 GRIL 能够重现该构造在解析上预测的行为和参数。此外，同一架构还可以扩展并在通用基准上进行训练，包括 Long Range Arena、语言建模和联想回忆。综上，这些结果确立了窗口化交叉乘积自注意力这一具体的归纳偏置，使 LRNN 能够通过类似梯度下降的更新进行上下文学习，同时仍可在通用任务上进行训练。