0 / 5 節読了

低ランク分解モデルと最適化の課題

機械学習モデル、特に大規模な行列やテンソルを扱う場合、低ランク分解はモデルのパラメータ数を削減し、計算効率を高め、過学習を抑制するための強力な手法です。例えば、行列 $W$ を $W = UV^ op$ のように分解するモデルが一般的です。このようなモデルの最適化において、勾配降下法（GD）は暗黙的に低ランク解にバイアスがかかることが、これまでの研究で示唆されてきました。しかし、Adamのような適応的学習率最適化手法が、GDと同じように振る舞うのか、あるいは異なる特性を示すのかは、これまで十分に解明されていませんでした。私の調査では、この点に明確な違いがあることが判明しました。

損失関数のゲージ対称性と最適化手法の分類

この違いの核心は、損失関数の「ゲージ対称性」と最適化手法の「ゲージ同変性」にあります。ゲージ対称性とは、モデルのパラメータ $(U, V)$ を $(UQ, VQ)$ のように変換しても、損失関数の値が変化しない性質を指します。ここで $Q$ は任意の直交行列です。勾配流が低ランク解に収束するメカニズムは、このゲージ対称性によって可能になります。しかし、この低ランクバイアスは、最適化手法がゲージ対称性を保つ、すなわち「ゲージ同変性」を持つ場合にのみ利用可能です。私の分析に基づくと、勾配降下法（GD）、モメンタム、Muon、Shampoo、そして「shared-scalar」Adamはゲージ同変性を満たします。これらの手法は、パラメータの各要素に一様なスケーリングを適用するため、ゲージ変換に対する相対的な関係が保たれます。対照的に、Adam、RMSProp、およびその他の座標ごとの適応的学習率手法はゲージ同変性を満たしません。これらの手法は、パラメータの各要素に対して異なるスケーリングを適用するため、ゲージ変換によって相対的な関係が歪められてしまうのです。

ゲージ同変性の理論的根拠と構造定理

ゲージ同変性を満たすメモリレスな更新ルールは、「Gram-determined left preconditioners」として特徴付けられます。これは、更新が $U$ のグラム行列 $U^ op U$ にのみ依存する前処理によって行われることを意味します。この構造定理は、ゲージ同変な最適化手法がどのように低ランクバイアスを維持できるかを理論的に裏付けます。さらに、「transfer theorem」により、ゲージ同変な「common-scalar flows」が勾配流のパスワイズな特性を引き継ぐことが示されています。これは、これらの手法が勾配降下法と同様に低ランク解へと収束する傾向を持つことを意味します。この理論的枠組みは、最適化手法の選択がモデルの学習パスと最終的な解の性質にどのように影響するかを理解するための重要な基盤となります。

実証実験が示す最適化手法の性能差

私は、underdetermined matrix sensingのタスクにおいて、9つの異なる更新ルールを、植え付けられた真の低ランク解に対する回復誤差で比較しました。結果として、ゲージ同変性を持つ手法が、より低い回復誤差を達成することが確認されました。特に、座標ごとの前処理から共有スカラーの前処理へと移行する1パラメータ族は、低ランクバイアスを単調に回復させ、異方性がその原因であることを明確に示しました。また、Muonに関する相反する報告については、「スペクトルスケジュール」を用いることで解決しました。Muonは、完全に低ランクなターゲットに対しては優れた回復能力を示しますが、スペクトルテールが成長するにつれてその優位性を失います。さらに、Transformerモデルでは、Adamが最初のステップで2つのゲージ同変な初期化を分離し、最終的にヘッドごとの不変量 $WQ^ op WK$ がフロベニウス距離で56%も乖離するという結果が出ました。これは、ヘッドごとの回転では埋められないギャップです。ハイパースペクトルデータセットでの実験では、同じ訓練損失に到達しても、勾配降下法がAdamに比べて、最低サンプリング密度で保持誤差を43〜44%削減し、より低い実効ランクでモデルを学習させることができました。

モデル設計における基底選択の重要性

これらの結果は、モデルの基底選択が単なるチューニングの細部ではなく、最適化が最終的にどのような補間関数を選択するかという、極めて重要な決定であることを明確に示しています。Adamのような適応的学習率最適化手法は、その利便性から広く利用されていますが、低ランク構造を持つモデルやデータに対しては、勾配降下法のようなゲージ同変な手法の方が、より望ましい低ランク解へと導き、優れた汎化性能を発揮する可能性があります。したがって、AIモデルを設計する際には、使用する最適化手法がモデルの構造やデータ特性とどのように相互作用するかを深く理解し、慎重に選択することが不可欠です。この知見は、今後のモデル開発と最適化戦略に大きな影響を与えるでしょう。

PR

文賢 →
柴亮太
柴亮太の視点

Adamは便利ですが、万能ではないです。特に低ランクモデルでは、勾配降下法の方が低ランク解を見つけやすい。最適化手法は、モデルの最終性能を左右する重要な選択です。自分のプロダクトでも、まずは一次情報として両方試して、最適な解を見つけます。