Unlocking Feature Learning in Gated Delta Networks at Scale

技術要約：Gated Delta Networkにおける特徴量学習のスケールアップを実現する

問題提起
大規模言語モデル（LLM）のスケーリングは、主に2つのボトルネックに直面している。一つは、より大きなモデルを訓練する際の膨大な計算コストであり、もう一つは、標準的なTransformerの自己注意機構（Self-Attention）が持つ二次的な計算複雑性である。線形回帰アーキテクチャ、例えばGated Delta Networks（GDN）は、長距離依存関係を捉えつつ、推論時の線形時間を維持できる効率的な代替案として浮上している。しかし、これらのモデルを効果的にスケールさせるための実装には決定的な空白が存在する。Maximal Update Parametrization（µP）は、標準的なTransformerや一部の線形モデルにおいてゼロショットのハイパーパラメータ転送（特に学習率）を可能にしてきたが、GDNへの適用については未開拓である。核心となる課題は、このアーキテクチャの回帰状態のダイナミクスが、既存のµPの導出で想定されているスカラー回帰や標準的なフィードフォワード構造ではなく、データ依存のゲーティングメカニズムとフル行列のデルタ則を介して更新される点にある。Mambaのような構造化状態空間モデル（SSM）に関する先行研究では、バニラのµPが対角SSMに対して失敗することが示されているが、GDNが用いるフル行列の状態更新は、それらの特定の修正を適用できない性質を持っている。

手法
著者らは、Gated Delta NetworksのためのµP定式化を厳密に導出することで、この空白に対処している。その手法は、投影された特徴量、潜在状態 $S_t$、およびリードアウトの座標サイズ推定値を、特有のゲーティングメカニズムと回帰状態遷移を含む全フォワードパスを通じて伝播させるものである。

主な分析ステップは以下の通りである：

1. フォワードパス解析： 著者らは、投影された特徴量、潜在状態 $S_t$、およびリードアウトの座標サイズを導出した。クエリとキーのL2正規化、およびランク1の書き込み更新が、特定の分散スケーリングをもたらすことを特定した。極めて重要な点として、正規化層への入力が $\Theta(1)$ の座標サイズを維持するように、RMSNorm層の前に $\sqrt{d}$ 倍率を挿入することを提案している。これにより、特徴量学習に必要な勾配スケーリングの崩壊を防ぐ。
2. ゲーティングメカニズム解析： 本論文では、非線形ゲーティングパラメータ $\alpha_t$（減衰）および $\beta_t$（書き込み強度）を分析している。これらは通常の重みとは異なり、学習可能な重みベクトル（$W_\alpha, W_\beta$）とスカラーパラメータ（$a_{log}, b$）を含むスカラー事前活性化から導出される。分析の結果、これらのゲーティング成分は、状態更新式における特定の役割 due to その役割により、標準的なµPのスケーリング則に従わないことが明らかになった。
3. バックプロパゲーションと学習率の導出：
   • AdamW: AdamWは勾配をその二次のモーメントで正規化するため、実効的な更新量は重みのクラス間で一貫しており、学習率の要件を単純化している。
   • SGD: 純粋なSGDの下では、生の勾配の大きさは重みのクラス間で大きく異なる。著者らは、ゲーティング重み行列（$W_\alpha, W_\beta$）には $\Theta(1/\sqrt{d})$ の学習率スケーリングが必要であり、スカラーゲーティングパラメータ（$a_{log}, b$）には $\Theta(\sqrt{d})$ のスケーリングが必要であることを導出した。これらのスケーリングは、学習率が通常一様であるか、あるいは単純な $1/n$ ルールに従う標準的なµPの設定とは異なる。

主要な貢献

1. 理論的導出： 本論文は、Gated Delta Networksに対する完全なµP定式化を初めて提供する。全フォワードパスを通じて座標サイズの推定値を導出し、すべての重みクラスに対して原理的な初期化分散とフォワード倍率を確立した。
2. 新規のスケーリング則： 著者らは、SGDの下でのゲーティング成分に非標準的な学習率スケーリングが必要であることを発見した：
   • ゲーティング重み行列（$W_\alpha, W_\beta$）には $\Theta(1/\sqrt{d})$ のスケーリングが必要。
   • スカラーゲーティングパラメータ（$a_{log}, b$）には $\Theta(\sqrt{d})$ のスケーリングが必要。
3. 実証的検証： 著者らは、AdamWおよびSGDオプティマイザの両方を用いて、複数の隠れ次元（256から1563まで）にわたるGDN言語モデルを事前学習した。
   • 結果： 提案されたµP構成の下では、最適な学習率は、両方のオプティマイザにおいてすべてのモデル幅に対してゼロショットで転送された。対照的に、標準的なパラメトリゼーション（SP）および元のµP構成は転送に失敗し、モデルの幅が増加するにつれて最適な学習率が大きくシフトした。
   • 安定性： 提案された構成は安定した学習率転送を可能にし、理論的導出の正当性を裏付けた。

意義と主張
本論文は、その分析がGated Delta Networksを特徴量学習のためにスケールさせるためのパラメトリゼーションという未解決の問題を解決することを主張している。導出されたµP定式化が、ゼロショットのハイパーパラメータ転送を可能にすることを実証することで、本研究は理論的導出の正当性と実用的な有用性を検証している。著者らは、これがGated Delta Networksに特化した、µP整合的なパラメトリゼーションを導出し検証した最初の研究であることを強調しており、これは対角SSMに関する先行研究とは一線を画すものである。得られた知見は、適切なパラメトリゼーションが線形回帰モデルをスケールさせるために不可欠であること、そしてGDNにおける特定のゲーティングメカニズムが、標準的なTransformerや他の線形アーキテクチャとは異なる独自のスケーリング則を必要とすることを示唆している。著者らは、これらの導出が、他の線形またはハイブリッドアーキテクチャのスケール則に関する将来の研究の指針となることを期待している。