Small RL Controller, Large Language Model: RL-Guided Adaptive Sampling for Test-Time Scaling

Technische Samenvatting: RL-gestuurde Adaptieve Sampling voor Test-Time Scaling

Probleemstelling
Test-time scaling is een cruciale techniek geworden voor het verbeteren van de redeneervermogens van Large Language Models (LLM's) zonder extra training. Methoden zoals Self-Consistency (SC), Tree-of-Thoughts en Best-of-N sampling verbeteren de kwaliteit van antwoorden door meer rekenkracht tijdens de inferentie toe te wijzen. Deze winst brengt echter aanzienlijke kosten met zich mee in zowel totale berekening als latentie. Bestaande adaptieve sampling-methoden, zoals Adaptive Self-Consistency (ASC) en Early-Stopping Self-Consistency (ESC), proberen deze kosten te beperken door dynamisch te beslissen wanneer het samplen moet stoppen. Deze benaderingen lijden echter aan significante beperkingen: ze vertrouwen vaak op door mensen ontworpen heuristieken of distributionele aannames, vereisen hulp-signalen (zoals interne modelvertrouwen, verborgen toestanden of vraagmoeilijkheid) die mogelijk niet beschikbaar zijn, en kunnen invasief zijn voor standaard inferentie-pipelines. Er is behoefte aan een lichtgewicht, principieel methode die het evenwicht bewaart tussen antwoordcorrectheid, latentie en computationele kosten zonder hulpkenmerken of modelinterventie te vereisen.

Methodologie
De auteurs stellen RL-Guided Sampling voor, een framework dat adaptieve sampling formuleert als een finite-horizon Markov Decision Process (MDP).

• MDP Formulering:

  • Toestand ($s_t$): De toestand wordt uitsluitend geconstrueerd uit de statistieken van de huidige pool van gesamplede antwoorden ($D_t$). Het bevat de tellingen van de top-$K$ meest frequente antwoordklassen, het totaal aantal gegenereerde samples en de entropie van deze tellingen. Cruciaal is dat de toestand niet afhankelijk is van query-specifieke informatie of interne modelsignalen.
  • Actie ($a_t$): In elke ronde selecteert de controller een actie uit een discrete set: ofwel stoppen met samplen (actie 0) of een specifiek aantal aanvullende kandidaat-antwoorden parallel genereren (acties $k_1, k_2, \dots$).
  • Beloningsfunctie: De beloning balanceert de uiteindelijke antwoordkwaliteit tegen de inferentiekosten.
    • Stapsgewijze Straf: Een straf wordt opgelegd voor elke extra samplingronde ($\lambda_{lat}$) en voor elk nieuw gegenereerd sample ($\lambda_{comp}$).
    • Terminale Beloning: Bij beëindiging (ofwel door de controller die stopt, ofwel door het bereiken van het maximale budget) ontvangt het model een beloning gebaseerd op de correctheid van het meerderheidsstem-antwoord. De doelcorrectheid ($y^\star$) wordt gedefinieerd als de meerderheidsstem die zou worden verkregen als het samplen zou doorgaan tot het maximale budget $N$, waardoor de beloning wordt losgekoppeld van ground-truth labels.
  • Optimalisatie: De controller wordt getraind met Proximal Policy Optimization (PPO) om de verwachte cumulatieve beloning te maximaliseren. De auteurs merken op dat dit objectief geïnterpreteerd kan worden als de Lagrangian-relaxatie van een geconstreind optimalisatieprobleem waarbij het doel is de nauwkeurigheid te maximaliseren onder strikte latentie- en computationele budgetten.

• Architectuur en Implementatie:
  De controller is een lichtgewicht vierlaagse Multi-Layer Perceptron (MLP). Het vereist alleen de definitieve gegenereerde antwoorden om te functioneren, wat het niet-invasief en inzetbaar maakt op CPU's. Het vereist geen finetuning van het onderliggende LLM.

Belangrijkste Bijdragen

1. MDP Formulering voor Adaptieve Sampling: Het artikel introduceert een nieuwe formulering van adaptieve sampling als een MDP, waarmee men beweegt van heuristische regels naar een geleerd beleid dat expliciet de prestatie-kostenafweging navigeert.
2. Lichtgewicht en Niet-Invasief Ontwerp: De methode vertrouwt puur op de statistieken van de gesamplede antwoordpool, vereist geen hulp-signalen (zoals confidence scores) en geen modificatie van het redeneerproces van het LLM. De controller is klein genoeg om getraind en ingezet te worden op CPU's.
3. Theoretische Interpretatie: De auteurs bieden een theoretische verbinding tussen hun RL-objectief en geconstreinde optimalisatie, waarbij zij aantonen dat de gewogen beloningsfunctie fungeert als een Lagrangian-relaxatie voor het maximaliseren van nauwkeurigheid onder strikte middelenrestricties.
4. Generalisatie: Het framework demonstreert sterke transferbaarheid, waarbij beleid getraind op één dataset of model (bijv. Qwen-3-0.6B) kan worden toegepast op verschillende benchmarks en zelfs op andere, meer capabele modellen (bijv. GPT-4.1-nano) met minimale prestatievermindering.

Experimentele Resultaten
De methode werd geëvalueerd op drie wiskundige redeneerbenchmarks (AIME24, AIME25, HMMT 2025) met diverse LLM samplers (Qwen-3 familie en GPT-4.1-nano).

• Prestaties vs. Baselines: Vergeleken met sterke baselines zoals ASC en ESC, bereikt RL-Guided Sampling superieure trade-offs:
  • vs. ASC: Vermindert het aantal samplingrondes met ongeveer 3–4$\times$ en het totaal aantal samples met ~30% terwijl de nauwkeurigheid vergelijkbaar blijft. Dit pakt de hoge latentie-bottleneck van de sequentiële sampling van ASC aan.
  • vs. ESC: Vermindert de totale sample-vereisten met ~33% en de samplingrondes met ~10%, waarbij een hogere nauwkeurigheid wordt bereikt door te voorkomen dat er te agressief vroegtijdig wordt gestopt.
  • vs. SC: Hoewel alle adaptieve methoden de totale aantal samples verminderen vergeleken met fixed-budget SC, bereikt RL-Guided Sampling dit met aanzienlijk minder rondes dan ASC en een betere efficiëntie dan ESC.
• Scaling Gedrag: Scaling curves geven aan dat RL-Guided Sampling consequent de baselines overtreft over diverse parameterinstellingen, waarbij het betere nauwkeurigheid biedt voor hetzelfde of minder samples en rondes.
• Resource Allocatie: Analyse laat zien dat het geleerde beleid intuïtief meer samples toewijst aan queries met een hogere antwoordentropie (diversiteit), wat aansluit bij de noodzaak om meer samples nodig te hebben om te convergeren naar een betrouwbare meerderheid.
• Ablatie-studies: De "Running Majority" target (meerderheidsstem bij het maximale budget $N$) bleek superieur aan het gebruik van ground-truth labels (Real Label), wat de prestaties verslechterde door ruis te introduceren en het leren van universele stopcriteria te verhinderen.

Betekenis en Claims
Het artikel claimt dat RL-Guided Sampling een principieel, lichtgewicht alternatief biedt voor bestaande adaptieve sampling-methoden. Door het probleem te formuleren als een MDP en een kleine controller te trainen via RL, balanceert de methode effectief antwoordcorrectheid, latentie en computationele kosten zonder invasieve modificaties aan het LLM of toegang tot de interne toestanden van het model te vereisen. De auteurs benadrukken dat de aanpak bijzonder waardevol is voor praktische implementatie, waarbij het getraind kan worden op goedkope open-source modellen en getransformeerd kan worden om de inferentiekosten van sterkere, propriëtaire of closed-source systemen te verlagen. Het werk vertegenwoordigt een eerste stap richting RL-gebaseerde adaptieve sampling, waarbij toekomstige richtingen worden gesuggereerd voor het integreren van rijkere state-signalen en real-world kostenmetrieken.