Online Causal Kalman Filtering for Stable and Effective Policy Optimization

Titel: Online Causale Kalman-filtering voor Stabiele en Effectieve Policy-Optimalisatie

1. Het Probleem: Instabiliteit bij RL voor LLM's

Reinforcement Learning (RL) is cruciaal voor het verbeteren van Large Language Models (LLM's) via methoden zoals GRPO (Group Relative Policy Optimization). Echter, bij schaalvergroting ondervinden deze methoden ernstige stabiliteitsproblemen, veroorzaakt door:

• Hoge variantie in Importance Sampling (IS) ratios: Op token-niveau vertonen de IS-ratio's (die het verschil tussen de oude en nieuwe policy weergeven) extreme fluctuaties.
• Structurele inconsistentie: Empirisch onderzoek toont aan dat "off-policy" afwijkingen (waar de nieuwe policy afwijkt van de oude) niet lokaal homogeen zijn. In plaats van stabiele patronen binnen een redeneringsstap, wisselen de token-toestanden (bijv. ratio > 1 of < 1) zeer frequent en willekeurig.
• Gevolgen: Deze ruis en snelle schommelingen vervormen de policy-gradient updates voor aangrenzende tokens, wat leidt tot training-collaps, entropie-verlies en instabiel leren.

Bestaande oplossingen proberen dit op te lossen door:

1. Een vast sequence-level IS-ratio te gebruiken (gemiddeld over de hele zin), wat lokale structuur volledig verwijdert.
2. Token-ratio's apart aan te passen, maar zonder rekening te houden met de temporale samenhang tussen tokens.

2. Methodologie: KPO (Kalman Policy Optimization)

De auteurs introduceren KPO, een methode die online causale Kalman-filtering toepast op token-wise IS-ratio's om ruis te filteren terwijl lokale structurele coherentie behouden blijft.

Kernconcepten:

• Modelleer IS-ratio's als een tijdsreeks: De auteurs behandelen de logaritmische IS-ratio's ($z_t = \log r_t$) als een waargenomen tijdsreeks die voortkomt uit een onderliggende, gladde "latente" toestand ($\rho_t$).
• State-Space Model:
  • Procesmodel: $\rho_t = \rho_{t-1} + \eta_t$ (een random walk met procesruis $Q$). Dit modelleert de verwachting dat off-policy afwijkingen zich geleidelijk ontwikkelen binnen een semantisch segment.
  • Observatiemodel: $z_t = \rho_t + \epsilon_t$ (met observatieruis $V$). Dit erkent dat de berekende ratio's ruis bevatten.
• Causale Kalman-filtering:
  • De filter werkt autoregressief: de schatting voor token $t$ gebruikt alleen informatie van voorgaande tokens en de huidige token, nooit toekomstige tokens.
  • Stappen:
    1. Predictie: Schat de latente toestand op basis van de vorige stap.
    2. Kalman-gain berekening: Bepaal een adaptieve stapgrootte ($K_t$) die weegt tussen de voorspelling en de nieuwe meting, afhankelijk van de onzekerheid ($P$) en ruis ($V$).
    3. Update: Corrigeer de voorspelling met de nieuwe meting om een gladde, gestructureerde schatting ($\hat{\rho}_t$) te krijgen.
• Integratie in de Loss-functie: De gefilterde ratio's ($\hat{r}_t = \exp(\hat{\rho}_t)$) vervangen de ruwe ratio's in de GRPO/PPO doelfunctie. Dit behoudt token-level granulariteit maar dempt de hoge-frequentie ruis.

3. Belangrijkste Bijdragen

1. Empirische Inzicht: Het paper onthult voor het eerst dat lokale off-policy afwijkingen op token-niveau structureel inconsistent zijn (hoge schakelfrequentie, korte "run-lengths"), wat de basis vormt voor training-instabiliteit.
2. Nieuwe Optimalisatiemethode (KPO): Een innovatieve toepassing van Kalman-filtering op RL voor LLM's. Het filtert ruis weg terwijl het de lokale semantische samenhang (coherentie) binnen een reeks behoudt.
3. Robuuste Resultaten: KPO levert consistent betere prestaties op dan state-of-the-art methoden (zoals GSPO, GMPO, en standaard GRPO) op uitdagende wiskundige redeneringsopgaven, met name bij lange redeneringsketens.

4. Resultaten en Experimenten

De methode werd getest op zes uitdagende wiskundige benchmarks (o.a. AIME'24, AIME'25, AMC'23, MATH500, OlympiadBench) met het Qwen3-4B model.

• Prestaties: KPO (zowel de "clipped" als "unclipped" variant) behaalde over het algemeen de beste resultaten.
  • Op AIME'24 en AIME'25 (zeer moeilijk, lange redeneringen) boekte KPO-clipped aanzienlijke verbeteringen in zowel avg@16 als pass@16 ten opzichte van de beste concurrenten (GSPO).
  • Bijvoorbeeld: Op AIME'25 steeg pass@16 van 50.00 (GSPO) naar 60.00 (KPO-clipped).
• Trainingstabiliteit:
  • Entropie: Terwijl standaard GRPO snel instort naar bijna nul entropie (deterministisch gedrag), behoudt KPO een gezonde entropie, wat betekent dat het model blijft exploreren.
  • Verlies en Ruis: De policy-gradient loss van KPO is veel stabieler en minder ruisig dan bij GRPO.
  • Dynamiek: KPO voorkomt de vroege saturatie die bij GSPO/GMPO optreedt en de divergentie die bij GRPO voorkomt.
• Verdeling van IS-ratio's: Na filtering veranderde de dynamiek drastisch:
  • De switch-frequentie daalde van 0.43 naar 0.01 (tokens wisselen veel minder vaak van toestand).
  • De run-length (aantal opeenvolgende tokens in dezelfde toestand) steeg van ~1.6 naar ~120-135 tokens.
  • Dit bevestigt dat KPO korte, ruisige fluctuaties omzet in lange, coherente segmenten.

5. Betekenis en Conclusie

Dit paper biedt een fundamenteel nieuw perspectief op hoe we importance sampling in RL voor taalmodellen moeten benaderen. In plaats van te kiezen tussen ruwe token-level data (te ruisig) of een globaal gemiddelde (te veel informatieverlies), introduceert KPO een temporeel gefilterde benadering.

• Stabiliteit: Het maakt grootschalige RL-training (zoals GRPO) veel stabieler door de variatie in de gradienten te reduceren zonder de nuance van token-level beloning te verliezen.
• Efficiëntie: Het is een lichtgewicht methode die compatibel is met bestaande pipelines en geen extra critic-model vereist.
• Toekomst: De auteurs wijzen erop dat de autoregressieve aard van de Kalman-filtering parallelisatie bemoeilijkt, wat een interessante richting is voor toekomstig onderzoek (bijv. parallelle Kalman-filtering).

Kortom, KPO lost het probleem van "training collapse" bij off-policy RL voor LLM's op door de inherente temporale structuur van taalredenering te respecteren en te benutten.