Pruning

プルーニングはモデル最適化の手法であり、学習済みのニューラルネットワークの出力への寄与が小さい重み、ニューロン、フィルター、アテンションヘッド、またはレイヤー全体を削除します。これにより、メモリ、計算量、エネルギーの消費を抑えられる小型モデルが得られます。ディープラーニングでは、この用語はモデルプルーニングと同義です。ほとんどのネットワークはパラメーター数が大幅に過剰です。Hanらは、学習後に値の小さい結合を削除し、残りを再学習することで、ImageNetにおいて精度を損なうことなく、AlexNetのパラメーター数を6,100万から670万に、VGG-16のパラメーター数を1億3,800万から1,030万に削減しました。プルーニングは、メモリ、電力、推論レイテンシの予算が限られるエッジAIへのデプロイ前によく行われます。

プルーニングの仕組み#

どのプルーニング手法でも、削除するパラメーターと削除するタイミングを決めます。重要度のスコアには通常、重みの大きさを使います。ゼロに近い重みほど出力への影響が小さいためです。一方、勾配ベースのスコアは、各パラメーターが損失関数に与える影響をより直接的に推定します。多数のパラメーターを一度に削除すると、反復的なプルーニングより精度が大きく低下する傾向があります。反復的なプルーニングでは、少量ずつ削除し、残った重みをファインチューニングして精度を回復させ、サイズや速度の目標を達成するまでこれを繰り返します。宝くじ仮説によると、密なネットワークには、元の初期値にリセットすると単独で学習でき、ネットワーク全体に匹敵する精度を得られる疎なサブネットワークが含まれています。

非構造化、構造化、半構造化プルーニング#

非構造化プルーニングでは、テンソル内の任意の位置にある個々の重みをゼロにします。同じ精度で最も高いスパース性を実現できますが、その結果は不規則な疎行列となります。疎なランタイムがゼロをスキップしない限り、密なGPUやCPUのカーネルは、通常どおり全計算を行います。構造化プルーニングでは、畳み込みニューラルネットワークやTransformerからフィルター、チャネル、アテンションヘッド、レイヤー全体を削除します。そのため、プルーニング後の各レイヤーはより小さな密なレイヤーとなり、どのハードウェアでも高速に動作しますが、削除したパラメーターあたりの精度低下は大きくなります。PyTorchのプルーニングユーティリティでは、既存のレイヤーにマスクを適用して、どちらの方式も実行できます。

import copy

from torch import nn
from torch.nn.utils import prune

# Two copies of one convolution layer with 32 filters of shape 16x3x3
unstructured = nn.Conv2d(16, 32, kernel_size=3)
structured = copy.deepcopy(unstructured)

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

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

for name, layer in [("Unstructured", unstructured), ("Structured", structured)]:
    zeros = (layer.weight == 0).float().mean()
    empty_filters = (layer.weight.flatten(1).abs().sum(dim=1) == 0).sum()
    print(f"{name}: {zeros:.0%} of weights zeroed, {empty_filters}/32 filters entirely empty")

どちらのレイヤーも重みの50%を失いますが、構造化レイヤーでは32個のフィルターのうち16個が完全に空になります。この16個は物理的に削除でき、16フィルターの密な畳み込みレイヤーを残せます。

半構造化プルーニングは、この2つの中間に位置します。NVIDIAの2:4スパース性パターンでは、連続する4つの重みのうち2つをゼロにします。Ampere以降のGPUに搭載されたSparse Tensor Coreはこれらをスキップし、同じ行列演算を半分の時間で実行します。NVIDIAは、A100上でTensorRTを使用したスパースResNeXt-101について、エンドツーエンドで最大20%の性能向上を測定しました。密なハードウェアでは、ゼロを散在させるだけでは速度は上がりません。prune()ユーティリティを使ってUltralytics YOLOv5xをプルーニングし、ファインチューニングなしでグローバルな非構造化スパース性を30%にすると、COCO mAP_50-95は0.507から0.489に低下し、V100の推論時間は画像あたり5.2 msのままでした。

実際のアプリケーション#

制約のあるハードウェア上でモデルを動作させる必要がある場合、プルーニングが特に重要です。チャネルプルーニングを行った物体検出ネットワークは、ドローン、スマートカメラ、スマートフォンのメモリと電力の予算内に収まります。これにより、クラウド接続なしでリアルタイム推論が可能になり、データのプライバシーを守るため画像をデバイス上に保持できます。製造工程の検査でも、同じ圧縮によって、低コストのエッジモジュール上で生産ライン脇の欠陥検出器を実行できます。

大規模言語モデルもプルーニングされます。SparseGPTは、再学習せずにOPT-175BとBLOOM-176Bを一度の処理で非構造化スパース性60%までプルーニングし、パープレキシティの増加をわずかに抑えました。NVIDIAのMinitronの研究では、レイヤーまたは隠れ次元を削除してLlama 3.1 8Bを4Bモデルにプルーニングし、その後、元のモデルからの知識蒸留によって品質を回復しました。

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

プルーニングはほかの圧縮手法と組み合わせて使われることが多く、ニューラルネットワーク以外では、この用語に古くからある別の意味もあります。

• モデル量子化：量子化ではすべての重みを保持し、数値精度を下げます。たとえば、32ビット浮動小数点数を8ビット整数に変換します。プルーニングしたモデルは、さらなる削減のためにその後量子化されることがよくあります。
• 知識蒸留：蒸留では、独立した小型の生徒モデルを学習させ、大型の教師モデルを模倣させます。大幅なプルーニングによって失われた精度を回復するためによく使われます。
• ドロップアウトレイヤー：ドロップアウトでは、正則化のために1回の学習ステップ中だけランダムな活性化をゼロにします。パラメーター数は変わりませんが、プルーニングではパラメーターを恒久的に削除します。
• 決定木のプルーニング：従来の機械学習では、学習済みの決定木から枝を削除して過学習を抑えます。scikit-learnのコスト複雑度プルーニングもその一例で、ccp_alphaパラメーターが大きくなるにつれて、影響の小さいサブツリーを削除します。