矩阵正交化提升循环模型的记忆能力

# 矩阵正交化提升循环模型记忆能力 来源：https://ayushtambde.com/blog/matrix-orthogonalization-improves-memory-in-recurrent-models/ 2026年6月30日 *本研究由Paradigm (https://paradigm.xyz/) 资助。* Transformer在关联记忆（AR）方面表现卓越：注意力机制使每个令牌能直接访问其前面的令牌，这是其他架构（如循环神经网络RNN）难以匹敌的能力。 但在某些领域，我们无法承受Transformer二次注意力机制的开销。例如长时序强化学习（Dreamer风格）。对于这类应用，我们需要让循环神经网络发挥作用，但又不想放弃关联记忆能力。 目前已知在关联记忆方面表现最佳的RNN是mLSTM (https://arxiv.org/abs/2405.04517)，它是LSTM的一个变体，维护一个矩阵记忆。在MQAR (https://arxiv.org/pdf/2312.04927) 基准测试中，mLSTM相比基线模型表现出显著的记忆提升。但纯记忆能力可能不足以衡量循环模型的性能。在环境转换可能包含噪声的领域，一个有用的代理测试是噪声关联记忆（NAR）。 由于MQAR不测量NAR，我们可以参考MAD的噪声AR任务套件 (https://arxiv.org/pdf/2403.17844)。下面是一个任务示例： `0 9 3 10 12 13 15 14 0 9 5 8 2 9` 这里，键`0`映射到值`9`，键`3`映射到值`10`，以此类推。MAD生成器对键、值和干扰项使用不同的令牌范围。因此，如果键范围是`0-5`，那么令牌`12-15`就是干扰项。擅长NAR的模型应该在位置10预测出`9`，因为它看到了开头的`0 -> 9`，同时忽略了穿插其中的干扰令牌。 那么，如何改进循环模型的NAR呢？我们可以借鉴Muon的一些思想，Muon是一种在语言建模中非常成功的优化器。Muon对其动量进行正交化，起到均衡表示方向的作用。它防止少数强方向主导更新，并提升弱方向。特别相关的是最近的研究表明，Muon在尾部关联记忆学习方面优于Adam (https://arxiv.org/abs/2509.26030)。其思想是，这种均衡可以防止弱记忆被挤压。 受此启发，我们决定测试在读取时对mLSTM记忆矩阵进行正交化，并配合这一额外过程训练，是否能提高NAR性能。 我们比较了mLSTM基线与其正交化变体在MAD噪声AR样本上的下一令牌预测任务。训练和评估使用MAD噪声召回，设置`frac_noise`为`0.8`，涵盖不同的词汇大小和序列长度。所有模型均使用AdamW进行训练（`betas = 0.9, 0.999`，`weight_decay = 0.01`），训练步数2000步，批次大小64。学习率通过为每个任务设置扫描`3e-4`、`1e-3`、`3e-3`和`1e-2`来确定。 我们在每一步生成新的训练批次，并为每个实验维护一个独立的固定验证集。对于正交化，我们使用Frobenius范数进行归一化（`eps = 1e-6`），并应用五次Newton-Schulz迭代。我们允许梯度流经该过程。关键是，我们不将正交化后的记忆写回，因为发现这会降低性能。我们只将其用于读出。实验的完全可复现代码可在此处 (https://github.com/at2005/mlstm-orthogonalize) 找到。 使用五次Newton-Schulz迭代的正交化mLSTM算法 图1. 正交化mLSTM算法。 正交化mLSTM与mLSTM基线的训练步数 vs 验证准确率 图2. 训练步数 vs 验证准确率。 $$ \small \begin{array}{lccc} \hline \text{场景} & \text{正交化} & \text{基线} & \Delta \\ \hline \text{vocab 80, len 512} & 87.5 \pm 12.4\ (20/24) & 69.1 \pm 17.8\ (17/24) & +18.4 \pm 18.1 \\ \text{vocab 80, len 768} & 91.7 \pm 11.4\ (22/24) & 75.9 \pm 12.0\ (13/24) & +15.7 \pm 16.8 \\ \text{vocab 80, len 1024} & 98.5 \pm 2.4\ (23/24) & 83.3 \pm 13.6\ (19/24) & +15.2 \pm 14.0 \\ \text{vocab 96, len 768} & 62.4 \pm 18.4\ (14/24) & 22.0 \pm 14.4\ (4/24) & +40.4 \pm 17.6 \\ \text{vocab 96, len 1024} & 68.5 \pm 18.3\ (16/24) & 23.1 \pm 15.3\ (4/24) & +45.4 \pm 18.6 \\ \hline \end{array} $$ 图3. MAD噪声召回结果。条目为2000步时的最终验证准确率，平均值 ± 95%置信区间（24个种子）；括号内显示准确率超过80%的种子数。Delta值按种子配对计算。参数数量：vocab 80为77,716，vocab 96为80,740。正交化运行使用学习率3e-3；基线仅在vocab 80、序列长度768时使用学习率1e-2，其余情况使用3e-3。 我们发现正交化在所有任务上都提升了成功率和平均准确率。有趣的是，在进入vocab-96场景时，差距似乎进一步拉大，表明正交化在最困难的NAR任务（原始mLSTM难以应对）中帮助最大。在最后两种情况下（vocab 96，序列长度768/1024），正交化将mLSTM从失败的边缘（4/24个解种子）提升到显著更可靠的表现（14-16个解种子）。对于我们意图进行的小干预来说，这一效果令人瞩目。Newton-Schulz在固定参数数量的情况下带来了额外收益，但以增加FLOPs和挂钟时间为代价。 我们应谨慎，不要过度解读这些结果。它们仅适用于小模型规模，且NAR是一项合成任务。值得研究的是，NAR的增益是否能转化为更大模型在真实世界基准上的增益。 感谢Dan Robinson、Alpin Yukseloglu和Glen Taggart在撰写本文时提供的反馈和建议。