Résumé technique : DISA - Échantillonnage d'importance hors ligne pour l'apprentissage par renforcement des LLM à correspondance de distribution

1. Énoncé du problème

Les agents de raisonnement modernes sont de plus en plus évalués sur leur capacité à générer plusieurs chemins de solution valides (par exemple, des chaînes de pensée diversifiées, des plans ou des traces d'utilisation d'outils) plutôt qu'une seule réponse correcte. Les objectifs standard d'apprentissage par renforcement (RL), tels que l'optimisation de politique proximale (PPO) ou l'optimisation de politique relative de groupe (GRPO), se concentrent sur la maximisation de la récompense. Bien qu'efficaces pour la précision d'une seule réponse, ces méthodes tendent à souffrir d'un effondrement de mode, où la politique concentre la masse de probabilité sur une seule trajectoire à haute récompense, éliminant ainsi la diversité des solutions valides requise pour une évaluation robuste (par exemple, les métriques pass@k).

L'apprentissage par renforcement à correspondance de distribution offre une alternative en entraînant la politique à s'ajuster à une distribution cible où la probabilité de trajectoire est proportionnelle à la récompense exponentiée :
$$\tilde{\pi}(o | q) = \frac{1}{Z(q)} \pi_{\text{ref}}(o | q) \exp(\beta \tilde{r}(q, o))$$
où $Z(q)$ est la fonction de partition (une constante de normalisation dépendante de l'invite $q$).

Le goulot d'étranglement central des méthodes existantes de correspondance de distribution (par exemple, FlowRL, GFlowNets) est le couplage en ligne entre la politique ($\pi_\theta$) et l'estimateur de la fonction de partition ($Z_\phi$). Ces méthodes entraînent conjointement les deux composants sur des déroulements en ligne partagés en utilisant une perte conjointe (par exemple, l'équilibre de trajectoire). Ce couplage crée deux problèmes critiques :

1. Propagation d'erreurs : Les erreurs d'étalonnage de l'estimateur de fonction de partition $Z_\phi$ se propagent directement dans le gradient de la politique, faussant les mises à jour.
2. Absence de diagnostics : Il n'existe aucun signal de qualité indépendant pour la fonction de partition ; ses performances sont entremêlées avec la perte de la politique, rendant impossible le diagnostic de savoir si les échecs proviennent de la politique ou de l'estimation de la fonction de partition.

2. Méthodologie : Ancrage par échantillonnage d'importance découplé (DISA)

Les auteurs proposent DISA, un pipeline en trois étapes qui découple l'estimation de la fonction de partition de la boucle d'optimisation de la politique. L'idée clé est que la fonction de partition $Z(q)$ dépend uniquement de la politique de référence $\pi_{\text{ref}}$, de la fonction de récompense et de l'inverse de la température $\beta$, mais pas de la politique entraînable $\pi_\theta$. Par conséquent, elle peut être estimée hors ligne avant le début de l'entraînement par RL.

Étape 1 : Estimation par échantillonnage d'importance (IS) hors ligne

Pour un ensemble d'invites hors ligne $\mathcal{D}_Z$, la méthode tire $N$ trajectoires de proposition à partir d'un modèle de proposition robuste $p_T(\cdot|q)$ (par exemple, un modèle enseignant plus grand).

• Poids d'importance : Pour chaque trajectoire $o_i$, un poids d'importance est calculé : $w(q, o_i) = \frac{\pi_{\text{ref}}(o_i|q)}{p_T(o_i|q)} \exp(\beta \tilde{r}(q, o_i))$.
• Agrégation dans l'espace des logarithmes : Pour assurer la stabilité numérique et réduire le biais, l'estimation de la fonction de partition pour l'invite $q$ est calculée dans l'espace des logarithmes en utilisant un agrégateur logsumexp :
  $$\log \hat{Z}_{\text{IS}}(q) = \text{logsumexp}_i \left( \log \pi_{\text{ref}}(o_i|q) - \log p_T(o_i|q) + \beta \tilde{r}(q, o_i) \right) - \log N$$
  Cela fournit une estimation non biaisée de $Z(q)$ à l'échelle linéaire, avec un biais contrôlé de $O(1/N)$ à l'échelle logarithmique.

Étape 2 : Amortissement par régression

Les estimations par invite $\log \hat{Z}_{\text{IS}}(q)$ servent d'étiquettes pour entraîner un régresseur conditionné par l'invite léger $g_\psi(q)$ (par exemple, un petit MLP).

• Le régresseur est entraîné par régression des moindres carrés pour minimiser l'erreur quadratique entre sa prédiction et les étiquettes IS hors ligne.
• Une fois entraîné, les paramètres du régresseur $\psi^*$ sont figés. Cela crée une estimation réutilisable et hors ligne de la fonction de partition qui se généralise aux invites en dehors de l'ensemble d'estimation initial.

Étape 3 : Apprentissage de politique en ligne ancré

Pendant la phase de RL en ligne, le régresseur figé $g_{\psi^*}(q)$ remplace la fonction de partition apprenable $Z_\phi(q)$ dans la perte d'équilibre de trajectoire (TB).

• Objectif : La politique $\pi_\theta$ est optimisée pour minimiser le résidu :
  $$\mathcal{L}_{\text{DISA}}(\theta) = \mathbb{E} \left[ \sum_i \omega_i \left( g_{\psi^*}(q) + \log \pi_\theta(o_i|q) - \log \pi_{\text{ref}}(o_i|q) - \beta \tilde{r}(q, o_i) \right)^2 \right]$$
• Découplage : Crucialement, les gradients sont arrêtés de traverser $g_{\psi^*}$. La fonction de partition agit comme une constante fixe, par invite (un « ancre »), pendant les mises à jour de la politique.

3. Contributions clés

1. Identification du couplage comme obstacle central : L'article identifie l'entrelacement de l'estimation de la fonction de partition et de l'apprentissage de la politique comme la source principale des erreurs d'étalonnage et de l'aveuglement diagnostique dans l'apprentissage par renforcement à correspondance de distribution.
2. Le pipeline DISA : Une approche novatrice en trois étapes qui déplace l'estimation de la fonction de partition hors ligne via l'échantillonnage d'importance, l'amortit par régression et la fige pour le RL en ligne.
3. Garanties théoriques :
   • Préservation exacte : Si la fonction de partition hors ligne est exacte, le minimiseur global de l'objectif DISA correspond à la distribution cible théorique inclinée par la récompense.
   • Stationnarité avec biais : Même avec un biais lié uniquement à l'invite (où l'erreur dépend uniquement de $q$ et non de la trajectoire), la distribution cible reste un point stationnaire du flux de gradient sur-politique.
   • Perturbation bornée : Les erreurs d'estimation résiduelles n'entrent dans l'objectif de la politique que par le biais d'une enveloppe de perturbation bornée (dérivée via Cauchy-Schwarz), assurant la stabilité.
4. Validation empirique : Évaluation complète sur des benchmarks de mathématiques et de code démontrant que DISA préserve les avantages de la correspondance de distribution tout en évitant les pièges du couplage en ligne.

4. Résultats expérimentaux

Les auteurs ont évalué DISA sur deux architectures à poids ouverts (Qwen2.5-7B et Qwen3-4B-Base) à travers six benchmarks de mathématiques (y compris AIME 2024/2025, AMC 2023, MATH-500) et trois benchmarks de code (LiveCodeBench, CodeForces, HumanEval+).

• Raisonnement mathématique :
  • DISA a égalé ou légèrement dépassé la ligne de base couplée en ligne FlowRL (par exemple, 65,5 contre 64,8 Mean@8 sur Qwen3-4B-Base).
  • Il a nettement surpassé les lignes de base de maximisation de récompense GRPO et GSPO (par exemple, +16,1 points par rapport à GRPO sur Qwen3-4B-Base).
• Génération de code :
  • Sur pass@16 (une métrique sensible à la diversité des solutions), DISA a obtenu les scores les plus élevés parmi toutes les méthodes (par exemple, 51,5 sur Qwen3-4B-Base), surpassant FlowRL, GSPO et GRPO.
  • Sur pass@1, DISA est resté compétitif, ne montrant aucune régression significative par rapport aux méthodes de maximisation de récompense.
• Analyse de la diversité :
  • Une évaluation par un LLM-juge sur la diversité des stratégies de solution (utilisant une échelle de Likert de 1 à 5) a montré que DISA a conservé la plus grande fraction de la diversité originale de l'architecture (3,72/5,0), contre 3,40 pour FlowRL et nettement mieux que GRPO/GSPO (~3,0).
• Études d'ablation :
  • Force de la proposition : L'utilisation d'un modèle de proposition plus faible (Qwen3-4B contre 235B) a augmenté la variance des étiquettes hors ligne, entraînant une baisse de performance, mais le régresseur a tout de même convergé, confirmant le compromis biais-variance prédit par la théorie.
  • Inverse de la température ($\beta$) : Les études de sensibilité ont confirmé que le sous-inclinaison ($\beta=10$) conduit à une perte de signal légère, tandis que la sur-inclinaison ($\beta=20$) provoque un « précipice de variance » et un effondrement de mode, validant l'analyse théorique des poids d'échantillonnage d'importance.
• Distillation vs RL : Comparé au fine-tuning (SFT) sur les mêmes trajectoires à haute récompense utilisées pour l'estimation hors ligne de DISA, DISA a fourni un gain supplémentaire allant jusqu'à 13,8 points Mean@8, démontrant que l'ajustement de la distribution douce et inclinée par la récompense récupère des informations de paysage perdues par les cibles SFT rigides.

5. Importance et revendications

L'article affirme que DISA représente un changement architectural significatif dans le post-entraînement des LLM pour le raisonnement. En découplant l'estimation de la fonction de partition de la boucle de politique, DISA :

• Élimine les erreurs d'étalonnage : Il empêche les erreurs d'estimation de la fonction de partition de fausser les gradients de la politique.
• Permet des diagnostics indépendants : La qualité de la fonction de partition peut être surveillée séparément de l'entraînement de la politique via la perte de validation du régresseur.
• Préserve la diversité : Il maintient avec succès la structure multimodale de l'espace de solution, ce qui est crucial pour les tâches de raisonnement complexes où plusieurs chemins valides existent.
• Solidité théorique : La méthode fournit des garanties formelles que le processus d'estimation hors ligne n'altère pas le point fixe théorique de l'objectif de correspondance de distribution, à condition que l'erreur d'estimation soit bornée.

Les auteurs concluent que DISA offre une alternative robuste et théoriquement fondée à la correspondance de distribution couplée en ligne, atteignant des performances de pointe tant en précision qu'en diversité, sans l'instabilité associée à l'entraînement conjoint de fonctions de partition et de politiques.