MITのミューオン最適化器、リプシッツ境界でTransformerを安定化

大規模なTransformerモデルを安定して訓練することは、特にそれらのモデルがサイズと複雑さを増し続けるにつれて、ディープラーニングにおける長年の大きな課題でした。MITの研究者たちは、根本的な問題、すなわち活性化値の制御不能な増大と、それに起因する訓練中の損失の急増に、しばしば制約のない重みと活性化のノルムが原因で、対処してきました。

彼らの革新的な解決策は、Transformerに「証明可能なリプシッツ境界」を強制適用することです。これは、活性化正規化、QK正規化、ロジットのソフトキャッピングなどの一般的な安定化手法に頼ることなく、モデルの重みのスペクトル特性を直接調整することによって達成されます。

リプシッツ境界とその重要性の理解

ニューラルネットワークにおけるリプシッツ境界は、ネットワークの出力が入力または内部重みの摂動に応じて変化できる最大レートを定量化します。簡単に言えば、リプシッツ境界が低いほど、ネットワークは小さな変化やノイズに対して感度が低く、よりロバストで予測可能になります。この特性は、訓練中の安定性を確保し、敵対的ロバストネス（悪意のある入力操作に対する耐性）を高め、プライバシーを向上させ、新しいデータへのより良い汎化を促進するために極めて重要です。

従来の安定化方法の問題点

歴史的に、大規模Transformerにおける安定性の達成には、レイヤー正規化、QK正規化、ロジットtanhソフトキャッピングなど、さまざまな「その場しのぎ」の解決策が関わってきました。これらの方法はある程度の安定性を提供しますが、不安定性の根本原因、すなわち重み行列内の「スペクトルノルム」（最大特異値）の制御不能な増大には直接対処していません。この無制限な増大は、特に非常に大規模なモデルにおいて、活性化の爆発と訓練の不安定性の主要な要因です。

MITチームの中心的な仮説は、重み自体のスペクトル特性を直接調整することによって、つまり最適化器や活性化を超えて、ネットワークのリプシッツ性を厳密に制御し、それによって不安定性をその根源から解決できるというものです。

主要なイノベーション：ミューオン最適化器と重みスペクトル正則化

研究者たちのアプローチは、「ミューオン（Muon）」最適化器に基づいています。この最適化器はすでに勾配をスペクトル的に正則化しており、各勾配ステップがスペクトルノルムを設定された制限を超えて増加させないようにします。MITの主要なイノベーションは、この正則化をモデルの重みに拡張したことです。各訓練ステップの後、彼らはすべての重み行列の特異値を制限する操作を適用します。特異値は、行列が入力値をどれだけ引き伸ばしたり縮めたりするかを記述する数学的要素であり、それらを制限することで重みの増幅係数を直接制御します。

この重み正則化の注目すべき成果は、活性化ノルム（ネットワーク層内の値の大きさ）が非常に小さく保たれることです。彼らのGPT-2スケールのTransformerでは、最大活性化エントリーは約100を超えることはありませんでした。これは、最大活性化が148,000を超える可能性のある制約なしのベースラインとは対照的です。重要なことに、この安定性は、従来のレイヤー正規化、QKノルム、ロジットtanhのトリックを一切使用せずに達成されました。活性化の大きさが小さいことで、これらのモデルはFP8のような低精度データ形式とも互換性があり、効率的なハードウェア展開に非常に有益です。

リプシッツ制約を強制適用する方法

研究者たちは、重みノルム制約を強制適用するためのさまざまな方法を検討・比較し、高い性能を維持し、リプシッツ境界を保証し、性能とリプシッツ性のトレードオフを最適化する能力を評価しました。

• 重み減衰（Weight Decay）： 標準的な正則化方法ですが、スペクトルノルムを常に正確に制御できるわけではありません。

• スペクトル正規化（Spectral Normalization）： 重み行列の最大特異値を制限しますが、すべての特異値に全体的に影響を与える可能性があります。

• スペクトルソフトキャップ（Spectral Soft Cap）： すべての特異値を並行してスムーズかつ効率的に制限する新しい技術です。この方法は、ミューオンの安定ランク更新と効果的に連携するように特別に共同設計されており、より厳密な境界を可能にします。

• スペクトルハンマー（Spectral Hammer）： 最大特異値のみを最大値に設定する方法で、AdamW最適化器との使用に最適です。

実験結果と洞察

この研究は、さまざまなモデルスケールで重要な発見を示しました。

• モデル評価： 小規模なTransformer（例えばShakespeare、証明可能なリプシッツ境界が2未満）の場合、この方法は60%の検証精度を達成し、検証損失において制約なしのベースラインを上回りました。NanoGPT（1億4500万パラメータ）のような大規模モデルの場合、10未満の厳密なリプシッツ境界で21.2%の検証精度が得られました。強力な制約なしのベースライン（39.4%の精度）の性能に合わせるには、はるかに大きな上限（例えば10^264）が必要でした。これは、非常に厳密なリプシッツ制約と、大規模でのピーク表現力の達成との間の現在のトレードオフを浮き彫りにしています。

• 制約方法の効率： ミューオン最適化器とスペクトルソフトキャップの組み合わせは、損失-リプシッツのトレードオフにおいて常に最前線を走り、重み減衰を用いたAdamWと比較して、同等かそれ以上の検証損失でより低いリプシッツ定数を達成しました。

• 安定性とロバストネス： 制約されたリプシッツ定数で訓練されたモデルは、敵対的ロバストネスが著しく向上し、敵対的攻撃下での精度低下が、制約なしのベースラインと比較してはるかに軽度でした。

• 活性化の大きさ： 既に述べたように、スペクトル重み正則化は、規模が拡大しても最大活性化を常に小さく保ちました。これにより、ハードウェアにおける「低精度訓練と推論」の新たな道が開かれ、活性化が小さいことで計算、メモリ、電力コストを劇的に削減できます。

限界と今後の方向性

これらの進歩にもかかわらず、この研究はいくつかの未解決の疑問と限界を指摘しています。

• 重みノルム、ロジットスケーリング、アテンションスケーリングの最適なトレードオフの選択は、原理的な方法よりも依然として経験的な探索に大きく依存しています。

• モデルに対して計算された現在のグローバルなリプシッツ境界は、実際の活性化ノルムが非常に小さい場合でも、天文学的に大きい（例えば10^264）可能性があります。これは、理論的な境界が観測された挙動よりもはるかに緩いことを示しています。

• モデルの規模が拡大し続ける中で、厳密に小さいリプシッツ境界で制約なしのベースラインの性能に合わせることが実現可能かどうかは依然として不明です。この分野ではさらなる研究が必要です。

結論

MITの研究者たちによるこの研究は、スペクトル重み正則化、特にミューオン最適化器と統合された場合に、強制されたリプシッツ境界を持つ大規模Transformerを安定して訓練するための強力な方法を提供することを示しています。このアプローチは、従来の活性化正規化やその他の場当たり的な安定化トリックの必要性を排除し、より深く、より根本的なレベルで不安定性に対処します。活性化をコンパクトで予測可能な範囲に保つことで、この方法は敵対的ロバストネスを大幅に向上させ、低精度AI展開を通じてハードウェア効率を向上させる大きな可能性を秘めています。この研究は、ニューラルネットワーク正則化のための新しい効率的な計算プリミティブへの道を開き、高度なAIシステムの安全性、プライバシー、および実用的な展開に広範な影響を与えます。