ストレージだけでは足りない

前回の記事でK-Meansベース量子化を見た。重みを少数の代表値にまとめてモデルサイズを大きく減らしたが、決定的な限界があった。縮むのはストレージ容量だけ で、推論時にはコードブックでデコードして 依然としてfloating-point演算 をするという点だ。

今回はその壁を越える。重みを整数で保存するのを超えて、乗算と加算まで全部整数演算で 推論する方法 — Linear Quantization(線形量子化) だ。2018年のJacob et al.の論文1が提案し、TensorFlow LiteのINT8量子化がまさにこの方式だ。


核心: 整数と実数をつなぐアフィン写像

Linear Quantizationは 整数を実数へ送るアフィン(affine)写像 だ。たった一行だ。

$$r = S \cdot (q - Z)$$
  • $r$ — 元の実数(real)値(floating-point)
  • $q$ — 量子化された 整数(quantized integer)
  • $Z$ — zero point。実数 $0$ にちょうど対応する整数だ。(整数)
  • $S$ — scale。整数一目盛りが実数でどれだけかを表す。(floating-point)

つまり整数 $q$ から zero point $Z$ を引いて scale $S$ を掛ければ、元の実数が復元される。逆に実数を整数にするときは $q = \text{round}(r / S + Z)$ で丸める。

$Z$ が別にある理由が重要だ。実数 $0$(例: ReLUの後の0、パディング)はニューラルネットで非常によく現れ、正確に表現されなければならない。zero pointは 実数0がどれかの整数に誤差なくぴったり合う ことを保証する仕掛けだ。

下のウィジェットで実際に確認してみよう。実数の重み行列が整数 $q$ に量子化され、再び $S(q-Z)$ で復元される。bitsを変えると整数範囲が、asymmetric/symmetricを押すと $Z$ の扱い方が変わり、下に scale $S$、zero point $Z$、量子化誤差 がリアルタイムで更新される。


SとZはどう決めるか

$N$ビット整数の範囲は2の補数を基準に決まる。

ビット幅$q_{\min}$$q_{\max}$
2−21
3−43
4−87
$N$$-2^{N-1}$$2^{N-1}-1$

ここで実数範囲 $[r_{\min}, r_{\max}]$ の両端が整数範囲の両端に対応するとしよう。

$$r_{\max} = S(q_{\max} - Z), \qquad r_{\min} = S(q_{\min} - Z)$$

2つの式を引くと $Z$ が消えて scale が出る。

$$S = \frac{r_{\max} - r_{\min}}{q_{\max} - q_{\min}}$$

ウィジェットのデフォルト(2ビット、実数範囲 $[-1.08, 2.12]$)で計算すると $S = \frac{2.12 - (-1.08)}{1 - (-2)} = \frac{3.20}{3} \approx 1.07$ だ。

zero point は同じ式から出る。$r_{\min} = S(q_{\min} - Z)$ を $Z$ について解き、整数でなければならないので丸める。

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

同じ例で $Z = \text{round}\left(-2 - \frac{-1.08}{1.07}\right) = \text{round}(-0.99) = -1$ だ。ウィジェットのstatsに出る $S$、$Z$ と合うか確認してみよう。


本当の核心: 整数行列積

ここまでは「整数で保存する」だ。いよいよ 演算も整数で 行う部分を見よう。ニューラルネットの基本演算である行列積 $Y = WX$ を考える。各値をアフィン写像に置き換えると:

$$S_Y(q_Y - Z_Y) = S_W(q_W - Z_W)\cdot S_X(q_X - Z_X)$$

$q_Y$ について整理すると:

$$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 q_X$、$Z_X q_W$、$Z_W Z_X$ の項はすべて整数同士の積・和だ。さらに $Z_W Z_X$ のように入力に無関係な項は 事前計算(precompute) しておける。つまり括弧全体が 整数算術 で終わる。

残るのは先頭の $\frac{S_W S_X}{S_Y}$ だけで、これはfloating-pointだ。ここがこの方法の妙技だ。経験的にこのスケール比は 常に $(0, 1)$ 区間 にあるので、こう書ける。

$$\frac{S_W S_X}{S_Y} = 2^{-n} M_0, \qquad M_0 \in [0.5, 1)$$

$M_0$ は 固定小数点(fixed-point)乗算 で、$2^{-n}$ は単なる ビットシフト(bit shift) で処理される。結局floating-pointユニットは全く要らない。すべての演算が整数の乗算・加算・シフトで終わる。


さらにシンプルに: 対称量子化とbias畳み込み

実戦では2つの簡略化を使う。

対称(symmetric)量子化 — 重みの分布は大抵0を中心に対称だ。そこで重みのzero pointを最初から $Z_W = 0$ に固定する。すると上の式で $Z_W q_X$ と $Z_W Z_X$ の項が消えて、ずっとすっきりする。

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

bias畳み込み — バイアスのある $Y = WX + b$ では、$Z_b = 0$、$S_b = S_W S_X$ と置けばバイアスまで自然に吸収される。入力に無関係な部分($q_b - Z_X q_W$)を一つの $q_{bias}$ にまとめておくと:

$$q_Y = \frac{S_W S_X}{S_Y}\big(q_W q_X + q_{bias}\big) + Z_Y$$

畳み込み(convolution)層も同じ構造だ。$q_W q_X$ の位置に $\text{Conv}(q_W, q_X)$ が入るだけだ。まとめると整数推論パイプラインはこうなる。

  1. 整数の入力・重みで 整数MAC(積和) → int32累算
  2. 事前計算した $q_{bias}$ を 整数加算
  3. $\frac{S_W S_X}{S_Y}$ を 固定小数点の積 + シフト でrescale → Nビット整数へ
  4. zero point $Z_Y$ を 整数加算 → 整数出力

floating-pointはどこにもない。


結果: どれだけ失い、どれだけ速くなるか

このようにINT8へ量子化しても、精度はほぼ保たれる。

モデルFloat精度INT8精度
ResNet-5076.4%74.9%
Inception-V378.4%75.4%

1–3%p前後の損失で、モデルは1/4サイズ、演算は整数 になる。モバイル(Snapdragon)ではfloatに対し整数専用推論が、同じ精度でずっと低い遅延を示す — 整数ユニットの方が安く速いからだ。(第1回で見た「低ビット整数演算はfloatより数十倍安い」という話が、ここで現実になる。)


まとめ

  • アフィン写像 $r = S(q - Z)$ 一つで整数 $q$ と実数 $r$ をつなぐ。$Z$(zero point)は実数0を整数に誤差なく合わせる仕掛け、$S$(scale)は一目盛りの大きさ。
  • $S = \frac{r_{\max}-r_{\min}}{q_{\max}-q_{\min}}$、$Z = \text{round}(q_{\min} - r_{\min}/S)$。
  • 整数行列積: $q_W q_X$ とprecompute項で括弧が整数算術、先頭のスケール比は固定小数点の積 + シフト。→ floating-pointなしで推論。
  • 対称量子化($Z_W=0$)とbias畳み込みで式がさらに単純になる。

最大の違いはここだ。K-Means量子化は ストレージだけ 整数だった(演算はfloat)。Linear Quantizationは ストレージも整数、演算も整数 だ。

方式ストレージ演算
FP重みFP算術
K-Means量子化整数インデックス + FPコードブックFP算術
Linear量子化整数重み整数算術

これでニューラルネット量子化の2つの大きな分岐 — ストレージを減らすK-Means、演算まで整数化するLinear — を両方見た。次は極端まで推し進める話、重みを $+1/-1$ の1ビットに減らす Binary・Ternary量子化 だ。


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