この記事はディープラーニング時代のデータ型編の続編です。前回の記事でINT/FPのデータ型がビットをどう数値として解釈するかを見たなら、この記事ではその低ビット型を使って実際に学習済みモデルを削り、結果を測定します。コードはすべて github.com/warpspaceinc/efficient-ml-practice のノートブックから持ってきたもので、以下のplotはそのコードをColab(T4)でそのまま回して出力したものです。

32ビットの数値は本当にすべて必要なのか

前回の記事の最後で観察したことが一つありました。ニューラルネットワークの重みは、たいてい0の近くに釣鐘型で密集していて、裾が長いということです。実際にMNIST用の小さなMLP(784→256→128→10)をわずか2エポックだけ学習させたあと、第1層の重みをヒストグラムに描いてみると、まさにその通りになります。

fc1の重み分布(float32)— 0付近に密集した釣鐘型

値がこれほど狭い範囲に密集しているのに、本当に一つ一つを32ビットの実数として個別に保存し、32ビットで掛け算する必要があるのでしょうか。量子化(quantization) は「いいえ」と答えます。この記事では、代表的な2つの流れをコードで直接回してみます。

  • K-Means量子化 — 近い値を代表値でまとめて保存を減らす(非均等)。
  • Linear量子化 — 整数格子に写像して保存も演算も整数にする(均等)。

どちらの方式も学習が終わったモデルに事後適用する PTQ(post-training quantization) で、ここでは重みだけを量子化します。


1. K-Means量子化 — 代表値の共有で保存を減らす

最初のアイデアは2016年のDeep Compression論文1から来ました。近い値の重みを一つの塊にまとめて、代表値一つで済ませようというものです。たとえば 2.09, 1.92, 1.87 はすべて「だいたい2.0」なので、これらの値を 2.00 一つに置き換え、各位置には「何番目の代表値を使うか」というインデックスだけを記録しておきます。

手順は3ステップです。

  1. クラスタリング — すべての重みを $k = 2^N$ 個のグループにK-Meansクラスタリング。
  2. コードブック(codebook) — 各グループの中心値(centroid)を代表値として保存。このリストがコードブック。
  3. インデックス — 各位置には32ビットの実数の代わりに「何番目のクラスタか」を指す $N$ ビットの整数だけを保存。

コードではscikit-learnの KMeans をそのまま使います。局所最適解を避けるため、min〜maxを均等分割した値で初期化します。

from sklearn.cluster import KMeans
import numpy as np

def kmeans_quantize(w, n_bits, fit_sample=20000):
    """重みをK-Meansで量子化。復元値、コードブック、インデックスを返す。"""
    shape = w.shape
    flat = w.reshape(-1, 1)
    k = 2 ** n_bits
    lo, hi = float(flat.min()), float(flat.max())
    init = np.linspace(lo, hi, k).reshape(-1, 1)          # 線形初期化
    idx = np.random.choice(len(flat), min(fit_sample, len(flat)), replace=False)
    km = KMeans(n_clusters=k, init=init, n_init=1, max_iter=50).fit(flat[idx])
    codebook = km.cluster_centers_.reshape(-1)            # 代表値リスト
    indices = km.predict(flat)                            # 各重み → クラスタインデックス
    recon = codebook[indices].reshape(shape)              # 復元
    return recon, codebook, indices

2ビット($k=4$)で回した結果を元の分布の上に重ねて描くと、代表値(縦線)が値の密集する0付近に細かく配置されるのが分かります。これが「非均等(non-uniform)」の意味です。

K-Means 2ビット量子化 — 代表値が0付近に密な非均等格子

どれくらい減るか

圧縮率はビットの会計で計算します。パラメータ $M$ 個、$N$ ビット量子化なら:

$$\text{元} = 32M \text{ bit}, \qquad \text{圧縮} = \underbrace{N \cdot M}_{\text{インデックス}} + \underbrace{32 \cdot 2^N}_{\text{コードブック}} \text{ bit}$$
num = w.size; k = 2 ** N_BITS
orig_bits = 32 * num
comp_bits = N_BITS * num + 32 * k     # インデックス + コードブック
print(f"compression: {orig_bits/comp_bits:.2f}x")

