この記事は軽量化シリーズの3本目（#3）です。データ型 → 量子化が「数値をより少ないビットで」だったとすれば、今回はいっそ接続そのものを刈り取るプルーニングです。コードはgithub.com/warpspaceinc/efficient-ml-practiceのpruning.ipynbノートブックにあり、以下のすべてのplotはそのコードをColab(A100)で実際に回して実測した私たちの数字です。

1989年の「脳損傷」

プルーニングを語るとき、外せない論文が一つあります。タイトルが物騒です — Optimal Brain Damage(最適脳損傷)。著者はYann LeCun、John Denker、Sara Solla、AT&T Bell Labs、発表はNeurIPS 1989です。

少し立ち止まって考えてみると驚きます。1989年。 今から36年前です。ワールド・ワイド・ウェブが世に出る前であり、LeCunがCNN(LeNet)で有名になる前であり、「ディープラーニング」という言葉すら存在しなかった時代です。そのLeCunは今やMetaのチーフAIサイエンティストであり、チューリング賞受賞者です。「ニューラルネットワークで重要度の低い接続を刈り取れば、より小さく、よりよく汎化するネットワークが得られる」 というプルーニングの核心的アイデアが、そんなに昔にすでに定式化されていたのです。

論文の発想は生物学的な比喩ともつながっています。人間の脳もシナプスをやみくもに増やしません。ニューロンあたりのシナプス数は、乳児期に約2,500個 → 2〜4歳頃に15,000個まで急増し → 成人になると約7,000個へと再び減少します。 役に立たない接続を切り落とすこの「刈り込み(pruning)」こそが、かえって効率的な脳をつくるというわけです。

この記事では、その36年の洞察が今日NVIDIA GPUの2:4 sparsityとしてどのように現実になるのかを、3つの問いに沿って追いかけます。

  1. 何を刈るのか (criterion)
  2. どれだけ刈れるのか (再学習と圧縮の驚き)
  3. どのように刈ればハードウェアが速くなるのか (granularityと2:4)

プルーニングとは & 5段階フレームワーク

プルーニングとは重みを0にして接続を削除することです。denseなネットワークをsparseにして、記憶・演算・電力を節約します。形式的には「非零の重みの数を$N$個以下に制限しつつ損失を最小化する」問題です。

$$\arg\min_{W_P} L(x; W_P) \quad \text{s.t.} \quad \lVert W_P\rVert_0 < N$$

Han et al.(NeurIPS 2015)以降、プルーニングは通常5段階に整理されます: ① 学習 → ② granularity(どのように) → ③ criterion(何を) → ④ ratio(どれだけ) → ⑤ fine-tune(回復). この記事ではそのうち核心である criterion → ratio/回復 → granularity の順に見ていきます。

動機は量子化編と同じです。重みはたいてい0付近に集中しており、メモリアクセスは演算より数百倍高価です(Horowitz: DRAM 640pJ vs int ADD 0.1pJ)。だから「多くの重みがほぼ0なのに、本当にすべて必要なのか?」という問いが自然に湧いてきます。

1. 何を刈るか — Criterion

まず最初に答えるべき問いは「どの重みを消すのか」です。候補は3つあります。

  • Magnitude(大きさ): 重要度 $= |w|$。小さければ消す。安くて単純。
  • OBD saliency(2次/ヘシアン): OBDの答えです。各重みを消したときに損失がどれだけ増えるかをテイラー展開で近似します。
  • Random: ランダム。比較用の下限。

OBDの導出はこうです。損失 $L$ を重みの変化 $\delta w$ について2次まで展開すると:

$$\delta L = \sum_i g_i\,\delta w_i + \frac{1}{2}\sum_i h_{ii}\,\delta w_i^2 + \frac{1}{2}\sum_{i\neq j} h_{ij}\,\delta w_i \delta w_j + O(\lVert\delta w\rVert^3)$$

ここで3つの仮定を入れます — 収束完了(1次項 $g_i \approx 0$)、対角近似(交差項 $h_{ij}$ を無視)、2次近似(3次以上を無視)。すると、重み一つを0にするとき($\delta w_i = -w_i$)の損失増加、すなわちsaliencyがきれいに残ります。

