技術要約：NVFP4 LLM蒸留における出力一致を超えて

1. 問題提起

大規模言語モデル（LLM）がレイテンシやメモリ制約のあるプロダクション環境にデプロイされるにつれ、低精度推論フォーマット、特にNVFP4（きめ細かなブロックスケーリングを備えた4ビット浮動小数点フォーマット）への需要が高まっています。**量子化認識蒸留（QAD）**は、KLダイバージェンスを通じて量子化された生徒モデルの出力分布をフル精度の教師モデルに整合させることで、事後訓練量子化（PTQ）による精度低下を回復する堅牢な手法として台頭してきました。しかし、本論文は既存のパラダイムにおける決定的な欠陥を指摘しています。

著者らは、出力レベルの整合性だけでは、忠実な能力回復には不十分であると主張しています。彼らの分析によれば、標準的なQADはKLダイバージェンスのみを最適化するため、深刻な内部表現の劣化を隠蔽してしまう可能性があります。具体的には、生徒モデルが教師の出力ロジットをうまく一致させることができても、その中間活性化の幾何学的構造（ジオメトリ）は意図しない再編成（ドリフト）を起こしています。この「ショートカット学習」は、強化学習（RL）による事後訓練を経たモデルにおいて特に深刻であり、高い出力精度を実現しているにもかかわらず、複雑な推論やコーディングタスクに必要な内部的な意味階層が損なわれてしまいます。

2. 手法：CKA-QAD

出力のマッチングと内部の幾何学的な乖離に対処するため、著者らは標準的なQADに**中心化カーネルアライメント（CKA）**正則化を追加した手法、CKA-QADを提案しています。

コアメカニズム

出力レベルの蒸留のみに依存するのではなく、CKA-QADは、生徒と教師の活性化の層ごとのグラム行列を整合させる軽量な正則化を導入します。

• CKA指標: この手法は線形CKAを利用しており、中心化された活性化行列（$\bar{H}_t$ および $\bar{H}_s$）の類似性を測定します。CKAは直交変換や等方的なスケーリングに対して不変であるため、NVFP4量子化に特有の非定常な活性化分布やスケールシフトに対して堅牢であることから採用されました。
• 目的関数: 全体の損失関数は、標準的なTop-K KLダイバージェンス損失（$L_{TopK-KL}$）と、動的なCKAアライメント項（$L_{CKA}$）を組み合わせたものです。
  $$L_{total} = L_{TopK-KL} + \text{sg}\left(\frac{L_{TopK-KL}}{L_{CKA} + \epsilon}\right) L_{CKA}$$
  ストップグラディエント（sg）演算子の使用と動的なスケーリングにより、グローバルな重み $\lambda$ の手動チューニングなしで、CKA項をロジット損失とバランスさせることができます。

実装の詳細

• 層の選択: CKAは、一致したデコーダーブロック（post-attention/MLP/Mamba更新後）の出力に対して計算されます。残差ストリームのスケールや自己回帰的な依存関係によるアーティファクトを避けるため、事前正規化（pre-normalization）の活性化、アテンション重み、およびKVキャッシュの状態は明示的に除外されています。
• エンジニアリングの最適化: 多くの層にわたってCKAを計算する際の計算オーバーヘッドを管理するために、以下の工夫を行っています：
  1. Top-K 蒸留: KL損失は、教師の上位8192個のロジットに対してのみ計算されます。
  2. 特徴空間CKA: CKAは、明示的なトークン間のグラム行列ではなく、特徴共分散行列（$\bar{H}^\top \bar{H}$）を介して計算され、メモリ複雑性を低減します。
  3. フックベースのキャプチャ: 活性化はフォワードフックを介してキャプチャされ、教師の活性化は計算グラフから切り離されます。

3. 主な貢献

1. QADの失敗モードの診断: 本論文は、KLのみの蒸留が、特にミッドデプス（中間層）のTransformer FFN層やMambaステートブロックにおいて、表現のドリフトを悪化させる可能性があるという、表現レベルの診断を提供しています。このドリフトは、推論やコーディングタスクのパフォーマンスのボトルネックと相関しています。
2. CKA誘導型NVFP4 QAD: 低ビット蒸留中に内部の表現幾何学を明示的に保存する手法であるCKA-QADを導入しました。これにより、最小限のオーバーヘッドでBF16レベルに近い整合性を回復します。
3. 実証的検証: 2つの異なるアーキテクチャで包括的な評価を実施しました：
   • Nemotron 3 Nano: 強化学習による事後訓練を経たハイブリッドMamba–Transformer MoEモデル。
   • Qwen3-4B-Thinking-2507: コンパクトなTransformer推論モデル。

4. 実験結果

著者らは、BF16の教師モデル、NVFP4 PTQベースライン、および標準的なKLのみのQADに対してCKA-QADを評価しました。

表現のアライメント

• Nemotron 3 Nano: 標準的なQAD（KLのみ）は、PTQ（0.993）と比較して平均CKA（0.958）が大幅に低下し、最終層では深刻な崩壊（0.740）が見られました。CKA-QADは、平均CKAを0.994に、最終層のCKAを0.985に回復させ、BF16の教師にほぼ一致させました。
• Qwen3-4B: 同様の傾向が見られ、KLのみのQADでは平均CKAが0.98、最終層のCKAが0.86まで低下しました。CKA-QADはこれらをそれぞれ0.99および0.97に回復させました。

ダウンストリームの精度

CKA-QADは、複雑な推論およびコーディングのベンチマークにおいて一貫した改善を示しました：

• AIME25 (数学的推論): Qwen3-4Bにおいて、68.5% (KLのみ) から 72.3% (CKA-QAD) へ向上。
• GPQA-D (専門知識): 59.5% から 61.1% へ向上。
• LiveCodeBench-v5 (コード生成): 57.9% から 59.8% へ向上。
• Nemotron 3 Nano: 同様の回復を示し、AIME25は87.9%から**88.3%**へ、LiveCodeBenchは68.9%から70.3%へと向上しました。

効率性のオーバーヘッド

本手法は、わずかなトレーニングコストを導入します：

• 実行時間（Wall-clock time）: ステップあたり +0.5%。
• ピークVRAM: +7.0%。
• スループット: -0.5%。
• 収束: 目標とするKLに達するまでのステップ数がわずかに減少（-3.3%）。

5. 意義と主張

本論文は、CKA誘導型の表現アライメントを、量子化LLMの回復のための実用的な補完手法として位置付けています。著者らは以下のように主張しています：

• 機能的整合性は必要条件であるが、十分条件ではない: 出力分布を一致させることは、深い推論に必要な内部的な意味トポロジーの保存を保証するものではありません。
• 幾何学的保存は低ビットの汎化を助ける: 生徒モデルが、損失関数を満たすものの内部構造を劣化させる「ショートカット」マッピングを学習することを防ぐことで、CKA正則化は、事後訓練（特にRL）中にエンコードされた微細な意味階層を、積極的な量子化下でも保持するのに役立ちます。
• 量子化アーティファクトに対する堅牢性: 点ごとの特徴量損失（例：MSE）とは異なり、CKAのスケールや回転に対する不変性は、ブロックごとのスケーリングが非定常な分布を誘発するNVFP4にとって、極めて適しています。

著者らは、数学的推論やコード合成のような、多段階の特徴伝播を必要とするタスクにおいては、内部の幾何学を保存することが極めて重要であり、CKA-QADが最小限のエンジニアリングオーバーヘッドで高忠実度な低ビットLLMデプロイへの実行可能な道筋を提供すると結論付けています。