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

arXiv: https://arxiv.org/abs/2610.17679 · HTML: https://mlsanity.com/paper/2610.17679/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Test-time compute and efficient thinking, one of 17 directions in LLM reasoning.
- Builds on 29 of its 36 references that are in the map.
- First posted on arXiv 13 Oct 2026.

## Placed in

- LLM reasoning / [Test-time compute and efficient thinking](https://mlsanity.com/lr/test-time-compute-and-efficient-thinking/)

## Builds on (most cited first, 29 of 29)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Supervised Contrastive Learning](https://mlsanity.com/paper/2004.11362/) (2020)
- [Are NLP Models really able to Solve Simple Math Word Problems?](https://mlsanity.com/paper/2103.07191/) (2021)
- [Training Large Language Models to Reason in a Continuous Latent Space](https://mlsanity.com/paper/2412.06769/) (2024)
- [PAL: Program-aided Language Models](https://mlsanity.com/paper/2211.10435/) (2022)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Dataset Cartography: Mapping and Diagnosing Datasets with Training Dynamics](https://mlsanity.com/paper/2009.10795/) (2020)
- [Solving General Arithmetic Word Problems](https://mlsanity.com/paper/1608.01413/) (2016)
- [Confident Adaptive Language Modeling](https://mlsanity.com/paper/2207.07061/) (2022)
- [Quiet-STaR: Language Models Can Teach Themselves to Think Before Speaking](https://mlsanity.com/paper/2403.09629/) (2024)
- [Think before you speak: Training Language Models With Pause Tokens](https://mlsanity.com/paper/2310.02226/) (2023)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [Depth-Adaptive Transformer](https://mlsanity.com/paper/1910.10073/) (2019)
- [From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step](https://mlsanity.com/paper/2405.14838/) (2024)
- [CODI: Compressing Chain-of-Thought into Continuous Space via Self-Distillation](https://mlsanity.com/paper/2502.21074/) (2025)
- [Implicit Chain of Thought Reasoning via Knowledge Distillation](https://mlsanity.com/paper/2311.01460/) (2023)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)
- [SoftCoT: Soft Chain-of-Thought for Efficient Reasoning with LLMs](https://mlsanity.com/paper/2502.12134/) (2025)
- [Parallel Continuous Chain-of-Thought with Jacobi Iteration](https://mlsanity.com/paper/2506.18582/) (2025)
- [Do Latent-CoT Models Think Step-by-Step? A Mechanistic Study on Sequential Reasoning Tasks](https://mlsanity.com/paper/2602.00449/) (2026)
- [Dynamics Within Latent Chain-of-Thought: An Empirical Study of Causal Structure](https://mlsanity.com/paper/2602.08783/) (2026)
- [Learning When to Stop: Adaptive Latent Reasoning via Reinforcement Learning](https://mlsanity.com/paper/2511.21581/) (2025)
- [Thinking in Latents: Adaptive Anchor Refinement for Implicit Reasoning in LLMs](https://mlsanity.com/paper/2603.15051/) (2026)

## Abstract (arXiv)

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.