Whatever Remains Must Be True: Filtering Drives Reasoning in LLMs, Shaping Diversity

En proposant une méthode qui approxime une distribution cible par filtrage des réponses incorrectes via la famille des divergences $\alpha$, cette étude surmonte la perte de diversité des modèles de langage entraînés par apprentissage par renforcement, permettant d'atteindre un état de l'art sur le front de Pareto précision-coverage pour la preuve de théorèmes.

L'essentiel

🧠 Le Dilemme du Génie : Précision ou Créativité ?

Imaginez que vous avez un génie très intelligent (c'est notre modèle d'IA) qui sait résoudre des problèmes de mathématiques complexes. Au début, ce génie est très curieux : il essaie des milliers de chemins différents pour trouver la solution. Certains chemins sont absurdes, d'autres sont brillants, et d'autres sont juste "corrects".

Le problème, c'est que pour le rendre encore plus performant, les chercheurs ont utilisé une méthode appelée Apprentissage par Renforcement (RL). C'est un peu comme si on donnait au génie une récompense (un bonbon) chaque fois qu'il trouve la bonne réponse, et qu'on le gronde s'il se trompe.

Ce qui s'est passé :
Le génie a compris le jeu. Il a arrêté d'essayer des choses nouvelles. Il a commencé à répéter exactement la même solution parfaite qu'il a découverte une fois, encore et encore.

• Avantage : Il est devenu ultra-précis sur cette solution.
• Inconvénient : Il a perdu sa créativité. Si la première solution qu'il trouve est bloquée ou si le problème change légèrement, il est perdu. Il a oublié toutes les autres façons de résoudre le problème. C'est ce qu'on appelle la "collapse de mode" (effondrement vers une seule option).

🎯 La Nouvelle Approche : "Tout ce qui reste doit être vrai"

Les auteurs de ce papier disent : "Attendez, on ne veut pas que le génie oublie ses autres idées !"

Ils proposent une nouvelle méthode, qu'ils appellent DMVR (Distributional Matching with Verifiable Rewards). Voici comment cela fonctionne avec une analogie simple :

1. Le Filtre Magique (Le Vérificateur)

Imaginez que vous avez un tas de réponses générées par le génie. Vous avez un filtre magique (un vérificateur mathématique) qui ne laisse passer que les réponses correctes.

• Les mauvaises réponses sont jetées à la poubelle.
• Les bonnes réponses sont gardées.

2. Le Choix du Chef d'Orchestre (La Divergence)

C'est ici que la magie opère. La question est : Comment le génie doit-il apprendre à garder ces bonnes réponses ?

• L'ancienne méthode (RL classique) : C'est comme un chef d'orchestre qui force tous les musiciens à jouer exactement la même note, la plus forte possible, pour être sûr qu'elle soit juste. Résultat : une note parfaite, mais une musique monotone. Le génie se concentre sur une seule "zone" de succès et oublie le reste.
• La nouvelle méthode (Alpha-DPG) : Les chercheurs proposent d'utiliser un réglage fin (le paramètre $\alpha$).
  • Si on règle le bouton vers la Précision, le génie se concentre sur les solutions les plus probables (comme avant).
  • Si on règle le bouton vers la Diversité, le génie garde toutes les solutions correctes, même celles qui sont un peu plus rares ou étranges, tant qu'elles sont justes.

🎨 L'Analogie du Peintre

Imaginez que vous demandez à un peintre de dessiner un pomme rouge.

• L'IA classique (RL) : Elle va dessiner une pomme rouge parfaite, mais elle va dessiner uniquement cette pomme-là, exactement au même endroit, avec le même éclairage, 1000 fois de suite. Si on lui demande une pomme verte, elle panique.
• La nouvelle méthode (Alpha-DPG) : Le peintre regarde toutes les pommes rouges qu'il a déjà dessinées dans sa vie. Il jette celles qui sont vertes ou pourries. Il garde toutes les pommes rouges : celles qui sont lisses, celles qui sont tachées, celles qui sont grandes, celles qui sont petites.
  • Il peut choisir de se concentrer sur les pommes les plus "parfaites" (haute précision).
  • Ou il peut choisir de garder toute la variété de pommes rouges (haute diversité), sachant que si l'une d'elles ne marche pas, il en a 100 autres sous la main.

🏆 Pourquoi c'est important ?

