Model Pruning

モデルプルーニングは、学習済みニューラルネットワークの出力への寄与が小さいパラメーターを削除します。対象は個々の重みから、フィルター、注意ヘッド、レイヤー全体まで及び、モデルを小型化し、実行コストを削減します。多くのネットワークは過剰にパラメーターを持つため、特に後からモデルをファインチューニングすると、精度をほとんど損なわずに多くの重みを削除できます。2015年のプルーニングと再学習の手法では、ImageNetの精度を損なうことなく、AlexNetのパラメーター数を6,100万から670万に、VGG-16を1億3,800万から1,030万に削減しました。プルーニングは、スマートフォン、カメラ、その他のエッジAIハードウェア向けにモデルを準備する標準的な工程であり、現在では大規模言語モデル（LLM）にも適用されています。

モデルプルーニングの仕組み#

プルーニングでは、候補となる各重みや構造の重要度をスコア化し、スコアが最も低い部分を一定の割合で削除します。その後、残ったモデルの重みをファインチューニングして精度を回復します。最も単純で、現在も最も一般的なスコアは大きさです。ゼロに近い重みやL1ノルム、L2ノルムが小さいフィルターは、削除しても出力への変化が最も小さくなります。ほかの基準では、勾配や活性化を使って、各部分を削除した場合に損失がどれだけ増えるかを推定します。一度に大部分を削除すると精度を損ないやすいため、通常はプルーニングとファインチューニングを交互に行い、モデルが目標サイズに達するか、精度の低下が許容範囲を超えるまで繰り返します。

プルーニングは、研究者がネットワークのサイズについて考える方法も変えました。宝くじ仮説では、密なネットワークには、元の初期重みから独立して学習させてもネットワーク全体と同じ精度に達する小さなサブネットワークが含まれるとされています。提唱者らはMNISTとCIFAR-10において、元のサイズの10～20%未満のサブネットワークを一貫して発見しました。

非構造化プルーニングと構造化プルーニング#

非構造化プルーニングでは、テンソル内の任意の位置にある個々の重みをゼロにします。精度をほとんど損なわずに高いスパース性を実現できますが、テンソルの形状は変わらないため、標準的な密カーネルでは同じ計算が必要です。高速化にはスパースカーネルまたはハードウェアのサポートが必要です。Ampere世代以降のNVIDIA GPUでは、連続する4つの重みごとに2つをゼロにする2:4スパース性パターンをサポートし、密な演算の最大2倍の速度で実行できます。構造化プルーニングでは、フィルター、チャネル、ヘッド、レイヤー全体を削除し、どのプロセッサーでも高速に動作する小型の密モデルを作ります。同じ数の重みを削除する場合、精度への影響はより大きくなります。フィルターを削除すると、次のレイヤーにある対応する入力チャネルも削除されます。Torch-Pruningなどのライブラリは、この依存関係を自動的に追跡します。PyTorchのプルーニングユーティリティでは、2つの方式の違いを確認できます。

import torch
from torch.nn.utils import prune

# Two identical convolution layers, one for each pruning style
torch.manual_seed(0)
a = torch.nn.Conv2d(16, 32, kernel_size=3)
b = torch.nn.Conv2d(16, 32, kernel_size=3)

# Unstructured: zero the 50% of individual weights with the smallest magnitude
prune.l1_unstructured(a, name="weight", amount=0.5)

# Structured: zero the 50% of output filters with the smallest L2 norm
prune.ln_structured(b, name="weight", amount=0.5, n=2, dim=0)

for name, layer in [("unstructured", a), ("structured", b)]:
    zeros = (layer.weight == 0).float().mean().item()
    empty = (layer.weight.flatten(1) == 0).all(dim=1).sum().item()
    print(f"{name}: {zeros:.0%} of weights zero, {empty} of 32 filters empty")

どちらのレイヤーも、重みの50%がゼロになります。空のフィルターができるのは構造化レイヤーのみで、32個中16個です。これらを削除するとフィルターが半分の密レイヤーとなり、一般的なハードウェアでの計算量が半分になります。スパース性に対応したハードウェアは、非構造化方式にも有効です。Ambarellaが自社のパッケージフォークで保守している、Ultralytics YOLO向けのCVflow統合では、YOLOの学習中に重みを段階的にプルーニングし、エッジAIチップへのコンパイルに向けたONNXエクスポートでもスパース性を維持します。

プルーニングと関連手法の比較#

• モデル量子化: たとえば32ビット浮動小数点数から8ビット整数へと、すべての重みの精度を下げますが、重み自体はすべて保持します。プルーニングは演算数を減らし、量子化は各演算のコストを下げるため、両者を組み合わせることがよくあります。
• 知識蒸留: 小さな生徒モデルに大きな教師モデルの動作を模倣させます。フィルターやレイヤー全体の削除によって低下した精度を回復するため、構造化プルーニングの後に行われることがよくあります。
• Dropout: 正則化の一種として、学習中にランダムな活性化を一時的にゼロにします。推論ではネットワーク全体を使用します。

大規模言語モデルのプルーニング#

プルーニングを行うたびにLLMを再学習するにはコストがかかりすぎるため、ワンショット手法では小規模なキャリブレーションセットを使い、1回の処理でプルーニングします。SparseGPTは、1,750億パラメーターのOPTと1,760億パラメーターのBLOOMを、パープレキシティの増加を無視できる程度に抑えながら、4.5時間未満で非構造化スパース性60%にプルーニングしました。また、GPUが高速化できる2:4パターンにも対応しています。構造化プルーニングと知識蒸留を組み合わせると、より小さな密モデルを作成できます。NVIDIAのMinitronは、事前学習済みの150億パラメータモデルから80億および40億パラメータのモデルを作成しました。埋め込みチャネル、MLPニューロン、注意ヘッドをプルーニングした後、ゼロから学習する場合より最大40分の1少ないトークン数で再学習しています。