Résumé technique : vOPD (Distillation On-Policy avec Baseline de Variable de Contrôle)

1. Énoncé du problème

La distillation On-Policy (OPD) est devenue un paradigme dominant pour l'entraînement postérieur des grands modèles de langage (LLM), en particulier dans les domaines du raisonnement. Contrairement à l'apprentissage par renforcement avec récompenses vérifiables (RLVR), qui repose sur des récompenses terminales clairsemées, l'OPD minimise la divergence de Kullback-Leibler (KL) inverse entre un étudiant et un enseignant en utilisant des signaux denses au niveau des tokens. Cela permet un entraînement plus rapide et une gestion efficace des tâches longues de chaîne de pensée (CoT).

Cependant, l'OPD rencontre un problème critique de stabilité en pratique. L'implémentation standard utilise un estimateur de Monte Carlo à échantillon unique pour approximer le gradient. Cette approche souffre d'une variance de gradient élevée, entraînant une dynamique d'entraînement instable. Les méthodes de stabilisation existantes tentent d'atténuer ce problème mais introduisent des compromis significatifs :

• OPD à vocabulaire complet ($OPD_{full-V}$) : Calcule la divergence KL sur l'ensemble du vocabulaire. Bien que cela élimine la variance, cela engendre une surcharge computationnelle prohibitive (nécessitant un passage arrière sur le vocabulaire complet, soit environ 150k tokens) à chaque étape.
• OPD Top-k ($OPD_{top-k}$) : Restreint le calcul KL aux $k$ tokens les plus probables. Bien que plus léger, cela introduit un biais dans le gradient en tronquant le support, et les résultats empiriques montrent qu'il n'offre que des gains de performance marginaux par rapport à l'OPD de base instable.

Le défi central consiste à stabiliser l'entraînement OPD sans sacrifier l'efficacité computationnelle de l'estimateur à échantillon unique ni introduire de biais dans l'objectif.

2. Méthodologie : vOPD

Les auteurs proposent vOPD (Distillation On-Policy avec une baseline de variable de contrôle), qui reformule l'OPD comme un problème d'apprentissage par renforcement (RL) par gradient de politique et applique une technique standard de réduction de variance : la soustraction d'une baseline de variable de contrôle.

2.1 Fondement théorique

Dans le RL par gradient de politique, soustraire une baseline $b_t(c_t)$ de la récompense $r_t$ produit une fonction d'avantage $a_t = r_t - b_t$. Crucialement, si la baseline est indépendante de l'action échantillonnée, le gradient reste non biaisé, mais sa variance est réduite. Le choix canonique pour une baseline est la fonction de valeur $V^{\pi}(c_t)$.

L'idée clé de vOPD est que pour l'OPD, la fonction de valeur admet une solution sous forme fermée qui ne nécessite ni modèle critique supplémentaire ni inférence :
$$V^{\pi_\theta}(c_t) = \mathbb{E}_{y_t \sim \pi_\theta} [r_t(c_t, y_t)] = -D_{KL}(\pi_\theta(\cdot|c_t) \parallel \pi_T(\cdot|c_t))$$
Ici, la récompense est définie comme $r_t = \log \pi_T - \log \pi_\theta$. La récompense attendue sous la distribution de l'étudiant est simplement la divergence KL inverse négative entre l'étudiant et l'enseignant à ce contexte.

2.2 Implémentation algorithmique

L'estimateur de gradient vOPD est formulé comme suit :
$$\nabla_\theta J_{vOPD} = \mathbb{E} \left[ \sum_{t} \left( r_t(c_t, y_t) + D_{KL}(\pi_\theta \parallel \pi_T) \right) \nabla_\theta \log \pi_\theta(y_t|c_t) \right]$$

• Non-biais : La baseline (le terme KL) est calculée à partir des distributions $\pi_\theta$ et $\pi_T$ mais est détachée du flux de gradient concernant le token échantillonné $y_t$. Ainsi, le gradient reste non biaisé.
• Efficacité : La baseline KL est calculée lors du passage avant (déjà requis pour le calcul de la récompense). Le passage arrière ne touche toujours que le token unique échantillonné, préservant l'efficacité computationnelle de l'OPD de base.
• Mécanisme de réduction de variance : La baseline agit comme un régularisateur. Lorsque les distributions de l'étudiant et de l'enseignant divergent fortement (KL élevé), la récompense $r_t$ est grande et négative. La baseline $D_{KL}$ est grande et positive, annulant efficacement la récompense négative et atténuant les gradients à queue lourde résultants qui causent l'instabilité.

2.3 Approximation Top-k ($vOPD_{top-k}$)

