💬 NLP

TraceNAS: Zero-shot LLM Pruning via Gradient Trace Correlation

TraceNAS est un cadre de recherche d'architecture neuronale hautement efficace et sans entraînement qui exploite la corrélation des traces de gradient pour identifier des modèles de langage de grande taille élagués de manière non uniforme et optimale en alignant leurs paysages de perte avec ceux des modèles pré-entraînés originaux, atteignant ainsi des performances compétitives avec les méthodes tenant compte de l'entraînement pour une fraction du coût computationnel.

Auteurs originaux : Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

Publié 2026-02-04

📖 5 min de lecture🧠 Analyse approfondie

CC BY 4.0

Auteurs originaux : Prajna G. Malettira, Manish Nagaraj, Arjun Roy, Shubham Negi, Kaushik Roy

Article original sous licence CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Ceci est une explication générée par l'IA de l'article ci-dessous. Elle n'a pas été rédigée ni approuvée par les auteurs. Pour une précision technique, consultez l'article original. Lire la clause de non-responsabilité complète

Imaginez que vous possédez une bibliothèque massive et incroyablement intelligente (un grand modèle de langage, ou LLM) qui contient la somme des connaissances humaines. Elle est si grande et lourde qu'il est impossible de la transporter ou de l'utiliser sur un téléphone ordinaire. Vous voulez la réduire à une version de poche sans perdre sa capacité à raconter des histoires, à résoudre des énigmes ou à comprendre les blagues.

C'est le problème que le papier TraceNAS tente de résoudre.

Le Problème : Couper les mauvaises choses

D'habitude, quand on essaie de rétrécir ces modèles géants, on agit comme un jardinier maladroit. On pourrait dire : « Coupons 50 % des branches de chaque arbre », ou « Retirons les feuilles les plus lourdes ». C'est ce qu'on appelle l'élagage uniforme (uniform pruning).

Mais voici le piège : toutes les parties du cerveau (ou de la bibliothèque) ne sont pas égales. Certaines parties sont des « organes vitaux » qui détiennent la logique complexe, tandis que d'autres ne sont que du « gras » qui peut être facilement taillé. Si vous coupez le mauvais organe vital, le modèle oublie tout. Si vous ne coupez que le gras, il reste trop lourd.

Les méthodes existantes tentent de déterminer ce qu'il faut couper soit en :

Regardant une partie à la fois : Elles vérifient si un neurone spécifique est important, mais elles ignorent comment ce neurone communique avec le reste du cerveau.
Entraînant le modèle pendant la coupe : Elles essaient de rétrécir le modèle et de lui réapprendre simultanément comment penser. C'est comme essayer d'apprendre à conduire une voiture tout en reconstruisant simultanément le moteur. Cela prend un temps infini et nécessite une quantité massive de carburant (puissance de calcul).

La Solution : TraceNAS (Le détective de la « Trace de Gradient »)

Les auteurs proposent une nouvelle méthode appelée TraceNAS. Considérez cela comme un scanner de haute technologie par rayons X capable de regarder l'« âme » du modèle sans réellement le toucher ou le réenseigner.

Voici comment cela fonctionne, en utilisant des analogies simples :

1. L'« Écho » de l'esprit original

Imaginez que le modèle original, géant, est un chef cuisinier expert qui a préparé un repas parfait. Lorsque vous prenez une bouchée, la saveur laisse une « trace » spécifique sur votre langue.

L'ancienne méthode : Pour voir si une nouvelle recette plus petite fonctionne, vous devriez cuisiner tout le plat, le goûter, et si c'est mauvais, tout recommencer.
La méthode TraceNAS : Au lieu de cuisiner tout le plat, TraceNAS examine les ingrédients et les étapes de la recette (les gradients) pour prédire si la saveur sera la bonne. Il vérifie si la « trace de saveur » de la version réduite et élaguée correspond à la « trace de saveur » du grand chef original.

2. Le test de l'« Ombre » (Corrélation de la trace de gradient)

Le papier utilise un concept mathématique appelé Corrélation de la trace de gradient.

Imaginez que le modèle original est un navire géant naviguant sur un océan calme. Il laisse un sillage spécifique (une traînée de vagues) derrière lui.
Lorsque vous essayez de construire un bateau plus petit (un modèle élagué), TraceNAS demande : « Est-ce que ce plus petit bateau laisse un sillage qui ressemble exactement au sillage du grand navire ? »
Si les sillages correspondent, cela signifie que le petit bateau suit le même chemin et atteindra probablement la même destination (performera bien) une fois qu'il aura un peu de pratique. Si les sillages sont chaotiques, le bateau va s'échouer.

3. Le raccourci « Low-Rank » (Bas Rang)

Calculer ces sillages pour un navire géant nécessite normalement un supercalculateur. TraceNAS est astucieux : il réalise que le mouvement du navire se produit principalement dans quelques directions principales. Il utilise une astuce de bas rang (Low-Rank) pour n'observer que les directions les plus importantes du sillage.

