Résumé Technique : Alignement de Trajectoire Multi-Granulaire (MTA) pour la Distillation de LLM

Énoncé du Problème

Les méthodes existantes de Distillation de Connaissances (KD) pour les Grands Modèles de Langage (LLM) se concentrent principalement sur l'alignement des distributions de probabilité de sortie (logits) ou des états cachés intermédiaires à des couches fixes et statiques. Bien qu'efficaces pour transférer des signaux locaux, ces approches emploient souvent une stratégie d'alignement uniforme au niveau du token à travers toutes les profondeurs du réseau.

L'article soutient que ce traitement uniforme est en décalage avec la nature hiérarchique du langage humain et l'organisation fonctionnelle des modèles Transformer. La recherche linguistique et les études d'interprétabilité des Transformer indiquent que :

1. Les couches inférieures privilégient des caractéristiques de surface, lexicales et de récupération factuelle.
2. Les couches supérieures encodent progressivement des informations sémantiques abstraites et un raisonnement compositionnel.

En traitant tous les tokens de manière isolée et en appliquant le même objectif d'alignement de façon uniforme, les méthodes actuelles échouent à capturer la trajectoire évolutive progressive des représentations alors qu'elles se transforment d'unités lexicales en structures sémantiques abstraites. Cette limitation restreint la capacité du modèle étudiant à se généraliser à des tâches en aval.

Méthodologie : Alignement de Trajectoire Multi-Granulaire (MTA)

Pour remédier à ces limitations, les auteurs proposent MTA, un cadre modulaire qui aligne la trajectoire représentative des modèles enseignant et étudiant en utilisant une stratégie adaptative aux couches. MTA opère sur l'hypothèse que la distillation doit respecter les rôles sémantiques dépendants de la profondeur des couches Transformer.

Composants Principaux

1. Spans Multi-Granulaires Adaptatifs aux Couches :

   • Couches Inférieures : Les alignements sont effectués au niveau du mot (Spans de Mots) pour préserver l'ancrage lexical et les dépendances fines.
   • Couches Supérieures : Les alignements passent au niveau de la phrase (Groupes Nominaux et Groupes Verbaux) pour capturer la sémantique compositionnelle et les relations structurelles d'ordre supérieur.
   • Construction des Spans : Les unités sémantiques sont extraites à l'aide d'un analyseur syntaxique (par exemple, spaCy). Les poids d'importance des tokens sont calculés via un mécanisme d'auto-attention bidirectionnelle (excluant les boucles de soi) pour agréger les représentations de tokens en représentations de spans.

2. Perte d'Alignement Structurel Dynamique ($L_{DSA}$) :

   • Cet objectif impose une cohérence géométrique entre l'enseignant et l'étudiant.
   • Au lieu de faire correspondre des vecteurs de tokens individuels, $L_{DSA}$ minimise la divergence dans les distances relationnelles par paires entre les spans sémantiques (mots ou phrases) au sein d'une couche spécifique.
   • Il privilégie les spans sémantiquement saillants en pondérant l'alignement sur la base de l'importance agrégée des tokens constituants.
   • Mathématiquement, il aligne la structure topologique de l'espace de représentation :
     $$L_{DSA} = \frac{1}{\|L_{key}\|} \sum_{l \in L_{key}} \sum_{i,j} w_{ij} \left( d(U^S_{i,l}, U^S_{j,l}) - d(U^T_{i,\phi(l)}, U^T_{j,\phi(l)}) \right)^2$$
     où $d(\cdot)$ est la distance cosinus, et $U$ représente la représentation de span pondérée.

3. Perte d'Alignement des Représentations Cachées ($L_{Hid}$) :

   • Pour compléter les contraintes structurelles, une perte au niveau des caractéristiques aligne des états cachés spécifiques.
   • Étant donné que les dimensions de l'étudiant sont généralement plus petites que celles de l'enseignant, un projecteur linéaire apprenable mappe les états de l'étudiant vers l'espace de l'enseignant.
   • La perte minimise la distance cosinus pondérée entre les états de l'étudiant projetés et les états de l'enseignant, en se concentrant sur les tokens les plus informatifs.

