技術要約: CLHA – 重み共有型Mixture-of-Experts Transformerの量子化のためのクロスレイヤー・ヘッセ行列集約

1. 問題提起

Megrez2に見られるような、最近の効率的なMixture-of-Experts (MoE) Transformerアーキテクチャは、クロスレイヤー重み共有 (Cross-Layer Weight Sharing: CLWS) を採用しています。この設計では、隣接するTransformerレイヤー（例：$\ell$ と $\ell+1$）が同じルーテッド・エキスパートの重みテンソルを再利用することで、トークンあたりの活性化パラメータ数を維持しつつ、パラメータのフットプリントを半分に削減しています。

しかし、GPTQスタイルのアプローチを含む既存のポストトレーニング量子化 (Post-Training Quantization: PTQ) 手法は、これらの共有された重みを単一のレイヤーに属するものとして扱います。これは、単一の「プライマリ」レイヤーの活性化統計量を用いてキャリブレーションを行うものです。これにより、根本的なミスマッチが生じます：

• 残差ストリームの表現はレイヤー間で進化するため、レイヤー $\ell_a$ で共有エキスパートを活性化させるトークンの分布は、レイロ $\ell_b$ での分布とは異なります。
• 既存の手法は二次レイヤーの活性化統計量を無視しているため、特に低ビット領域（例：2ビット）において、再構成誤差の最小化が最適に行われません。

2. 手法

2.1 理論的基礎: クロスレイヤー・ヘッセ行列集約

本論文は、複数のレイヤー $\ell_1, \dots, \ell_n$ で共有される重み $W$ を量子化するための目的関数を定式化しています。正しい目的関数は、その重みを使用するすべてのレイヤーにわたる総再構成誤差を最小化することです：
$$L(W_q) = \sum_{i} \alpha_i \text{tr}[(W_q - W)H_{\ell_i}(W_q - W)^\top]$$
ここで、$H_{\ell_i}$ はレイヤー $\ell_i$ のヘッセ行列（活性化共分散）であり、$\alpha_i$ はルーティング頻度（各レイヤーへのトークンの割合）を表します。

命題1は、この結合された誤差を最小化することが、結合ヘッセ行列を用いたGPTQを適用することと数学的に等価であることを確立しています：
$$H_{comb} = \sum_{i} \alpha_i H_{\ell_i}$$
これは、単一のレイヤーのヘッセ行列（$H_{\ell_a}$）のみを使用して誤差を最小化するパーレイヤー独立キャリブレーション (Per-Layer Independent Calibration: PLIC) とは対照的です。

2.2 CLHAアルゴリズム

著者らは、3パスのキャリブレーション・アルゴリズムであるCLHA (Cross-Layer Hessian Aggregation) を提案しています。

1. 収集 (Collection): 各キャリブレーション・シーケンスに対して、すべての活性化エキスパートについて、レイヤーごとのヘッセ行列とトークン数を累積します。
2. 集約 (Aggregation): 各共有エキスパートのペア $(\ell_a, \ell_b)$ に対して、結合ヘッセ行列を計算します。論文では2つの重み付け戦略を評価しています：
   • ルーティング頻度による重み付け ($\alpha = n_a / (n_a + n_b)$): 各レイヤーに受信された実際のトークン数によってヘッセ行列を重み付けします。
   • 等価重み付け ($\alpha = 0.5$): 正則化として機能し、キャリブレーションデータが限られており、二次レイヤーのヘッセ行列の推定値にノイズが多い場合に有用です。
3. 量子化 (Quantization): 集約された $H_{comb}$ を用いて、標準的なGPTQの列逐次更新を適用します。

2.3 CLHA-MP (混合精度)

常に活性化する共有エキスパート (always-active shared experts) は、ルーテッド・エキスパート（$K/E$ の割合を受ける）とは異なり、すべての $N$ 個のトークンを受け取ることを認識し、その量子化感度は $E/K$ 倍高いことが指摘されています。CLHA-MP は、これらの共有エキスパートに $b+1$ ビット を割り当て、ルーテッド・エキスパートを $b$ ビットに保つことで、最小限のストレージオーバーヘッド（エキスパート・パラメータの10〜20%）で対応します。

2.4 実装上の修正: フックの曖昧性 (Hook Ambiguity)

