Résumé technique : Au-delà de l'effondrement de mode : l'appariement de distributions pour un raisonnement diversifié

1. Énoncé du problème

L'article identifie l'effondrement de mode comme une limitation fondamentale des méthodes d'apprentissage par renforcement (RL) on-policy, en particulier celles utilisées pour entraîner des Modèles de Raisonnement à Grande Échelle (LRM) via l'apprentissage par renforcement avec récompenses vérifiables (RLVR). Alors que des méthodes comme l'optimisation de politique relative de groupe (GRPO) ont connu du succès dans la résolution de problèmes mathématiques et la génération de code, elles présentent une tendance à concentrer la masse de probabilité sur une seule trajectoire à haute récompense une fois celle-ci découverte. Ce comportement supprime les stratégies de résolution alternatives, arrête l'exploration et conduit à un raisonnement fragile qui ne parvient pas à généraliser.

Les auteurs soutiennent que ce phénomène n'est pas simplement un artefact d'implémentation, mais une conséquence directe de l'objectif d'optimisation : le RL on-policy standard minimise implicitement la divergence de Kullback-Leibler (KL) inverse ($D_{KL}(\pi_\theta \| p^*)$). Le KL inverse est intrinsèquement « chercheur de mode », amplifiant la première trajectoire réussie rencontrée plutôt que de maintenir une distribution sur plusieurs solutions à haute récompense et diversifiées. En revanche, la divergence de KL directe ($D_{KL}(p^* \| \pi_\theta)$) présente un comportement « couvrant les modes », encourageant la politique à correspondre au support complet de la distribution cible, mais elle est généralement intraitable car elle nécessite un échantillonnage à partir de la distribution globale pondérée par la récompense.

2. Méthodologie : DMPO

Pour remédier à l'effondrement de mode sans nécessiter un échantillonnage global intractable, les auteurs proposent l'Optimisation de Politique d'Appariement de Distributions (DMPO).

Insight central

Le DMPO approxime la minimisation du KL direct au niveau du groupe plutôt que globalement. Au lieu d'échantillonner à partir d'une distribution cible globale, l'algorithme construit une distribution cible sur un groupe spécifique de trajectoires échantillonnées pour un prompt donné.

Formulation algorithmique

1. Distribution cible au niveau du groupe : Pour un groupe de $G$ trajectoires $\mathcal{O} = \{o_1, \dots, o_G\}$ échantillonnées à partir de la politique actuelle $\pi_{\theta_{old}}$, une distribution cible $p(o_i | \mathcal{O})$ est construite proportionnelle à l'exponentielle de leurs récompenses (une distribution de Boltzmann) :
   $$p(o_i | \mathcal{O}) = \frac{\exp(r(o_i)/\alpha)}{\sum_{j=1}^G \exp(r(o_j)/\alpha)}$$
   Ceci est calculable car cela ne nécessite que la fonction de partition au niveau du groupe.

2. Distribution de politique au niveau du groupe : Pour aligner la politique sur cette cible, la méthode définit une distribution de politique $q_\theta(o_i | \mathcal{O})$ basée sur les log-vraisemblances normalisées par la longueur pour tenir compte du biais lié à la longueur des trajectoires :
   $$q_\theta(o_i | \mathcal{O}) = \frac{\exp(\varphi(o_i))}{\sum_{j=1}^G \exp(\varphi(o_j))}$$
   où $\varphi(o_i)$ est la log-probabilité moyenne par token.

3. Perte d'appariement de distributions : La politique est alignée sur la cible en minimisant la divergence entre $q_\theta$ et $p$. Les auteurs utilisent l'Erreur Quadratique Moyenne (MSE) au lieu de la divergence KL pour la stabilité numérique, empêchant l'explosion du gradient lorsque la probabilité de la politique tend vers zéro :
   $$\mathcal{L}_{DM}(\theta) = \frac{1}{G} \sum_{i=1}^G (p(o_i | \mathcal{O}) - q_\theta(o_i | \mathcal{O}))^2$$

4. Objectif unifié : L'objectif final combine le terme de maximisation de récompense standard GRPO avec la régularisation d'appariement de distributions :
   $$\mathcal{L}_{DMPO}(\theta) = \mathcal{L}_{GRPO}(\theta) + \lambda \mathcal{L}_{DM}(\theta)$$
   Ici, $\mathcal{L}_{GRPO}$ favorise l'exploitation (maximisation de la récompense), tandis que $\mathcal{L}_{DM}$ impose l'exploration en empêchant la politique de s'effondrer vers un seul mode au sein de la région à haute récompense.

