拆解Shampoo的启发式设计,提出更简洁的优化算法
Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
- 从Frobenius范数逼近角度分析,分离预条件矩阵的特征值与特征基更新
- 证明直接修正特征值可消除学习率嫁接的必要性,减少超参数调优
- 提出自适应特征基计算频率准则,提升训练稳定性与可解释性
Shampoo在AlgoPerf竞赛中的成功重新引发了对基于Kronecker分解的优化算法的兴趣。尽管表现优异,其依赖学习率嫁接和过时预条件等启发式策略,增加算法复杂度并需额外调参,缺乏理论支持。本文从Frobenius范数逼近全矩阵Adam的角度出发,解耦预条件矩阵的特征值与特征基更新。结果表明,来自Adam的嫁接可缓解特征值的过时与错标定问题;直接校正特征值则可消除学习率嫁接需求。为控制特征基计算频率带来的误差,提出基于冷启动QR算法终止的自适应准则,实现不同预条件矩阵更新频率解耦,便于研究近似误差对收敛的影响。这些技术为去除Shampoo的启发式设计、构建更优的Kronecker分解优化算法提供了理论基础。
原文摘要 · Abstract (English)
The recent success of Shampoo in the AlgoPerf contest has sparked renewed interest in Kronecker-factorization-based optimization algorithms for training neural networks. Despite its success, Shampoo relies heavily on several heuristics such as learning rate grafting and stale preconditioning to achieve performance at-scale. These heuristics increase algorithmic complexity, necessitate further hyperparameter tuning, and lack theoretical justification. This paper investigates these heuristics from the angle of Frobenius norm approximation to full-matrix Adam and decouples the preconditioner's eigenvalues and eigenbasis updates. We show that grafting from Adam mitigates the staleness and mis-scaling of the preconditioner's eigenvalues and how correcting the eigenvalues directly eliminates the need for learning rate grafting. To manage the error induced by infrequent eigenbasis computations, we propose an adaptive criterion for determining the eigenbasis computation frequency motivated by terminating a warm-started QR algorithm. This criterion decouples the update frequency of different preconditioner matrices and enables us to investigate the impact of approximation error on convergence. These practical techniques offer a principled angle towards removing Shampoo's heuristics and developing improved Kronecker-factorization-based training algorithms.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。