Pruning

L’élagage est une technique d’optimisation de modèle qui supprime les poids, neurones, filtres, têtes d’attention ou couches entières qui contribuent peu à la sortie d’un réseau de neurones entraîné, pour obtenir un modèle plus petit nécessitant moins de mémoire, de calcul et d’énergie. En apprentissage profond, le terme est synonyme d’élagage de modèle. La plupart des réseaux sont fortement surparamétrés : Han et al. ont réduit AlexNet de 61 millions à 6,7 millions de paramètres et VGG-16 de 138 millions à 10,3 millions sur ImageNet, sans perte de précision, en entraînant le modèle, en supprimant les connexions de faible amplitude, puis en réentraînant le reste. L’élagage est une étape courante avant le déploiement de l’IA en périphérie, où les budgets de mémoire, d’énergie et de latence d’inférence sont serrés.

Fonctionnement de l’élagage#

Chaque méthode d’élagage décide quels paramètres supprimer et à quel moment. Le score d’importance correspond généralement à l’amplitude du poids, car les poids proches de zéro modifient le moins la sortie, tandis que les scores fondés sur le gradient estiment plus directement l’effet de chaque paramètre sur la fonction de perte. Supprimer une grande fraction d’un coup coûte généralement plus en précision que l’élagage itératif, qui retire une petite fraction, ajuste finement les poids restants pour récupérer les performances, puis répète l’opération jusqu’à atteindre l’objectif de taille ou de vitesse. L’hypothèse du ticket de loterie ajoute que les réseaux denses contiennent des sous-réseaux clairsemés qui, réinitialisés à leur initialisation d’origine, s’entraînent seuls pour atteindre une précision comparable à celle du réseau complet.

Élagage non structuré, structuré et semi-structuré#

L’élagage non structuré met à zéro des poids individuels n’importe où dans un tenseur. Il atteint la plus forte parcimonie pour une précision donnée, mais produit une matrice creuse irrégulière que les noyaux GPU et CPU denses traitent toujours au coût maximal, sauf si un moteur d’exécution spécialisé saute les zéros. L’élagage structuré supprime des filtres, canaux, têtes d’attention ou couches entiers d’un réseau neuronal convolutif ou d’un transformer ; chaque couche élaguée devient ainsi une couche dense plus petite, plus rapide sur n’importe quel matériel, au prix d’une perte de précision plus importante par paramètre supprimé. Les outils d’élagage de PyTorch appliquent les deux types sous forme de masques sur une couche existante :

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

Les deux couches perdent 50 % de leurs poids, mais seule la couche structurée se retrouve avec 16 de ses 32 filtres entièrement vides, qui peuvent être physiquement supprimés pour obtenir une convolution dense à 16 filtres.

L’élagage semi-structuré se situe entre les deux. Le motif de parcimonie 2:4 de NVIDIA met à zéro deux poids sur quatre consécutifs, que les Sparse Tensor Cores des GPU Ampere et ultérieurs ignorent pour effectuer le même calcul matriciel en deux fois moins de temps ; NVIDIA a mesuré des gains de bout en bout allant jusqu’à 20 % pour un ResNeXt-101 parcimonieux dans TensorRT sur un A100. Sur du matériel dense, des zéros dispersés ne procurent à eux seuls aucun gain de vitesse : l’élagage d’Ultralytics YOLOv5x jusqu’à une parcimonie globale non structurée de 30 % avec son utilitaire prune(), sans réglage fin, a fait passer le mAP_50-95 COCO de 0,507 à 0,489 et laissé le temps d’inférence sur V100 inchangé, à 5,2 ms par image.

Applications concrètes#

L’élagage est particulièrement important lorsqu’un modèle doit s’exécuter sur du matériel aux ressources limitées. Les réseaux de détection d’objets élagués par canal respectent les budgets de mémoire et d’énergie des drones, des caméras intelligentes et des téléphones, ce qui permet une inférence en temps réel sans connexion au cloud et garde les images sur l’appareil afin de préserver la confidentialité des données. Dans l’inspection en fabrication, la même compression permet à un détecteur de défauts de fonctionner sur un module de périphérie peu coûteux, à côté de la ligne de production.

Les grands modèles de langage sont eux aussi élagués. SparseGPT élague OPT-175B et BLOOM-176B jusqu’à 60 % de parcimonie non structurée en une seule opération, sans réentraînement et avec une hausse négligeable de la perplexité. Les travaux Minitron de NVIDIA ont réduit Llama 3.1 8B à un modèle 4B en supprimant des couches ou des dimensions cachées, puis ont récupéré les performances par distillation des connaissances à partir du modèle d’origine.

Élagage et techniques apparentées#

L’élagage est souvent combiné à d’autres méthodes de compression, et le terme a aussi un sens plus ancien en dehors des réseaux de neurones.

• Quantification de modèle : la quantification conserve tous les poids et réduit leur précision numérique, par exemple en passant des nombres à virgule flottante sur 32 bits à des entiers sur 8 bits. Les modèles élagués sont souvent quantifiés ensuite pour réaliser des économies supplémentaires.
• Distillation des connaissances : la distillation entraîne un étudiant compact distinct à imiter un enseignant plus grand ; c’est une méthode courante pour récupérer la précision perdue par un modèle fortement élagué.
• Couche de dropout : le dropout met à zéro des activations aléatoires pendant une seule étape d’entraînement, à des fins de régularisation, et ne modifie pas le nombre de paramètres, tandis que l’élagage supprime définitivement des paramètres.
• Élagage d’arbre de décision : en apprentissage automatique classique, l’élagage supprime des branches d’un arbre de décision ajusté pour réduire le surapprentissage, comme dans l’élagage par complexité-coût de scikit-learn, qui supprime les sous-arbres les plus faibles à mesure que son paramètre ccp_alpha augmente.