Sphere Retraction Normalizations

技術要約：Sphere Retraction Normalizations（球面退縮正規化）

問題提起
深層ニューラルネットワークの安定的な学習は、依然として大きな課題である。既存の手法は、主に残差接続と正規化（例：TransformerにおけるPre-LNおよびPost-LN）に依存しているが、これらは固有の限界を抱えている。Pre-LNアーキテクチャは、層が深くなるにつれて学習への寄与が減少する「深さの呪い」に苦しみ、また、膨大な活性化が蓄積してアテンション・シンク（attention sinks）を引き起こす、境界のない残差ストリームを許容してしまう。対照的に、Post-LNアーキテクチャは初期化時に出力層付近で大きな勾配に直面するため、最適化を遅らせる慎重な学習率ウォームアップ・スケジュールを必要とする。近年の研究であるGeodesic Normalization (GeoNorm) は、リーマン幾何学を用いて隠れ状態を超球面上に閉じ込めることで学習の安定化を図っているが、これは更新において指数写像（exponential map）のみに依存している。著者らは、指数写像はより広範な退縮写像（retraction maps）のファミリーにおける一構成要素に過ぎないと主張し、超球面上の代替となる退縮写像が、より優れた性能や安定性を提供できるのではないかという問いを投げかけている。

手法
本論文は、超球面上の**退縮写像（retraction maps）**に基づく、深層接続のための統一されたフレームワーク、Sphere Retraction Normalizations (SpheretNorms) を提案する。

1. 統一フレームワーク: 著者らは、標準的な残差接続 (RC) と GeoNorm を、多様体 $\mathcal{M}$ 上で定義される単一の計算スケルトンへと再構成した。更新規則は次のように抽象化される：
   $$h_l = R_{h_{l-1}}\left[ \alpha_l \cdot \text{proj}_{T_{h_{l-1}}\mathcal{M}}(F_l(h_{l-1})) \right]$$
   ここで $R$ は退縮写像である。GeoNorm が超球面 ($S^{d-1}$) と指数写像を選択するのに対し、このフレームワークでは他の退痴も可能となる。

2. 具体的なインスタンス化:

   • Proj-SpheretNorm: このフレームワークを計量射影退縮（metric projection retraction）（別名：gnomonic retraction）を用いて実装したもの。これは、現在の状態と更新量のベクトル和を正規化して超球面上に位置させるプロセスを含む。
   • Cay-SpheretNorm: 現在の状態と更新ベクトルから生成される反対称行列のケイリー変換（Cayley transform）を介して導出される、ケイリー退縮（Cayley retraction）（ステレオグラフィック退縮）を使用する。これにより、閉形式の代数的な更新が得られる。

3. 一般化 (p-SpheretNorm): 著者らは、Proj-SpheretNorm と Cay-SpheretNorm の両方が、更新の大きさを回転角に変換する平面回転構造を共有していることを観察した。これに基づき、p-SpheretNorm と呼ばれる、一つのパラメータを持つファミリーを導入する。ここで、回転角は $\theta_l = p \arctan(\beta_l / p)$ と定義され、$\beta_l$ は正規化されたステップサイズである。

   • $p=1$ のとき、Proj-SpheretNorm に回帰する。
   • $p=2$ のとき、Cay-SpheretNorm に回帰する。
   • $p \to \infty$ のとき、GeoNorm（指数写像）に収束する。
   • $p \to 0^+$ のとき、恒等写像に収束する。
   • 決定的なことに、$p$ が有限である場合、回転角は飽和するため、指数写像とは異なり、無制限な増大を防ぐことができる。

4. 実装の詳細: 安定性を確保するため、著者らは学習可能なステップサイズ $\alpha_l$ に対して減衰因子 $\lambda_l = 1/\sqrt{l}$ を採用し、最初の層に修正されたScaleNormを適用して隠れ状態の半径を初期化している。

主な貢献

1. 理論的統一: 本論文は、標準的なユークリッド残差接続とGeoNormを、指数写像によって駆動されるリーマン残差接続の一種としての退縮ベースの更新として統一した。
2. 新規の正規化層: 指数写像に必要な三角関数の計算コストを回避し、代数演算のみに依存する、二つの新しいノルム保存型正規化層（Proj-SpheretNorm および Cay-SpheretNorm）を導入した。
3. パラメトリック・ファミリー: これらの手法を、恒等写像、代数的退縮、および指数写像の間の隔たりを埋める単一パラメータ・ファミリーである p-SpheretNorm へと一般化した。著者らは、このファミリーのすべてのメンバーが、同一の一階接線力学（first-order tangent dynamics）を持つ二次の退縮であることを証明している。
4. 実証的な優位性: nanoGPT（最大36層、6.55Bトークン）を用いた広範な実験により、1-SpheretNorm（$p=1$ に対応）が、OpenWebText および FineWeb-Edu データセットにおいて、Pre-LN、Pre-DyT、Peri-LN、Keel、および GeoNorm を凌駕し、最も低い事前学習および検証損失、および最高のゼロショット精度を達成した。特に、1-SomeretNorm は Pre-LN や Keel が発散する設定でも収束した。

結果

• 事前学習: 中規模（24層）および大規模（36層）モデルにおいて、1-SpheretNorm は最高の検証パープレキシティを達成した。例えば、FineWeb-Edu の24層において、検証パープレキシティを Peri-LN の 19.09 から 16.81 へと減少させた。
• 安定性: Pre-LN（中規模の OpenWebText で発散）や Keel（小規模の FineWeb-Edu で発散）とは異なり、すべての SpheretNorm 変種はテストされたすべての設定において収束した。
• ダウンストリーム・タスク: FineWeb-Edu で訓練された大規模モデルにおいて、SpheretNorm 変種は12指標中11指標で最高スコアを記録した。OpenWebText においても、12指標中10指標で最高スコアを達成した。この利得はパープレキシティ指標（例：LAMBADA および WikiText-2）で最も顕著であり、1-SpheretNorm は GeoNorm や他のベースラインを大幅に上回った。
• パラメータ感度: 結果は、最適な性能が有限の $p$（具体的には $p=1$）において得られることを示しており、指数写像（$p \to \infty$ の極限）は、球面残差ストリームにとって好ましい退縮ではなく、単なるスペクトラムの一端であることを示唆している。

意義と主張
本論文は、退縮写像の選択が、超球面上の深層接続における重要な設計上の選択であることを主張している。代数的退縮（特に計量射影）を用いることで、ノルム保存と学習の安定性を達成しながら、より優れた性能を実現できることを著者らは示している。この研究は、「指数写像」が多様体上の安定した深層学習のための根本的な必然性ではなく、より広いスペクトラムにおける特定の点に過ぎないことを示唆している。著者らは、深層接続メカニズムの設計は、グラフ理論、最適化、および力学系の交差点にある数学的問題であると結論付けており、情報の伝播とヤコビ行列の条件付けの間の最適なトレードオフを特定することが、今後の研究における未解決の方向性であると認めている。また、特殊なアーキテクチャ（例：等変ニューラルネットワーク）に関する制限や、現在の範囲を超えた大規模モデルでの評価の必要性についても言及している。