Dynamic K-Head: Difficulty-Supervised Budget Selection for Continuous Latent Reasoning

Ying Li; Yuqi Chen; Sixuan Li; Yunfeng Cai

cs.LGOct 13, 2026

Dynamic K-Head: Difficulty-Supervised Budget Selection for Continuous Latent Reasoning

Authors: Ying Li, Yuqi Chen, Sixuan Li, Yunfeng Cai

Organizations: Beijing Institute of Mathematical Sciences and Applications, Beijing, China · Tsinghua University, Beijing, China · The University of Hong Kong, Hong Kong, China · Peking University, Beijing, China

Abstract

Chain-of-thought (CoT) improves reasoning but incurs substantial autoregressive rationale cost. Continuous CoT moves intermediate computation into recurrent hidden states, yet typically uses a fixed number of updates even though the best budget varies across questions and accuracy need not improve monotonically with $K$ . We introduce Dynamic K-Head, a question-conditioned, pre-rollout budget selector. It is jointly trained with the reasoner under selected budgets, using difficulty supervision based on scores assigned from the question alone. We also show that latent-update count can misrepresent end-to-end cost because a short latent rollout can shift computation into a longer explicit continuation. We therefore count latent updates, implementation-specific special tokens, and all decoded text. Instantiated with CoDI on Llama 3.2 1B and GSM8K, Dynamic K-Head reaches 54.44% exact-match accuracy at 12.03 average reasoning/output tokens, compared with 52.31% at 15.21 tokens for fixed- $K=6$ CoDI---a 2.12-point gain with 20.9% lower token cost. Across three out-of-distribution arithmetic datasets, it reduces cost by 19.3--25.7% while improving accuracy by 0.30--2.22 points. These results demonstrate that difficulty-supervised budget selection can improve both accuracy and efficiency, while stable budget selection remains a central challenge.

Figures & tables

Figure 1: Fixed-budget accuracy is non-monotonic and conceals instance-level recoveries. Panel (a) sweeps inference budget

k

for CoDI trained with fixed

K=6

. Panel (b) partitions the 155

K=6

errors recovered by another swept budget: 125 are recovered only below six, 17 only above six, and 13 on both sides. Thus the smaller- and larger-budget recovery totals are 138 and 30, respectively, with the 13 “both sides” problems included in each.

Figure 2: Schematic of Dynamic K-Head training and execution. Anchor, positive, and negative questions are prefilled by the LLM, and the K-head produces a distribution over latent budgets. Learned prototypes map this distribution to

u_{i}

, while per-question difficulty scores make nearby questions positive pairs and distant questions negative pairs. A sampled

K

controls the illustrated shared train/inference path, where the anchor question is concatenated with exactly

K

latent tokens before answer decoding.

GPT-2
	GSM8K		GSM8K-NL
Method	Accuracy (%)	Avg. token cost	Accuracy (%)	Avg. token cost
No-CoT SFT	14.25	6.34	14.25	6.34
CoT SFT	48.90	32.27	33.59	67.60
Coconut	33.13	11.21	–	–
CoDI	39.80	14.24	37.38	14.26

Table 1: Accuracy and average token cost on the official GSM8K test set across two backbones and the GSM8K and GSM8K-NL training variants. AdaLR denotes Adaptive Latent Reasoning. All Dynamic K-Head rows use DeepSeek difficulty supervision. Token cost follows Eq. 8 . Within each column, bold and underlined values indicate the best and second-best results, respectively; CoT SFT is excluded from the accuracy ranking, and No-CoT SFT is excluded from the cost ranking.

Figure 3: Accuracy versus average token cost for the Table 1 results, with GPT-2 on the left and Llama 3.2 1B on the right. GSM8K and GSM8K-NL are overlaid using the same method-specific markers, and the horizontal axis uses a logarithmic scale. Dashed lines indicate the Pareto frontiers.

Figure 4: Accuracy (top) and average token cost (bottom) on three OOD datasets. Within each column, methods are ordered as No-CoT SFT, CoDI (

K=6

), Dynamic K-Head, and CoT SFT. Each panel uses a dataset-specific vertical scale.

GPT-2
Method	Accuracy (%)	Avg. token cost
CoDI	39.80	14.24
Dynamic K-Head (NLL)	38.29	14.37
Dynamic K-Head (difficulty)	37.68	11.56
Llama 3.2 1B
CoDI	52.31	15.21

