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

Asmita Bhardwaj; Yuya Jeremy Ong; Eelaaf Zahid; Basel Shbita

cs.CLMar 19, 2026

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

Authors: Asmita Bhardwaj, Yuya Jeremy Ong, Eelaaf Zahid, Basel Shbita

Organizations: UC San Diego San Diego, CA · Plastic Labs San Jose, CA · IBM Research Cambridge, MA · IBM Research San Jose, CA

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.

Figures & tables

Figure 1: Overview of the RL-based decoder sampler. The agent observes the decoding state from a frozen LLM, selects decoding parameters (e.g., temperature and top- $p$ ), and receives rewards guiding adaptive sampling.

Figure 2: Reward over time for WikiHow + Granite-3.3 under the Proposed Reward. The upward trend in the moving average demonstrates PPO learns a useful sampling policy.

	Granite-3.3-2B			Qwen-2.5-0.5B
	WikiHow	arXiv	BookSum	WikiHow	arXiv	BookSum
Greedy	0.333	0.220	0.147	0.114	0.154	0.014
Static	0.369	0.275	0.212	0.140	0.174	0.083
RL (Proposed)	0.390	0.353	0.277	0.204	0.226	0.147
$\Delta$ vs Greedy (%)	+17.17	+60.45	+88.44	+78.95	+46.75	+950.00
$\Delta$ vs Static (%)	+5.69	+28.36	+30.66	+45.71	+29.89	+77.11

Table 1: Main evaluation results with the Proposed Reward. RL achieves consistent improvements over greedy and static baselines. Scores are reported alongside percentage gains and PPO early

\to

late improvements.

Reward Variant	PPO	Greedy	Static	$\Delta$ vs Greedy (%)	$\Delta$ vs Static (%)	Early $\to$ Late (%)
ROUGE-only	0.169	0.158	0.170	+6.96	-0.59	+15.38
Core shaping	0.448	0.335	0.400	+33.73	+12.00	-4.56
No coverage bonus	0.349	0.219	0.275	+59.36	+26.91	+24.91
Soft repetition penalty	0.503	0.431	0.462	+16.72	+8.88	+6.81
Sigmoid scaling	0.382	0.302	0.323	+26.49	+18.27	-4.27
Proposed (ours)	0.353	0.220	0.275	+60.45	+28.36	-4.38

Table 2: Reward ablations on arXiv ( Granite-3.3-2B ) . We report absolute rewards, percentage improvements over baselines, and PPO early

\to

late change.

Explore similar work

Learning Adaptive LLM Decoding
Mar 10, 2026Chloe H. Su, Zhe Ye, Samuel Tenka +3Language Model Decoding RL for Language Models
Department of Computer Science, Harvard University, Boston, USA · Kempner Institute, Harvard University, Cambridge, MA, USA · UC Berkeley, CA, USA +1
Sampling for Quality: Training-Free Reward-Guided LLM Decoding via Sequential Monte Carlo
Apr 7, 2026Jelena Markovic-Voronov, Wenhui Zhu, Bo Long +5Language Model Decoding LLM Code Generation
LinkedIn
Entropy-informed Decoding: Adaptive Information-Driven Branching
May 10, 2026Benjamin Patrick Evans, Sumitra Ganesh, Leo ArdonLanguage Model Decoding LLM Inference Acceleration
JP Morgan AI Research, London, UK · JP Morgan AI Research, New York, USA

Mar 10, 2026J/K move · Enter open · S save

Learning Adaptive LLM Decoding

Chloe H. Su, Zhe Ye, Samuel Tenka, Aidan Yang+2

Department of Computer Science, Harvard University, Boston, USA · Kempner Institute, Harvard University, Cambridge, MA, USA · UC Berkeley, CA, USA+1

Language Model Decoding RL for Language Models

Open paper AskPDF arXiv

Decoding from large language models (LLMs) typically relies on fixed sampling hyperparameters (e.g., temperature, top-p), despite substantial variation in task difficulty and uncertainty across prompts and individual decoding steps. We propose to learn adaptive decoding policies that dynamically select sampling strategies at inference time, conditioned on available compute resources. Rather than fine-tuning the language model itself, we introduce lightweight decoding adapters trained with reinforcement learning and verifiable terminal rewards (e.g. correctness on math and coding tasks). At the sequence level, we frame decoding as a contextual bandit problem: a policy selects a decoding strategy (e.g. greedy, top-k, min-p) for each prompt, conditioned on the prompt embedding and a parallel sampling budget. At the token level, we model decoding as a partially observable Markov decision process (POMDP), where a policy selects sampling actions at each token step based on internal model features and the remaining token budget. Experiments on the MATH and CodeContests benchmarks show that the learned adapters improve the accuracy-budget tradeoff: on MATH, the token-level adapter improves Pass@1 accuracy by up to 10.2% over the best static baseline under a fixed token budget, while the sequence-level adapter yields 2-3% gains under fixed parallel sampling. Ablation analyses support the contribution of both sequence- and token-level adaptation.

Abstract

Explore similar work

Learning Adaptive LLM Decoding

Sampling for Quality: Training-Free Reward-Guided LLM Decoding via Sequential Monte Carlo

Entropy-informed Decoding: Adaptive Information-Driven Branching