Analogie : Au lieu de mesurer chaque ride à la surface de l'océan, il ne mesure que les trois plus grosses vagues. Cela leur permet de réaliser le test sur une seule carte graphique (GPU) en seulement 8,5 heures, alors que les anciennes méthodes auraient pris des jours ou des semaines sur un cluster entier d'ordinateurs.

Les Résultats : Rapide, Économique et Intelligent

Le papier a testé cela sur des modèles célèbres comme Llama et Qwen.

Vitesse : Ils ont trouvé la « coupe » optimale en 8,5 heures sur un seul ordinateur. Les anciennes méthodes prenaient 10 fois plus de temps et consommaient 10 fois plus d'énergie.
Précision : Les petits modèles résultants sont aussi performants que les modèles qui ont été entraînés pendant des semaines pour trouver les meilleures coupes.
Non-uniformité : Contrairement au « jardinier maladroit » qui coupe tout de manière égale, TraceNAS a trouvé un « ajustement sur mesure ». Il a préservé les parties lourdes et complexes du cerveau tout en ne taillant que le gras, créant ainsi un modèle efficace mais toujours très intelligent.

L'essentiel à retenir

TraceNAS est un outil qui vous permet de réduire un cerveau d'IA géant en une version de poche sans avoir besoin de le réenseigner ou de dépenser une fortune en électricité. Il y parvient en vérifiant si l'« ombre » du petit cerveau correspond à l'« ombre » du grand cerveau. Si les ombres s'alignent, le petit cerveau est prêt à fonctionner.

Cela rend possible l'exécution d'IA puissantes sur des appareils ordinaires sans avoir besoin d'un centre de données massif, tout en économisant une quantité énorme de temps et d'énergie.

Résumé Technique : TraceNAS

Définition du Problème

Le déploiement des modèles de langage de grande taille (LLM) est contraint par des nombres massifs de paramètres et des exigences de calcul élevées. Bien que l'élagage structuré offre une voie vers un déploiement efficace en supprimant des unités architecturales (par exemple, les têtes d'attention, les blocs de transformateurs), les méthodes existantes font face à des limitations significatives :

