蒸留・量子化・枝刈りの使い分け — 小さくする3つの道

モデルを小さくする道は3本あります。数値の目盛りを粗くする量子化、重みそのものを間引く枝刈り、小さいモデルに作り直す蒸留。圧縮率・精度・実装コストの三軸で正面から比べ、「訓練を伴うものを先に、量子化を最後に」という組み合わせ順序の理由まで、前提知識ゼロから。

Distilling the Knowledge in a Neural Network

荷物を軽くするには、三つのやり方がある

引っ越しでダンボールの数を減らしたいとします。手は3つあります。

1つめは、記録の精度を落とすこと。アルバムの写真を1枚ずつ粗い解像度で保存し直せば、枚数はそのままで容量だけ減ります。2つめは、使っていない物を捨てること。3年触っていない鍋を処分すれば、その分だけ物理的に減ります。3つめは、中身を見た上で作り直すこと。全部の荷物を並べて「結局この暮らしに要るものはこれだけだ」と判断し、最初から小さい一式を組み直す。

ニューラルネットワークを小さくする技術も、驚くほどこの3つに対応しています。順に量子化、枝刈り、蒸留です。この記事は、この3つを同じ土俵に載せて比べます。どれか1つを深掘りするのではなく、「いま自分の状況ではどれを使うべきか」「複数使うならどの順で並べるか」に答えを出すのが目的です。

3つの正体を、一行ずつ

まず言葉を揃えます。3つは削っているものが違います。

量子化（quantization） は、重みや活性値を表す数値のビット幅を下げます。1個の重みを16ビット（fp16）で持っていたのを4ビットで持つ。パラメータの個数は1つも減らず、モデルの形も変わりません。減るのは1個あたりのバイト数だけです。

枝刈り（pruning） は、重みの一部を消して0にするか、ヘッド・チャネル・層といった単位ごと取り除きます。残った重みの表現精度はそのままで、本数が減ります。

蒸留（distillation） は、大きいモデル（教師）の出力を手本にして、別の小さいモデル（生徒）を訓練し直すことです。層の数も隠れ次元も語彙も、なんならアーキテクチャそのものも自由に変えられます。元のモデルを加工するのではなく、作り直すのが蒸留です。

一行でまとめると、量子化は「精度を削る」、枝刈りは「本数を削る」、蒸留は「作り直す」。この違いが、そのまま次の三軸の差になります。

三軸で正面から比べる

この表の中で、最初に目を止めるべき行は「速くなるか」です。「小さくなる」と「速くなる」は同じ意味ではありません。 量子化はメモリに載る量が減り、読み出すバイト数も減るので、メモリ帯域が支配的な生成フェーズでは素直に速くなります。枝刈りは、後で述べるとおり、消した重みを「消したまま計算をスキップできる」形になっていなければ1秒も速くなりません。ここを取り違えたまま提案書を書くと、実装が終わってから何も改善していないことに気づきます。

まず、いちばん想像しにくい「目盛りを粗くすると何が起きるか」を、画像で体感しておきましょう。仕組みはJPEGの量子化とほぼ同じ発想です。

滑らかに劣化するのではなく、ある点まで粘って、そこから急に落ちる。この形を覚えておいてください。3つの手法すべてに共通する挙動です。

量子化 — 一番安く、一番先に試すべき手

量子化の中身は、実数の範囲を等間隔の格子に写す割り算と丸め算です。もっとも単純な対称量子化なら、こう書けます。

$w_i$ が元の重み、$b$ が使えるビット数、$s$ が「目盛りの幅（スケール）」です。式(1)は「そのグループで一番大きい重みが、ちょうど表せる最大の整数になるように目盛り幅を決め、あとは全部その目盛りに丸める」と言っているだけです。保存するのは丸めた整数と、グループごとの $s$ ひとつ。