重みを共有してモデルを小さくする — K-Meansベース量子化

32ビットの数値は本当に全部必要か

前回の記事でデータ型を分解しながら、一つ気づいたことがある。ニューラルネットの重みは大抵 0付近に密に集まり、裾が長い。だとすれば、こんな疑問が自然に浮かぶ。

似た値を持つ重みがこれほど多いのに、本当に一つ一つを32ビットの実数として別々に保存する必要があるのか?

似た値を 一つの塊にまとめて代表値一つで済ませたら どうだろう? たとえば 2.09, 2.12, 1.92, 1.87 はどれも「だいたい2.0」だ。この4つを 2.00 一つで置き換え、各位置には「何番目の代表値を使うか」というインデックスだけを記録する。

これが K-Meansベースの重み量子化(weight quantization) の核心的なアイデアだ。2016年のDeep Compression論文¹が提案し、精度をほとんど落とさずにモデルを数十倍小さくした。

---

アイデア: クラスタリングで重みを共有する

手順は3ステップだ。

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

$k=4$個のクラスタなら、インデックスは2ビット($\log_2 4$)で十分だ。元の32ビットの数値が 2ビットのインデックス + 小さなコードブック に縮む。

下のウィジェットで実際に動かしてみよう。4×4の重み行列があり、runを押すとK-Meansが反復されるたびに 中心値(コードブック)が収束 し、各セルが属するクラスタの色で塗られる。bitsを変えるとクラスタの個数が変わり、下の 量子化誤差(MSE)・ストレージ・圧縮率 がリアルタイムで更新される。

デフォルト状態(2-bit)でrunを最後まで回すと、コードブックはおよそ -1.00, 0.00, 1.50, 2.00 に収束する。各重みはこの4つの値のいずれかに復元(reconstruct)される。元の値と復元値の差が 量子化誤差 だ。

• bitsを 3 に上げてみよう → クラスタが8個に増えて誤差は減るが、インデックスが大きくなりストレージは増える。
• bitsを 1 に下げてみよう → 代表値2個にまとめられて誤差が大きくなる。

精度(誤差)とサイズ(ストレージ)の間の トレードオフ が手に取るように分かる。

---

どれだけ縮むか

さっきのウィジェットの数字を追ってみよう。4×4 = 16個の重み、2-bit量子化:

• 元: $16 \times 32\text{bit} = 512\text{bit} = 64\text{B}$
• 圧縮: インデックス $16 \times 2\text{bit} = 32\text{bit} = 4\text{B}$  +  コードブック $4 \times 32\text{bit} = 128\text{bit} = 16\text{B}$  =  $20\text{B}$
• $64 / 20 = \textbf{3.2倍}$ 小さくなった。

この例はパラメータが16個しかないので、コードブック(16B)が相対的に大きく見える。しかし実際のネットワークはパラメータ数$M$が数百万〜数十億なので $M \gg 2^N$ となり、コードブックのコストは無視できる。一般化すると:

$$\text{元} = 32M \text{ bit}, \qquad \text{圧縮} = \underbrace{N \cdot M}_{\text{インデックス}} + \underbrace{32 \cdot 2^N}_{\text{コードブック}} \text{ bit}$$

$M \gg 2^N$ のとき、コードブックの項は消えて圧縮率は $32M / NM = \mathbf{32/N}$ 倍に収束する。つまり2ビットなら約16倍、4ビットなら8倍だ。下でパラメータ数$M$とビット数$N$を動かすと、コードブックがいつ無視できるようになるかが見える。

---

精度は? — コードブックを再学習する

重みを代表値にまとめれば当然、誤差が生じる。そのままだと精度が落ちる。Deep Compressionの巧妙な点は、量子化した コードブック自体を再学習(fine-tuning) させることだ。

方法はこうだ。逆伝播で各重みの勾配を求めた後、

1. 同じクラスタに属する重みの勾配を集めて(group by index)、
2. 足し合わせ(reduce)、
3. 学習率を掛けて そのクラスタの中心値(centroid)を更新する。

個々の重みではなく 代表値をわずか数個だけ 学習するわけだ。こうすると量子化で生じた誤差をかなり取り戻せる。

