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

TL;DR

This paper presents a reinforcement learning-based test-time decoding policy that dynamically adjusts sampling parameters for large language models, significantly improving generation quality across diverse tasks without retraining.

Contribution

It introduces a lightweight RL-based decoder policy that adapts sampling strategies at test-time, enhancing output quality and consistency across domains.

Findings

Outperforms greedy and static decoding baselines by up to 88%

Structured shaping rewards lead to more stable improvements

Reinforcement learning enables domain-aware, user-controllable generation

Abstract

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.