Technische Samenvatting: Representatie-gebaseerde Exploratie voor Taalmodellen

Probleemstelling

Reinforcement learning (RL) heeft veelbelovende resultaten laten zien bij het verbeteren van taalmodellen (LM's) voor taken met verifieerbare beloningen, zoals wiskundig redeneren en codegeneratie. Echter, een kritieke beperking van huidige RL post-training technieken is de onzekerheid over de vraag of ze werkelijk nieuwe gedragingen ontdekken of slechts bestaande capaciteiten in het basismodel (die met een bescheiden waarschijnlijkheid aanwezig zijn) "verscherpen" (versterken). Bestaand bewijs suggereert dat standaard RL-recepten vaak er niet in slagen om capaciteiten te ontsluiten die afwezig zijn in het voorgetrainde model, maar in plaats daarvan convergeren op bekende oplossingen. Bovendien worden datacuratie en algoritmische interventies om dit te mitigeren een bottleneck, met name in complexe, open eind domeinen.

Het artikel stelt dat bewuste exploratie—het expliciet stimuleren van het model om werkelijk nieuwe en diverse gedragen te ontdekken—essentieel is om het volledige potentieel van RL voor taalmodelredeneren te realiseren. De centrale uitdaging is het kwantificeren van nieuwheid en diversiteit binnen de combinatorisch grote ruimte van taal op een schaalbare manier, terwijl de voorkennis die in voorgetrainde modellen is ingebed wordt gebruikt om deze zoektocht efficiënt te sturen.

Methodologie

De auteurs stellen een representatie-gebaseerde exploratiestrategie voor, die gebruikmaakt van elliptische bonussen afgeleid van de verborgen toestanden (hidden states) van het voorgetrainde taalmodel. De methodologie wordt geëvalueerd via een tweeledige aanpak:

1. Selectie tijdens de Inferentie

Om de rol van diversiteit te isoleren van optimalisatie en generalisatie, definiëren de auteurs eerst een selectieprobleem tijdens de inferentie. Gegeven een prompt $x$ en een grote pool van $N$ kandidaat-responsen gegenereerd door een vast model $\pi$, is het doel om een subset van $k$ responsen te selecteren die maximaal divers is en een hoge waarschijnlijkheid heeft om een correct antwoord te bevatten.

• Mechanisme (RepExp): Het algoritme selecteert iteratief responsen om een elliptische bonus te maximaliseren. Voor een respons $y$ wordt de featurevector $\bar{h}_\theta(x, y)$ geconstrueerd door de laatste laag verborgen toestanden van het model te middelen over de tokens van de respons.
• Diversiteitsmetriek: De nieuwheid van een kandidaat $y$ gegeven een set eerder geselecteerde responsen wordt gedefinieerd als $div(y) = \bar{h}_\theta(x, y)^\top \Sigma^{-1} \bar{h}_\theta(x, y)$, waarbij $\Sigma$ de covariantie-matrix is van de geselecteerde features. De inverse covariantie wordt efficiënt bijgewerkt met behulp van de Woodbury-identiteit.
• Doel: Deze aanpak selecteert een "coreset" van diverse responsen om een verifier te bevragen, met als doel de efficiëntie van de verifier te verbeteren (het aantal benodigde samples om een correct antwoord te vinden te verminderen).

2. RL Post-Training

De auteurs integreren deze exploratiestrategie in de GRPO (Group Relative Policy Optimization) post-training pipeline.

• Beloningsaugmentatie: In plaats van een subset te selecteren, wordt de elliptische bonus direct toegevoegd aan de extrinsieke beloning. Voor een rollout $y_i$ wordt de beloning $r(x, y_i) + \beta \cdot \bar{h}_\theta(x, y_i)^\top \Sigma^{-1} \bar{h}_\theta(x, y_i)$.
• Implementatie: De covariantie-matrix $\Sigma$ wordt voor elke batch van trainingsrollouts opnieuw geïnitialiseerd om de nieuwheid te meten ten opzichte van de huidige batch. De featurevectoren worden naar een lagere dimensie geprojecteerd (bijv. 512) via sparse random projectie om schaalbaarheid te garanderen.

3. Token-niveau Extensie (Voorlopig)

Als proof of concept verkennen de auteurs ook het modificeren van het autoregressieve generatieproces zelf. Door de token-niveau logits te verstoren met een elliptische bonus bij elke generatiestap, beogen zij diversiteit te sturen tijdens het generatieproces in plaats van post-hoc selectie.

Belangrijkste Bijdragen en Bevindingen

1. Verbeterde Verifier Efficiëntie (Inference-Time)

Het artikel demonstreert dat representatie-gebaseerde exploratie de verifier efficiëntie aanzienlijk verbetert over een breed scala aan modellen (Phi, Llama, Qwen) en taken (MATH, GSM8K, MBPP+, Game of 24).

• Resultaat: Voor modellen zoals Qwen-2.5-14b-Instruct bereikt de methode een verbetering van meer dan 50% in het aantal benodigde samples om een correct antwoord te vinden vergeleken met willekeurige sampling.
• Correlatie met Modelsterkte: De voordelen van exploratie schalen met de kracht van het model; sterkere modellen (bijv. Qwen-32B) profiteren aanzienlijk, terwijl zwakkere modellen geen verbetering of zelfs degradatie kunnen zien.
• Gevoeligheid voor Moeilijkheidsgraad: De methode levert de grootste efficiëntiewinst op de moeilijkste vragen (bijv. de bovenste 20% moeilijkste in MATH).

2. Verbeterde Post-Training Prestaties en Diversiteit

Het integreren van representatie-gebaseerde exploratie in de RL post-training levert superieure prestaties op vergeleken met standaard GRPO en "Unlikeliness" baselines.

• Pass@k Verbetering: De methode behoudt of verbetert de pass@k ratio's voor grote waarden van $k$, terwijl standaard GRPO vaak de prestaties ten opzichte van het basismodel degradeert voor grote $k$ (een fenomeen bekend als "diversity collapse").
• Sample Efficiëntie: Op de AIME 2024 benchmark bereikt de post-getrainde Qwen-2.5-7b-Instruct met gebruik van RepExp een pass@80 die overeenkomt met de pass@256 van standaard GRPO, wat een 3x verbetering betekent in test-time sample efficiëntie.
• Anti-Sharpening: Analyse van de log-likelihood onder het basismodel onthult dat RepExp responsen genereert die aanzienlijk nieuwer zijn (lagere likelihood) dan die van standaard GRPO, dat de verdeling van het basismodel de neiging heeft te verscherpen. Dit suggereert dat de methode succesvol ontsnapt aan het "sharpening regime".

3. Toepassing op Hoge Kosten Verificatie

In een voorlopig experiment op eiwitsequentiegeneratie (waar verificatie via nat-lab werk duur is), verbeterde de methode de verifier efficiëntie met wel 56,9x vergeleken met willekeurige sampling, wat de potentiële bruikbaarheid aantoont in domeinen waar verificatie kostbaar is.

Betekenis en Claims

Het artikel beweert dat bewuste exploratie, gestuurd door een principiële, op representatie gebaseerde notie van diversiteit, een praktische weg is naar het ontdekken van nieuwe gedragingen in taalmodellen die verder gaan dan louter verscherping.

• Nieuwigheid van Aanpak: In tegen tegenstelling tot eerder werk dat vertrouwt op hulpnetwerken (bijv. Random Network Distillation) of complexe datacuratie, maakt deze methode gebruik van de eigen verborgen toestanden van het voorgetrainde model, waardoor geen extra leerbediening nodig is.
• Validatiestrategie: De auteurs benadrukken de waarde van hun tweestaps-evaluatie (inferentie-tijd selectie gevolgd door post-training) als een manier om de voordelen van diversiteit te ontkoppelen van andere RL-mechanismen.
• Bescheiden Omvang: De auteurs erkennen dat hun huidige implementatie een "eenvoudig, principieel" schema is. Ze beweren niet dat ze alle exploratie-uitdagingen hebben opgelost, maar suggereren dat deze aanpak een sterke baseline biedt en een kader vormt voor het begrijpen van de trade-offs van meer geavanceerde strategieën. Ze merken op dat hoewel token-niveau exploratie veelbelovend is, het momenteel hogere computationele kosten met zich meebrengt die optimalisatie vereisen.

Samenvattend betogen de auteurs dat door expliciet de ontdekking van diverse gedragingen te stimuleren via een op representatie gebaseerde bonus, taalmodellen kunnen worden gestuurd om nieuwe oplossingen te vinden die standaard RL post-training mist, waardoor hun redeneervermogen effectiever wordt uitgebreid.