Learning Adaptive LLM Decoding

Chloe H. Su; Zhe Ye; Samuel Tenka; Aidan Yang; Soonho Kong; Udaya Ghai

cs.LGMar 10, 2026

Learning Adaptive LLM Decoding

Authors: Chloe H. Su, Zhe Ye, Samuel Tenka, Aidan Yang, Soonho Kong, Udaya Ghai

Organizations: Department of Computer Science, Harvard University, Boston, USA · Kempner Institute, Harvard University, Cambridge, MA, USA · UC Berkeley, CA, USA · Amazon, Boston, MA 02210, USA

Abstract

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.

Figures & tables

Panel of Figure 1 — Figure 1: Overview of the proposed decoding adapter (DA) for a frozen language model (LM). Blue blocks denote input tokens $x_{i}$ ; green blocks denote generated tokens $y_{i}$ .

		Static sampling		Sequence-level adapter
Metric	Setting	Best	Mixed	w/o budget	w/ budget
Pass@1	TURN ( Du et al., 2025 )	72.72
	w/o CoT	71.70 $\pm$ 1.25	71.20 $\pm$ 1.26	72.60 $\pm$ 1.24	72.90 $\pm$ 1.23
	$\Delta$ Abs. (Rel.)			$\uparrow$ 0.90 (+1.26%)	$\uparrow$ 1.20 (+1.67%)
	mix CoT	72.10 $\pm$ 1.24	71.93 $\pm$ 1.25	73.60 $\pm$ 1.22	74.20 $\pm$ 1.21
	$\Delta$ Abs. (Rel.)			$\uparrow$ 1.50 (+2.08%)	$\uparrow$ 2.10 (+2.91%)

Table 1: MATH. Comparison of static sampling baselines and the sequence-level adapter under different settings. For CoT mix , we apply CoT prompting to 30% of prompts (ratio

=0.3

). Static sampling reports the best single strategy ( Best ) and a fixed mixture ( Mixed ). The adapter is evaluated with a budget provided in feature set and without. Values are percentages (mean

\pm

confidence interval).

		Static sampling		Sequence-level adapter
Metric	Setting	Best	Mixed	w/o budget	w/ budget
Pass@1	w/o CoT	11.43 $\pm$ 2.28	10.53 $\pm$ 2.19	14.53 $\pm$ 2.52	14.50 $\pm$ 2.52
	$\Delta$ Abs. (Rel.)			$\uparrow$ 3.10 (+27.12%)	$\uparrow$ 3.07 (+26.85%)
	mix CoT	13.97 $\pm$ 2.48	14.80 $\pm$ 2.54	17.06 $\pm$ 2.69	19.70 $\pm$ 2.85
	$\Delta$ Abs. (Rel.)			$\uparrow$ 2.26 (+15.27%)	$\uparrow$ 4.90 (+33.11%)
Pass@8	w/o CoT	22.80 $\pm$ 3.00	22.23 $\pm$ 2.97	26.08 $\pm$ 3.14	23.10 $\pm$ 3.02

Table 2: CodeContests : comparison of static sampling strategies vs. the sequence -level sampling adapter under different settings.

Metric	Setting	Reported	Adapter (w/ budget)
Pass@1	non-thinking	19.1	20.1 $\pm$ 2.62
Pass@1	thinking	65.6	71.1 $\pm$ 2.96

Table 3: AIME’25 performance comparison of a sequence -level adapter trained on MATH-train vs. numbers reported in the Qwen3-4B technical report ( Team, 2025a ) .

		Static sampling				Sequence-level adapter
Metric	Setting	Best		Mixed		w/o budget		w/ budget
		MATH	Code	MATH	Code	MATH	Code	MATH	Code
Pass@1	w/o CoT	72.00 $\pm$ 3.18	11.80 $\pm$ 2.39	71.10 $\pm$ 3.17	10.30 $\pm$ 2.18	73.55 $\pm$ 3.12	13.70 $\pm$ 2.46	74.20 $\pm$ 3.13	13.70 $\pm$ 2.46
Pass@1	$\Delta$ Abs. (Rel.)					$\uparrow$ 1.55 (2.15%)	$\uparrow$ 1.90 (16.10%)	$\uparrow$ 2.20 (3.06%)	$\uparrow$ 1.90 (16.10%)
	mix CoT	73.80 $\pm$ 3.15	17.10 $\pm$ 2.69	73.30 $\pm$ 3.17	14.50 $\pm$ 2.52	74.80 $\pm$ 3.11	18.50 $\pm$ 2.89	74.70 $\pm$ 3.11	18.20 $\pm$ 2.76
	$\Delta$ Abs. (Rel.)					$\uparrow$ 1.00 (1.36%)	$\uparrow$ 1.40 (8.19%)	$\uparrow$ 0.90 (1.22%)	$\uparrow$ 1.10 (6.43%)

