arXiv:cs.AI· Ying Li, Yuqi Chen, Sixuan Li, Yunfeng Cai·· 23 小时前
Dynamic K-Head:面向连续潜在推理的难度监督预算选择
Dynamic K-Head: Difficulty-Supervised Budget Selection for Continuous Latent Reasoning
AI 导读
Dynamic K-Head 是一种以问题为条件的预展开预算选择器,与推理器在选定预算下联合训练,难度监督仅来自问题本身的评分。
正文
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.
| Subjects: | Machine Learning (cs.LG); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.17679 [cs.LG] |
| (or arXiv:2610.17679v1 [cs.LG] for this version) | |
| https://doi.org/10.48550/arXiv.2610.17679 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Ying Li [view email]
[v1]
Tue, 13 Oct 2026 04:02:46 UTC (137 KB)
来源:arXiv:cs.AI · arxiv.org