$$\text{saliency}_{w_i} = \frac{1}{2}\, h_{ii}\, w_i^2$$

大きさの大きい重み($w_i^2$)ほど、そして損失曲率が急なところ($h_{ii}$)ほど重要だという意味です。問題はヘシアン $h_{ii}$ が高価だということ — そこで実務では対角ヘシアンを**empirical Fisher(勾配の二乗の平均)**で近似します。私たちの実験もそうしました。

実測: 3つの基準を比較

MNIST MLP(784→512→512→10)を学習したのち、3つの基準でそれぞれsparsityを上げていきながら(再学習なしで)精度を測りました。

One-shotプルーニング: OBD・magnitude・randomの精度 vs sparsity

読み取るべき核心:

  • randomは50%を超えただけで崩壊します。 何を刈るかが決定的だという証拠です。
  • magnitudeとOBDはほぼ重なります。 どちらも90%まで精度を保ちます。極端(98%)ではOBDの2次の洞察が少し先行しますが(44.6% vs 29.3%)、その手前まではmagnitudeがむしろわずかに勝ります。

下で直接触ってみてください。スライダーでsparsityを上げると、切られたシナプスは右のニューラルネットから消え、接続が1つも残らないニューロンは除去されます。同時に演算量(MAC・FLOPs)も減ります — ただしこの削減は、ハードウェアが実際にゼロをスキップする場合にのみ実現します。基準をmagnitude ↔ randomに変えると、同じsparsity(同じMAC数)でも、どの接続・ニューロンが先に消えるかが変わります。

この結果こそが**「なぜ実務はOBDではなくmagnitudeを使うのか」**への答えです。OBDは理論的に優雅ですが、ヘシアンが高価です。一方でmagnitudeはタダなのに、これだけ強力です。だから36年経った今でも、ほとんどのプルーニングはmagnitudeベースです。OBDが開いた扉(「何を刈るかを原理的に選ぼう」)を、実際には最も単純な基準が実用的にくぐり抜けたわけです。

2. どれだけ刈れるか — 再学習と圧縮の驚き

一度に大量に刈れば、当然精度は落ちます。ですが**刈って → 再学習(fine-tune)**を繰り返すと、話はまったく変わってきます。再学習の学習率は通常、元の1/10〜1/100と低く設定します。

プルーニング+再学習: 驚くほど高いsparsityまで精度を維持

  • prune only(グレー): 再学習なしで刈ると95%で88%まで落ちます。
  • iterative prune+retrain(グリーン): 少しずつ刈りながら毎回再学習すれば、**重みを95%刈り取っても(=20分の1のパラメータ)、denseと事実上同じ精度(97.3% vs 97.6%)**を保ちます。

重みの20分の1しか残さなくても性能がそのままだという事実こそが、プルーニングが今なお研究され続ける理由です。プルーニング後の重み分布を見れば、何が起きたのか一目でわかります — ちょうど0に巨大なスパイクが立ちます。

90%プルーニング後の重み分布: 0にスパイク

この「刈って再び学習」というアイデアを極限まで推し進めたのがDeep Compression(Han et al., 2016)です。iterative pruningでAlexNetのプルーニング限界を5倍から9倍まで引き上げ、量子化・ハフマン符号化まで載せて最大数十倍を達成しました。

3. どう刈れば速くなるか — Granularityとハードウェア

ここで必ず押さえておくべき落とし穴があります。プルーニングでsparsityを得たからといって、勝手に速くなるわけではありません。

じっくり考えてみると、重みを「刈る」というのは、実はその場所に0を入れることと変わりありません。そして0をそのまま乗算-累算器に入れて計算すると — $0 \times a = 0$ — 演算回数は一つも減りません。 上で精度を測った実験も、実は重みを0にしたのちそのままdenseカーネルで回したもので、精度は測れましたが、記憶も速度もまったく減りませんでした(量子化編の「fake quantization」とまったく同じ状況です)。0を掛けようが0.3を掛けようが、乗算器の立場では同じ一回の乗算です。

