💬 NLP

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

This paper introduces a reinforcement learning-based decoder sampler that learns a lightweight, test-time policy to dynamically adjust sampling parameters for Large Language Models, significantly outperforming static decoding strategies across summarization tasks without requiring model retraining.

Original authors: Asmita Bhardwaj, Yuya Jeremy Ong, Eelaaf Zahid, Basel Shbita

Published 2026-03-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Asmita Bhardwaj, Yuya Jeremy Ong, Eelaaf Zahid, Basel Shbita

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, encyclopedic chef (the Large Language Model or LLM) who knows how to cook every dish in the world. This chef has a massive cookbook of knowledge, but they are a bit rigid. They follow a strict recipe book that never changes.

Usually, when you ask this chef to write a story or summarize a news article, they use a standard "cooking method" (called decoding).

Greedy Decoding is like the chef always picking the single most obvious ingredient for the next step. It's safe, but the meal might taste boring and repetitive.
Static Sampling is like the chef adding a fixed amount of "spice" (randomness) to every dish, regardless of whether it's a soup or a cake. Sometimes it works, but often the flavor is off.

The problem? One size does not fit all. A news summary needs to be crisp and factual, while a story needs to be wild and creative. But the chef's standard method doesn't know the difference.

The Solution: The "Smart Sous-Chef"

This paper introduces a new idea: instead of trying to retrain the whole chef (which is expensive and slow), we hire a tiny, super-smart Sous-Chef (an AI Agent) to stand next to the main chef and guide them while they are cooking.

Here is how this "Sous-Chef" works, using the paper's Reinforcement Learning approach:

The Job: The Sous-Chef watches the main chef cook. At every step, the Sous-Chef decides how much "randomness" (temperature) or "variety" (top-p) to add to the next ingredient.
The Feedback Loop: As the dish is being made, the Sous-Chef tastes it.
- If the summary is too long, the Sous-Chef says, "Cut back the spice!"
- If the story is too boring, the Sous-Chef says, "Add more flavor!"
- If the summary repeats the same words, the Sous-Chef says, "Stop that!"
Learning by Doing: The Sous-Chef isn't a human; it's a small computer program. It learns from these taste tests. If it makes a good decision, it gets a "point" (Reward). If it makes a bad one, it gets a "thumbs down." Over time, it learns the perfect strategy for every type of dish without ever changing the main chef's recipe book.

The Experiment: Testing the Sous-Chef

The researchers tested this system on three very different "kitchens" (datasets):

BookSum: Long, complex story chapters (like a novel).
arXiv: Dry, scientific papers (like a textbook).
WikiHow: Step-by-step instructions (like a DIY guide).

They used two different "chefs" (models): a small one (Qwen-0.5B) and a medium one (Granite-2B).

The Results:
The Sous-Chef was a game-changer.

On the BookSum (stories), the new method produced summaries that were 88% better than the standard method.
On WikiHow (instructions), it was 79% better.
Even on the scientific papers, it did better than the old ways.

The Secret Sauce: The "Reward Menu"

One of the most important findings was about how the Sous-Chef was taught.

The Bad Teacher: If you only told the Sous-Chef, "Make sure the words match the original text" (a metric called ROUGE), the Sous-Chef learned to just copy-paste. The results were boring and didn't improve much.
The Good Teacher: The researchers gave the Sous-Chef a balanced menu of goals:
- "Match the important words."
- "Keep the length just right."
- "Don't repeat yourself."
- "Make sure it sounds complete."

When the Sous-Chef had this balanced menu, it learned to cook perfectly for the specific task. It learned that a story needs different "spices" than a science paper.

Why This Matters

No Heavy Lifting: You don't need to rebuild the whole kitchen (retrain the giant AI model). You just add a small, smart guide. This saves massive amounts of money and computing power.
Instant Adaptation: The Sous-Chef can switch styles instantly. One minute it's writing a poem, the next it's writing a legal contract, adjusting its strategy on the fly.
Self-Improvement: This is a step toward "Recursive Self-Improvement." The system watches itself, judges its own work, and gets better in real-time, all without changing its core brain.

The Takeaway

Think of this paper as inventing a smart remote control for AI. Instead of trying to rewrite the AI's brain to make it smarter, we just gave it a remote that lets it adjust its own volume, contrast, and brightness depending on what it's watching. The result? A much better show, every single time.

1. Problem Statement

Large Language Models (LLMs) rely on decoding strategies to select output tokens from predicted probability distributions. Current standard methods (e.g., greedy search, fixed temperature, top-k, top-p) are static and task-agnostic. They use fixed heuristics that do not adapt to:

The evolving context of the generation.
Specific domain requirements (e.g., factual accuracy in scientific papers vs. stylistic flexibility in fiction).
Trade-offs between fluency, diversity, and length.