Table 4: Mixed training of math and coding: comparison of static sampling strategies vs. the sequence -level sampling adapter. We use Code to refer to CodeContests in this table.

Figure 2: Action distributions and validation reward on MATH for sequence-level adapter strategies. Actions are defined in table 11 . The x-axis denotes training progress; the left y-axis reports action probability; the right y-axis reports validation reward.

		Static sampling		Adapter
Metric	Setting	Greedy	Mixed	w/o budget	$\Delta$	budget	$\Delta$
Pass@1	w/o CoT	71.33 $\pm$ 1.25	71.60 $\pm$ 1.25	78.28 $\pm$ 1.14	+6.68	80.82 $\pm$ 1.07	+9.49
Pass@1	mix CoT	72.10 $\pm$ 1.24	72.67 $\pm$ 1.24	78.52 $\pm$ 1.14	+5.85	82.33 $\pm$ 1.03	+10.23

Table 5: Comparison of sampling strategies on Math : performance comparison of token -level sampling adapter versus static sampling strategies. Here budget is the remaining token budget within a sequence. Values are percentages.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

		Static sampling		Sequence-level adapter
Metric	Setting	Best	Mixed	w/o budget	w/ budget
Pass@1	w/o CoT	72.86 $\pm$ 1.23	72.38 $\pm$ 1.24	73.20 $\pm$ 1.23	73.67 $\pm$ 1.22
	mix CoT	73.02 $\pm$ 1.23	72.92 $\pm$ 1.23	73.80 $\pm$ 1.22	74.07 $\pm$ 1.21
Pass@8	w/o CoT	78.38 $\pm$ 1.14	77.10 $\pm$ 1.16	78.20 $\pm$ 1.14	78.80 $\pm$ 1.13
Pass@8	mix CoT	78.98 $\pm$ 1.13	78.58 $\pm$ 1.14	79.00 $\pm$ 1.13	79.13 $\pm$ 1.13

Appendix

Table 6: Comparison of sampling strategies on MATH. against the proposed sequence -level sampling adapter, using Qwen3-8B. For the ‘w/ CoT mix’ configuration, a CoT ratio of 0.3 is applied. All reported values are percentages.

		Static sampling		Sequence-level adapter
Metric	Setting	Best	Mixed	w/o budget	w/ budget
Pass@1	w/o CoT	12.80 $\pm$ 2.39	12.00 $\pm$ 2.33	18.80 $\pm$ 2.80	17.90 $\pm$ 2.74
Pass@1	mix CoT	19.70 $\pm$ 2.85	14.50 $\pm$ 2.52	22.20 $\pm$ 2.97	23.10 $\pm$ 3.02
Pass@8	w/o CoT	23.10 $\pm$ 3.02	22.23 $\pm$ 2.98	26.50 $\pm$ 3.16	27.40 $\pm$ 3.19
Pass@8	mix CoT	29.10 $\pm$ 3.25	23.90 $\pm$ 3.05	32.50 $\pm$ 3.35	34.20 $\pm$ 3.40

Appendix

Table 7: Comparison of sampling strategies on CodeContests. : vs. the sequence -level sampling adapter under two inference settings (without/with a mix of CoT prompts), using Qwen3-8B. For ‘w/ CoT mix’, we use a CoT ratio of 0.3.

		Static sampling		Sequence-level adapter
Metric	Setting	Best	Mixed	Just 4 actions	w/ mix of inference strategy
Pass@1	w/o CoT	61.52 $\pm$ 1.35	56.42 $\pm$ 1.37	61.52 $\pm$ 1.35	62.62 $\pm$ 1.34
Pass@1	mix CoT	62.10 $\pm$ 1.34	59.00 $\pm$ 1.36	62.76 $\pm$ 1.34	63.20 $\pm$ 1.34
Pass@8	w/o CoT	61.52 $\pm$ 1.35	66.74 $\pm$ 1.31	68.25 $\pm$ 1.29	69.66 $\pm$ 1.27
Pass@8	mix CoT	63.00 $\pm$ 1.34	67.10 $\pm$ 1.30	69.32 $\pm$ 1.28	70.88 $\pm$ 1.26

Appendix

