量子化アウェア・ヒーリング：4ビット圧縮モデルがフル精度を超える新手法

発表の概要

Hugging Faceのブログに、量子化アウェア・ヒーリング（QAH）という新しいモデル修復手法に関する論文が紹介されました。QAHは、構造圧縮と4ビット量子化を経たモデルの性能を、圧縮前の元モデルを教師とした蒸留で回復させる手法です。論文では、GPT-OSS 120Bを60Bに圧縮し、MXFP4形式に量子化したモデルにQAHを適用した結果が報告されています。

その結果、9つのベンチマークのうち7つで、同じ60Bアーキテクチャのbfloat16版（フル精度）を上回る精度を達成しました。つまり、4ビット化したモデルが、その元となった16ビット版のチェックポイントより小さくて正確だという逆転現象が起きています。残る2つのベンチマーク（MMLU-ProとSciCode）でも、差は1.5ポイント未満に収まっています。

従来手法の限界

従来の効率化パイプラインは、アーキテクチャ圧縮、重みの量子化、その後の修復という3段階で構成されます。修復手法の主流は量子化アウェア訓練（QAT）ですが、これは量子化を模擬した演算を入れながらタスク損失で再学習するためコストが高く、訓練が長期化すると不安定になるという課題があります。もう1つの量子化アウェア蒸留（QAD）は、フル精度の教師からの蒸留で済みますが、構造圧縮されたモデルには独立したフル精度版が存在しないという問題がありました。

このため、構造圧縮と量子化の両方を経たモデルをどう修復するかは、これまであまり研究が進んでいませんでした。QAHは、圧縮後の回復チェックポイントではなく、圧縮前の元モデルを直接教師として使うことでこの問題を解決します。教師と生徒でアーキテクチャが異なっていても、出力分布を使って蒸留できることがポイントです。

QATとの比較

論文では、QAHとQATを同じ条件で比較するため、GPT-OSS 9BをMXFP4に量子化し、平均精度と収束速度を調べています。ピーク精度はQAHが54.9、QATが54.6とほぼ同じですが、収束までの速さはQAHが約100ステップ、QATが約700ステップと大きく異なります。

さらに、QATはピークを過ぎると精度が急落し、1,200ステップまでに約19ポイント低下しました。一方QAHは、ピークに達した後もその精度を維持します。これは、固定された教師分布へのKLダイバージェンスによる蒸留が、生徒モデルにそれ以上動くインセンティブを与えないためと考えられています。

実務への影響と注意点

QAHを適用したモデルは、bfloat16の生徒モデルと比べて重みメモリが約4分の1になります。パラメータ数も半分になるため、トークンあたりの計算量は約半分になり、より小さなGPUでも実行しやすくなる利点があります。モデルファミリー全体をbfloat16ではなく4ビットで提供する場合、計算量は約8分の1に削減できると見込まれています。

一方で、実験はGPT-OSS系モデルに限定されており、他のアーキテクチャでも同様の効果が得られるかはまだ不明です。また、QAHは圧縮前の元モデルを教師として使うため、そのモデルへのアクセスが前提になります。論文の詳細な設定や、再現可能な実装が公開されるかどうかが、次の判断材料になりそうです。