Pruning via Causal Attribution Preserves Reasoning Performance in Large Language Models

Résumé Technique : Causal Attribution Pruning (CAP)

Énoncé du Problème
Les grands modèles de langage (LLM) font preuve de solides capacités de raisonnement multi-étapes mais engendrent des coûts d'inférence substantiels, ce qui entrave leur déploiement dans le monde réel. Bien que l'élagage (pruning) des modèles offre une solution en supprimant les paramètres redondants sans réentraînement, les méthodes existantes dégradent souvent les performances sur les tâches de raisonnement complexes. Les approches actuelles reposent sur des signaux corrélationnels, tels que la magnitude des poids (par exemple, Wanda) ou les statistiques de gradient locales. Ces méthodes supposent que les petits poids sont sans importance ; cependant, des paramètres de faible magnitude peuvent être critiques pour les circuits de raisonnement multi-étapes. De plus, les critères purement corrélationnels ne parviennent pas à capturer les interactions inter-couches essentielles au raisonnement de type "chaîne de pensée" (Chain-of-Thought - CoT), menant fréquemment à un "effondrement de couche" où des modules entiers perdent leur fonction.

Méthodologie : Causal Attribution Pruning (CAP)
Les auteurs introduisent le Causal Attribution Pruning (CAP), une méthode d'élagage sans réentraînement et en une seule étape (one-shot), conçue pour préserver la performance de raisonnement. Contra-rément aux méthodes qui classent les poids selon des propriétés statiques, CAP mesure l'impact causal de composants spécifiques du modèle sur la performance de la tâche. Le cadre opère en trois étapes :

1. Calibration et Score Causal (Étape 1) :

   • Un petit jeu de données de calibration ($D_{cal}$) est construit en échantillonnant à partir de benchmarks centrés sur le raisonnement : GSM8K (mathématiques), StrategyQA (sens commun multi-étapes) et ARC-Challenge (sciences).
   • Pour chaque tête d'attention $h_\ell$ du modèle, CAP effectue une expérience de masquage interventionnel. La tête est temporairement masquée lors de la passe directe (forward pass), et l'augmentation résultante de la perte de cross-entropy par jeton ($\Delta(h_\ell)$) est mesurée par rapport au modèle intact.
   • Pour garantir la robustesse face à la variance due aux petits échantillons, le jeu de données est partitionné en $K$ sous-ensembles, et l'augmentation médiane de la perte est calculée.
   • Ce processus identifie les têtes qui sont fonctionnellement critiques pour le raisonnement, les distinguant de celles qui sont simplement numériquement grandes mais redondantes.

2. Conversion de l'Importance au Niveau des Poids (Étape 2) :

   • Les scores causaux au niveau des têtes sont convertis en facteurs d'importance normalisés $I(h_\ell)$, mappés sur une plage de $[1.0, 10.0]$. Les têtes à fort impact reçoivent des facteurs plus élevés.
   • Ces facteurs sont propagés aux poids individuels au sein des matrices de projection correspondantes ($W_q, W_k, W_v, W_o$). Le score d'importance pour un poids spécifique $W_{ij}$ est calculé comme le produit de sa magnitude et du facteur d'importance de sa région de tête associée : $S_{ij} = |W_{ij}| \cdot I(h_{region})$.
   • Pour les couches MLP, qui ne possèdent pas de structure naturelle partitionnée par têtes, la méthode assigne un importance moyenne par couche dérivée des têtes d'attention colocalisées. Cela est noté comme une approximation.

3. Élagage Global (Étape 3) :

   • Tous les poids à travers le modèle sont classés par leurs scores de magnitude pondérés par l'importance ($S_{ij}$).
   • Un seuil global est déterminé via une recherche binaire pour atteindre la parcimonie cible (ex: 10%, 20%, 50%). Les poids en dessous du seuil sont mis à zéro, créant un motif de parcimonie non structuré compatible avec les infrastructures d'inférence standards.