Dans le monde réel, surtout pour les mathématiques ou la programmation, il ne suffit pas de trouver une solution. Parfois, la première solution est trop complexe, ou le contexte change.

Grâce à cette méthode, les chercheurs ont créé une courbe de performance idéale (la frontière de Pareto) :

1. Ils peuvent avoir un modèle très précis (qui trouve la solution du premier coup).
2. OU un modèle très diversifié (qui trouve la solution en essayant 256 fois, mais qui explore des chemins très différents).
3. OU, le plus important, ils peuvent ajuster le curseur pour avoir un peu des deux, selon les besoins.

En résumé

Cette recherche nous dit : "Ne forcez pas l'IA à être un robot qui répète une seule bonne réponse. Donnez-lui un filtre pour éliminer les erreurs, mais laissez-lui la liberté de garder toutes les façons différentes de réussir."

C'est comme passer d'un éléphant qui marche sur une seule ligne droite (précis mais fragile) à une fourmilière intelligente qui explore tous les chemins possibles pour trouver la nourriture (divers et robuste).

Résumé technique

1. Problématique : La perte de diversité dans le RLVR

Les modèles de langage (LLM) sont de plus en plus affinés pour des tâches de raisonnement complexe (comme la preuve de théorèmes) grâce à l'apprentissage par renforcement à partir de récompenses vérifiables (RLVR). Des méthodes comme PPO ou GRPO optimisent une politique pour maximiser la probabilité de générer des réponses correctes, vérifiées par un outil externe (ex: Lean, un assistant de preuve).

Cependant, l'article identifie un problème majeur : la perte de diversité (ou "effondrement de mode"). Les modèles entraînés par RLVR tendent à se concentrer sur un sous-ensemble restreint de solutions correctes, ignorant d'autres solutions valides présentes dans la distribution de base.

L'hypothèse centrale des auteurs :
Ce phénomène n'est pas dû à la nature des récompenses, mais à la fonction de divergence implicite optimisée par le RLVR. Le RLVR minimise la divergence de Kullback-Leibler (KL) inverse ($D_{KL}(\pi || p)$), qui est une divergence "chercheuse de mode" (mode-seeking).

• Comportement : Elle pénalise sévèrement le modèle s'il attribue de la probabilité à des régions où la distribution cible $p$ est nulle, mais elle est indifférente si le modèle ignore des modes (solutions) de $p$.
• Conséquence : Le modèle "écrase" la distribution pour se concentrer sur quelques pics de haute probabilité, sacrifiant la couverture de l'espace des solutions.

2. Méthodologie : DMVR et $\alpha$-DPG

Les auteurs proposent un cadre unifié appelé DMVR (Distributional Matching with Verifiable Rewards) et une méthode d'optimisation flexible basée sur les divergences $\alpha$.

A. Définition de la distribution cible explicite

Au lieu d'optimiser une récompense pseudo-stochastique, les auteurs définissent explicitement la distribution cible idéale $p_x$ pour un problème $x$ :
$$p_x(y) \propto \pi_{base}(y|x) \cdot v(y, x)$$
Où :

• $\pi_{base}$ est le modèle de base.
• $v(y, x)$ est le vérificateur binaire (1 si correct, 0 sinon).

Cette distribution filtre toutes les réponses incorrectes tout en préservant les probabilités relatives des réponses correctes du modèle de base. Elle garantit à la fois la correction et la diversité maximale possible.

B. L'approche $\alpha$-DPG

Pour approximer cette cible $p_x$ avec une politique $\pi_\theta$, les auteurs utilisent l'algorithme Distributional Policy Gradient (DPG) généralisé aux divergences $\alpha$ ($\alpha$-DPG).

La famille des divergences $\alpha$ permet d'interpoler continûment entre deux comportements extrêmes :

1. $\alpha \to 0$ (KL Forward) : $D_{KL}(p || \pi)$. C'est une divergence "couvrante" (mass-covering). Elle pénalise le modèle s'il ne couvre pas les modes de la cible, favorisant la diversité mais risquant d'inclure des régions de faible récompense.
2. $\alpha \to 1$ (KL Inverse) : $D_{KL}(\pi || p)$. C'est la divergence utilisée par le RLVR classique. Elle favorise la précision (concentration sur les modes de haute probabilité) mais sacrifie la diversité.
3. $\alpha \in (0, 1)$ : Un compromis contrôlable.

