Last-Iterate Stationarity of Muon with Finitely Many Newton-Schulz Iterations
Abstractの概要
本論文では、正確な極分解因子の計算ではなく固定回数の5次Newton–Schulz反復によって更新方向を生成する、実践的なMuonの最終反復（last-iterate）における収束挙動を解析しています。対象とする設定は滑らかで下に有界な非凸最適化であり、決定論的勾配と持続的なノイズを伴う確率的勾配を区別して解析を行っています。著者らはまず、入力の安定化を行わない有限ステップのMuonは滑らかな凸2次関数上でさえ周期軌道に陥る可能性があり、最終反復の定常性が自明には成り立たないことを示します。その上で、正の安定化を導入することで決定論的ケースでは十分に小さい一定のステップサイズによる最終反復の勾配収束が可能になる一方、確率的ケースでは学習率の減少とモメンタムの増加を組み合わせる必要があることを証明します。さらに、一般的な点ごとの多項式レートは必ずしも成立しないものの、ほぼすべての後半の反復に対して密度1のレートが確立できることも示しています。
新規性
主な新規性は、先行研究の大半を占める平均反復や最良反復の保証ではなく、実践的な有限ステップMuonに対する最終反復定常性理論を構築した点にあります。また、安定化がない場合の周期化や固定された確率的ハイパーパラメータ下での持続ノイズフロアという2つの具体的な失敗モードを特定し、点ごとの多項式レートが成り立たない場合に密度1の収束レートを確立した点でも独自性を有しています。
成果
厳密な勾配に対しては、正の入力安定化と十分に小さい一定学習率により、勾配の二乗可総性と1未満の任意の固定モメンタム係数における勾配列全体のゼロへの収束が証明されます。確率的勾配に対しては、固定の学習率およびモメンタムでは有界な不偏ノイズ下で収束に失敗し得ることを示しつつ、減少するステップサイズと増加するモメンタムを用いることで概収束および平均二乗最終反復収束が回復することを明らかにしています。さらに密度1の勾配減衰レートを証明するとともに、一般的な点ごとの多項式レートが保証されないことを示す滑らかな反例も提示しています。
論文の注目点
- 入力の安定化がない場合、実践的な有限ステップMuonは滑らかな凸2次関数上でも2周期振動を示す可能性があり、単純な決定論的設定でも最終反復収束が失敗し得ます。
- 正の安定化を導入することで、著者らは連立降下追従リアプノフ解析を導出し、小さな一定ステップサイズのもとで決定論的な滑らか非凸目的関数に対する全列勾配収束を保証します。
- 確率的設定では、持続的な不偏かつ有界分散のノイズにより固定ハイパーパラメータでの収束が妨げられますが、協調的な学習率の減少とモメンタムの増加によって概収束および平均二乗の最終反復定常性が回復し、密度1のレートが支持されます。