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

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.

Machine Learning (cs.LG); Artificial Intelligence (cs.AI)

arXiv:2610.17679 [cs.LG]

(or arXiv:2610.17679v1 [cs.LG] for this version)

https://doi.org/10.48550/arXiv.2610.17679 Focus to learn more arXiv-issued DOI via DataCite (pending registration)

arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.

Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.

Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.