本論文は、重大なエンジニアリング上の落とし穴を特定しています。共有の ExpertFFN モジュールにアタッチされた標準的なPyTorchのフォワード・フックは、そのモジュールを呼び出す任意のレイヤーに対して実行されます。これにより、異なるレイヤーからの活性化が単一のヘッセ行列累積器に混ざり合い、レイヤーごとの属性付けが破損してしまいます。

• 解決策: 著者らは、MoELayer.forward() のシグネチャを通じて渡されるレイヤー認識コールバック・メカニズムを提案しています。これにより、レイヤーインデックスをコールバックに明示的に提供し、基盤となる重みモジュールが共有されている場合でも、$H_a$ と $H_b$ に対して個別のヘッセ行列を確実に蓄積できるようにします。

3. 主な貢献

1. 理論的命題: 共有された重みの最適なヘッセ行列は、レイヤーごとのヘッセ行列のルーティング頻度による加重和であるということを定式化した。
2. CLHAアルゴリズム: GPTQフレームワーク内における結合ヘッセ行列キャリブレーションの実用的な実装。
3. CLHA-MP: 高感度な共有エキスパートに余剰の1ビットを割り当てる混合精度拡張。
4. 実装の修正: 共有モジュール・アーキテクチャにおいてキャリブレーションを密かに破損させる「フックの曖昧性」を特定し、解決した。
5. 実証的検証: CLHAによる改善が、クロスレイヤー間のルーティングの乖離（divergence）と相関していることを示した。

4. 実験結果

WikiText-2でゼロから学習された 34.5MパラメータのCLWS-MoE を用いて実験が行われました。

• 2ビットにおける性能: これは誤ったキャリブレーションの影響が蓄積される重要な領域です。
  • CLHA-rt (ルーティング重み付け) は、PLICと比較してパープレキシティ (PPL) を 4.3% 低減しました。
  • CLHA-eq (等価重み付け) は、PPLを 5.4% 低減しました。
  • CLHA-MP は、PLICに対して 18.6% の改善を達成しました。
  • 4ビットおよび3ビットでは、差はノイズの範囲内（1%未満）でした。
• キャリブレーション誤差 vs テスト誤差: 興味深いことに、PLICは訓練セットに対するキャリブレーション・ヘッセ行列誤差 (CHE) が低くなりました。これは、PLICがその分布に対して直接最適化しているためです。しかし、CLHAはテストセットに対してより優れた汎化性能を示しました。これは、結合ヘッセ行列が単一レイヤーの視点よりも、共有重みの真の分布をより良く捉えていることを示唆しています。
• ルーティングの鋭さ（Sharpness）実験: ルーターのロジットを人工的に鋭くし（クロスレイヤーの乖離を増加させ）、以下の観察を行いました：
  • PLICの性能は、乖離が増すにつれて 5% 低下しました。
  • CLHAは安定していました（劣化は1%未満）。
  • CLHAとPLICの性能差は、自然なルーティング時の +4.3% から、高乖離時の +8.0% へと拡大しました。これは、CLHAの利点がクロスレイヤーの乖離に応じて拡大するという理論的予測を裏付けています。

5. 意義と主張

本論文は、MoEモデルにおけるクロスレイヤー重み共有に特化した最初のPTQ分析を提供すると主張しています。その意義は以下の通りです：

• 目的関数の修正: 共有された重みを単一レイヤーのエンティティとして扱うことは数学的に最適ではなく、ヘッセ行列の集約が必要不可欠であることを示したこと。
• 利点のスケーラビリティ: モデルが大型化し、ルーティングパターンがより専門化（クロスレイヤーの乖離が増大）するにつれ、CLHAの恩恵は大きくなる可能性が高く、将来のプロダクション向けCLWS-MoE（例：Megrez2, DeepSeek, Qwen）にとって不可欠であることを示したこと。
• 実用的な必要性: 既存のPTQツールをCLWSアーキテクチャにそのまま適用することは、フックの曖昧性のために根本的に不完全であることを指摘し、この領域における将来のあらゆるキャリブレーション手法に必要な修正を提供したこと。

著者らは、モデルが小さく（34.5M）、学習ステップも限定的であることを踏まえ、絶対的な性能については控えめな姿勢を保っていますが、標準的なキャリブレーションに対する相対的な改善は堅牢であり、理論的に裏付けられていることを強調しています。