Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models

1. Problématique

Les grands modèles de langage (LLM) modernes sont généralement entraînés via des techniques d'apprentissage par renforcement (RL) pour produire une seule réponse optimale par requête. Cette approche, souvent basée sur des récompenses binaires de correction (RLVR), tend à faire "effondrer" la distribution de probabilité du modèle vers un mode unique dominant.

Bien que cela fonctionne bien pour les benchmarks à réponse unique, cela pose problème dans des scénarios réels complexes (diagnostic médical, questions ambiguës, codage) où :

• L'incertitude est irréductible : Plusieurs réponses peuvent être valides simultanément (ex: plusieurs diagnostics possibles pour un ensemble de symptômes).
• L'information est incomplète : Le modèle doit explorer plusieurs hypothèses plausibles plutôt que de parier sur une seule.
• La diversité est cruciale : Les méthodes d'inférence actuelles (comme l'échantillonnage répété "best-of-k") sont coûteuses en calcul et inefficaces, car le modèle est entraîné pour se concentrer sur une seule réponse, ce qui entraîne une redondance dans les chaînes de raisonnement générées.

L'objectif est donc d'entraîner les modèles à générer directement une distribution de réponses plausibles (un ensemble de candidats avec des estimations de confiance) en une seule passe forward, sans recourir à un échantillonnage itératif coûteux.

2. Méthodologie : Multi-Answer Reinforcement Learning (Multi-Answer RL)

Les auteurs proposent une nouvelle approche d'apprentissage par renforcement qui optimise explicitement la génération d'ensembles de réponses plutôt que d'une seule.

A. Objectif RLVR Multi-Answer (RLVR Multi)

Au lieu de maximiser la récompense pour une seule réponse $y$, le modèle est entraîné à générer un ensemble structuré de $K$ réponses distinctes $A = \{a_1, ..., a_K\}$.
La fonction de récompense est définie au niveau de l'ensemble :
$$R_{RLVR}^{multi}(A, Y^*) = \sum_{i=1}^{K} \mathbb{1}[a_i \in Y^*]$$
où $Y^*$ est l'ensemble des réponses correctes (ground-truth). Cela encourage le modèle à couvrir un maximum de réponses valides dans un seul passage. Une récompense de format est ajoutée pour garantir l'unicité des réponses générées.

B. Objectif RLCR Multi-Answer (RLCR Multi)

Pour intégrer l'estimation de l'incertitude, les auteurs étendent l'approche en demandant au modèle de générer non seulement les réponses, mais aussi un score de confiance $q_i \in [0, 1]$ pour chaque réponse.
La récompense combine la correction et l'étalonnage (calibration) via une règle de score propre (le Multi-Brier Score) :
$$R_{RLCR}^{multi}(A, Q, Y^*) = R_{RLVR}^{multi}(A, Y^*) - \frac{1}{K} \sum_{i=1}^{K} (q_i - \mathbb{1}[a_i \in Y^*])^2$$
Cela force le modèle à associer une probabilité de justesse précise à chaque hypothèse générée, produisant ainsi une distribution explicite $P$ sur les réponses.

C. Architecture et Entraînement

• Base : Utilisation de modèles Qwen3-8B.
• Algorithme : GRPO (Group Relative Policy Optimization) modifié.
• Format : Le modèle génère une chaîne de pensée (CoT) unique suivie de balises structurées (<answer1>, <confidence1>, etc.) contenant les $K$ candidats.
• Contrainte d'unicité : Si les réponses extraites ne sont pas uniques, la récompense est annulée pour forcer la diversité.

3. Contributions Clés

1. Changement de paradigme d'entraînement : Passage d'un objectif "mode unique" à un objectif "distributionnel" intégré directement dans la phase d'entraînement.
2. Efficacité computationnelle : Démontre que générer un ensemble de réponses en une seule passe est nettement plus efficace (moins de tokens) que l'échantillonnage répété de modèles entraînés pour une seule réponse.
3. Étalonnage (Calibration) : Introduction d'une méthode permettant d'obtenir des scores de confiance calibrés pour chaque élément d'un ensemble de réponses, crucial pour les applications à haut risque.
4. Validation sur des tâches complexes : Application réussie sur des tâches où la vérité terrain est un ensemble (DDXPlus médical), ambiguë (HotPotQA modifié) ou à solutions multiples (MBPP pour le codage).

4. Résultats Expérimentaux

Les expériences ont été menées sur trois jeux de données : DDXPlus (diagnostic médical), HotPotQA-Modified (QA ambiguë) et MBPP (codage).

• Couverture et Diversité :

  • Les modèles Multi-Answer RL surpassent significativement les baselines "Single-Answer" (même avec échantillonnage répété) en termes de couverture (nombre moyen de réponses correctes par ensemble).
  • Sur DDXPlus, la couverture augmente de 0,62 (Single) à 0,79 (Multi).
  • La diversité des réponses est nettement supérieure : les modèles Single tendent à répéter les mêmes réponses (effondrement de mode), tandis que les modèles Multi génèrent des hypothèses structurellement différentes.

• Efficacité (Tokens) :

  • Pour générer 3 réponses, le modèle Multi-Answer utilise moins de la moitié des tokens par rapport à un modèle Single-Answer échantillonné 3 fois.
  • Sur MBPP (codage), l'approche Multi-Answer réduit l'utilisation de tokens de plus de 50% tout en augmentant la précision Top-1 de plus de 50% (de 0,29 à 0,49).

• Étalonnage (Calibration) :

  • L'approche RLCR Multi produit des scores de confiance bien mieux étalonnés que RLVR Multi.
  • Les courbes de calibration montrent que RLVR Multi est systématiquement trop confiant (overconfident), tandis que RLCR Multi suit plus fidèlement la ligne d'identité, même si une légère déviation persiste aux confiances élevées.

• Stabilité : L'entraînement reste stable même lorsque le nombre de réponses générées ($K$) augmente (de 2 à 5), permettant au modèle de découvrir progressivement plus d'hypothèses correctes.

5. Signification et Impact

Ce travail remet en question la pratique standard de l'optimisation des LLM vers une réponse unique. Il démontre que :

1. L'exploration est inhérente au raisonnement : Dans des domaines comme la médecine ou le codage, la capacité à maintenir plusieurs hypothèses actives est supérieure à la capacité de choisir une seule réponse "parfaite" trop tôt.
2. L'efficacité par l'entraînement : L'inférence par échantillonnage répété (inference-time scaling) est souvent une solution de contournement coûteuse. L'entraînement explicite à la génération d'ensembles permet d'internaliser cette recherche, réduisant drastiquement le coût computationnel.
3. Fiabilité accrue : En fournissant non seulement des réponses alternatives mais aussi des estimations de probabilité calibrées, les modèles deviennent plus adaptés aux applications critiques où l'incertitude doit être quantifiée et communiquée.

En conclusion, le Multi-Answer RL offre une alternative principielle et efficace aux méthodes d'inférence actuelles, permettant aux modèles de langage de raisonner de manière distributionnelle et de mieux refléter la complexité du monde réel.