Résumé technique : TRIM (Saliences dérivées de l'attention au niveau des tokens pour un réglage des instructions efficace en données)

1. Énoncé du problème

Le réglage des instructions est crucial pour aligner les grands modèles de langage (LLM) avec l'intention de l'utilisateur, mais il repose traditionnellement sur des corpus massifs et diversifiés. Bien que des recherches récentes suggèrent que de petits sous-ensembles de haute qualité (co-ensembles) puissent égaler ou dépasser les performances obtenues avec des données complètes, la curation de ces sous-ensembles reste un défi majeur.

Les méthodes existantes de sélection de données font face à deux limitations principales :

1. Coût computationnel : Les approches les plus avancées reposent souvent sur des fonctions d'influence ou des signaux basés sur les gradients (par exemple, LESS, TAGCOS), qui nécessitent des passes arrière coûteuses pour chaque échantillon candidat, les rendant irréalisables pour des corpus d'instructions massifs.
2. Granularité et biais : Les méthodes actuelles opèrent généralement au niveau de l'échantillon, attribuant un score d'importance unique à une séquence entière. Cette approche grossière introduit des biais systématiques :
   • Biais de longueur : L'agrégation de signaux au niveau des tokens (comme la perte ou les magnitudes de gradient) corrèle souvent les scores avec la longueur de la séquence, favorisant les exemples plus courts même s'ils contiennent moins de raisonnement.
   • Biais centré sur la perte : La prédiction du token suivant répartit le crédit uniformément sur tous les tokens, ignorant le fait que les preuves pertinentes pour la tâche sont souvent rares et localisées sur des tokens spécifiques (par exemple, opérateurs mathématiques, connecteurs logiques).

La question centrale abordée par ce travail est la suivante : Étant donné seulement quelques exemples cibles, comment pouvons-nous identifier efficacement des données d'instruction à fort impact en passant d'un score au niveau de l'échantillon à des signaux au niveau des tokens ?

2. Méthodologie : TRIM

Les auteurs proposent TRIM (Token Relevance via Interpretable Multi-layer Attention), un cadre centré sur les tokens et ne nécessitant que le passage avant, évitant totalement le calcul des gradients. La méthode opère en deux étapes en utilisant un modèle « notateur » ($M_{warmup}$) qui a été brièvement réchauffé sur un petit sous-ensemble aléatoire du pool de candidats.

Étape I : Construction d'empreintes digitales de tokens pondérées par la saillance

En utilisant un petit ensemble de validation cible ($T_{val}$), TRIM construit un dictionnaire d'« empreintes digitales » représentant l'apparition des tokens définissant la tâche dans leur contexte.

• Calcul de la saillance des tokens : Au lieu d'utiliser des gradients, TRIM dérive l'importance des tokens à partir des mécanismes d'attention. Il calcule deux signaux complémentaires pour chaque token $i$ :
  • Saillance de ligne ($Q_i$) : Mesure la « netteté » de la distribution d'attention d'un token (une faible entropie indique que le token se concentre sur des clés spécifiques).
  • Saillance de colonne ($K_j$) : Mesure la quantité d'attention qu'un token reçoit d'autres tokens (agissant comme un « hub »).
  • Ces signaux sont agrégés à travers les couches et les têtes, normalisés, puis combinés via une moyenne pondérée ($\alpha_i$) pour produire un score de saillance final.
• Construction de l'empreinte digitale : Pour chaque classe de tokens (token de vocabulaire unique), un vecteur d'empreinte digitale ($f_t$) est créé. Il s'agit de la moyenne pondérée par la saillance des états cachés de la dernière couche de toutes les occurrences de ce token dans $T_{val}$.
• Repondération par fréquence : Pour empêcher les tokens de formatage courants de dominer, les poids de saillance sont mis à l'échelle par l'Inverse Document Frequency (IDF) calculé sur le pool de candidats.

Étape II : Notation et sélection des candidats

TRIM note chaque échantillon candidat dans le grand corpus ($S$) en faisant correspondre ses représentations de tokens aux empreintes digitales.

• Appariement au niveau des tokens : Pour chaque token d'un échantillon candidat, le modèle calcule la similarité cosinus entre son état caché et l'empreinte digitale de la classe correspondante.
• Gestion des tokens non vus : Si une classe de tokens n'apparaît pas dans $T_{val}$, elle est mappée à la classe empreinte la plus proche dans l'espace d'incrustation, son score de similarité étant pénalisé par un facteur $\lambda$.
• Regroupement robuste : Les scores au niveau des tokens sont agrégés en un score d'échantillon unique en utilisant une combinaison de :
  • Moyenne : Informativité globale.
  • Maximum : Présence de tokens hautement cruciaux.
  • Couverture : Un terme récompensant les échantillons où les tokens notés sont moins rares, atténuant le biais de longueur.
• Sélection : Les échantillons les mieux classés forment le co-ensemble final.

3. Contributions clés

1. Paradigme centré sur les tokens : TRIM passe d'une notation grossière au niveau de l'échantillon à un appariement de représentations fin et granulaire au niveau des tokens, capturant les caractéristiques structurelles (syntaxe, opérateurs) qui définissent une tâche.
2. Efficacité par passage avant uniquement : En s'appuyant uniquement sur la saillance dérivée de l'attention et les états cachés, TRIM élimine le besoin de passes arrière lors de la phase de sélection, réalisant des accélérations d'un ordre de grandeur par rapport aux méthodes basées sur les gradients.
3. Atténuation des biais : L'agrégation au niveau des tokens atténue intrinsèquement le biais de longueur, permettant la sélection d'exemples de raisonnement plus longs et plus complexes que les méthodes au niveau de l'échantillon rejettent souvent.
4. Transférabilité : La méthode démontre que les co-ensembles sélectionnés par un petit modèle notateur (par exemple, 7B) se transfèrent efficacement vers des modèles plus grands (par exemple, 13B) et différentes architectures (par exemple, Mistral), capturant des motifs structurels agnostiques au modèle.

4. Résultats expérimentaux

Les auteurs ont évalué TRIM sur LLAMA-2-7B (et transfert vers 13B/Mistral) à travers des benchmarks de raisonnement général (MMLU, BBH, TYDIQA) et une tâche mathématique difficile (GSM8K) en utilisant un pool de candidats de 270 000 exemples d'instructions.

• Performance : Avec un co-ensemble de 5 %, TRIM a atteint une précision moyenne de 48,56 % sur les benchmarks généraux, surpassant la meilleure ligne de base (LESS, 48,27 %) et dépassant le réglage fin avec des données complètes sur TYDIQA.
• Tâches difficiles (GSM8K) : Dans un scénario de faible chevauchement (sélection de données mathématiques à partir d'un corpus général), un co-ensemble TRIM de 5 % a atteint une précision de 29,52 %, égalant presque le réglage fin avec des données complètes (30,25 %) et surpassant LESS d'environ 8,8 %. Même avec un budget de 1 %, TRIM (22,33 %) a nettement surpassé les alternatives.
• Efficacité : TRIM a réduit le temps d'exécution de la sélection à 4,40 heures contre 35,76 heures pour LESS et 87,95 heures pour S2L, principalement grâce à l'élimination des passes arrière et de la collecte de trajectoires.
• Biais de longueur : L'analyse a montré que TRIM sélectionnait des co-ensembles d'une longueur moyenne de 446 tokens, significativement plus longs que LESS (323) et TAGCOS (259), confirmant sa capacité à sélectionner des exemples de raisonnement complexes.
• Sélection intra-domaine : TRIM a été adapté pour sélectionner des données sans exemples cibles externes, atteignant des performances comparables à S2L et LESS tout en étant plus efficace en termes de calcul.

5. Importance et affirmations

L'article présente TRIM comme une alternative évolutive et efficace pour construire des ensembles de données de réglage des instructions de haute qualité. Les auteurs affirment qu'en exploitant des « empreintes digitales » dérivées de l'attention, TRIM peut identifier des données d'entraînement qui correspondent aux motifs structurels définissant une tâche, plutôt que de s'appuyer sur des calculs de gradients coûteux ou une similarité de surface.

L'ouvrage affirme que :

• TRIM offre une voie pratique vers un réglage efficace en données, permettant aux praticiens de curer des ensembles de données plus petits et de meilleure qualité avec une fraction du coût computationnel.
• La notation au niveau des tokens est un signal supérieur pour le réglage des instructions, en particulier pour les tâches nécessitant un raisonnement structuré, car elle évite les biais systématiques inhérents aux métriques au niveau de l'échantillon.
• La méthode est robuste à travers les échelles de modèles et les architectures, suggérant que les motifs structurels définissant une tâche sont transférables entre différentes familles de LLM.

Les auteurs adoptent une position modeste concernant la sécurité, notant que si TRIM améliore l'efficacité, il ne détecte pas intrinsèquement les contenus nuisibles ; par conséquent, il doit être utilisé parallèlement à des pratiques établies de gouvernance des ensembles de données (filtrage, audit) pour atténuer les risques d'amplification des biais.