重みの分布で見ると直感的だ。量子化前は連続的な釣鐘型の分布だったのが、離散的ないくつかのスパイクに変わり、再学習を経てそのスパイクがデータに合わせて微調整される。代表値の「位置」が損失を減らす方向へ移っていくのだ。

その結果、何ビットまで下げても精度が保たれるか という問いへの答えはかなり驚きだ。Deep Compressionの実験ではAlexNetで:

• 畳み込み(Conv)層: 4ビットまで精度損失なし
• 全結合(FC)層: 2ビットまで精度損失なし

つまり層ごとに必要なビット数が異なり、パラメータの多いFC層ほどより積極的に削れる。

---

最後の一絞り: ハフマン符号化

量子化まで終えると、各重みはいくつかのインデックス(例: 2ビットなら0–3)のいずれかになる。ところがこのインデックスは 均等には現れない。重みが0付近に集まっているので、0に対応するインデックスは非常に頻繁に現れ、両端のインデックスはまれにしか現れない。実際のネットワークのインデックスのヒストグラムを描くと、片側に大きく偏っている。

ここでロスレス圧縮をもう一絞りできる。ハフマン符号化(Huffman coding) はシンボルの頻度に応じてビット長を変えて割り当てる。

• よく現れる値 → 短い符号
• まれに現れる値 → 長い符号

すべてのインデックスに同じ2ビットを使う代わりに、頻度に合わせて符号長を分ければ、平均ビット数が減る。たとえば分布が片側に偏った場合:

代表値	頻度	固定2ビット	ハフマン符号
0.00	50%	01	0 (1ビット)
−1.00	25%	00	10 (2ビット)
1.50	15%	10	110 (3ビット)
2.00	10%	11	111 (3ビット)

このとき平均ビット数は $0.5\times1 + 0.25\times2 + 0.15\times3 + 0.1\times3 = 1.75$ ビットで、固定2ビットより短い。情報理論的に最適な符号長は分布の エントロピー に近づくので、偏った分布ほど利得が大きくなる。

核心は ロスレス だという点だ。値を一つも変えず、純粋に表現の無駄だけを取り除く。Deep Compressionではこの段階がパイプラインの最後に付いて、圧縮率をもう一段引き上げる(量子化まで27–31倍 → ハフマンまで35–49倍)。

---

重要な限界: これは「ストレージ」しか減らさない

ここで必ず押さえるべき点がある。K-Means量子化が減らすのは ディスク・メモリに保存されるモデルサイズだけ だ。

推論(inference)の時点ではどうなるか? 保存されているのは整数インデックス + 実数コードブックだ。実際に乗算・加算を行うには、インデックスをコードブックで デコードして元の実数重みを復元 しなければならない。つまり:

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

メモリ帯域は節約できたが、計算そのものは全く速くなっていない。積和演算(MAC)ユニットは相変わらずfloatを回す。

本当に 整数演算で推論まで速く するには別のアプローチが必要だ。重みを整数で保存し、整数算術で乗算まで行う Linear Quantization(線形量子化) がそれだ。その話は次回に。

---

まとめ

• アイデア: 似た重みをK-Meansでまとめて代表値(コードブック)を共有し、各位置には小さなインデックスだけを保存する。
• 圧縮: 32ビット → $N$ビットインデックス。$M \gg 2^N$ なら約 $32/N$ 倍小さくなる。
• 精度: コードブック(centroid)を再学習して誤差を取り戻す。Conv 4ビット・FC 2ビットまでロスレス。
• 追加圧縮: インデックス分布が偏っているので、ハフマン符号化でロスレス圧縮をもう一絞りできる。
• 限界: ストレージしか減らず、推論の演算は依然としてfloat。これはDeep Compressionパイプライン(Pruning → Quantization → Huffman)の真ん中のピースで、整数演算まで高速化するにはLinear Quantizationが必要だ。

参考までに、原論文のDeep Compressionはここに 枝刈り(pruning) まで加えた3段階パイプライン(枝刈り → 量子化 → ハフマン)で、AlexNetを35倍、VGGを49倍まで小さくした(精度を維持)。SqueezeNetに適用すると0.47MB、510倍の圧縮まで到達した。枝刈りは別途扱うテーマだ。