Model Pruning

L’élagage de modèles supprime les paramètres qui contribuent peu aux sorties d’un réseau neuronal entraîné, qu’il s’agisse de poids individuels, de filtres entiers, de têtes d’attention ou de couches, afin de rendre le modèle plus petit et moins coûteux à exécuter. La plupart des réseaux sont surparamétrés ; une grande partie de leurs poids peut donc être supprimée avec une faible perte d’exactitude, surtout si le modèle est ensuite ajusté finement. Une méthode d’élagage et de réentraînement de 2015 a réduit AlexNet de 61 millions à 6,7 millions de paramètres et VGG-16 de 138 millions à 10,3 millions, sans perte d’exactitude sur ImageNet. L’élagage est une étape standard de la préparation des modèles pour les téléphones, les caméras et d’autres matériels d’IA de périphérie, et il est désormais également appliqué aux grands modèles de langage (LLM).

Fonctionnement de l’élagage des modèles#

L’élagage évalue l’importance de chaque poids ou structure candidate, supprime la fraction ayant le score le plus faible et ajuste finement les poids du modèle restants pour retrouver de l’exactitude. Le critère le plus simple, et toujours le plus courant, est la magnitude : les poids proches de zéro ou les filtres ayant une petite norme L1 ou L2 modifient le moins la sortie lorsqu’on les supprime. D’autres critères utilisent les gradients ou les activations pour estimer l’augmentation de la perte causée par la suppression de chaque composant. La suppression d’une grande fraction en une fois tend à nuire à l’exactitude ; l’élagage se fait donc généralement par étapes, chacune suivie d’un ajustement fin, jusqu’à ce que le modèle atteigne la taille cible ou que l’exactitude baisse au-delà de la tolérance :

L’élagage a également changé la façon dont les chercheurs envisagent la taille des réseaux. L’hypothèse du ticket de loterie avance qu’un réseau dense contient de petits sous-réseaux qui, entraînés seuls à partir de leurs poids initiaux d’origine, atteignent l’exactitude du réseau complet. Ses auteurs ont systématiquement trouvé de tels sous-réseaux dont la taille était inférieure à 10–20 % de celle du réseau d’origine sur MNIST et CIFAR-10.

Élagage non structuré et structuré#

L’élagage non structuré met à zéro des poids individuels n’importe où dans un tenseur. Il permet d’obtenir une forte parcimonie avec une faible perte d’exactitude, mais les tenseurs conservent leur forme : les noyaux denses standard effectuent donc le même travail, et les gains de vitesse nécessitent des noyaux creux ou une prise en charge matérielle. Depuis la génération Ampere, les GPU NVIDIA exécutent un motif de ce type, la parcimonie 2:4, qui place deux zéros sur chaque groupe de quatre poids consécutifs, à une vitesse de calcul dense pouvant être doublée. L’élagage structuré supprime des filtres, canaux, têtes ou couches entiers et produit un modèle dense plus petit, plus rapide sur n’importe quel processeur, mais avec une perte d’exactitude supérieure pour le même nombre de poids supprimés. La suppression d’un filtre entraîne aussi celle du canal d’entrée correspondant dans la couche suivante, et des bibliothèques comme Torch-Pruning suivent automatiquement ces dépendances. Les outils d’élagage de PyTorch illustrent la différence entre les deux méthodes :

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

Les deux couches se retrouvent avec 50 % de leurs poids à zéro. Seule la couche structurée comporte des filtres vides : 16 sur 32. Leur suppression laisse une couche dense avec moitié moins de filtres, ce qui divise par deux le calcul sur du matériel ordinaire. Le matériel conçu pour la parcimonie bénéficie également de l’élagage non structuré : l’intégration CVflow d’Ambarella pour Ultralytics YOLO, maintenue par Ambarella dans sa propre branche du paquet, élague progressivement les poids pendant l’entraînement de YOLO et conserve cette parcimonie lors de l’exportation ONNX pour la compilation sur ses puces d’IA de périphérie.

Élagage et techniques apparentées#

• Quantification de modèle : Réduit la précision de chaque poids, par exemple en passant de nombres flottants 32 bits à des entiers 8 bits, tout en les conservant tous. Élagage et quantification sont souvent combinés : l’élagage réduit le nombre d’opérations et la quantification rend chacune d’elles moins coûteuse.
• Distillation de connaissances : Entraîne un modèle étudiant plus petit à imiter un modèle enseignant plus grand ; elle suit souvent l’élagage structuré pour rétablir l’exactitude perdue lors de la suppression de filtres ou de couches entiers.
• Abandon : Met temporairement à zéro des activations aléatoires pendant l’entraînement à des fins de régularisation ; le réseau complet est utilisé lors de l’inférence.

Élagage des grands modèles de langage#

Le réentraînement d’un LLM après chaque étape d’élagage est trop coûteux ; les méthodes en une seule passe l’élaguent donc d’un seul coup à l’aide d’un petit jeu de calibration. SparseGPT a réduit les modèles OPT de 175 milliards de paramètres et BLOOM de 176 milliards à une parcimonie non structurée de 60 % en moins de 4,5 heures, avec une augmentation négligeable de la perplexité ; il prend également en charge le motif 2:4 accéléré par les GPU. L’élagage structuré combiné à la distillation de connaissances produit des modèles denses plus petits : le Minitron de NVIDIA a dérivé des modèles 8B et 4B d’un modèle préentraîné 15B en élaguant les canaux d’incorporation, les neurones MLP et les têtes d’attention, puis en les réentraînant avec jusqu’à 40 fois moins de jetons qu’un entraînement depuis zéro.