技術的概要：プレノーマルトランスフォーマーにおけるゲート付き正規化除去とスケールアンカリング

問題提起

正規化層（LayerNorm、RMSNorm など）はトランスフォーマーアーキテクチャの標準的な構成要素であり、活性化を安定させ、深層ネットワークの学習を可能にします。しかし、これらの層は推論時に大きなオーバーヘッドを招く、サンプル依存の非線形計算を導入します。特にプレノーマルトランスフォーマーにおいて、これらのトークンごとの統計量が推論フェーズ全体で厳密に必要かどうかは依然として不明確です。以前の研究では、初期化の工夫や逐次的な微調整を通じて正規化の除去を試みてきましたが、内部正規化層と最終正規化層の機能的役割の違い、特にそれらがプレログイト表現のスケールをどのように制御するかについての理解は不足しています。

手法

著者らは、プレノーマルトランスフォーマー向けに設計された、ゲート付き正規化除去アプローチであるTaperNormを提案します。この手法は、モデルが標準的な正規化から、学習されたサンプル非依存の線形またはアフィン写像へ学習中に遷移することを可能にし、推論時に正規化操作の除去を実現します。

中核メカニズム：TaperNorm

TaperNorm は、グローバルゲート $g \in [0, 1]$ を用いて、正規化ブランチと固定スケーリングブランチの間を動的に補間します。トークンベクトル $h$ に対して：
$$\text{TaperNorm}(h; g) = g \cdot \frac{h}{r(h)} D\gamma + (1 - g) \cdot c \cdot h D\tilde{\gamma}$$
ここで：

• $r(h)$ は RMSNorm スケーリング因子です。
• $D\gamma$ と $D\tilde{\gamma}$ は学習可能な対角スケーリング行列です。
• $c$ は層ごとのスカラーです。
• $g$ は 1（標準的な正規化）で始まり、学習中にコサイン減衰して 0 になるグローバルゲートです。

学習手順：

1. ウォームアップ： ゲート $g$ を 1 に保持します。モデルは標準的なプレノーマルトランスフォーマーとして学習します。
2. 較正： ウォームアップの終了時に、正規化ブランチとスケーリングブランチ間の最小二乗差を最小化するようにスカラー $c$ を計算し、分布シフトなしに滑らかな遷移を確保します。
3. テーパリング： ゲート $g$ を 0 まで減衰させます。スケーリングブランチのパラメータ（$\tilde{\gamma}$）は学習可能なままです。
4. 推論折り込み： $g=0$ になると、TaperNorm 層は単純な線形/アフィン写像（$c \cdot h D\tilde{\gamma}$）となります。この操作は数学的に隣接する線形投影（例：QKV 投影または MLP 重み）に折り込むことができ、推論グラフから明示的な正規化操作を除去します。

スケールアンカリング分析

本論文は、内部正規化層と最終正規化層の間の決定的な違いを特定しています：

• アンカーとしての最終正規化： 最終正規化層は暗黙的なスケールアンカーとして機能します。数学的には、これは 0 次同次であり、出力における放射状勾配成分を除去することを意味します。このアンカーがなければ、交差エントロピーの最小化はモデルに損失を減らすためにログイトの大きさを無限に増大させる（「ログイト・チェイシング」）よう駆動し、不安定さを招きます。
• 明示的な代替案： 最終正規化が除去された場合（「All-Taper」設定）、著者らは固定ターゲットスケール損失（$L_{aux}$）を導入します。この補助損失は、トークンスケールが固定ターゲット $s_{tgt}$ から逸脱することを罰し、ログイト・チェイシングの傾向に対抗する明示的な放射状復元力を提供します。

2 つの実験設定がテストされました：

1. Internal-Taper： 内部ブロック正規化器は線形写像へテーパリングされますが、最終正規化はアクティブなままです（暗黙的アンカー）。
2. All-Taper： 最終のものを含むすべての正規化器がテーパリングされます。これには安定性を維持するための固定ターゲットスケール損失が必要です（明示的アンカー）。