Consequently, these static methods often yield suboptimal or inconsistent results across different domains. Furthermore, improving decoding behavior typically requires fine-tuning the LLM itself, which is computationally expensive and impractical for large-scale deployment. The paper addresses the need for a lightweight, adaptive decoding mechanism that allows a model to self-improve during inference without modifying its internal weights.

2. Methodology

The authors propose a Reinforcement Learning (RL) based Decoder Sampler that treats decoding as a sequential decision-making problem. The core innovation is a "frozen LLM + learnable policy" architecture.

A. Formulation as a Markov Decision Process (MDP)

The decoding process is framed as an MDP where an agent (the policy sampler) interacts with a frozen LLM:

State ( $S_t$ ): Encapsulates the decoding context, including:
- Prompt and generated prefix.
- Model logits (top-50).
- Hidden representations (mean-pooled final hidden states).
- Prefix length and normalized entropy of the next-token distribution.
Action ( $A_t$ ): The agent selects decoding hyperparameters dynamically. In this work, the action space is 2D, controlling:
- Temperature ( $T$ ): Mapped to $[0.2, 1.2]$ .
- Top-p ( $p$ ): Mapped to $[0.8, 1.0]$ .
- Note: The policy outputs Gaussian actions, squashed via sigmoid to fit these ranges.
Reward ( $R_t$ ): Task-specific signals derived from the generated output to guide the agent. The proposed reward is a composite of:
- ROUGE-L F1 (50%): Overlap with reference summaries.
- Length Penalty (20%): Deviation from ideal source-derived length.
- Coverage Bonus (10%): Overlap with important source tokens.
- Repetition Penalty: Applied if >30% of words are repeated.
- Completeness Penalty: Applied if sentence-ending punctuation is missing.

B. Policy Learning

Algorithm: Proximal Policy Optimization (PPO) is used for its stability and sample efficiency.
Architecture: The policy is a lightweight 2-layer MLP (hidden size 256, GELU activations) that maps state embeddings to the action distribution.
Training Loop:
1. The agent observes the state from the frozen LLM.
2. It selects decoding parameters (Temperature/Top-p).
3. The LLM generates tokens based on these parameters.
4. The reward is computed based on the final output.
5. The policy is updated to maximize the cumulative reward.
Key Constraint: The LLM weights remain frozen throughout; only the lightweight policy network is trained.

3. Key Contributions

Test-Time Adaptive Decoding: Introduces a formal framework for decoding as a test-time decision process with explicit state, action, and reward spaces, enabling adaptation without retraining the base model.
Empirical Validation: Demonstrates consistent performance gains over static baselines (greedy and fixed temperature) across two models (Granite-3.3-2B, Qwen-2.5-0.5B) and three diverse datasets (BookSum, arXiv, WikiHow).
Reward Design Analysis: Conducts extensive ablation studies showing that composite reward shaping (combining overlap, length, coverage, and repetition metrics) is critical for stable learning, whereas single-metric rewards (e.g., ROUGE-only) often fail or underperform.

4. Experimental Results

The study evaluated the method on summarization tasks, comparing the RL policy against Greedy Decoding and Static Sampling (fixed $T=0.3$ ).

Performance Gains:
- The RL policy consistently outperformed static baselines.
- BookSum (Granite-3.3): +88.44% relative gain over greedy decoding.
- WikiHow (Qwen-2.5): +78.95% relative gain over greedy decoding.
- arXiv: Smaller but consistent gains (+60.45% over greedy), indicating robustness even in formulaic domains.
Ablation Insights:
- ROUGE-only: Yielded negligible or negative gains, proving overlap metrics alone are insufficient.
- Composite Rewards: The "Proposed Reward" (multi-component) achieved the best balance of high reward and training stability.
- Coverage Term: Removing the coverage bonus actually improved stability and gains in some cases, suggesting it can introduce noise.
- Repetition Penalty: Strict penalties were necessary; softer penalties dampened improvements.
Model Size: Significant improvements were observed even on the smaller 0.5B parameter model, suggesting the approach is effective for resource-constrained settings where fine-tuning is impossible.

5. Significance and Implications

Practical Self-Improvement: This work provides a concrete mechanism for "recursive self-improvement" at test-time. It allows a system to monitor its own generation, evaluate outcomes, and adjust future actions within a single episode without altering the base model's knowledge.
Decoupling Control from Capacity: By separating the decoding strategy from the model weights, the approach enables domain-aware and user-controllable generation without the computational cost of fine-tuning.
Safety and Reversibility: Since the policy is external and lightweight, it can be easily disabled or rolled back if it causes degradation, offering a lower-risk pathway to self-improving systems compared to weight-level updates.
Future Directions: The authors suggest extending this to richer action spaces (top-k, penalties), using LLM-as-a-judge for rewards, and applying the framework to dialogue, code generation, and reasoning tasks.

In summary, the paper demonstrates that Reinforcement Learning can effectively learn a lightweight policy to dynamically tune decoding parameters, significantly improving generation quality across diverse domains while keeping the underlying LLM frozen.