2026年6月30日 本工作的资金来源于Paradigm。变换器展现出令人惊讶的联想记忆（AR）能力：注意力机制使每个标记能够直接访问其之前的标记，这是其他架构（如递归神经网络（RNN））难以做到的一种机制。但对于某些领域，我们无法承受变换器的二次注意力开销。一个例子是以Dreamer风格的长范围强化学习（RL）。对于这类应用，我们需要使递归神经网络有效，但又不想放弃联想记忆。当前已知的在联想记忆中表现最佳的RNN是mLSTM，这是一种保持矩阵记忆的LSTM变体。mLSTM在一个基准MQAR上显示出显著改善的记忆回忆能力。但单纯的回忆可能不足以衡量递归性能。在环境转变可能会有噪声的领域中，一个有用的代理测试是噪声联想记忆（NAR）。因为MQAR不测量NAR，我们可以查看MAD的噪声AR任务套件。以下是一个任务的示例：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。这个思想是，这种平衡防止弱记忆被挤出。受到此启发，我们决定测试在读取过程中对mLSTM记忆矩阵进行正交化，并通过这个附加过程进行训练是否提高NAR性能。我们将mLSTM基线与其正交化变体进行比较，使用MAD的噪声AR样本进行下一标记预测。在训练和评估中，我们使用MAD噪声回忆，frac_noise设置为0.8，涵盖一系列词汇大小和序列长度。所有模型使用AdamW（betas = 0.9，0.999，weight_decay = 0.01）训练2k步，批量大小为64。学习率通过在3e-4、1e-3、3e-3和1e-2之间选择每个任务设置进行选择。我们在每一步生成一个新的训练批次，并在每个实验中保持一个固定的验证集。对于正交化，我们通过Frobenius范数进行归一化（eps = 1e-6），并进行五次Newton-Schulz迭代。我们允许梯度通过该过程流动。至关重要的是，我们不将正交化的记忆写回，因为我们发现这会降低性能。我们仅将其用于读取。我们实验的完全可复现代码可以在这里找到。图1. 正交化的mLSTM算法。图2. 训练步骤与验证准确率。$$ iny egin{array}{lccc} ext{制度} & ext{正交化} & ext{基线} & ext{差异} \ ext{词汇80，长度512} & 87.5 ext{±} 12.4 ext{（20/24）} & 69.1 ext{±} 17.8 ext{（17/24）} & +18.4 ext{±} 18.1 \ ext{词汇80，长度768} & 91.7 ext{±} 11.4 ext{（22/24）} & 75.9 ext{±} 12.0 ext{（13/24）} & +15.7 ext{±} 16.8 \ ext{词汇80，长度1024} & 98.5 ext{±} 2.4 ext{（23/24）} & 83.3 ext{±} 13.6 ext{（19/24）} & +15.2 ext{±} 14.0 \ ext{词汇96，长度768} & 62.4 ext{±} 18.4 ext{（14/24）} & 22.0 ext{±} 14.4 ext{（4/24）} & +40.4 ext{±} 17.6 \ ext{词汇96，长度1024} & 68.5 ext{±} 18.3 ext{（16/24）} & 23.1 ext{±} 15.3 ext{（4/24）} & +45.4 ext{±} 18.6 \ ext{差异} \ ext{pairs by seed}\ ext{Params are 77,716 for vocab 80 and 80,740 for vocab 96.}\ ext{正交化运行使用LR 3e-3;基线仅在词汇80，序列长度768上使用LR 1e-2，其他情况使用LR 3e-3。} \ ext{我们发现正交化在各方面提高了成功率和平均准确性。} \ ext{有趣的是，当我们进入词汇-96制度时，间隔似乎在扩大，这表明正交化最能帮助在原始mLSTM经历困难的NAR任务。} \ ext{在后两种情况下（词汇96，序列长度768/1024），正交化将mLSTM从失败的边缘（4/24个解决种子）带到了更加可靠的表现（14-16个解决种子）。} \ ext{这对我们希望进行的小干预而言是非常显著的。} \ ext{Newton-Schulz以固定的参数数量为我们带来了额外的收益，权衡了额外的FLOPs和墙钟时间。} \ ext{我们应谨慎解读这些结果。} \ ext{它们在小模型制度中成立，且NAR是一个合成任务。} \ ext{值得探讨的是，NAR的收益是否能转化为大模型在真实世界基准中的收益。} \ ext{感谢Dan Robinson, Alpin Yukselog}