4. Optimisation :

   • MTA est conçu comme un module plug-in. La perte totale combine la perte de distillation de base ($L_{Base}$) des cadres existants (par exemple, DistiLLM, FDD) avec les pertes structurelles et cachées proposées :
     $$L_{Total} = L_{Base} + \lambda_{DSA}L_{DSA} + \lambda_{Hid}L_{Hid}$$

Contributions Clés

1. Identification des Limitations : Les auteurs identifient que l'alignement intermédiaire uniforme ignore l'évolution hiérarchique des représentations, conduisant à un transfert de connaissances sous-optimal.
2. Cadre Adaptatif aux Couches : Ils introduisent MTA, qui transitionne dynamiquement de l'alignement au niveau du mot dans les couches inférieures à l'alignement structurel au niveau de la phrase dans les couches supérieures, reflétant la nature compositionnelle ascendante du langage.
3. Module Généralisé : MTA est démontré comme une couche d'optimisation généralisée compatible avec les méthodes de distillation de pointe (DistiLLM, DistiLLM-2, FDD) partageant la tokenisation.
4. Validation Empirique : Des expériences extensives montrent que MTA améliore constamment les performances sur diverses architectures de modèles (GPT-2, Qwen1.5, OPT) et benchmarks.

Résultats Expérimentaux

Les auteurs ont évalué MTA sur des tâches de suivi d'instructions en utilisant trois paires étudiant-enseignant :

• GPT-2 1.5B $\to$ 120M
• Qwen1.5 1.8B $\to$ 0.5B
• OPT 6.7B $\to$ 1.3B

Métriques de Performance :

• ROUGE-L : MTA a constamment surpassé les lignes de base (FDD, DistiLLM, DistiLLM-2) sur tous les ensembles de données (Dolly, SelfInst, Vicuna, S-NI). Par exemple, dans le réglage GPT-2 1.5B $\to$ 120M, l'ajout de MTA à DistiLLM a amélioré le ROUGE-L moyen de 20,21 à 21,45.
• LLM-as-a-Judge : En utilisant GPT-4o-mini pour évaluer la précision factuelle, la cohérence et l'utilité, les modèles améliorés par MTA ont montré des améliorations constantes des scores par rapport aux lignes de base.

Études d'Ablation :

• Composantes de Perte : À la fois $L_{DSA}$ (structurelle) et $L_{Hid}$ (caractéristique) ont contribué positivement, leur combinaison produisant les meilleurs résultats.
• Granularité : La stratégie « Full-level » (adaptative aux couches) a surpassé les stratégies statiques (Mot uniquement ou Phrase uniquement), confirmant que l'alignement de l'ancrage lexical dans les couches inférieures et de la sémantique compositionnelle dans les couches supérieures est crucial.
• Sélection de Couches : La méthode est robuste face à différents calendriers de couches, bien qu'une allocation hybride (1 couche de niveau Mot : 2 couches de niveau Phrase) ait été jugée optimale pour les configurations testées.

Efficacité :

• MTA introduit une surcharge d'entraînement principalement due à l'extraction de spans syntaxiques (réduite à ~0,27s/étape via optimisation) et à l'appariement géométrique.
• Les coûts d'inférence restent inchangés car l'extraction de spans n'est pas utilisée au moment du test.
• Des comparaisons appariées dans le temps ont confirmé que les gains de performance sont dus à l'objectif d'alignement et non à un budget d'entraînement accru.

Importance et Revendications

L'article revendique que MTA offre une approche fidèle sémantiquement à la distillation de connaissances en modélisant explicitement la trajectoire représentative hiérarchique des LLM. En allant au-delà de l'appariement statique au niveau du token vers un alignement structurel dynamique et multi-granulaire, MTA permet aux modèles étudiants de mieux reproduire la dynamique de traitement de l'information interne de l'enseignant.

Les auteurs positionnent MTA non pas comme un remplacement des cadres de distillation existants, mais comme un module complémentaire, plug-and-play qui améliore la généralisation et la robustesse des LLM compressés. Ce travail souligne l'importance de respecter la hiérarchie linguistique et l'évolution des caractéristiques dépendante de la profondeur dans la conception des futures techniques de compression de modèles.