什么是矩阵正交化提升循环模型内存性能研究？
2026年6月30日的一项研究由Paradigm资助。Transformer展现出卓越的关联回忆（AR）能力，其注意力机制让每个token能直接访问前面的token，这是循环神经网络（RNN）等其他架构难以企及的。然而，在某些领域，如以Dreamer风格进行的长时程强化学习（long - horizon RL），无法承受Transformer二次注意力机制带来的开销。对于这类应用，需要让循环神经网络发挥作用，同时又不想放弃关联回忆能力。
目前关联回忆表现最佳的RNN是什么？
目前已知在关联回忆方面表现最佳的RNN是mLSTM，它是LSTM的一种变体，能够维护矩阵内存。在MQAR这一基准测试中，mLSTM的回忆性能相较于基线模型有显著提升。不过，单纯的回忆能力可能不足以衡量循环模型的性能，在环境转换存在噪声的领域，噪声关联回忆（NAR）是一项很有用的替代测试。由于MQAR不测试NAR，可以参考MAD的噪声AR任务套件。
MAD的噪声AR任务套件示例是怎样的？
例如`0 9 3 10 12 13 15 14 0 9 5 8 2 9`，在这个示例中，键`0`对应值`9`，键`3`对应值`10`，依此类推。MAD生成器会为键、值和干扰项使用不同的token范围。如果键的范围是`0 - 5`，那么token`12 - 15`就是干扰项。一个擅长NAR的模型，在看到开头的`0 -> 9`后，应该能在第10个位置预测出`9`，同时忽略穿插其中的干扰项token。
如何提升循环模型的NAR性能？
可以借鉴Muon的一些思路，它是一种在语言建模方面非常成功的优化器。Muon会对其动量进行正交化处理，像一个方向均衡器，能防止少数强方向主导更新过程，同时提升较弱方向的影响力。近期研究表明Muon在尾端关联记忆学习方面优于Adam，其原理是这种均衡处理能防止较弱的记忆被挤出。受此启发，决定测试在读取mLSTM内存矩阵时进行正交化处理，并在训练中加入这一额外步骤，是否能提升NAR性能。
实验是如何进行的？
使用MAD噪声AR样本，对mLSTM基线模型和正交化变体在预测下一个token方面进行了比较。在训练和评估过程中，使用MAD噪声回忆任务，将`frac_noise`设置为`0.8`，并测试了不同的词汇表大小和序列长度。所有模型都使用AdamW进行2000步训练（`betas = 0.9, 0.999`，`weight_decay = 0.01`），批量大小为64。学习率是通过对每个任务设置分别测试`3e - 4`、`1e - 3`、`3e - 3`和`1e - 2`来确定的。在每一步训练时都会生成一个新的批次，并为每个实验维护一个独立的固定验证集。对于正交化处理，使用Frobenius范数进行归一化（`eps = 1e - 6`），并应用五次牛顿 - 舒尔茨迭代，允许梯度在这个过程中流动，且不会将正交化后的内存写回，只在读取时使用它。实验的完全可复现代码可找到。
实验结果如何？
实验有相关的算法图和训练步数与验证准确率对比图。MAD噪声回忆结果显示，数据为2000步训练后的最终验证准确率，均值 ± 95%置信区间，基于24个随机种子；括号内显示准确率超过80%的种子数量。Delta是按种子配对计算的。词汇表大小为80时参数为77716，词汇表大小为96时参数为80740。正交化模型使用学习率3e - 3；基线模型仅在词汇表80、序列长度768时使用学习率1e - 2，其他情况使用3e - 3。发现正交化处理全面提升了成功率和平均准确率。当进入词汇表大小为96的模式时，性能差距似乎会扩大，表明正交化在原始mLSTM难以应对的困难NAR任务中帮助最大。在后面两种情况（词汇表96，序列长度768/1024）下，正交化让mLSTM从几乎失败的边缘（24个种子中只有4个成功）提升到了更可靠的性能水平（14 - 16个成功种子）。牛顿 - 舒尔茨迭代在固定参数数量的情况下带来了额外的性能提升，但也会增加浮点运算次数和训练时间。
实验结果应如何解读？
在解读这些结果时应保持谨慎。这些结果是在小模型范围内得出的，而且NAR是一个合成任务。值得进一步研究的是，NAR性能的提升是否能转化为大型模型在实际基准测试中的性能提升。最后感谢Dan Robinson、Alpin Yukseloglu和Glen Taggart在撰写本文时提供的反馈和建议。
)
