Résumé Technique : ARKD – Distillation de Divergence KL Bidirectionnelle Guidée par l'Apprentissage par Renforcement Adaptatif

1. Énoncé du Problème

La Distillation de Connaissance (KD - Knowledge Distillation) est une technique primaire pour compresser les Grands Modèles de Langage (LLM) afin de réduire les coûts d'inférence tout en maintenant les performances. Cependant, les méthodes traditionnelles de KD reposant sur un objectif unique de divergence de Kullback-Leibler (KL) font face à un compromis fondamental dans la génération de texte ouvert :

• KL Direct (FKL - Forward KL) : Encourage un comportement de « couverture de mode » (mode-covering), où l'étudiant tente de couvrir tous les modes de la distribution de l'enseignant. Cela conduit souvent à une surestimation des probabilités dans les régions de faible confiance et de longue traîne, dégradant la qualité de la génération.
• KL Inverse (RKL - Reverse KL) : Encourage un comportement de « recherche de mode » (mode-seeking), se concentrant sur les modes dominants et ignorant les régions de faible probabilité. Bien que cela empêche la surestimation de la longue traîne, cela compromet la diversité et la couverture, pouvant mener à un effondrement de mode (mode collapse).

Les méthodes d'interpolation statiques existantes (par exemple, un pondération fixe de 0,5:0,5) ne parviennent pas à capturer la nature dynamique de l'écart entre l'étudiant et l'enseignant au cours de l'évolution de l'entraînement. Un coefficient statique ne peut pas s'adapter aux exigences changeantes de l'alignement des distributions, qui nécessitent des équilibres différents entre FKL et RKL à différentes étapes de l'entraînement.

2. Méthodologie : ARKD

Les auteurs proposent ARKD (Adaptive Reinforcement Learning-Guided Bidirectional KL Divergence Distillation), un cadre qui assigne dynamiquement des poids au FKL et au RKL en utilisant un agent d'apprentissage par renforcement (RL).

Fondement Théorique

L'article établit d'abord la cohérence théorique du FKL et du RKL. L'analyse montre que les deux objectifs convergent vers la même solution ($q_\theta = p$) lorsque les gradients sont nuls. Cela confirme que les deux divergences sont valides pour l'appariement de distribution mais présentent des comportements complémentaires pendant la trajectoire d'optimisation, justifiant leur utilisation conjointe.

Cadre basé sur le RL

ARKD reformule le problème d'assignation de poids comme un processus de décision séquentiel (Processus de Décision Markovien) :

• État ($s$) : Un vecteur à 6 dimensions capturant la dynamique actuelle de l'entraînement, incluant les entropies de l'enseignant et de l'étudiant ($H_T, H_S$), les variances ($\sigma^2_T, \sigma^2_S$), et les valeurs de perte FKL/RKL actuelles ($L_{FKL}, L_{RKL}$).
• Action ($\alpha$) : Un réseau de politique ($\pi_\phi$) produit un poids scalaire $\alpha \in (0, 1)$, où $\alpha$ pondère le terme FKL et $(1-\alpha)$ pondère le terme RKL.
• Récompense ($r$) : Définie comme la négative de la perte de distillation ($r = -L_{distil}$) pour le batch actuel. L'objectif est de maximiser la récompense cumulée (minimiser la perte totale) sur la trajectoire d'entraînement.
• Optimisation :
  • Mise à jour de l'étudiant : Minimise la perte de distillation pondérée ainsi qu'une perte de modélisation de langage.
  • Mise à jour de la politique : Utilise l'algorithme REINFORCE avec une moyenne mobile exponentielle (EMA) de base et une régularisation de l'entropie pour stabiliser l'entraînement et éviter la saturation des limites.

Dynamique d'Entraînement

Le réseau de politique apprend à naviguer de manière adaptative dans le compromis :