主要な貢献

1. ゲート付き正規化除去： RMSNorm/LayerNorm モジュールを固定された線形/アフィン再スケーリングへテーパリングする手法。これにより、重み折り込みを通じて推論時にトークンごとの統計量を排除できます。
2. 出力スケールアンカリング理論： 最終正規化層の主な固有の役割は放射状勾配の除去にあることを示す分析。このアンカーがなければ交差エントロピーが無限大のログイト成長を誘発し、固定ターゲットスケール損失がこの安定化役割を効果的に代替できることを示しています。
3. 実証的検証： TinyStories 事前学習と GPT-2 微調整の実験により、内部正規化は検証損失のわずかな増加で除去可能であることが示されました。このアプローチは、内部スケーリングを折り込んだ後、KV キャッシュ付き自己回帰デコードにおいて最大1.18 倍のスループット向上を達成しました。

結果

• 事前学習（TinyStories）：
  • Internal-Taper (+aux)： モデルサイズ（100 万から 3000 万パラメータ）全体で、ベースラインの検証損失の 0.7〜1.5% 以内を達成しました。
  • All-Taper (+aux)： ベースラインの損失の 0.8〜1.8% 以内を達成し、明示的なスケールアンカリングにより完全にノーマルフリーなモデルが実現可能であることを示しました。
  • ダイナミクス： スケールアンカーを持たないモデルは、ログイトの大きさの急激な成長（「ログイト・チェイシング」）を示しました。固定ターゲットスケール損失はこの挙動を正常に抑制し、勾配の大きさをアンカー付きベースラインと整合させました。
• 推論スループット：
  • NVIDIA H100 において、Internal-Taper (fused) 設定は、標準的な RMSNorm ベースラインと比較して1.15〜1.18 倍のトークン/秒スループットを達成しました。Internal-Taper (unfused) 設定（明示的スケーリング操作）は1.10〜1.14 倍の高速化を達成しました。
• 微調整（GPT-2）：
  • 提案されたInternal-TaperLNおよびAll-TaperLN手法は、OpenWebText および The Pile データセットにおいて、「FakeLN」アプローチ（Baroni et al., 2025）および標準的な Vanilla モデルと競合する交差エントロピー損失を達成しました。
  • 逐次的除去手法とは異なり、TaperNorm は単一のグローバルゲートを使用して、選択されたすべての層を同時に除去します。

意義と主張

本論文は、プレノーマルトランスフォーマーにおいて内部正規化層と最終正規化層が異なる役割を果たすと主張しています。内部正規化器は、性能低下を無視できる程度に学習されたサンプル非依存の線形写像に置き換え可能ですが、最終正規化層は出力スケールアンカーとして独自の機能を果たします。

この研究の意義は以下の点にあります：

1. 効率性： 推論グラフから計算コストの高い正規化操作を除去する実用的な道筋を提供し、ゼロから再学習することなく測定可能なスループット向上（最大 1.18 倍）をもたらします。
2. 理論的洞察： 最終正規化層が提供する安定性は、主に放射状勾配の除去によるものであることを明確にします。この洞察により、暗黙的アンカーを明示的固定ターゲットスケール損失に置き換えることで、「完全にノーマルフリー」なモデルを設計することが可能になります。
3. ロバスト性： このアプローチは学習初期には標準的なプレノーマルダイナミクスを維持し、後になってから線形写像へ遷移するため、最初から正規化なしで学習を試みる手法よりも安定しているように見えます。

著者らは範囲に関して慎重であり、実験は最大 3000 万パラメータのモデルと特定のハードウェアスタックに限定されていると指摘しています。また、数十億パラメータ領域における活性化外れ値に対するロバスト性や、交差エントロピー以外の目的関数への適用可能性は、今後の研究における未解決の課題であると認めています。