Adaptive Decoding via Test-Time Policy Learning for Self-Improving Generation
This work addresses the need for domain-aware and user-controllable generation in LLMs without retraining, though it is incremental as it builds on existing reinforcement learning methods for decoding.
The paper tackled the problem of static decoding strategies in large language models leading to suboptimal generation quality across domains by introducing a reinforcement learning-based decoder sampler that adjusts sampling parameters at test-time, achieving relative gains of up to +88% on summarization datasets.
Decoding strategies largely determine the quality of Large Language Model (LLM) outputs, yet widely used heuristics such as greedy or fixed temperature/top-p decoding are static and often task-agnostic, leading to suboptimal or inconsistent generation quality across domains that demand stylistic or structural flexibility. We introduce a reinforcement learning-based decoder sampler that treats decoding as sequential decision-making and learns a lightweight policy to adjust sampling parameters at test-time while keeping LLM weights frozen. We evaluated summarization datasets including BookSum, arXiv, and WikiHow using Granite-3.3-2B and Qwen-2.5-0.5B. Our policy sampler consistently outperforms greedy and static baselines, achieving relative gains of up to +88% (BookSum, Granite) and +79% (WikiHow, Qwen). Reward ablations show that overlap-only objectives underperform compared to composite rewards, while structured shaping terms (length, coverage, repetition, completeness) enable stable and sustained improvements. These findings highlight reinforcement learning as a practical mechanism for test-time adaptation in decoding, enabling domain-aware and user-controllable generation without retraining large models.