1. Phase d'Exploration : Au début de l'entraînement, la politique explore une gamme plus large de $\alpha$ (ex: 0,3–0,6), mettant l'accent sur le FKL pour assurer une couverture large et prévenir l'effondrement de mode.
2. Phase de Convergence : À mesure que la distribution de l'étudiant se rapproche de celle de l'enseignant, la politique bascule vers des valeurs de $\alpha$ plus faibles (dominance RKL) pour se concentrer sur un alignement précis et éviter la surestimation de la longue traîne.
3. Phase de Stabilité : La politique se stabilise à un $\alpha$ spécifique (ex: $\approx 0,18$) pour le raffinement final.

3. Principales Contributions

1. Analyse Théorique : L'article fournit une analyse théorique établissant la complémentarité et la cohérence de convergence du FKL et du RKL, formant la base de la distillation adaptative.
2. Cadre ARKD : Il introduit un nouveau cadre d'optimisation conjointe basé sur le RL où un réseau de politique pondère dynamiquement le FKL et le RKL en fonction des caractéristiques dépendantes de l'état, allant au-delà des heuristiques statiques.
3. Validation Empirique : Des expériences approfondies démontrent qu'ARKD améliore systématiquement la qualité de la génération, l'ajustement de la distribution et la généralisation à travers de multiples benchmarks et échelles de modèles.

4. Résultats Expérimentaux

Les auteurs ont évalué ARKD sur des tâches de réglage d'instructions (instruction tuning) en utilisant des configurations GPT-2 (1,5B $\to$ 120M/340M) et LLaMA (13B $\to$ 7B).

• Performance In-Distribution (ID) : Sur le benchmark DollyEval, ARKD a obtenu le BertScore le plus élevé et des scores Rouge-L compétitifs ou supérieurs par rapport aux lignes de base (SFT, SeqKD, FKL, RKL, FKL+RKL et MiniLLM). Par exemple, sur GPT-2 1,5B $\to$ 120M, ARKD a atteint un Rouge-L de 24,5 et un BertScore de 0,815, surpassant la forte ligne de base MiniLLM.
• Généralisation Out-of-Distribution (OOD) : ARKD a démontré une généralisation supérieure sur des jeux de données non vus (SelfInst, VicunaEval, S-NI, UnNI), surpassant systématiquement toutes les lignes de base à travers les deux familles de modèles.
• Études d'Ablation :
  • vs Programmation/Planification Linéaire Statique : ARKD a nettement surpassé le pondération statique ( $\alpha$ fixe) et les stratégies de planification linéaire (décroissance/augmentation basée sur le temps), qui échouent à capturer les dynamiques d'alignement complexes.
  • vs Heuristiques Gloutonnes : ARKD a surpassé les heuristiques gloutonnes (sélectionnant la divergence à perte plus faible à chaque étape) de 0,4 à 0,6 point. L'approche gloutonne s'est avérée myope, dérivant souvent vers un RKL pur, tandis qu'ARKD maintient une composante FKL bénéfique par un mélange continu.
  • Analyse de Trajectoire : La trajectoire de $\alpha$ apprise a confirmé une stratégie intelligente : un accent précoce sur le FKL pour la couverture, transitant vers le RKL pour la précision, un schéma non atteignable par conception manuelle.

5. Signification et Revendications

L'article affirme qu'ARKD fait progresser l'état de l'art de la distillation de connaissance adaptative en abordant les limites des méthodes basées sur la divergence KL statique. En exploitant l'apprentissage par renforcement pour équilibrer dynamiquement les comportements de couverture de mode et de recherche de mode, ARKD parvient à un meilleur alignement de distribution qui améliore à la fois la qualité de la génération et la généralisation.

Les auteurs positionnent ARKD comme une approche fondée sur des principes qui valide la nécessité d'un pondération dynamique et sensible à l'état plutôt que des heuristiques fixes. Ils notent que bien que la politique actuelle repose sur un vecteur d'état conçu à la main et des récompenses de batch immédiates, le cadre est compatible avec des familles de $f$-divergence plus larges et des paradigmes on-policy. Le travail fournit une base solide pour une compression efficace des LLM en démontrant qu'une optimisation adaptative guidée par le RL peut systématiquement surpasser les lignes de base statiques et gloutonnes à travers diverses tâches et échelles de modèles.