Pruning

Pruning ist eine Modelloptimierung, bei der Gewichte, Neuronen, Filter, Attention-Heads oder ganze Schichten entfernt werden, die wenig zur Ausgabe eines trainierten neuronalen Netzes beitragen. So entsteht ein kleineres Modell, das weniger Speicher, Rechenleistung und Energie benötigt. Im Deep Learning ist der Begriff gleichbedeutend mit Modell-Pruning. Die meisten Netze sind stark überparametrisiert: Han et al. reduzierten AlexNet von 61 Millionen auf 6,7 Millionen Parameter und VGG-16 von 138 Millionen auf 10,3 Millionen auf ImageNet, ohne die Genauigkeit zu verringern. Dazu trainierten sie die Modelle, entfernten Verbindungen mit geringer Gewichtsstärke und trainierten die übrigen erneut. Pruning ist ein gängiger Schritt vor der Bereitstellung von Edge-KI, bei der die Budgets für Speicher, Leistung und Inferenzlatenz knapp bemessen sind.

So funktioniert Pruning#

Jedes Pruning-Verfahren entscheidet, welche Parameter entfernt werden und wann das geschieht. Als Maß für die Wichtigkeit dient meist der Betrag des Gewichts, da Gewichte nahe null die Ausgabe am wenigsten verändern. Gradientenbasierte Maße schätzen den Einfluss jedes Parameters auf die Verlustfunktion direkter. Das gleichzeitige Entfernen eines großen Anteils beeinträchtigt die Genauigkeit meist stärker als iteratives Pruning: Dabei wird ein kleiner Anteil entfernt, die verbleibenden Gewichte werden zur Wiederherstellung der Genauigkeit feinabgestimmt, und der Vorgang wird bis zum Erreichen des Größen- oder Geschwindigkeitsziels wiederholt. Die Lottery-Ticket-Hypothese besagt außerdem, dass dichte Netze dünn besetzte Teilnetze enthalten, die nach dem Zurücksetzen auf ihre ursprüngliche Initialisierung eigenständig eine Genauigkeit erreichen können, die mit dem vollständigen Netz vergleichbar ist.

Unstrukturiertes, strukturiertes und halbstrukturiertes Pruning#

Unstrukturiertes Pruning setzt einzelne Gewichte an beliebigen Stellen eines Tensors auf null. Bei gleicher Genauigkeit erreicht es die höchste Sparsity, doch das Ergebnis ist eine unregelmäßige Sparse-Matrix, die dichte GPU- und CPU-Kernels weiterhin mit vollem Aufwand verarbeiten, sofern eine Sparse-Laufzeitumgebung die Nullen nicht überspringt. Strukturiertes Pruning entfernt ganze Filter, Kanäle, Attention-Heads oder Schichten aus einem Convolutional Neural Network oder Transformer. Dadurch wird jede beschnittene Schicht zu einer kleineren dichten Schicht, die auf jeder Hardware schneller läuft, allerdings um den Preis eines größeren Genauigkeitsverlusts pro entferntem Parameter. Die Pruning-Werkzeuge von PyTorch wenden beide Arten als Masken auf eine bestehende Schicht an:

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")

Beide Schichten verlieren 50 % ihrer Gewichte. Doch nur bei der strukturierten Schicht sind anschließend 16 ihrer 32 Filter vollständig leer. Diese Filter können physisch entfernt werden, sodass eine dichte Faltungsschicht mit 16 Filtern übrig bleibt.

Strukturiertes Pruning liegt zwischen den beiden Ansätzen. Beim 2:4-Sparsitätsmuster von NVIDIA werden zwei von jeweils vier aufeinanderfolgenden Gewichten auf null gesetzt. Sparse Tensor Cores in Ampere-GPUs und neueren Modellen überspringen diese Gewichte und führen dieselbe Matrixberechnung in der halben Zeit aus. NVIDIA maß bei einem dünnbesetzten ResNeXt-101 in TensorRT auf einer A100 eine Ende-zu-Ende-Beschleunigung von bis zu 20 %. Auf dichter Hardware bringt das bloße Verteilen von Nullen keinen Geschwindigkeitsvorteil: Das Pruning von Ultralytics YOLOv5x auf 30 % globale unstrukturierte Sparsität mit dem Hilfsprogramm prune() senkte ohne weiteres Training den COCO-mAP_50-95-Wert von 0,507 auf 0,489. Die Inferenzzeit auf der V100 blieb mit 5,2 ms pro Bild unverändert.

Anwendungen in der Praxis#

Pruning ist besonders wichtig, wenn ein Modell auf eingeschränkter Hardware ausgeführt werden muss. Kanalbeschnittene Objekterkennungs netze passen in die Speicher- und Leistungsbudgets von Drohnen, intelligenten Kameras und Smartphones. Das ermöglicht Echtzeitinferenz ohne Cloud-Verbindung und sorgt dafür, dass Bilddaten zum Schutz der Datenschutz auf dem Gerät bleiben. Bei der Inspektion in der Fertigung ermöglicht dieselbe Komprimierung den Betrieb eines Fehlerdetektors auf einem kostengünstigen Edge-Modul neben der Produktionslinie.

Auch große Sprachmodelle werden beschnitten. SparseGPT reduziert OPT-175B und BLOOM-176B in einem Schritt ohne erneutes Training auf 60 % unstrukturierte Sparsity, wobei die Perplexität nur vernachlässigbar steigt. Bei NVIDIAs Minitron-Arbeit wurde Llama 3.1 8B durch Entfernen von Schichten oder verborgenen Dimensionen auf ein 4B-Modell reduziert. Anschließend wurde die Qualität mithilfe von Wissensdestillation aus dem ursprünglichen Modell wiederhergestellt.

Pruning im Vergleich zu verwandten Verfahren#

Pruning wird häufig mit anderen Komprimierungsverfahren kombiniert. Der Begriff hat außerdem außerhalb neuronaler Netze eine ältere Bedeutung.

• Modellquantisierung: Bei der Quantisierung bleiben alle Gewichte erhalten, aber ihre numerische Genauigkeit wird verringert, zum Beispiel von 32-Bit-Gleitkommazahlen auf 8-Bit-Ganzzahlen. Beschnittene Modelle werden oft anschließend quantisiert, um weitere Einsparungen zu erzielen.
• Wissensdestillation: Bei der Destillation wird ein separates, kompaktes Schülermodell darauf trainiert, ein größeres Lehrermodell nachzuahmen. Sie ist eine gängige Methode, um die Genauigkeit wiederherzustellen, die ein stark beschnittenes Modell verliert.
• Dropout-Schicht: Dropout setzt zur Regularisierung während eines einzelnen Trainingsschritts zufällige Aktivierungen auf null und lässt die Parameteranzahl unverändert. Pruning entfernt dagegen Parameter dauerhaft.
• Pruning von Entscheidungsbäumen: Beim klassischen maschinellen Lernen werden Zweige eines angepassten Entscheidungsbaums entfernt, um Overfitting zu verringern. Ein Beispiel ist das Cost-Complexity-Pruning in scikit-learn, das mit steigendem Wert des Parameters ccp_alpha die schwächsten Teilbäume entfernt.