Pour réduire davantage le coût du passage avant du calcul du KL à vocabulaire complet, les auteurs proposent d'approximer la baseline en utilisant uniquement les $k$ premiers tokens de l'étudiant.

• Contrairement à $OPD_{top-k}$, qui utilise le support top-$k$ pour remplacer l'objectif de perte (introduisant un biais), $vOPD_{top-k}$ utilise l'approximation top-$k$ uniquement pour la baseline.
• Puisque la baseline est détachée du token échantillonné, le gradient reste non biaisé quelle que soit la valeur de $k$.
• Empiriquement, même de petites valeurs de $k$ (par exemple $k=5$) fournissent une réduction de variance suffisante pour stabiliser l'entraînement.

3. Contributions clés

1. Baseline de valeur sous forme fermée : L'article démontre que la fonction de valeur de l'OPD est exactement la divergence KL inverse négative, la rendant directement disponible à partir du passage avant sans modèles supplémentaires.
2. Stabilisation non biaisée : vOPD stabilise l'entraînement en réduisant la variance du gradient tout en maintenant la nature non biaisée de l'estimateur de Monte Carlo à échantillon unique, contrairement aux variantes top-k qui biaisent l'objectif.
3. Efficacité computationnelle : vOPD conserve le passage arrière à un seul token de l'OPD de base. La variante $vOPD_{top-k}$ réduit davantage les coûts du passage avant, offrant un compromis précision-computation supérieur par rapport aux méthodes à vocabulaire complet.
4. Analyse théorique : Les auteurs fournissent une dérivation montrant que la réduction de variance est proportionnelle au carré de la divergence KL, confirmant que la méthode cible spécifiquement les tokens à fort désaccord responsables de l'instabilité.

4. Résultats expérimentaux

Les auteurs ont évalué vOPD sur quatre modèles (Qwen3-1.7B, 4B et Olmo-3-7B) à travers six benchmarks de raisonnement (MATH500, Minerva Math, AMC23, AIME24/25, SciKnowEval, GPQA-Diamond).

• Gains de performance : vOPD surpasse systématiquement l'OPD de base.
  • Sur MATH500, vOPD a réalisé des gains absolus allant jusqu'à +6,2 % par rapport à l'OPD de base.
  • Sur l'ensemble des benchmarks, les variantes vOPD ont montré une amélioration moyenne de +3 % par rapport à l'OPD de base.
  • Les performances de vOPD ont égalé ou légèrement dépassé celles de $OPD_{full-V}$ (la méthode la plus coûteuse) tout en étant significativement plus rapides.
• Comparaison avec les variantes de stabilisation :
  • $OPD_{top-k}$ n'a produit que des gains marginaux (par exemple, +0,4 % en moyenne) par rapport à l'OPD de base, confirmant le problème de biais.
  • Les variantes vOPD ont nettement surpassé $OPD_{top-k}$ et égalé $OPD_{full-V}$.
• Efficacité :
  • vOPD a réduit le temps réel d'exécution jusqu'à 57,7 % par rapport à $OPD_{full-V}$.
  • $vOPD_{top-k}$ a été la plus efficace, offrant le meilleur compromis précision-computation.
• Métriques de stabilité :
  • Normes de gradient : vOPD a maintenu des normes de gradient 1 à 2 ordres de grandeur inférieurs à ceux de l'OPD de base, confirmant la suppression des gradients bruyants et déstabilisateurs.
  • Sensibilité aux hyperparamètres : Les performances sont restées stables sur différentes valeurs de $k$ (de 5 à 100), avec même des approximations grossières ($k=5$) fournissant une stabilisation suffisante.

5. Importance et affirmations

L'article affirme que vOPD comble le fossé entre l'apprentissage par renforcement et la distillation de connaissances en fournissant une approche principielle et efficace pour une distillation On-Policy stable.

• Mécanisme : Il identifie la variance élevée des gradients provenant des estimateurs à échantillon unique comme la cause racine de l'instabilité de l'OPD et la résout en utilisant une baseline de variable de contrôle dérivée de la structure même du problème (la divergence KL).
• Impact pratique : Il offre un remplacement direct pour l'OPD de base ne nécessitant ni modèles supplémentaires (critiques) ni passages arrière coûteux, le rendant adapté à l'entraînement postérieur à l'échelle industrielle où l'efficacité computationnelle est primordiale.
• Généralisation : La méthode s'est révélée robuste à travers différentes échelles de modèles, architectures (Qwen, Olmo) et domaines de raisonnement (mathématiques et sciences).

Les auteurs concluent que vOPD agit comme un régularisateur principiel sur les tokens de récompense déstabilisateurs, transformant l'OPD d'une heuristique instable en un paradigme d'entraînement fiable sans compromettre l'efficacité qui l'avait rendu attrayant à l'origine.