3. Contributions clés

A. Algorithme DMPO

L'article présente le DMPO comme une solution principielle et pratique à l'effondrement de mode. En approximating le KL direct au niveau du groupe, il atteint un comportement couvrant les modes sans nécessiter de fonction de partition apprise ni d'échantillonnage global. La méthode est compatible avec les cadres GRPO existants, ne nécessitant que l'ajout d'un seul terme de régularisation.

B. MM-NP-Bench

Les auteurs introduisent MM-NP-Bench, une référence multimodale conçue pour évaluer rigoureusement l'exploration dans les modèles de raisonnement.

• Portée : Elle étend le NP-Bench basé sur le texte pour inclure des représentations visuelles de 10 tâches d'optimisation combinatoire NP-difficiles (par exemple, TSP, coloration de graphes, couverture de sommets).
• Évaluation à double métrique : Pour rendre l'effondrement de mode observable, la référence utilise deux métriques :
  • Taux de réussite (SR) : Mesure la satisfaction des contraintes (faisabilité).
  • Ratio de qualité (QR) : Mesure la qualité de l'optimisation par rapport à une base heuristique.
  • Signification : L'effondrement de mode se manifeste par un SR élevé mais un QR faible (trouver des solutions valides mais sous-optimales).
• Infrastructure : Comprend des générateurs paramétriques, des vérificateurs basés sur des règles et des solveurs heuristiques pour fournir des bases quasi optimales.

C. Validation empirique

L'article fournit des expériences extensives démontrant que le DMPO surpasse des bases solides (GRPO, GSPO, GPG, FlowRL, etc.) dans plusieurs domaines.

4. Résultats

Performance sur les tâches d'optimisation

• MM-NP-Bench (Multimodal) : Le DMPO a atteint un Ratio de Qualité de 43,1 %, soit une amélioration relative de 12 % par rapport au GRPO (38,4 %). Il a également obtenu le taux de réussite le plus élevé (61,9 % contre 55,7 %).
• NP-Bench (Texte uniquement) : Le DMPO a atteint un Ratio de Qualité de 43,9 %, soit une amélioration relative de 9 % par rapport au GRPO (40,1 %).
• Ablation : La méthode a montré une robustesse aux choix d'hyperparamètres ($\lambda \in [0,3, 0,7]$ et $\alpha \in [0,8, 1,2]$).

Généralisation et transfert

• Raisonnement mathématique : L'entraînement sur des tâches d'optimisation avec DMPO a amélioré les performances sur les références de raisonnement mathématique (AIME, AMC, etc.) de 2,0 % par rapport au GRPO.
• Tâches hors domaine : Le DMPO a démontré une amélioration de 2,3 % sur les références de raisonnement visuel hors domaine (LogicVista, MathVista, etc.), suggérant que l'entraînement préservant la diversité améliore les capacités de raisonnement général.

Dynamique d'entraînement

L'analyse des journaux d'entraînement a révélé que, tandis que le GRPO souffre d'un « effondrement de longueur » (dégénérescence en réponses courtes et triviales) et de plateaux de récompense prématurés, le DMPO maintient une amélioration continue et préserve des chaînes de raisonnement robustes (~400 tokens). Le DMPO maintient une entropie plus faible que le GRPO, indiquant une « diversité structurée » (exécution confiante de différentes stratégies logiques) plutôt qu'une confusion stochastique.

5. Signification et affirmations

L'article affirme établir l'appariement de distributions comme une approche pratique et principielle pour prévenir l'effondrement de mode dans le RL on-policy. Les auteurs soutiennent que :

1. Identification de la cause racine : L'effondrement de mode dans le RLVR est fondamentalement causé par la nature chercheuse de modes de la minimisation du KL inverse.
2. Solution réalisable : La minimisation du KL direct, traditionnellement considérée comme intractable, peut être efficacement approximée au niveau du groupe pour préserver la diversité des solutions.
3. Bénéfice général : Préserver la diversité lors de l'entraînement sur des tâches d'optimisation spécifiques se transfère en des capacités de raisonnement général améliorées, suggérant que l'exploration est un composant critique d'un raisonnement robuste dans les Modèles de Raisonnement à Grande Échelle.

Le travail positionne le DMPO non pas simplement comme un régularisateur de diversité, mais comme un mécanisme permettant la découverte de solutions de meilleure qualité que les objectifs chercheurs de modes manquent, améliorant ainsi à la fois les performances d'optimisation et la robustesse du raisonnement à travers les modalités.