$M \gg 2^N$ ならコードブック項は無視できるようになり、圧縮率は $32M/NM = \mathbf{32/N}$ 倍へ収束します。2ビットなら約16倍、4ビットなら8倍です。(第1層fc1を基準にした実測圧縮率: 15.99×

重要な限界: 保存しか減らない

ここは必ず押さえておくべきです。K-Means量子化が減らすのはディスク・メモリに保存されるサイズだけです。推論するときは、インデックスをコードブックでデコードして元の実数の重みを復元したうえで掛け算しなければなりません。

  • 保存: 整数インデックス(小さい)✓
  • 演算: 復元された32ビット実数で依然としてfloating-point演算

メモリ帯域は節約しましたが、乗算器(MAC)は依然としてfloatを回します。本当に整数演算まで速くするには、次の方式が必要です。


2. Linear量子化 — 整数演算だけで推論する

Linear(線形)量子化は、整数と実数をアフィン(affine)写像一行でつなぎます。2018年にJacob et al.2が提案し、TensorFlow LiteのINT8量子化がまさにこの方式です。

$$r = S \cdot (q - Z)$$
  • $r$ — 元の実数(real)値
  • $q$ — 量子化された整数
  • $Z$ — zero point。実数 $0$ にちょうど対応する整数。ReLUの後の0やパディングのようなよくある値を誤差なく表現するための仕組み。
  • $S$ — scale。整数一目盛りが実数でいくらか(floating-point)。

$S$ と $Z$ は、実数の範囲 $[r_{\min}, r_{\max}]$ の両端を整数の範囲 $[q_{\min}, q_{\max}]$ の両端に合わせるところから出てきます。

$$S = \frac{r_{\max} - r_{\min}}{q_{\max} - q_{\min}}, \qquad Z = \text{round}\left(q_{\min} - \frac{r_{\min}}{S}\right)$$

コードでは、この $S$、$Z$ さえ計算してやれば、量子化(round・clamp)と復元(dequant)はPyTorch組み込みの fake_quantize_per_tensor_affine が処理します。任意のビット幅は quant_min/quant_max で指定します。

import torch

def linear_quantize(w, n_bits, symmetric=False):
    """PyTorch組み込みのfake_quantizeで線形(アフィン)量子化。復元値とS, Zを返す。"""
    t = torch.as_tensor(w, dtype=torch.float32)
    if symmetric:                                   # Z=0固定(演算が単純)
        qmax = 2 ** (n_bits - 1) - 1; qmin = -qmax
        S = float(t.abs().max()) / qmax; Z = 0
    else:                                           # asymmetric: 範囲を正確に合わせる
        qmin = -2 ** (n_bits - 1); qmax = 2 ** (n_bits - 1) - 1
        S = (float(t.max()) - float(t.min())) / (qmax - qmin)
        Z = int(round(qmin - float(t.min()) / S))
    recon = torch.fake_quantize_per_tensor_affine(t, S, Z, qmin, qmax)
    return recon.numpy(), S, Z

2ビットで回すと、代表値(縦線)が一定の間隔 $S$ で均等に配置されます。K-Meansとちょうど対比される点です — こちらはデータがどこに密集していようと格子が均一です。

Linear 2ビット量子化 — 均等間隔の格子

なぜ「整数演算まで」可能なのか

Linear量子化の本当の値打ちは保存ではなく演算にあります。行列積 $Y = WX$ の各値をアフィン写像に置き換えると:

$$q_Y = \frac{S_W S_X}{S_Y}\big(q_W q_X - Z_W q_X - Z_X q_W + Z_W Z_X\big) + Z_Y$$

括弧の中は $q_W q_X$ を含めてすべて整数の掛け算・足し算で、$Z_W Z_X$ のように入力に依存しない項はあらかじめ計算しておけます。前のスケール比 $\frac{S_W S_X}{S_Y}$ だけが実数ですが、この値は経験的に常に $(0,1)$ の区間にあるので、$2^{-n} M_0$ の形の固定小数点の掛け算 + ビットシフトで処理されます。結局floating-pointユニットはまったく必要ありません — 第1編のエネルギー表で見た「整数演算はfloatより数十倍安い」がここで現実になります。

重みはたいてい0対称なので、実戦では $Z_W = 0$ である対称(symmetric)量子化をよく使います。そうすると $Z_W$ 関連の項がまるごと消え、式がさらにすっきりします。


3. で、どれくらい失うのか — 実際に測定

言葉だけでは掴めないので、3つの方式(Linear asymmetric / Linear symmetric / K-Means)をビット幅を変えながらMNISTのテスト精度で比較しました。モデルの3つのLinear層の重みをすべて量子化したうえで(復元値に置き換え)精度を測ったものです。

def eval_linear(model, n_bits, symmetric):
    m = copy.deepcopy(model)
    with torch.no_grad():
        for layer in [m.fc1, m.fc2, m.fc3]:
            wl = layer.weight.detach().cpu().numpy()
            recon, *_ = linear_quantize(wl, n_bits, symmetric=symmetric)
            layer.weight.copy_(torch.tensor(recon, dtype=torch.float32, device=device))
    return accuracy(m)

bits_list = [2, 3, 4, 8]
acc_asym = [eval_linear(model, b, False) for b in bits_list]
acc_sym  = [eval_linear(model, b, True)  for b in bits_list]
# K-Means も同じ構造の eval_kmeans(model, b) で比較

ビット幅によるMNISTテスト精度 — Linear(asym/sym)vs K-Means

bitsLinear (asym)Linear (sym)K-Meansfloat32
897.2%97.2%97.2%97.2%
497.1%97.0%97.1%97.2%
396.7%94.4%96.8%97.2%
258.1%18.8%93.9%97.2%

読み取るべき大きな絵はこうです。

  • 8ビットなら事実上タダ。 3つの方式すべてがfloat32と区別がつかないほど精度を保ちます。
  • 低ビットに行くほどK-Meansが有利。 代表値をデータ分布に合わせて(非均等)配置するため、同じ2ビットでも均等格子のLinearより誤差が小さくなります。下のMSE曲線でより顕著です。

fc1量子化誤差(MSE)vs ビット幅 — 非均等(K-Means)が低ビットで有利

しかし精度がすべてではありません。K-Meansが2ビットで誤差が小さくても、推論の演算は依然としてfloatです。Linearは誤差を少し多めに引き受ける代わりに、演算まで整数にして実際のハードウェアで速くなります。2つの方式は競合ではなく、目的の異なる道具です。

方式保存演算格子
FP重みFP算術
K-Means量子化整数インデックス + FPコードブックFP算術非均等
Linear量子化整数重み整数算術均等

実際に回してみる

上のplotと表はすべて、以下の2つのノートブックをColabで実行して出てきたものです。ランタイム → すべて実行 で終わり、自分のモデルの .safetensors をアップロードしてパスだけ変えればそのまま適用できます。

まとめ

  • 量子化 = 表現可能な値の個数を減らすこと。 重みは0付近に密集しているので、32ビットを丸ごと使うのは無駄です。
  • K-Means(非均等) は代表値を分布に合わせて配置し保存を減らすが、演算はfloat。低ビットで誤差が小さいです。
  • Linear(均等) はアフィン写像 $r=S(q-Z)$ で保存も演算も整数化。ハードウェア高速化の基盤です。
  • 8ビットはほぼタダ、2ビットまで押し込むと方式の選択が精度を分けます。

次の実験ノートでは、ここで扱いきれなかった活性値(activation)量子化QAT(量子化を意識した学習) で、整数推論パイプラインを最後まで組み上げてみます。


  1. Han, Mao, Dally. Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. ICLR 2016. ↩︎

  2. Jacob et al. Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference. CVPR 2018. ↩︎