Dépendances Locales vs Globales : La plupart des méthodes d'élagage évaluent l'importance des couches, des têtes ou des canaux de manière isolée en utilisant des heuristiques locales (par exemple, les magnitudes pondérées par l'activation ou la reconstruction basée sur le Hessien). Ces approches ignorent les dépendies structurelles globales complexes et la propagation des erreurs induites par l'élagage à travers le modèle.
Coût de Calcul de la Recherche Sensible à l'Entraînement : Les méthodes qui tentent de capturer les dépendances globales reposent souvent sur une recherche d'architecture neuronale (NAS) sensible à l'entraînement. Ces approches nécessitent des mises à jour de gradients itératives et des flux continus de jetons (tokens) inédits pour éviter le surapprentissage, rendant le processus de recherche aussi coûteux que l'entraînement de récupération post-élagage lui-même.
Inadéquation des Proxys Zero-Shot Existants : Les proxys actuels sans entraînement sont conçnés pour identifier des architectures entraînables et initialisées aléatoirement en se concentrant sur la capacité d'apprentissage (par exemple, la stabilité du gradient, la lissité du paysage de perte). Ils échouent à évaluer l'« héritage fonctionnel » — la capacité d'un modèle élagué à récupérer ses performances en préservant le paysage de perte préentraîné.

Méthodologie

TraceNAS est proposé comme un cadre de NAS sans entraînement conçu pour optimiser conjointement la profondeur et la largeur des LLM sans entraînement lors de la recherche.

1. Espace de Recherche et Encodage

Le cadre traite le LLM préentraîné comme un super-réseau. L'espace de recherche $G$ est défini comme une distribution conjointe sur :

Profondeur ( $d$ ) : Un masque discret sur $L$ blocs de transformateurs, permettant une réduction de profondeur non uniforme tout en maintenant la continuité du flux résiduel.
Largeur ( $\kappa$ ) : Des ratios de rétention de paramètres pour les sous-blocs d'attention et de MLP, permettant une parcimonie hétérogène à travers le modèle.

2. Réalisation Architecturale In-Place

Pour éviter la surcharge de mémoire liée à l'instanciation de milliers de modèles candidats, TraceNAS emploie une stratégie de masquage in-place :

Masquage de Largeur : Utilise un produit pondéré par l'activation (similaire à Wanda) pour générer des masques binaires pour les canaux de poids. Cette heuristique en $O(d^2)$ identifie les canaux fonctionnellement critiques tout en préservant les caractéristiques aberrantes (outliers).
Masquage de Profondeur : Contourne des blocs entiers s'ils sont désactivés, acheminant les activations via les connexions résiduelles.
Exécution : Le cadre modifie temporairement les matrices de poids du modèle de base pour calculer les gradients sur un ensemble de calibration minimal, puis restaure les poids originaux, préservant ainsi le modèle de base pour les itérations suivantes.

3. Le Proxy Zero-Shot : Corrélation de la Trace de Gradient

L'innovation centrale est un proxy invariant d'échelle $\Phi$ qui mesure l'alignement fonctionnel entre un candidat élagué ( $M_{sub}$ ) et le modèle de base préentraîné ( $M_{base}$ ).

Traces de Gradient : Les gradients sont calculés sur un ensemble de calibration. Pour garantir l'extensibilité, ceux-ci sont projetés dans un sous-espace de faible rang (LoRA), exploitant la faible dimensionnalité intrinsèque des LLM pour réduire la surcharge de mémoire de plusieurs ordres de grandeur.
Métrique d'Alignement : Le proxy calcule le coefficient de corrélation de Pearson entre les traces de gradient standardisées du modèle élagué et du modèle de base pour chaque couche. Cela garantit l'invariance d'échelle, en se concentrant sur l'alignement directionnel plutôt que sur les changements de magnitude.
Agrégation Pondérée par la Parcimonie : Le score final $\Phi$ est une somme pondérée des corrélations par couche, où les poids correspondent aux ratios de rétention de paramètres. Cela ancre le score global dans les régions à haute capacité, empêissant le bruit provenant de sous-blocs agressivement élagués de fausser le classement.

4. Recherche Évolutive

TraceNAS emploie une stratégie évolutive pour naviguer dans l'espace de recherche discret :

Initialisation : Une population est générée basée sur des priors d'importance globale.
Évolution : La population subit des croisements et des mutations. L'évolution de la profondeur utilise des mutations par inversion de bit et des croisements de plages contiguës, tandis que l'évolution de la largeur utilise l'interpolation et le jitter gaussien.
Sélection : Les candidats sont classés uniquement par le proxy $\Phi$ , sélectionnant les modèles ayant le plus haut potentiel de récupération de performance lors de l'entraînement pré-continué (CPT) post-élagage.

Contributions Clés

NAS Sans Entraînement et Scalable : TraceNAS fournit un cadre unifié pour l'élagage conjoint de la profondeur et de la largeur, permettant la découverte d'architectures non uniformes sans entraînement lors de la recherche.
Proxy d'Héritage Fonctionnel : Introduit un nouveau proxy basé sur le gradient et invariant d'échelle qui évalue l'alignement des traces de gradient entre les modèles élagués et préentraînés. Il atteint une haute corrélation avec la précision en aval (Spearman $\rho = 0,94$ , Kendall $\tau = 0,82$ ).
Efficacité : Le cadre réduit la surcharge de recherche en utilisant des sous-espaces de gradient de faible rang et un masquage in-place, permettant la découverte de modèles de haute fidélité sur un seul GPU.

Résultats Expérimentaux

TraceNAS a été évalué sur Llama-2-7B, Llama-3.1-8B et Qwen-2.5-14B, en les élaguant respectivement à 2,7B, 4,6B et 8,4B paramètres.

Efficacité de Recherche : TraceNAS a identifié les architectures optimales en 8,5 heures GPU (sur un A100), représentant une réduction de 10× du temps de recherche par rapport aux bases sensibles à l'entraînement comme DarwinLM. Il a nécessité seulement 98M de jetons de recherche, soit une réduction de 10× par rapport à DarwinLM.
Performance :
- Llama-2-7B : TraceNAS a atteint la précision moyenne la plus élevée (62,81 %) sur huit benchmarks de raisonnement, surpassant DarwinLM (62,57 %) et ShearedLLaMA (62,63 %). Notamment, la performance de ShearedLLaMA a chuté de manière significative (à 60,86 %) lorsqu'elle était restreinte au même budget de jetons utilisé par TraceNAS.
- Llama-3.1-8B & Qwen-2.5-14B : TraceNAS a surpassé les bases d'élagage uniforme et les méthodes sensibles à l'entraînement (DarwinLM, E3-Pruner) sur des tâches de raisonnement complexes (ex: ARC-Challenge, BoolQ) tout en nécessitant beaucoup moins de jetons de récupération (20B contre 94B pour Minitron).
Validation du Proxy : Le proxy a démontré une grande stabilité à travers diverses hyperparamètres (taille de l'échantillon, longueur du contexte et rang LoRA), confirmant que l'alignement de la trace de gradient capture efficacement l'héritage du modèle, même dans des contextes contraints.

Signification et Revendications

L'article affirme que TraceNAS répond au goulot d'étranglement critique du coût de calcul dans l'élagage structuré. En déplaçant l'accent de la « capacité d'apprentissage » (apprendre à partir de zéro) vers l'« héritage » (préserver le paysage de perte préentraîné), TraceNAS permet la découverte de modèles élagués non uniformes de haute performance avec un minimum de dépenses de ressources.

Les auteurs soulignent que leur méthode :

Élimine le coût prohibitif de l'entraînement lors de la recherche, rendant la compression de modèles à grande échelle accessible.
Capture les dépendances structurelles globales que les heuristiques locales manquent, assurant la préservation des capacités de raisonnement complexes.
Réduit l'empreinte carbone de la compression de modèles en diminuant drastiquement les heures-GPU requises pour la recherche d'architecture.

Ce travail positionne TraceNAS comme une fondation scalable pour le déploiement de LLM efficaces dans des environnements à ressources limitées, tout en maintenant l'intégrité fonctionnelle des modèles denses préentraînés.