$\mathcal{X}$-KD: General Experiential Knowledge Distillation for Large Language Models

1. Problématique

La distillation de connaissances (KD) pour les Modèles de Langage de Grande Taille (LLM) vise à transférer les capacités d'un modèle « enseignant » (grand, complexe) vers un modèle « étudiant » (plus petit, efficace). Cependant, les approches existantes souffrent de limitations majeures :

• Imitation comportementale vs apprentissage environnemental : La plupart des méthodes actuelles se concentrent sur l'imitation directe du comportement de l'enseignant (par exemple, via l'apprentissage par clonage comportemental ou la minimisation de la divergence entre les distributions de tokens). Elles négligent l'environnement d'apprentissage original (la fonction de récompense sous-jacente) qui a façonné les connaissances de l'enseignant.
• Limites des approches récentes : Des travaux comme MiniLLM ou GKD tentent d'aborder ces problèmes, mais ils s'éloignent souvent du paradigme d'apprentissage supervisé (utilisant des gradients de politique complexes) ou ne s'adaptent pas facilement à toutes les fonctions de divergence ou aux modèles en boîte noire.
• Inefficacité et manque de généralisation : Sans comprendre la « raison » pour laquelle l'enseignant a pris certaines décisions (la récompense), l'étudiant peut apprendre des motifs superficiels, ce qui nuit à l'efficacité des données et au compromis performance-diversité.

2. Méthodologie : X-KD

Les auteurs proposent X-KD (Experiential Knowledge Distillation), un cadre général qui permet à l'étudiant d'apprendre dans l'environnement d'apprentissage original de l'enseignant, inspiré par la Théorie de l'Apprentissage Expérientiel et l'Apprentissage par Renforcement Inverse (IRL).

A. Fondements Théoriques

X-KD modélise la génération de langage naturel comme un Processus de Décision Markovien (MDP). Au lieu de simplement copier les actions de l'enseignant, X-KD cherche à inférer la fonction de récompense $R_\pi$ qui a guidé l'enseignant, puis à entraîner l'étudiant pour maximiser cette récompense.

Le cadre s'appuie sur AVRIL (Approximated Variational Reward Imitation Learning), une méthode d'IRL bayésienne qui permet d'apprendre conjointement une fonction de récompense et une politique.

B. Architecture et Objectifs d'Entraînement

L'objectif global de X-KD est décomposé en deux parties :

1. Objectif de Distillation de Connaissances ($L_{seq}$ ou $L_{kd}$) : Minimiser la divergence entre la politique de l'enseignant et celle de l'étudiant (comme dans la KD classique).
2. Objectif Expérientiel ($L_{ex}$) : Un terme de régularisation qui assure que la politique de l'étudiant est cohérente avec la fonction de récompense estimée de l'enseignant.

La perte totale s'écrit :
$$L_{XKD} = L_{distillation} + \lambda \cdot L_{expérientiel}$$
Où $\lambda$ est un hyperparamètre contrôlant l'importance de l'apprentissage expérientiel.

Le terme expérientiel utilise un encodeur de récompense $q_\phi$ et un décodeur de valeurs Q $Q_\theta$ (déduit de la politique de l'étudiant) pour minimiser l'erreur Temporelle Différentielle (TD-error) par rapport à la distribution de récompense estimée.

C. Généralité du Cadre

X-KD est conçu pour être flexible et s'adapter à plusieurs scénarios :

• X-KD Séquentiel : Pour la distillation au niveau de la séquence (maximisation de la vraisemblance des séquences générées par l'enseignant).
• X-KD Supervisé : Pour la distillation basée sur la divergence (KL-divergence) au niveau des tokens.
• X-KD Généralisé (G-XKD) : Combine les approches précédentes et utilise des divergences généralisées (comme la divergence JS pondérée $\beta$) pour éviter les hallucinations et les générations de faible qualité, en intégrant des données « on-policy » (générées par l'étudiant) et « offline ».
• Boîte Noire : Une version adaptée permet d'utiliser X-KD même lorsque l'accès aux probabilités de sortie de l'enseignant est impossible (seules les réponses textuelles sont accessibles), en utilisant des échantillons de l'enseignant comme données d'entraînement.

3. Contributions Clés

1. Nouveau Paradigme : Introduction de X-KD, le premier cadre de distillation qui intègre explicitement l'apprentissage de la fonction de récompense de l'environnement original de l'enseignant, passant d'une simple imitation comportementale à un apprentissage expérientiel.
2. Flexibilité et Simplicité : Le cadre s'intègre naturellement dans le paradigme d'apprentissage supervisé (contrairement à MiniLLM qui utilise des gradients de politique complexes) et est compatible avec diverses fonctions de divergence et types de distillation (séquentielle, supervisée, boîte noire).
3. Preuve Théorique : Démonstration que X-KD peut être dérivé comme une extension naturelle des méthodes de KD existantes (SeqKD, SKD, GKD) via le cadre AVRIL.

4. Résultats Expérimentaux

Les auteurs ont évalué X-KD sur trois tâches de génération de langage naturel (NLG) : Résumé Abstrait (XSum), Traduction Automatique (WMT14 en-de) et Raisonnement Arithmétique (GSM8k).

• Performance Globale : X-KD surpasse systématiquement les méthodes de base (SeqKD, SKD, GKD) et le concurrent MiniLLM sur la plupart des tâches et configurations de modèles (de grands modèles vers petits ou moyens).
  • Exemple : Sur XSum, G-XKD obtient un score ROUGE-2 supérieur à GKD et MiniLLM.
• Compromis Performance-Diversité : X-KD maintient une meilleure qualité de génération même lorsque la diversité des sorties augmente (mesurée par le SelfBLEU). Contrairement aux méthodes de base dont la qualité chute rapidement avec la diversité, X-KD apprend des stratégies de résumé plus généralisables.
• Efficacité des Données : X-KD atteint des performances comparables aux méthodes de base en utilisant moins de données d'entraînement (par exemple, 75% des données suffisent pour égaler GKD entraîné sur 100% des données). Cela suggère que l'apprentissage expérientiel permet de capturer les principes fondamentaux de la tâche plus efficacement.
• Distillation en Boîte Noire : Dans les scénarios où l'enseignant est inaccessible (Llama-2-7b), la version SeqXKD obtient un taux de victoire supérieur face à SeqKD lors d'évaluations par GPT-4.

5. Signification et Impact

• Avancement Théorique : X-KD comble le fossé entre la distillation de connaissances et l'apprentissage par renforcement inverse, offrant une justification théorique solide pour l'importance de l'environnement d'apprentissage dans le transfert de connaissances.
• Praticité : En restant dans le cadre de l'apprentissage supervisé, X-KD est plus facile à intégrer dans les pipelines d'entraînement existants des LLM que les méthodes basées sur le gradient de politique.
• Limites et Perspectives : Les auteurs notent que la sensibilité au paramètre de poids expérientiel ($\lambda$) nécessite un réglage fin. De plus, bien que le cadre soit applicable aux boîtes noires, l'étude se concentre principalement sur les scénarios en boîte blanche. Des travaux futurs visent à systématiser l'application de X-KD aux modèles fermés.

En résumé, X-KD représente une avancée significative en permettant aux modèles étudiants de « comprendre » pourquoi un enseignant prend certaines décisions, plutôt que de simplement copier ces décisions, conduisant à des modèles plus robustes, efficaces et diversifiés.