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

Probleemstelling

Moderne grote taalmodellen (LLM's) worden doorgaans getraind met versterkende leerling (Reinforcement Learning, RL) om per query één "beste" antwoord te genereren. Dit proces neigt ertoe om de onderliggende kansverdeling over mogelijke antwoorden te laten instorten naar één dominante modus (mode collapse). Hoewel dit effectief is voor benchmarks met één correct antwoord, is het fundamenteel ongeschikt voor real-world scenario's waar:

1. Meerdere juiste antwoorden mogelijk zijn: Bijvoorbeeld in medische diagnosestelling (waar meerdere aandoeningen tegelijkertijd plausibel zijn) of bij het coderen (waar verschillende algoritmen tot dezelfde oplossing leiden).
2. Onzekerheid inherent is: Door incomplete of ambiguë informatie (bijv. epistemische onzekerheid).

Bestaande methoden om hiermee om te gaan, zoals het herhaaldelijk afnemen van samples tijdens de inferentie (inference-time scaling zoals best-of-k), zijn computatie-intensief en leiden vaak tot redundante redeneringen. Daarnaast missen deze modellen vaak goed gekalibreerde onzekerheidsschattingen voor de gegenereerde antwoorden.

Methodologie: Multi-Answer Reinforcement Learning

De auteurs stellen Multi-Answer RL voor, een aanpak die taalmodellen traint om expliciet een verdeling over meerdere plausibele antwoorden te genereren in één enkele forward pass, in plaats van één enkel antwoord.

Kerncomponenten:

1. Structuur van de Output: Het model wordt geprompt om een gestructureerde set van $K$ distincte kandidaat-antwoorden ($A = \{a_1, ..., a_K\}$) te genereren binnen één denkproces (Chain-of-Thought).
2. Multi-Answer RLVR (Verifiable Rewards):
   • Voor taken met verifieerbare antwoorden (zoals wiskunde of coderen) wordt een set-niveau beloning gebruikt.
   • De beloning $R_{multi}^{RLVR}$ telt hoeveel van de gegenereerde antwoorden in de set van grondwaarheid ($Y^*$) zitten.
   • Dit generaliseert de standaard binary correctness naar set-level correctness. Het moedigt het model aan om een zo breed mogelijk bereik van correcte antwoorden te dekken.
3. Multi-Answer RLCR (Calibration Rewards):
   • Voor taken met onzekerheid wordt het model getraind om naast de antwoorden ook een set van betrouwbaarheidsscores ($Q = \{q_1, ..., q_K\}$) te genereren.
   • De beloning combineert correctheid met een proper scoring rule (specifiek de Multi-Brier score).
   • De formule is: $R_{multi}^{RLCR} = R_{multi}^{RLVR} - R_{multi}^{Brier}$.
   • Dit traint het model om de kans dat een specifiek antwoord correct is, nauwkeurig te schatten, wat leidt tot goed gekalibreerde onzekerheid.
4. Uniekheid: Er wordt een format-beloning toegevoegd om te garanderen dat de gegenereerde antwoorden uniek zijn, wat voorkomt dat het model dezelfde oplossing herhaalt.

Belangrijkste Bijdragen

• Inherent Distributioneel Redeneren: In plaats van post-hoc sampling, internaliseert het model het zoeken naar meerdere hypotheses tijdens het generatieproces.
• Efficiëntie: Het genereren van een set antwoorden in één keer is aanzienlijk efficiënter dan het herhaaldelijk afnemen van samples van een single-answer model, omdat het redundantie in het denkproces elimineert.
• Uitgebreide Kalibratie: De methode introduceert een manier om niet alleen antwoorden, maar ook de bijbehorende onzekerheidsschattingen per antwoord in de set te kalibreren.
• Flexibiliteit: De aanpak werkt zowel voor scenario's met meerdere grondwaarheden (bijv. medische diagnose) als voor scenario's met één grondwaarheid maar hoge ambiguïteit (bijv. HotPotQA).

Resultaten

De methode is geëvalueerd op drie datasets: DDXPlus (medische diagnose), HotPotQA-Modified (ambiguë vragen), en MBPP (coderen).

• Dekking en Diversiteit: Multi-Answer RL (zowel RLVR als RLCR) presteert aanzienlijk beter dan single-answer baselines op metrics voor dekking (coverage) en diversiteit. Op de coderingsdataset (MBPP) steeg de top-1 nauwkeurigheid met meer dan 50% terwijl het aantal tokens met meer dan de helft werd verminderd.
• Token-efficiëntie: Het genereren van 3 antwoorden via Multi-Answer RL vereist aanzienlijk minder tokens dan het genereren van 3 onafhankelijke samples van een single-answer model. Dit komt doordat het single-answer model vaak dezelfde redeneringsketens herhaalt voordat het een ander antwoord vindt.
• Kalibratie: Modellen getraind met Multi-Answer RLCR tonen aanzienlijk betere kalibratie (lagere Brier scores en ECE) dan modellen getraind met RLVR. Ze zijn minder overconfident en kunnen betekenisvolle onzekerheidsschattingen geven voor meerdere mogelijke diagnoses.
• Stabiliteit: De prestaties blijven stabiel naarmate het aantal gevraagde antwoorden ($K$) toeneemt, waarbij het model steeds meer unieke correcte hypotheses blijft vinden zonder te degenereren.

Betekenis en Conclusie

Dit onderzoek positioneert Multi-Answer RL als een principieel en rekenkundig efficiënt alternatief voor externe inferentie-schaaltechnieken. Het lost het probleem van "mode collapse" op door diversiteit en onzekerheid expliciet te optimaliseren tijdens het trainingsproces.

De resultaten suggereren dat taalmodellen beter kunnen worden ingezet voor complexe, real-world taken (zoals medische besluitvorming of creatief coderen) waar het vinden van een reeks plausibele opties en het begrijpen van de bijbehorende onzekerheid cruciaal is, zonder de kosten van uitgebreide inferentie-time sampling. Het werk markeert een verschuiving van het trainen van modellen voor "het ene juiste antwoord" naar het trainen voor "een goed onderbouwde verdeling van mogelijke antwoorden".