Technische Samenvatting: Kalman-Guided Prompt Selectie (KGPS) voor Adaptieve RL Finetuning

1. Probleemstelling

Reinforcement Learning (RL) finetuning verbetert de redeneervermogens van Large Language Models (LLM's) aanzienlijk, maar de effectiviteit ervan is sterk afhankelijk van het selecteren van trainingsprompts met een geschikte moeilijkheidsgraad voor de huidige policy. Een kritieke uitdaging in dit domein is niet-stationariteit: naarmate de policy verbetert, evolueert de moeilijkheidsgraad van de prompts. Prompts die voorheen informatief waren, kunnen triviaal worden, terwijl prompts die voorheen te moeilijk waren, nu leerbaar kunnen zijn.

Bestaande online promptselectiemethoden kampen met een afweging:

• Evaluatie-gebaseerde benaderingen (bijv. Dynamic Sampling) zijn accuraat maar computationeel duur, omdat ze enorme hoeveelheden extra rollouts vereisen om de succesratio's van prompts te schatten, wat vaak resulteert in een 7× hogere inferentie-overhead dan standaard training.
• Predictie-gebaseerde benaderingen (bijv. MoPPS, GRESO) zijn efficiënt, maar modelleren de moeilijkheid van een prompt doorgaans als een stationaire latente parameter. Ze vertrouwen op heuristische vervalfactoren om niet-stationariteit te behandelen, wat faalt in het nauwkeurig volgen van het verschuivende moeilijkheidslandschap dat wordt veroorzaakt door continue policy-updates, wat leidt tot persistent hoge schattingsfouten.

Het artikel betoogt dat een principiële verschuiving van statische parameterschatting naar dynamische staatsschatting vereist is om afstemming met de evoluerende policy te behouden zonder extra inferentiekosten te veroorzaken.

2. Methodologie: Kalman-Guided Prompt Selectie (KGPS)

KGPS herformuleert promptselectie als een dynamisch staatsschattingprobleem. In plaats van de succesratio van een prompt als een vaste parameter te behandelen, modelleert KGPS deze als een tijdvariërende latente staat binnen een lineair-Gaussiaans toestandsruimtemodel (SSM).

Kerncomponenten:

1. Toestandsruimteformulering in de logit-ruimte:

   • De latente succesratio $\phi^t_\tau$ van een prompt $\tau$ op stap $t$ wordt getransformeerd naar de logit-ruimte: $\psi^t_\tau = \text{logit}(\phi^t_\tau)$.
   • Toestandsevolutie: De latente staat evolueert als een random walk: $\psi^t_\tau = \psi^{t-1}_\tau + w^t_\tau$, waarbij de procesruisvariantie $Q_t$ expliciet gekoppeld is aan de omvang van de policy-update: $Q_t = \gamma \cdot \|\theta^t - \theta^{t-1}\|^2$. Dit zorgt ervoor dat de onzekerheid toeneemt wanneer de policy substantieel verandert, wat de potentiële verschuiving in de moeilijkheidsgraad van de prompt weerspiegelt.
   • Observatiemodel: Wanneer een prompt wordt geselecteerd en $k$ keer wordt uitgevoerd (rollout), wordt de empirische succesratio getransformeerd via de delta-methode om een Gaussische observatie in de logit-ruimte te benaderen.

2. Kalman Filter voor het onderhouden van de posterior:

   • KGPS gebruikt een Kalman filter om een Gaussische posterior $N(\hat{\psi}^t_\tau, P^t_\tau)$ te onderhouden voor elke prompt in de pool.
   • Predictiestap: Bij elke trainingsstap wordt de variantie $P^t_\tau$ van alle prompts (geselecteerd of niet) opgezwollen met $Q_t$. Dit mechanisme verhoogt op natuurlijke wijze de onzekerheidsscore voor prompts die recentelijk niet zijn geëvalueerd, waardoor wordt voorkomen dat ze permanent worden genegeerd.
   • Update stap: Voor prompts die in de vorige stap zijn geselecteerd, incorporeert de filter de nieuwe rollout-observaties om de gemiddelde waarde te updaten en de variantie te verkleinen.

3. Posterior-Expected Utility Selectie:

   • In plaats van prompts te selecteren op basis van een puntinschatting van de moeilijkheidsgraad, selecteert KGPS prompts door een posterior-verwachte trainingsutiliteit te maximaliseren: $\tilde{A}(\tau, t) = \mathbb{E}_{\psi \sim N(\hat{\psi}^t_\tau, P^t_\tau)}[h(\psi)]$, waarbij $h(\psi) = \sigma(\psi)(1-\sigma(\psi))$ piekt bij een gemiddelde moeilijkheidsgraad.
   • Deze verwachting integreert over de volledige posterior distributie. Bijgevolg krijgen prompts met een hoge onzekerheid (grote $P^t_\tau$) van nature hogere scores, wat een balans vindt tussen het exploiteren van huidige schattingen en het exploreren van onzekere prompts.
   • De verwachting wordt efficiënt berekend met behulp van Gauss–Hermite kwadratuur, wat geen stochastische sampling vereist.

Belangrijk efficiëntiekenmerk: KGPS werkt volledig op de resultaten van de rollouts die al zijn gegenereerd voor de standaard policy training. Het vereist geen aanvullende rollouts of LLM-inferenties voor het selectieproces zelf.

3. Belangrijkste Bijdragen

Het artikel identificeert drie primaire bijdragen:

1. Herformulering van het volgen van moeilijkheidsgraad: De auteurs stellen voor om de moeilijkheidsgraad van een prompt te modelleren als een dynamisch staatsschattingprobleem met behulp van een lineair-Gaussiaans SSM, waarbij de procesruis gekoppeld is aan de omvang van de policy-updates, in plaats van het als een stationaire parameter te behandelen.
2. Kalman Filter Afleiding & Utility Score: Zij leiden een Kalman filter af voor online moeilijkheidsgraad-tracking en introduceren een posterior-verwachte utility score die exploitatie en exploratie balanceert zonder handmatig afgestemde decay-schema's.
3. Empirische Validatie: Uit uitgebreide experimenten blijkt dat KGPS consequent beter presteert dan bestaande baselines in zowel uiteindelijke nauwkeurigheid als rollout-efficiëntie over taken op het gebied van wiskunde, planning en geometrisch redeneren, met een verwaarloosbare computationele overhead.

4. Experimentele Resultaten

De auteurs hebben KGPS geëvalueerd op drie reasoning benchmarks (Wiskunde, Planning/Countdown, Visuele Geometrie) met verschillende model-schalen (Qwen3-0.6B tot DeepSeek-R1-Distill-7B) en RL-algoritmen (GRPO, PPO, Reinforce++).

• Prestaties vs. Rollout Kosten: Op het DeepSeek-R1-Distill-7B model bereikte KGPS een vergelijkbare of licht superieure prestatie ten opzichte van de evaluatie-gebaseerde Oracle (Dynamic Sampling), terwijl het 83% minder rollouts gebruikte. Vergeleken met de predictie-gebaseerde baseline MoPPS verbeterde KGPS de gemiddelde prestatie met 0,12 punten over zes wiskunde benchmarks met hetzelfde rollout-budget.
• Nauwkeurigheidswinst: Over alle model-schalen en taken heen behaalde KGPS consistent de hoogste testnauwkeurigheid onder online selectiemethoden, en overtrof vaak de Oracle (DS) die enorme computationele kosten met zich meebrengt.
• Schattingsnauwkeurigheid: KGPS behield een significant lagere Mean Absolute Error (MAE ≈ 0,15) tussen voorspelde en empirische succesratio's vergeleken met MoPPS (MAE ≈ 0,40), wat duidt op superieure tracking van niet-stationaire moeilijkheidsgraad.
• Ablatie studies:
  • Het verwijderen van de posterior verwachting (het gebruik van puntinschattingen in plaats daarvan) verminderde de prestaties met ~1,09 punten.
  • Het ontkoppelen van procesruis van de policy update magnitude (het gebruik van vaste ruis) verminderde de prestaties met ~1,21 punten, wat het belang van adaptieve onzekerheid bevestigt.
  • Het verwijderen van de warmup initialisatie veroorzaakte een lichte daling in prestaties, wat de noodzaak van stabiele initiële posteriors valideert.

5. Betekenis en Claims

Het artikel claimt dat KGPS een nieuwe state-of-the-art vestigt voor online promptselectie in RL finetuning. De betekenis ligt in het oplossen van de spanning tussen efficiëntie en adaptiviteit:

• Het bereikt de nauwkeurigheid van dure evaluatie-gebaseerde methoden zonder hun computationele overhead.
• Het overwint de beperkingen van statische predictie-gebaseerde methoden door de niet-stationaire aard van RL-trainingsdynamiek rigoureus te modelleren.
• Door moeilijkheidsgraad te behandelen als een dynamische staat in plaats van een statische parameter, biedt KGPS een principieel mechanisme voor onzekerheidsbeheer dat van nature de verwaarloosde prompts opnieuw bezoekt, waardoor het trainingscurriculum in lijn blijft met de evoluerende competentie van het model.

De auteurs concluderen dat dynamische staatsschatting een praktisch en superieur alternatief is voor statische moeilijkheidsgraad-approximaties voor adaptieve RL finetuning.