Résumé Technique : SRA (Alignement des Représentations de Séquences pour la Distillation de Modèles de Langage de Grande Taille)

Énoncé du Problème

L'article aborde le défi de la Distillation de Connaissances Cross-Tokenizers (CTKD), où un grand modèle de langage « enseignant » et un modèle « étudiant » plus petit utilisent des tokenizers et des vocabulaires différents. Les méthodes CTKD existantes reposent principalement sur des stratégies d'alignement au niveau des tokens ou des logits (par exemple, en utilisant la distance d'édition, le transport optimal ou des couches de projection). Les auteurs soutiennent que ces approches sont souvent fragiles et sensibles aux écarts entre les tokenizers, car elles tentent d'aligner des unités discrètes et atomiques qui peuvent ne pas correspondre sémantiquement entre différents systèmes de vocabulaire. De plus, les méthodes actuelles échouent souvent à capturer la dynamique sous-jacente de l'évolution des caractéristiques au sein de l'architecture Transformer lorsque les tokenizers diffèrent.

Méthodologie : Cadre SRA

Le cadre SRA proposé reformule la CTKD en déplaçant l'unité fondamentale d'alignement des tokens individuels vers des séquences (spans) robustes et indépendantes du tokenizer. Cette approche est fondée sur l'analogie physique des Systèmes Dynamiques Multi-Particules (MPDS), où le Transformer est considéré comme un système de particules interagissantes (tokens) évoluant à travers les couches.

1. Construction et Mappage des Séquences

• Définition de la Séquence : Au lieu d'aligner directement les tokens, SRA regroupe les tokens en « séquences » qui correspondent aux mêmes unités textuelles dans l'entrée brute.
• Alignement via LCS : Pour faire correspondre les séquences entre les modèles enseignant et étudiant, les auteurs calculent la Plus Longue Sous-Séquence Commune (LCS) des décalages de caractères produits par les deux tokenizers. Cela identifie les segments correspondants du texte d'entrée, définissant les indices de début et de fin pour les séquences alignées dans les deux modèles sans nécessiter de mappage direct token-à-token.

2. Représentation de Séquence via le Centre de Masse (CoM)

S'inspirant de la perspective MPDS, chaque séquence est traitée comme un cluster de particules. L'état de ce cluster est représenté par son Centre de Masse (CoM) :

• Pooling Pondéré : La représentation d'une séquence est calculée comme une moyenne pondérée par l'attention de ses états cachés de tokens constitutifs.
• Poids d'Attention : Le poids $w_i$ d'un token est dérivé des scores d'attention reçus par le dernier token de la séquence à travers toutes les têtes d'attention. Cela imite le concept physique où les particules « plus lourdes » (plus saillantes sémantiquement) dominent le centre du cluster.
• Formule : $C_i = \sum w_t H_t$, où $H_t$ est l'état caché et $w_t$ est le poids dérivé de l'attention.

3. Objectifs de Distillation

SRA utilise une fonction de perte combinée pour transférer les connaissances au niveau des séquences :

• Transfert d'États Cachés au Niveau des Séquences ($L_{HS}^{Span}$) :

  • Aligne les représentations CoM des séquences étudiant et enseignant en utilisant une perte de similarité cosinus.
  • Régularisateur Géométrique ($L_{Geo}$) : Pour préserver l'intégrité structurelle de l'espace de représentation, un régularisateur est ajouté afin de maintenir les distances relatives (distance cosinus) entre les paires de séquences dans le modèle étudiant, en les faisant correspondre à celles du modèle enseignant.
  • Perte Pondérée : La perte est pondérée par la saillance sémantique des séquences (dérivée de la somme des poids des tokens au sein de la séquence), assurant que l'étudiant se concentre sur les régions les plus informatives.

• Distillation des Logits de Séquences Alignées ($L_{KD}^{Span}$) :

  • Traite le décalage de vocabulaire ($V_{tea} \neq V_{stu}$) en restreignant le transfert de connaissances au sous-espace partagé $\tilde{V} = V_{tea} \cap V_{stu}$.
  • Les logits au niveau des séquences sont agrégés à partir des logits de tokens au sein du vocabulaire partagé, et une divergence de Kullback-Leibler (KL) est calculée entre les distributions de l'enseignant et de l'étudiant sur cet espace partagé.

• Objectif Total : La perte finale combine la perte d'entropie croisée standard sur les étiquettes de vérité terrain avec les pertes de distillation au niveau des séquences (états cachés et logits), équilibrées par un hyperparamètre $\alpha$.

Contributions Clés

1. Cadre Conceptuel Novel : L'article introduit un cadre CTKD fondé sur la vision des Systèmes Dynamiques Multi-Particules des Transformers, modélisant les séquences comme des clusters de particules et utilisant leur Centre de Masse pour des représentations stables et additives.
2. Distillation au Niveau des Séquences : Il propose une méthode pour distiller les connaissances au niveau des séquences, englobant à la fois les représentations cachées et les logits. Cela inclut un schéma de pondération dérivé de l'attention pour privilégier les séquences saillantes et un régularisateur géométrique pour préserver la géométrie relative de l'espace de représentation.
3. Validation Empirique : Les auteurs démontrent que SRA surpasse significativement les références CTKD de l'état de l'art (incluant ULD, DSKD, MinED et MultilevelOT) sur plusieurs paires enseignant-étudiant et benchmarks.

Résultats Expérimentaux

Les auteurs ont évalué SRA sur diverses paires inter-architectures et inter-tokenizers (par exemple, Qwen1.5/2.5, Mistral-7B, GPT-2, OPT, TinyLLaMA) en utilisant des jeux de données de suivi d'instructions (Dolly, Vicuna, SelfInst, S-NI, DialogSum).

• Performance : SRA a constamment obtenu les scores ROUGE-L les plus élevés et les meilleures performances moyennes dans tous les paramètres par rapport à des références solides.
• Qualité Sémantique : Les évaluations sémantiques utilisant GPT-4o-mini comme juge ont confirmé que SRA produit des réponses plus fidèles sémantiquement et mieux alignées contextuellement.
• Études d'Ablation :
  • Le retrait du régularisateur géométrique ou du mécanisme de pooling pondéré a entraîné une dégradation des performances, confirmant la nécessité de préserver les relations structurelles et de se concentrer sur les séquences saillantes.
  • La méthode a montré une robustesse même dans des scénarios de même tokenizer, surpassant les méthodes basées sur la distance et autres méthodes CTKD.
• Efficacité : Bien que SRA nécessite légèrement plus de mémoire GPU que certaines références en raison de l'appariement et de l'agrégation des séquences, elle a obtenu le temps d'entraînement moyen par étape le plus faible parmi les méthodes comparées.

Importance et Revendications

L'article revendique que SRA fournit une approche robuste et ancrée dans la physique pour la distillation de connaissances, surmontant la fragilité de l'alignement au niveau des tokens. En déplaçant l'unité d'alignement vers les séquences et en exploitant le concept de Centre de Masse, SRA soutient naturellement le transfert de l'évolution dynamique des caractéristiques de couche en couche. Les auteurs affirment que leur méthode valide le principe que la distillation doit se concentrer sur la dynamique de l'évolution géométrique du système plutôt que sur les seules sorties finales, offrant une solution évolutive pour déployer des modèles efficaces à travers des environnements de tokenizers hétérogènes.

Le travail est présenté comme une étape vers un transfert de connaissances plus stable et complet, bien que les auteurs reconnaissent des limitations telles que la nature statique du mappage actuel des logits et le coût computationnel de l'inférence de l'enseignant à la volée requise pour l'alignement des séquences.