arXivDaily arXiv每日学术速递 周一至周五更新
arXiv 2603.26554cs.LGstat.ML

谱优化器在关联记忆学习中的容量扩展分析

Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory

UC Berkeley(加州大学伯克利分校);Princeton University(普林斯顿大学);New York University(纽约大学);Flatiron Institute(Flatiron研究所)

Juno Kim, Eshaan Nichani, Denny Wu, Alberto Bietti, Jason D. Lee

AI总结:

本文研究了谱优化器在关联记忆问题中的性能,发现Muon的存储容量显著超过SGD,且在仅使用一阶信息时可与牛顿法匹配。通过实验验证了预测的扩展规律,为更实际的语言建模任务提供了理论基础。

摘要:

谱优化器如Muon在大规模语言模型训练中展现了强大的经验性能,但其优势来源及程度仍不明确。本文通过线性关联记忆问题研究这一问题,这是一个适用于transformer模型事实回忆的可解模型。特别是,我们超越正交嵌入,考虑高斯输入和输出,这使得存储的关联数量可以远超嵌入维度。我们的主要结果精确刻画了Muon、SGD和牛顿法在逻辑回归损失下,在幂律频率分布下的恢复率。我们证明Muon的存储容量显著超过SGD,且仅使用一阶信息时可与牛顿法匹配。此外,Muon在更大的临界批大小下饱和。我们进一步分析了阈值梯度近似的多步动态,并显示Muon在初始恢复率上显著快于SGD,而两种方法最终都收敛到信息论极限,速度相当。合成任务的实验验证了预测的扩展规律。我们的分析为谱预条件器信号放大提供了定量理解,并为建立更实际的语言建模任务和优化器的扩展规律奠定了基础。

英文摘要:

Spectral optimizers such as Muon have recently shown strong empirical performance in large-scale language model training, but the source and extent of their advantage remain poorly understood. We study this question through the linear associative memory problem, a tractable model for factual recall in transformer-based models. In particular, we go beyond orthogonal embeddings and consider Gaussian inputs and outputs, which allows the number of stored associations to greatly exceed the embedding dimension. Our main result sharply characterizes the recovery rates of one step of Muon, SGD, and Newton's method on the logistic regression loss under a power law frequency distribution. We show that the storage capacity of Muon significantly exceeds that of SGD, and even matches Newton's method while only using first-order information. Moreover, Muon saturates at a larger critical batch size. We further analyze the multi-step dynamics under a thresholded gradient approximation and show that Muon achieves a substantially faster initial recovery rate than SGD, while both methods eventually converge to the information-theoretic limit at comparable speeds. Experiments on synthetic tasks validate the predicted scaling laws. Our analysis provides a quantitative understanding of the signal amplification of spectral preconditioners and lays the groundwork for establishing scaling laws across more practical language modeling tasks and optimizers.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