谱优化器在学习关联记忆中的容量 scaling 定性分析
机器学习
2026-04-29 v2 机器学习
摘要
谱优化器如Muon在大规模语言模型训练中显示出强大的经验性能，但其优势的来源与程度仍未充分理解。我们通过线性关联记忆问题深入探讨此问题，该问题是Transformer-based模型中事实记忆的可解析模型。特别是，我们超越正交嵌入，考虑高斯输入与输出，这使得存储关联数大幅超过嵌入维度。我们的主要结果对针对幂律频率分布下的Logistic回归损失，对Muon、SGD与Newton方法单步恢复率进行了精确刻画。我们表明，Muon的存储容量显著超过SGD，甚至与Newton方法相当，仅使用一阶信息。此外，Muon在更大的临界批量大小处于饱和。我们进一步分析在阈值梯度近似下的多步动力学，表明Muon在初始恢复速率上显著优于SGD，而两者最终均收敛至信息论极限，收敛速度相当。合成任务实验验证了预测的规模定律。我们的分析为谱预条件器的信号放大提供了量化理解，为在更实用的语言建模任务中建立规模定律奠定了基础。
引用
@article{arxiv.2603.26554,
title = {Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory},
author = {Juno Kim and Eshaan Nichani and Denny Wu and Alberto Bietti and Jason D. Lee},
journal= {arXiv preprint arXiv:2603.26554},
year = {2026}
}
备注
84 pages, 9 figures