用于长上下文线性注意力的 Kronecker Delta Rule

我们提出了 KronLSTM，这是一种基于 Kronecker Delta Rule 的线性注意力架构。Kronecker Delta Rule 是一种用于联想矩阵记忆的新型更新规则。KronLSTM 通过引入记忆修订机制扩展了 xLSTM 系列，该机制能够校正指定的关联，同时仅沿矩阵空间中的单一方向修改记忆。与 delta-rule 更新相比，这种最小侵入式的记忆修订可减少对已存储信息的干扰，因此尤其适用于长上下文语言建模。由此产生的记忆动态可以数学描述为一个收敛至稳态分布的马尔可夫链，同时，在整个输入序列中，记忆矩阵的范数期望始终保持一致有界。因此，无需为记忆状态额外引入归一化和稳定化机制，KronLSTM 也能确保稳定的记忆动态。实验表明，KronLSTM 具备与现有线性注意力架构相当的竞争力，并且经常优于包括（Gated）DeltaNet 系列模型在内的现有架构，尤其在长上下文基准测试中表现出色。这些结果证明，Kronecker Delta Rule 是循环线性注意力架构中稳定联想记忆的一种有效且可解析的基础。我们的代码（https://github.com/NX-AI/KronLSTM）和模型（https://huggingface.co/collections/NX-AI/kronlstm）均已开源。