Depth Registers Unlock W4A4 on SwiGLU: A Reader/Generator Decomposition

この論文「Depth Registers Unlock W4A4 on SwiGLU: A Reader/Generator Decomposition」は、SwiGLU アーキテクチャを持つ大規模言語モデル（LLM）における、4 ビット重み・4 ビット活性化（W4A4）のポストトレーニング量子化（PTQ）が直面する課題と、その解決策に関する研究です。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細にまとめます。

1. 研究の背景と問題定義

• W4A4 量子化の課題: 推論時のメモリ使用量削減と高速化のために、W4A4（重みと活性化の両方を 4 ビットに量子化）は有望なアプローチですが、SwiGLU ベースのモデル（LLaMA, Mistral など）では、単純な丸め（Round-to-Nearest）を行うと、検証ペルプレキシティ（PPL）が FP16 の 23.6 から 1727 へと劇的に悪化（100 倍以上の損失）することが知られています。
• 既存手法の限界: SmoothQuant や QuaRot などの既存の補正手法は、特定の線形層を再スケーリングしたり回転させたりすることで性能を回復させますが、**「どの入力活性化の箇所（input-activation sites）が誤差の主要因となっているか」**を体系的に分解・特定したわけではありませんでした。
• 核心となる疑問: SwiGLU ブロック内のどの線形層の入力が、量子化エラーに対して最も脆弱（ロードベアリング）であり、残差軸（residual axis）からの介入で制御可能なのか？

2. 主要な理論的枠組み：Reader/Generator 分解

著者は SwiGLU ブロック内の 5 つの学習可能な線形層（qkv, w1, w3, o_proj, w2）を、入力源に基づいて 2 つのクラスに分解しました。

1. Reader（読者）: 残差ストリーム（RMSNorm 後）を入力とする層。
   • 対象：qkv, w1, w3
   • 特性：残差ストリームのノルムに直接依存するため、残差軸の介入（ノルム制御）によって入力値の分布を厳密に制御可能。
2. Generator（生成者）: ブロック内部の計算結果を入力とする層。
   • 対象：o_proj（アテンション重み付き和）、w2（SwiGLU の積演算 $silu(w1(x)) \odot w3(x)$）
   • 特性：
     • o_proj: 残差のノルムに依存するが、アテンション重みと $W_V$ に依存するため、制御は緩やか。
     • w2: **二項積（bilinear product）**であるため、個々の因子のノルムが制御されても、稀に大きな値を持つ因子同士が掛け合わさると、積の分布に重いテール（outlier）が生じる。理論的に、残差軸のノルム制御だけではこのテールを抑制できない（積の上限は因子の上限の積でしか評価できない）。

3. 提案手法：Depth Registers with Sink Loss (DR+sink)

この分解に基づき、トレーニング時に残差軸のノルムを制御する新しい介入手法 DR+sink を提案しました。

• Depth Registers (DR): 残差チャネル軸を「意味領域（semantic）」と「レジスタ領域（register）」に分割し、レジスタ領域に追加のチャネル（例：64 次元）を割り当てます。
• Sink Loss: レジスタ領域の最大ノルムが特定の閾値（$\tau=3.0$）を超えないようにするヒンジ損失（hinge loss）を課します。
• メカニズム: 本来は意味のあるチャネルに現れるはずの巨大な活性化値（outliers）を、この「レジスタ（吸い込み口）」に誘導・吸収させ、意味領域の活性化分布をガウス分布に近い形に保ちます。
• 特徴: これはポストトレーニングの補正（Calibration）や回転（Rotation）ではなく、トレーニング時のアーキテクチャ的介入であり、ゼロショット性能や FP16 の PPL を維持したまま行われます。

4. 実験結果

3 億パラメータ、50 億トークン（FineWeb-Edu）で学習したモデルを用いた実験結果は以下の通りです。

• W4A4 性能の劇的改善:
  • 単純な W4A4（Naive RTN）: PPL 1727
  • DR+sink 適用後: PPL 119（約 14 倍の改善）。
  • SmoothQuant と組み合わせると PPL 39.9 までさらに改善。
• Reader/Generator 分解の検証:
  • Reader の改善: DR+sink により、Reader 層（qkv, w1, w3）の入力活性化の尖度（kurtosis）が劇的に低下（例：qkv は 35.9 → 2.85）。
  • Generator の不変性: Generator 層（特に w2）の尖度はほとんど変化せず（1921 → 1860）、理論予測通り残差軸の介入では制御できないことが確認されました。
• Skip-Ablation による誤差分解:
  • Reader 層のみを FP16 に残して他を量子化すると、PPL は約 31 まで低下します。これは「Reader による誤差」を除去した後の残存誤差（Residue）です。
  • この残存誤差（約 0.28 nats）は、DR+sink を使ったかどうかに関わらず一定であり、w2（Generator）由来の誤差が支配的であることを示しています。
  • o_proj の寄与は小さく、w2 の寄与が圧倒的に大きいことが確認されました。

5. ポストホック手法との比較と理論的検証

• Per-Linear QuaRot: 各線形層の入力を独立に回転させる手法は、Reader 軸において DR+sink と同等の効果（PPL 約 25.4）を示しました。
• Full QuaRot: 残差軸の回転に加え、w2 入力に対するオンライン直交回転（Hadamard 変換）を追加した手法を試しましたが、PPL は 26.0 程度で改善しませんでした。
• 結論: 直交回転（Orthogonal Rotation）は、二項積（bilinear product）の重いテールを抑制できないという理論的予測（Theorem 3.4）が実証されました。w2 の問題解決には、回転ではなく混合精度（mixed-precision）や QAT（Quantization-Aware Training）などの別のアプローチが必要である可能性が示唆されました。

6. 論文の意義と貢献

1. メカニズムの解明: SwiGLU ブロック内の誤差源を「Reader（残差依存）」と「Generator（内部積依存）」に分解し、なぜ W4A4 量子化が失敗するかの根本原因を特定しました。
2. 新しいトレーニング介入: DR+sink を提案し、トレーニング時に残差軸のノルムを制御することで、Reader 側の量子化エラーを劇的に低減できることを実証しました。
3. 評価指標の確立: 「Skip-Ablation」を用いた誤差分解（$\Delta NLL$）により、どの層がどの程度の誤差を生んでいるかを定量的に評価する枠組みを提供しました。
4. 理論と実験の一致: 直交回転が二項積のテールを制御できないという理論的限界を、Full QuaRot の実験結果によって実証的に裏付けました。

まとめ

この論文は、W4A4 量子化における性能劣化が単一の要因ではなく、**「残差軸から制御可能な Reader 部分」と「内部積によって増幅される Generator 部分（特に w2）」**に分解可能であることを示しました。DR+sink は Reader 部分の誤差を大幅に削減しますが、w2 由来の誤差は残存します。これは、将来の量子化手法において、w2 に対しては回転ではなく、より強力な非線形な制御や混合精度の導入が必要であることを示唆する重要な知見です。