From Layers to Submodules: Rethinking Granularity in Replacement-Based LLM Compression

Résumé technique : De la couche au sous-module (SUBFIT)

Énoncé du problème
La compression post-entraînement des grands modèles de langage (LLM) vise à réduire les coûts de déploiement en supprimant des composants architecturaux. Bien que les méthodes qui suppriment des couches entières ou des blocs contigus soient efficaces, elles perturbent le flux résiduel, provoquant souvent une dégradation brutale de la qualité. Les approches récentes basées sur le « remplacement » traitent ce problème en substituant les composants supprimés par des modules légers ajustés. Cependant, les méthodes existantes partagent deux contraintes de conception restrictives :

1. Granularité à l'échelle de la couche complète : Elles traitent l'ensemble de la couche Transformer comme l'unité atomique de suppression.
2. Sélection contiguë : Elles suppriment des composants dans des plages de profondeur consécutives.

Les auteurs soutiennent que ces contraintes sont trop restrictives. La redondance dans les transformeurs pré-entraînés n'est pas confinée à des régions contiguës, et elle n'est pas répartie uniformément entre les sorties de l'Attention et du FeedForward (FFN). Par conséquent, différentes stratégies sont nécessaires pour approximer différents types de sous-modules, et les composants supprimables n'ont pas besoin de se regrouper dans des plages de profondeur contigues.

Méthodologie : SUBFIT
Le papier présente SUBFIT (Submodule-level Fitted residual replacement), une méthode de compression post-entraînement qui opère à la granularité des sous-modules individuels d'Attention et de FFN. Le processus comprend deux étapes principales :

1. Sélection non contiguë :

   • SUBFIT sélectionne un nombre cible de sous-modules d'Attention et de FFN de manière indépendante à travers la profondeur du modèle en fonction d'un budget de parcimonie $s$.
   • La sélection est non contiguë ; les modules sont choisis indépendamment par score sans former un bloc consécutif.
   • Stratégie de scoring :
     • Attention : Les modules sont classés par un « score d'impact » calculé sur les activations de calibration, mesurant à quel point un sous-module fait pivoter le flux résiduel et sa magnitude relative.
     • FFN : Pour les modèles de base, un score prenant en compte le remplacement est utilisé, combinant les effets de suppression et l'erreur de substitution. Pour les modèles ajustés par instructions (instruction-tuned), un score basé uniquement sur le cosinus est utilisé pour éviter le bruit provenant des magnitudes de flux résiduel mal calibrées.
   • Ordre : Les sous-modules d'Attention sont sélectionnés et remplacés en premier. Les sous-modules FFN sont ensuite scorés et sélectionnés sur le modèle partiellement compressé pour tenir compte du décalage du flux résiduel induit par le remplacement de l'Attention.

2. Remplacement résiduel ajusté :

   • Chaque sous-module sélectionné est remplacé par un substitut paramétrique léger $\hat{F}_\ell$ qui approxime la contribution résiduelle du sous-module original ($y = F_\ell(x)$) plutôt que son calcul interne.
   • Structure du substitut : Le substitut utilise une correction affine (porte élément par élément, biais et centrage) plus un terme d'interaction structuré défini par un opérateur de rang faible $U_\ell = V_\ell^\top W_\ell$.
   • Paramétrage asymétrique :
     • Attention : Chaque sous-module d'Attention supprimé reçoit sa propre base de rang faible ($V_\ell$) et ses propres paramètres, reflétant des rôles hétérogènes à travers les profondeurs.
     • FFN : Une base d'entrée partagée $V$ est utilisée à travers tous les couches FFN sélectionnées. Seuls les facteurs de sortie ($W_\ell$) et les termes affines restent spécifiques à chaque couche. Ce partage de paramètres inter-couches réduit considérablement le coût de déploiement.
   • Ajustement : Les paramètres sont récupérés de manière analytique via un ajustement en forme fermée sur les statistiques de calibration (moindres carrés régularisés par ridge et décomposition en valeurs propres), ne nécessant aucune optimisation basée sur le gradient ou fine-tuning complet du modèle.

Principales contributions

1. Identification des contraintes : Les auteurs identifient la granularité de la couche complète et la sélection contiguë comme des facteurs limitants dans la compression par remplacement existante.
2. Cadre SUBFIT : Une nouvelle méthode qui effectue une sélection et un remplacement conjoints au niveau du sous-module (Attention et FFN) avec une sélection non contiguë.
3. Conception asymétrique : Une architecture spécifique où les bypass d'Attention utilisent des bases indépendantes, tandis que les bypass de FFN partagent une base d'entrée commune pour minimiser la surcharge de paramètres.
4. Évaluation exhaustive : Tests approfondis sur 10 LLM (5 modèles de base, 5 modèles ajustés par instructions) et 5 niveaux de parcimonie (12,5 % à 37,5 %), démontrant de meilleurs compromis par rapport à quatre bases de comparaison (Streamline, ReplaceMe).

Résultats expérimentaux

• Perplexité (Modèles de base) : SUBFIT atteint la plus faible dégradation de la perplexité dans 9 cas sur 10 des configurations modèle-parcimonie. L'écart avec la meilleure base de comparaison s'élargit considérablement sous une compression agressive (ex: à 25 % de parcimonie, SUBFIT montre une dégradation de la perplexité de 2,42× contre 4,34× pour la meilleure base).
• Précision en aval (Modèles ajustés par instructions) : À 25 % de parcimonie, SUBFIT conserve en moyenne 84,6 % de la précision descendante du modèle dense, surpassant la meilleure base (81,6 %). C'est la seule méthode capable de maintenir plus de 80 % de précision à 25 % de parcimonie et plus de 73 % à 37,5 %.
• Stabilité : SUBFIT présente une variance de la dégradation de la perplexité nettement plus faible que les bases de comparaison, qui font preuve d'instabilité (ex: l'écart-type augmentant de plusieurs ordres de grandeur) à haute parcimonie.
• Efficacité : La méthode produit des accélérations de l'inférence mesurables (1,18×–1,40× pour le temps du premier jet/TTFT) et des économies de cache KV (proportionnelles à la parcimonie) sans nécessiter de noyaux creux (sparse kernels) ou de matériel personnalisé.
• Études d'ablation :
  • Compensation résiduelle : Le remplacement des FFN est le facteur dominant de la stabilité ; remplacer à la fois l'Attention et le FFN donne les meilleurs résultats.
  • Base partagée : La base FFF partagée atteint des performances comparables aux bases par couche, mais avec nettement moins de paramètres.
  • Sensibilité au rang : La qualité du FFN s'améliore jusqu'au rang 4096, tandis que l'Attention sature à un rang de 256, justifiant les réglages de rang asymétriques.

Signification et affirmations
Le papier affirme que la granularité est un axe sous-exploré dans la compression post-entraînement. En passant d'un remplacement au niveau de la couche à un remplacement au niveau du sous-module, SUBFIT récupère davantage le comportement du modèle dense qu'un remplacement contigu de la couche complète sous le même budget de paramètres. Les résultats suggèrent qu'une récupération non contiguë et spécifique au sous-module offre un meilleur compromis global qualité-efficacité, particulièrement sous une compression agressive. La méthode est conçue pour être peu coûteuse à déployer (ajustement en forme fermée) et compatible avec les piles d'inférence existantes. Les auteurs suggèrent modestement que cette approche peut se généraliser à d'autres unités de suppression, telles que les têtes d'attention ou les neurones individuels, mais ne revendiquent pas de généralisation immédiate à ces unités spécifiques dans ce travail.