利益を得るには、0をそもそも記憶もせず、掛けもせずスキップしなければなりません。 ところがそれが難しいのです。ここまでのプルーニングは0がどこにでも散らばるunstructured(非構造)方式なので、圧縮率は最高でもGPUの立場では0の位置が不規則でスキップできません。 「次の0がどこにあるか」を毎回探すコストが、乗算一回を省くより高くつくので、結局denseで回すほうがマシになります。

そこでプルーニングにはこの原則が付いてまわります:

プルーニングは、ハードウェアがそのsparsityパターンを支援して初めて実利益が出る。使うハードウェアが何を加速するのかを知って刈らなければならない。

これまでsparse演算のための専用ハードウェア(EIE, ISCA 2016など)が研究されてきましたが、今日もっとも広く使われている答えがNVIDIAの2:4 structured sparsityです。

2:4プルーニング

規則は単純です。連続する4個の重みごとに、ちょうど2個を0にします(50% sparsity)。0の位置が「4個中2個」と規則的なため、NVIDIA Ampere(A100)以降のSparse Tensor Coreがこのパターンをハードウェアで加速します。記憶も効率的です — 非零値の半分 + 各値の位置を指す2-bitインデックスさえあれば済みます。

magnitudeで4列ごとに小さい2個を刈るマスクは、こう作ります。

def two_four_mask(weight):
    """4列単位で|w|が大きい2個だけ残す2:4マスク。"""
    w = weight.detach().abs(); R, C = w.shape
    mask = torch.zeros_like(weight)
    for j in range(0, C - C % 4, 4):
        idx = w[:, j:j+4].topk(2, dim=1).indices   # 各4個のうち大きい2個
        mask[:, j:j+4].scatter_(1, idx, 1.0)
    return mask

私たちのMLPに2:4を適用した結果: dense 97.6% → 2:4(再学習なし)96.8% → 2:4 + 再学習 98.2%。半分を規則的に刈っても、再学習すれば完全に回復します。

Pythonで2:4を活用する

PyTorchは2:4パターンを圧縮保存し、Sparse Tensor Coreで掛けるAPIを提供します(fp16、Ampere+が必要)。

from torch.sparse import to_sparse_semi_structured

mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).bool()
W = (torch.rand(64, 64) * mask).half().cuda()
Ws = to_sparse_semi_structured(W)          # 非零値 + 2-bitインデックスで圧縮
x  = torch.rand(64, 64).half().cuda()

torch.mm(Ws, x)                            # Sparse Tensor Coreを使用

本当に速くなるか — 直接ベンチマーク

A100でdense vs 2:4 sparseの行列積の速度をサイズ別に測りました。結果は正直で、示唆に富んでいます。

A100での2:4 sparse vs dense fp16行列積の速度

  • 小さい行列(N < 4096)では2:4がむしろ遅いです。 圧縮・インデックシングのオーバーヘッドが利益より大きいのです。
  • 大きい行列でのみ利益が出ます。 N=8192で約1.72倍速くなります。

これが「ハードウェアとワークロードを知らなければならない」の実体です。同じ2:4プルーニングでも、Ampere以上のGPUで、十分に大きい行列でのみ速くなります。Turing(T4)や小さい行列では利益がないか、むしろ損です。プルーニングはアルゴリズムだけの問題ではなく、最初から最後までハードウェアとともに設計しなければならない問題です。

4. プルーニング × 量子化は直交する — Deep Compression

ここまでプルーニングを見てきましたが、先の量子化編ではビットを削りました。核心はこの二つが互いに直交(orthogonal)するということです。プルーニングは「どの重みを消すか」を、量子化は「残った重みを何ビットで表現するか」を決めます。軸が違うので一緒に使えば利益が掛け算されます。

これを最もよく示したのがDeep Compression(Han et al., 2016)の3段階パイプラインです。

  1. プルーニング + 再学習 — 弱い接続を消し、再び学習して回復。
  2. 量子化 + 再学習 — 生き残った重みをK-Meansコードブック数個にまとめ(コードブックを再学習)、インデックスだけを保存。
  3. ハフマン符号化 — インデックス分布が片側に偏っているので、無損失でもう一度圧縮。

