Adaptive Decoding via Test-Time Policy Learning for Self-Improving Generation

Probleemstelling

De kwaliteit van de output van Large Language Models (LLM's) wordt grotendeels bepaald door de decodestrategie (de methode waarmee tokens uit de voorspelde kansverdeling worden geselecteerd). Traditionele methoden, zoals greedy sampling, top-k, nucleus sampling (top-p) of beam search, maken gebruik van statische heuristieken. Deze heuristieken zijn vaak:

• Task-agnostisch: Ze passen niet aan op de specifieke eisen van verschillende taken of domeinen (bijv. een nieuwsartikel samenvatten versus een fictief hoofdstuk).
• Statisch: Parameters zoals temperatuur of top-p worden vooraf vastgesteld en veranderen niet tijdens het generatieproces.
• Beperkt: Ze kunnen moeilijk een balans vinden tussen vloeiendheid, diversiteit, feitelijke juistheid en stilistische controle.

Het aanpassen van het gedrag van een LLM vereist doorgaans het opnieuw trainen of fine-tunen van het model, wat computatief zeer duur en onpraktisch is. Er is behoefte aan een methode die het generatiegedrag tijdens het testen (test-time) kan aanpassen zonder de modelgewichten te wijzigen, wat past binnen het concept van "recursieve zelfverbetering".

Methodologie

De auteurs stellen een reinforcement learning (RL) gebaseerde decoder-sampler voor. In plaats van het LLM zelf te trainen, wordt een lichtgewicht agent (policy) getraind om de decodingsparameters dynamisch aan te passen.

1. Formulering als Markov Decision Process (MDP):

   • State (Toestand): De agent observeert de context van de generatie, bestaande uit de prompt, de gegenereerde prefix, de model-logits en de verborgen representaties. Concreet worden hierin gemiddelde hidden states, de top-k logits, de prefix-lengte en de genormaliseerde entropie van de volgende-token-verdeling samengevoegd.
   • Action (Actie): De agent kiest decodingsparameters. In dit werk focussen ze op het aanpassen van de temperatuur (T) en top-p (p). De policy outputt een Gaussische verdeling die wordt afgebeeld op de bereikwaarden T ∈ [0.2, 1.2] en p ∈ [0.8, 1.0].
   • Reward (Beloning): Een taakspecifieke beloningsfunctie evalueert de gegenereerde output en stuurt de agent. Voor samenvattingen wordt een samengestelde reward gebruikt (zie hieronder).

2. Training met PPO:

   • Het systeem gebruikt Proximal Policy Optimization (PPO) om de policy te trainen.
   • Het LLM blijft volledig bevroren (frozen); alleen de kleine policy-network (een 2-laags MLP) wordt getraind.
   • Dit creëert een gesloten feedbacklus: generatie → evaluatie via reward → aanpassing van de policy → volgende generatie binnen dezelfde sessie.

3. Reward Functie (Samenvattingstaken):
   De voorgestelde reward is een gewogen combinatie van:

   • ROUGE-L F1: Overlap met referentiesamenvattingen.
   • Lengtestraf: Afwijking van een ideale lengte gebaseerd op de bron.
   • Coverage bonus: Overlap met belangrijke bron-tokens.
   • Herhalingsstraf: Toepassing als meer dan 30% van de woorden wordt herhaald.
   • Volledigheidsstraf: Straf als er geen zin-afsluitende leestekens zijn.

Belangrijkste Bijdragen

1. Formalisatie van Decoding als Test-Time Besluitvorming: Het introduceren van een RL-framework waar decoding wordt gezien als een sequentieel beslissingsproces met expliciete ruimtes voor state, actie en reward, zonder LLM-gewichten te wijzigen.
2. Empirische Validatie: Uitgebreide evaluaties op twee modellen (Granite-3.3-2B en Qwen-2.5-0.5B) en drie verschillende datasets (BookSum, arXiv, WikiHow) die aantonen dat adaptieve sampling superieur is aan statische baselines.
3. Ablatie Studies: Diepgaande analyse van de invloed van verschillende reward-componenten, wat inzicht geeft in hoe reward-shaping het leerproces en de uiteindelijke prestaties beïnvloedt.

Resultaten

De experimenten tonen aan dat de RL-based policy consistent beter presteert dan greedy decoding en statische sampling:

• Prestatieverbetering: De RL-policy behaalde relatieve winsten van maximaal +88% op de BookSum-dataset (met Granite) en +79% op WikiHow (met Qwen). Zelfs op formelere datasets zoals arXiv werden verbeteringen geboekt.
• Lerend Vermogen: De "early-to-late" analyse (vergelijking van prestaties aan het begin versus het einde van de training) toont aan dat de policy daadwerkelijk leert en verbetert in plaats van vast te lopen in een statische strategie.
• Impact van Reward Design:
  • Alleen ROUGE als reward bleek ontoereikend (slechte prestaties).
  • Het toevoegen van gestructureerde "shaping terms" (lengte, herhaling, etc.) was cruciaal voor stabiele en duurzame verbeteringen.
  • Specifiek bleek dat het verwijderen van de "coverage bonus" soms zelfs leidde tot grotere, stabielere winsten, wat suggereert dat niet alle componenten in elke context even nuttig zijn.
• Efficiëntie: De policy is zeer lichtgewicht (256 hidden units), wat betekent dat de rekentijd per token verwaarloosbaar is in vergelijking met de LLM-forward pass.

Betekenis en Conclusie

Dit werk demonstreert dat reinforcement learning een praktische methode is voor test-time adaptatie in decodestrategieën.

• Zelfverbetering zonder hertraining: Het biedt een route voor systemen om zichzelf te verbeteren tijdens het generatieproces door feedback te gebruiken, zonder de dure en riskante stap van het opnieuw trainen van het basismodel.
• Domein-afhankelijkheid: De methode is bijzonder effectief in domeinen waar statische heuristieken falen door de noodzaak van stilistische flexibiliteit (zoals fictie of instructies).
• Toekomstperspectief: Hoewel de studie zich beperkt tot samenvattingen, opent dit pad de weg voor toepassing in dialoog, code-generatie en redenering. Het benadrukt dat het ontwerp van de reward-functie minstens zo belangrijk is als het RL-algoritme zelf.

Kortom, de paper presenteert een robuust, lichtgewicht framework dat de kloof overbrugt tussen starre decodingsheuristieken en volledig gefinetuned modellen, waardoor generatie flexibeler, domeinbewuster en controleerbaar wordt.