Quantization-Aware Training (QAT)

量子化対応学習（QAT）は、学習中に量子化をシミュレーションすることで、低精度推論に向けてモデルを準備する学習手法です。通常、モデルは32ビット浮動小数点数（FP32）で学習され、重みを保存します。モデル量子化では、デプロイ向けに8ビット整数（INT8）などの形式へ変換します。これにより重みのサイズは4分の1になり、CPU、GPU、ニューラルプロセッシングユニット（NPU）で高速な整数演算を利用できます。変換では各値を丸め、外れ値をクリップするため、精度が低下します。QATでは、モデルがまだ学習中の順伝播に丸め処理を挿入するため、勾配降下法によって量子化後も機能する重みを見つけられます。通常、学習済みモデルに対する短いファインチューニングの段階として実施します。Googleの整数演算のみを用いる推論方式では、モバイルCPU向けのINT8推論にこの種の学習手順を組み合わせました。それ以来、QATはエッジAIのデプロイで標準的な手法になっています。

量子化対応学習の仕組み#

QATでは、疑似量子化演算（量子化・逆量子化（Q/DQ）ペアとも呼ばれます）をモデルの重みと活性化に追加します。それぞれの演算は、スケールとゼロ点を使って浮動小数点数を最も近いINT8レベルに対応付けた後、ただちに浮動小数点数に戻します。そのため、ネットワークは浮動小数点演算を続けながら、INT8エンジンが使うものと同じ丸め値を扱います。スケールは、少数のデータバッチで各テンソルの一般的な範囲を記録するキャリブレーションから取得します。手法によっては、学習中にスケールを学習するものもあります。

丸め処理の勾配はほとんどの箇所でゼロとなり、学習が止まってしまいます。QATではストレートスルー推定器（STE）を使用し、誤差逆伝播中に丸め処理を恒等関数として扱うことで、勾配が基礎となるFP32の重みに伝わるようにします。

import torch

# FP32 weights that training keeps and updates
w = torch.tensor([0.3137, -0.7214, 0.0521, 1.1800], requires_grad=True)

# Symmetric INT8 fake quantization: snap each weight to one of 255 levels, then map back to float
scale = w.detach().abs().max().item() / 127
w_q = torch.fake_quantize_per_tensor_affine(w, scale, 0, -127, 127)
print(f"Fake-quantized: {[round(v, 4) for v in w_q.tolist()]}")

# The forward pass sees the rounded weights; the straight-through estimator still delivers gradients to w
x = torch.tensor([1.0, 2.0, 3.0, 4.0])
loss = (w_q * x).sum()
loss.backward()
print(f"Gradient on FP32 weights: {w.grad.tolist()}")

順伝播では、丸められた重み0.3159、−0.7247、0.0557、1.18を使用します。各値はFP32値から量子化ステップの半分（0.0046）以内ですが、FP32の重みに届く勾配は、丸め処理がなかった場合と同じく、入力である1.0、2.0、3.0、4.0と完全に一致します。多くのステップを経るうちに、オプティマイザーは丸め誤差の影響が最小になる位置へ重みを移動させ、学習済みモデルは学習されたスケールとともにINT8形式へエクスポートされます。

QATと学習後量子化の比較#

学習後量子化（PTQ）では、完成したFP32モデルを直接変換し、数百枚のキャリブレーション画像を使って範囲を設定します。数分で完了し、学習パイプラインも不要ですが、重みが丸め処理に適応する機会がないため、繊細なモデルでは精度が大きく低下することがあります。QATは全学習データを使ったファインチューニングが必要ですが、その代わりに精度低下の大部分を取り戻せます。通常、チームはまずPTQを試し、INT8による精度低下が大きすぎる場合に限ってQATへ切り替えます。

Ultralytics YOLOは、学習時にquantize=8を指定するとQATを実行します。NVIDIA Model Optimizerを通じてINT8の疑似量子化をキャリブレーションし、処理を組み込んでファインチューニングします。また、ONNXおよびTensorRTへのエクスポートで、別途キャリブレーションデータを使わずにQ/DQノードとして書き出せる範囲を保存します。COCO val2017でのTensorRTエンジンでは、PTQによって中型YOLO26モデルのmAP_50-95は0.527から0.470に低下しましたが、QATでは0.514を維持しました。一方、nanoモデルではPTQによる低下がわずか0.010だったため、QATを追加しても改善はありませんでした。

from ultralytics import YOLO

# 学習済みのFP32 YOLO26チェックポイントから開始します
model = YOLO("yolo26n.pt")

# 低い学習率を使い、INT8の疑似量子化を組み込んでファインチューニングします
model.train(data="coco8.yaml", epochs=5, quantize=8, optimizer="AdamW", lr0=0.00001, cos_lr=True)

# ONNXにエクスポートします。学習されたINT8の範囲がQ/DQノードになるため、キャリブレーションデータは不要です
model.export(format="onnx", quantize=8)

QATの用途#

• エッジビジョン: モバイル端末、スマートカメラ、ドローン、NPUで動作するINT8検出器です。整数演算のみのアクセラレーターでは、浮動小数点モデルを実行できない場合があります。
• 推論サーバー: TensorRTなどのGPUエンジンはINT8 Tensor Coreカーネルを実行してスループットを向上させます。QATを使うと、精度をFP32に近い状態に維持できます。
• 言語モデル: QATは現在、8ビット未満の精度にも対応しています。Googleが公開した4ビットQ4_0形式のGemma 3 QATチェックポイントは、メモリ使用量を大幅に抑えながら、bfloat16に近い品質を維持します。

主要なフレームワークには、PyTorchのtorchao、TensorFlowのModel Optimization Toolkit、OpenVINO向けのIntel NNCFなど、QATツールが含まれています。

QATと関連用語#

• モデル量子化: 数値精度を下げる手法の総称です。PTQとQATが主なワークフローです。
• 混合精度学習: メモリと時間を節約するため、FP16またはBF16で学習します。INT8推論に向けたモデルの準備は行いません。
• 半精度: FP16推論のことです。ほとんどのモデルは再学習なしで対応できるため、通常QATは必要ありません。
• モデルプルーニング: 精度を下げるのではなく、重みやチャネルを削除します。量子化とプルーニングは併用できます。
• 知識蒸留: 大規模なアーキテクチャを模倣する小規模なアーキテクチャを学習させます。QATと併用されることも多く、その場合、全精度モデルが自身の量子化コピーの教師モデルとして機能します。