Table 8: Comparison of sampling strategies on MATH. the sequence -level sampling adapter under two inference settings (without/with a mix of CoT prompts). 0.3.

		Static sampling		Adapter
Metric	Setting	Best	Mixed	w/o budget
Pass@1	w/o CoT	11.43 $\pm$ 2.28	10.53 $\pm$ 2.20	12.80 $\pm$ 2.39
Pass@1	mix CoT	13.97 $\pm$ 2.48	14.80 $\pm$ 2.54	17.90 $\pm$ 2.74

Appendix

Table 9: Coding: performance comparison of sequence -level adapter trained on MATH train versus baseline numbers.

Figure 3: Entropy modulation under token-level control. Token entropy distributions for a representative static strategy versus the learned token-level adapter, using the temperature-only action set from Table 14 . Action 0 is greedy; larger action indices correspond to higher temperature. Qualitatively, the adapter more often preserves stochasticity on higher-entropy tokens, while collapsing many low-entropy tokens to near-deterministic behavior.

Panel of Figure 4 — Figure 4: Generation length distributions. Generation length statistics for a representative static strategy and the learned adapters. Length shifts are present but not extreme, suggesting the token-level gains are not explained solely by systematic truncation/verbosity changes.

Figure 5: Word clouds by action. Most frequent tokens under each temperature-based action from Table 14 . This is a lightweight diagnostic and does not by itself explain the policy’s performance improvements.

Parameter	Value
Temperature	$[0.3,0.5,0.75,1.0,1.25]$
Top- $k$	$[5,10,50,\text{``off"}]$
Top- $p$	$[0.9,0.95,\text{``off"}]$
Min- $p$	$[0.1,0.2,\text{``off"}]$

Appendix

Table 10: Sweep ranges of as the action space for our sampling adapter.

Figure 6: Best of $N$ accuracy vs number of selected actions for greedy selection (our approach) vs the baseline of taking the top N actions.

Strategy	Parameters
	do_sample	temperature	top_k	top_p	min_p
1	true	0.75	–	0.9	0.1
2	true	1.0	–	0.9	–
3	true	1.25	10	0.9	–
4	true	1.0	5	0.9	–
5	true	1.25	50	0.9	–

Appendix

Table 11: Selected decoding strategies for math tasks.

Strategy	Parameters
	do_sample	temperature	top_k	top_p	min_p
1	true	1.25	–	0.95	0.1
2	true	1.25	–	0.95	0.2
3	true	1.0	–	0.95	0.2
4	true	1.25	50	0.9	0.1
5	true	1.25	5	0.9	–

Appendix

Table 12: Selected decoding strategies for coding tasks.

Strategy	Parameters
	do_sample	temperature	top_k	top_p	min_p
0	true	0.5	–	–	–
1	true	1.0	–	–	–
2	true	1.25	–	–	–
3	true	0.75	10	–	–
4	true	0.75	10	0.95	0.1

Appendix

Table 13: Selected decoding strategies for a mix of math and coding tasks.

Strategy	Parameters
	do_sample	temperature
0	false	0
1	true	0.5
2	true	1.0
3	true	1.25

Appendix

Table 14: Selected decoding strategies for token-level math tasks.

Parameter	Value
Architecture	MLP
Number of Layers	3
Activation	SiLU
Dropout	0.1
Optimizer	Adam
format	bf16

Appendix

Table 15: Parameters of token -level sampling adapter.

Panel of Figure 7 — Figure 7: Action distributions and val reward on CodeContests for sampling strategies with sequence-level adapter. The actions are defined in table 12 . Axes are same as in Figure 2 .

Panel of Figure 8 — Figure 8: Action distributions and val reward on MATH for sampling strategies with the token -level adapter. Actions are defined in table 11 . Axes match fig. 2 . Note that a CoT action probability of 0 indicates that no CoT prompt was sampled in that batch.

Explore similar work

Adaptive Decoding via Test-Time Policy Learning for Self-Improving Generation
Mar 19, 2026Asmita Bhardwaj, Yuya Jeremy Ong, Eelaaf Zahid +1Language Model Decoding Test-Time Adaptation
UC San Diego San Diego, CA · Plastic Labs San Jose, CA · IBM Research Cambridge, MA +1
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
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

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

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

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

UC San Diego San Diego, CA · Plastic Labs San Jose, CA · IBM Research Cambridge, MA+1

Language Model Decoding Test-Time Adaptation RL for Language Models

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

Entropy-informed Decoding: Adaptive Information-Driven Branching

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