La fonction de récompense pseudo ($\hat{R}_\theta$) pour $\alpha$-DPG est dérivée de la dérivée de la fonction génératrice de la divergence $\alpha$ :
$$\hat{R}_\theta(y, x) = \min\left( \left(\frac{p_x(y)}{\pi_\theta(y|x)}\right)^{1-\alpha} - 1, M \right)$$
(Note : Un clipping $M$ est appliqué pour stabiliser l'entraînement lorsque $\alpha$ est faible).

3. Contributions Clés

1. Cadre DMVR : Introduction d'un cadre théorique qui reformule l'entraînement RLVR comme une tâche d'appariement de distributions (Distributional Matching) vers une cible filtrée explicite.
2. Diagnostic théorique : Démonstration que le RLVR optimise implicitement une divergence KL inverse vers une distribution lissée, expliquant mathématiquement la perte de diversité observée empiriquement.
3. Méthode $\alpha$-DPG : Proposition d'une méthode unifiée qui englobe le RLVR (KL inverse), le KL-DPG (KL forward) et le Rejection Sampling Fine-Tuning (RS-FT) comme cas limites, permettant de naviguer sur le front de Pareto entre précision et diversité via le paramètre $\alpha$.
4. Validation empirique : Résultats state-of-the-art sur le benchmark Lean (preuve de théorèmes), montrant que l'on peut obtenir à la fois une haute précision (pass@1) et une haute couverture (pass@256) en ajustant $\alpha$.

4. Résultats Expérimentaux

Les expériences ont été menées sur le modèle DeepSeek-Prover-V1.5-SFT (7B paramètres) sur le dataset Lean Workbook.

• Front de Pareto Précision-Couverture :
  • Les modèles $\alpha$-DPG occupent le front de Pareto optimal.
  • Les valeurs faibles de $\alpha$ (ex: 0.25) maximisent la couverture (pass@256), surpassant toutes les méthodes précédentes (GRPO, RLOO, ReMax) tout en améliorant la précision par rapport au modèle de base.
  • Les valeurs élevées de $\alpha$ (ex: 0.999) atteignent une précision (pass@1) équivalente ou supérieure aux méthodes RL classiques, tout en conservant une meilleure couverture.
• Analyse de la difficulté des problèmes :
  • Les méthodes RL classiques (GRPO) améliorent la résolution des problèmes "moyens" en les rendant "faciles", mais dégradent souvent la résolution des problèmes "difficiles" (les rendant insolubles), signe d'un effondrement de mode.
  • $\alpha$-DPG (avec $\alpha$ faible) est plus conservateur : il améliore l'efficacité d'échantillonnage sur moins de problèmes, mais préserve la capacité à résoudre les problèmes difficiles.
• Analyse de la diversité :
  • Une corrélation positive est observée entre la diversité des tactiques et des prémisses utilisées dans les preuves et la performance en pass@256.
  • Les modèles $\alpha$-DPG à faible $\alpha$ maintiennent une entropie de Shannon et un indice de Simpson élevés, contrairement aux modèles GRPO qui s'effondrent rapidement.
• Analyse de Perplexité :
  • Les solutions générées par les modèles affinés (y compris GRPO) sont déjà très probables sous le modèle de base. Cela suggère que le RL n'invente pas de nouvelles capacités de raisonnement, mais réattribue simplement les probabilités des solutions existantes.

5. Signification et Conclusion

Cet article remet en question le dogme selon lequel le RL est nécessaire pour "découvrir" de nouvelles capacités de raisonnement. Il démontre que :

1. Le modèle de base contient déjà les solutions, mais avec une distribution de probabilité trop diffuse.
2. Le problème du RLVR n'est pas la cible (qui est correcte), mais la méthode d'approximation (divergence KL inverse) qui force une concentration excessive.
3. En utilisant des divergences $\alpha$ via le cadre DMVR, on peut contrôler explicitement le compromis entre la fiabilité (précision) et l'exploration (diversité).

Impact : Cette approche est particulièrement cruciale pour les tâches où la diversité est essentielle (comme la découverte mathématique ou la génération de code varié), car elle permet d'éviter l'effondrement de mode tout en garantissant la correction formelle des sorties. Elle offre une alternative plus stable et interprétable aux méthodes RL classiques pour l'alignement des LLMs.