私たちのMLPで実際に回してみました。各段階でfc2の重み分布がどう変わるかを見ると、プルーニングと量子化が分布の異なる部分を触っていることが一目でわかります。(0になった重みは隠し、生き残ったものだけを描きました。)

重み分布の進化: 正規 → プルーニング(中間を除去) → 再学習(2モード) → 量子化(離散レベル)

  • ① dense — おなじみの釣鐘型の正規分布。
  • ② プルーニング(80%) — 0近傍の小さな重みが丸ごと消えます。真ん中が吹き飛び、両側の裾だけが残ります。
  • ③ 再学習 — 生き残った重みがデータに再び合わせられ、はっきりとした2モード(bimodal)分布になります。
  • ④ 量子化(4-bit) — その連続分布が16個の離散レベルに畳み込まれます(ここが[量子化編]のK-Meansがやること)。

プルーニングは横軸の真ん中(小さい値)を空にし、量子化は残った値をいくつかの縦線に離散化します。触る場所が違うので重ならずに積み上がります。その結果、モデルサイズは段階ごとに掛け算で縮みます。

モデルサイズ縮小: dense → プルーニング(4.3×) → +量子化(17.8×) → +ハフマン(18.6×)

  • dense 2,612 KB → プルーニング 604 KB (4.3×) → +量子化 147 KB (17.8×) → +ハフマン 141 KB (18.6×)
  • その間、精度はむしろ 97.5% → 98.3% へと維持・改善されました(再学習のおかげ)。

つまりプルーニングだけで4.3×、量子化を載せて17.8×、ハフマンまでで18.6×。各手法の利益は足されるのではなく掛け算されます。 原論文はこのパイプラインでAlexNetを35×、VGG-16を49×縮めました(精度を落とさずに)。データ型 → 量子化 → プルーニング、このシリーズが扱った3つの軸は、結局一つのモデルに同時に載せられる直交したツールなのです。

自分で回してみる

上のすべてのplotは、以下のノートブックをColabで実行して出たものです。ランタイム → すべて実行でOKです(2:4速度測定はA100などAmpere GPUを選択)。

まとめ

  • 何を: randomはすぐに崩壊する。magnitude(|w|)は安くて強力で、OBDの2次saliencyに匹敵する — だから実務の標準になった。
  • どれだけ: prune→retrainを繰り返せば、95%(20倍)を刈っても精度が維持される。
  • どのように: unstructuredは圧縮率だけが高い。2:4 structuredでこそAmpere Sparse Tensor Coreが実際に加速する — ただし、大きい行列でのみ。
  • 一緒に: プルーニングと量子化は直交する。Deep Compressionのように積み重ねれば利益が掛け算される(私たちの実測で18.6×、精度維持)。

Optimal Brain Damageが1989年に蒔いた洞察は、今なお有効です。ただしそれが「実利益」になるためには、36年後のハードウェアがそのsparsityを受け止めてくれる必要がありました。プルーニングの歴史は、すなわちアルゴリズムとハードウェアが互いを待ち続けてきた歴史でもあるのです。

次回は、sparsityを層ごとに自動的に決める方法(sensitivity analysis・AMC・NetAdapt)を扱っていきます。

References

  • Y. LeCun, J. S. Denker, S. A. Solla. Optimal Brain Damage. Advances in Neural Information Processing Systems (NeurIPS), 1989.
  • S. Han, J. Pool, J. Tran, W. J. Dally. Learning both Weights and Connections for Efficient Neural Networks. NeurIPS, 2015.
  • S. Han, H. Mao, W. J. Dally. Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding. ICLR, 2016.
  • A. Mishra et al. Accelerating Sparse Deep Neural Networks (2:4 Structured Sparsity). arXiv:2104.08378, 2021.
  • S. Han et al. EIE: Efficient Inference Engine on Compressed Deep Neural Network. ISCA, 2016.
  • M. Horowitz. Computing’s Energy Problem (and what we can do about it). ISSCC, 2014.
  • 講義参考: MIT 6.5940 TinyML and Efficient Deep Learning Computing (Song Han), Lec 3–4 Pruning & Sparsity.