Table 2: Fixed-budget CoDI and Dynamic K-Head with two signals for difficulty supervision on GSM8K. NLL denotes online continuation NLL recomputed at each training epoch; difficulty denotes fixed DeepSeek scores computed from questions alone. Within each backbone block, bold and underlined values indicate the best and second-best results, respectively.

$\lambda_{K}$	Accuracy (%)	Avg. cost	$H(K)$	$H(K\mid Q)$	$I(Q;K)$
0.00	50.87	13.11	1.781	1.715	0.066
0.01	54.51	13.33	0.441	0.131	0.310
0.03	54.44	12.03	0.540	0.097	0.443
0.05	52.62	10.27	0.122	0.006	0.115
0.10	49.05	10.89	0.724	0.300	0.424

Table 3: Accuracy, average token cost, marginal entropy

H(K)

, predictive entropy

H(K\mid Q)

, and question–budget mutual information

I(Q;K)

for the loss-weight ablation on GSM8K. All information quantities are measured in nats.

Appendix figures & tables1 asset

Supplementary material from the paper’s appendix.

Appendix

Figure 5: DeepSeek difficulty scores, computed from questions alone, across training, validation, and test. The lower panel reports validation-minus-training and test-minus-training differences in percentage points.

Efficient Language Model Reasoning

Efficient language model reasoning studies how to improve problem-solving accuracy while reducing inference, memory, latency, and training costs.

See all 371 papers

Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs
Mar 29, 2026Amirul Rahman, Aisha Karim, Kenji Nakamura +1Efficient Language Model Reasoning Inference-Time Compute Allocation
University of Malaya
ROI-Reasoning: Rational Optimization for Inference via Pre-Computation Meta-Cognition
Jan 7, 2026Muyang Zhao, Chunhao Tian, Qi Qi +1Cost-Aware Inference Metacognition in Language Models
Gaoling School of Artificial Intelligence, Renmin University of China, Beijing, China
Draft-Thinking: Learning Efficient Reasoning in Long Chain-of-Thought LLMs
Feb 28, 2026Jie Cao, Tianwei Lin, Juekai Lin +7Efficient Language Model Reasoning CoT Reasoning
Zhejiang University · Tencent

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

Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs

Amirul Rahman, Aisha Karim, Kenji Nakamura, Yi-Fan Ng

University of Malaya

Efficient Language Model Reasoning Inference-Time Compute Allocation Overthinking in Language Models RL for Language Model Reasoning Numerical Reasoning in Language Models

Open paper AskPDF arXiv

Scaling test-time compute via extended reasoning has become a key paradigm for improving the capabilities of large language models (LLMs). However, existing approaches optimize reasoning under fixed or uniformly sampled token budgets, ignoring the fundamental mismatch between problem difficulty and allocated compute. This leads to overthinking on easy problems and underthinking on hard ones, resulting in suboptimal token efficiency across diverse reasoning scenarios. In this paper, we propose Budget-Adaptive Curriculum Reasoning (BCAE), a unified framework that jointly optimizes reasoning quality and token efficiency through three synergistic components: (1) a budget-conditioned unified policy that embeds the token budget as a continuous conditioning signal, eliminating the need for decoupled thinking and summarization strategies; (2) a curriculum-aware budget scheduler that adaptively shifts the training budget distribution from easy to hard problems based on real-time learning progress; and (3) a truncation-aware dense reward mechanism that provides fine-grained credit assignment at intermediate reasoning steps via process-level verification. We further introduce Budget-Conditioned Advantage Estimation (BCAE), a novel variance reduction technique that conditions the advantage baseline on the sampled budget, yielding more stable policy gradients. Experiments on mathematical reasoning benchmarks (MATH, GSM8K, AIME, and Minerva Math) demonstrate that BACR consistently outperforms other strong baselines across all token budgets, achieving up to 8.3% accuracy improvement under tight budgets while reducing average token consumption by 34% compared to unconstrained reasoning.

Abstract

Efficient Language Model Reasoning

Explore similar work

Avoiding Overthinking and Underthinking: Curriculum-Aware Budget Scheduling for LLMs

ROI-Reasoning: Rational Optimization for Inference via Pre-Computation Meta-Cognition

Draft-Thinking: Learning Efficient Reasoning in Long Chain-of-Thought LLMs