Contributions Clés

• Formalisation de l'Attribution Causale pour l'Élagage : Les auteurs formalisent une méthode pour mesurer l'augmentation attendue de la perte spécifique à la tâche lorsque des têtes d'attention individuelles sont masquées, allant au-delà des proxys corrélationnels pour quantifier la criticité fonctionnelle.
• Élagage au Niveau des Poids Sensible au Budget : Ils introduisent une procédure qui convertit les scores causaux au niveau des têtes en valeurs d'importance au niveau des poids. Cela permet un élagage de magnitude fin sous une contrainte de parcimonie globale tout en préservant la structure des circuits de raisonnement critiques.
• Validation Empirique : La méthode est évaluée sur Llama-3-8B-Instruct et Mistral-7B-Instruct à travers les benchmarks GSM8K, StrategyQA et ARC-Challenge.

Résultats

• Parcimonie Modérée (10–20%) : CAP surpasse généralement la ligne de base Wanda pour préserver la performance de raisonnement.
  • Sur Llama-3-8B-Instruct à 20% de parcimonie, CAP obtient une amélioration relative de 61% par rapport à Wanda sur ARC-Challenge (70,8% vs 43,9% de précision).
  • Sur GSM8K à 10% de parcimonie, CAP montre une amélioration relative de 20% par rapport à Wanda (71,5% vs 59,5%).
  • Les auteurs notent que bien que CAP soit parfois en retrait par rapport à Wanda sur des tâches spécifiques (ex: Llama-3 GSM8K à 20%), il préserve systématiquement la performance sur les tâches d'inférence multi-étapes là où les critères corrélationnels de Wanda échouent.
• Parcimonie Élevée (50%) : La performance se dégrade considérablement pour les deux méthodes. CAP subit un échec catastrophique sur Llama-3 (la perplexité bondit à 428,2), attribué à son traitement grossier des couches MLP. À ce niveau, le seuil global supprime des poids MLP essentiels qui ont été assignés une faible importance basée sur l'approximation des têtes d'attention.
• Perplexité vs Raisonnement : À une parcimonie modérée, CAP atteint souvent une meilleure précision de raisonnement malgré une perplexité légèrement plus élevée que Wanda, suggérant un découplage entre la qualité de modélisation du langage général et la préservation du raisonnement spécifique à la tâche.
• Limitations Architecturales : Des expériences préliminaires sur les modèles Mixture-of-Experts (MoE) suggèrent que CAP est moins efficace, car le routage dynamique rend l'importance des têtes non stationnaire.

Signification et Revendications
L'article affirme que l'attribution causale au niveau de la tête d'attention est un critère supérieur pour préserver la performance de raisonnement par rapport aux méthodes d'élagage corrélationnelles à niveaux de parcimonie équivalents. En mesurant explicitement la contribution causale des têtes aux tâches de raisonnement, CAP peut identifier et protéger les paramètres "numériquement petits mais fonctionnellement critiques" que les méthodes basées sur la magnitude rejettent.

Les auteurs soulignent que leur travail est de portée modeste :

• C'est une méthode sans réentraînement (training-free), ce qui la rend efficace en termes de calcul par rapport aux approches nécessitant des mises à jour de poids (comme SparseGPT).
• Elle est alignée sur la tâche, calibrée spécifiquement sur des données de raisonnement plutôt que sur la modélisation du langage générique.
• Elle reconnaît des limitations, particulièrement l'absence d'attribution causale fine pour les neurones MLP et la difficulté d'appliquer la méthode aux architectures à routage dynamique (MoE).

La contribution centrale est la démonstration que l'élagage interventionnel et aligné sur la tâche peut mieux préserver des capacités spécifiques (le raisonnement) que les méthodes optimisées pour la perplexité générale ou les statistiques de poids statiques, à condition que le niveau de parcimonie reste modéré.