PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ying Li, Yuqi Chen, Sixuan Li, Yunfeng Cai

arXiv:2610.17679v1cs.LGcs.AI

TL;DR

Fixed-budget continuous reasoning can miss questions recoverable at other budgets, while latent-update counts alone may misstate end-to-end cost. Dynamic K-Head selects a question-conditioned budget before rollout and uses difficulty supervision. On GSM8K, it improves accuracy and lowers token cost relative to fixed-K=6 CoDI, with joint gains also reported on three out-of-distribution datasets.

  • Problem

    Continuous latent reasoning typically uses a fixed update count even though the best budget varies by question and accuracy need not improve monotonically with K.

    Source: §1 · pp. 1–2

  • Method

    Dynamic K-Head selects a question-conditioned budget before rollout and jointly trains the selector and reasoner under selected budgets using difficulty supervision.

    Source: §1 · p. 3, §4.2 · p. 7

  • Results

    54.44% GSM8K accuracy at 12.03 average tokens compares with fixed-K=6 CoDI’s 52.31% at 15.21 tokens, a 2.12-point gain and 20.9% lower cost; joint accuracy-cost gains also hold across three out-of-distribution datasets.

    Source: §1 · p. 3, §8 · p. 14

  • Takeaways & Limitations

    Difficulty-supervised budget selection can improve both accuracy and efficiency, while stable budget selection remains a central challenge.

    Source: §8 · p. 14

    The contrastive objective supplies weak question-level supervision, without explicit per-example credit assignment for which candidate budgets yield correct answers or counterfactual evidence for unexecuted budgets.

    Source: §7 · p. 13

Abstract

from arXiv · show

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.

1 Introduction

Dynamic K-Head selects a question-conditioned latent-update budget before rollout, addressing the non-monotonic behavior of fixed-budget reasoning. On GSM8K, it improves CoDI’s accuracy while reducing reasoning/output token cost.

1 Introduction

  • Of 629 GSM8K questions missed at K=6, 138 are recovered by a smaller budget and 30 by a larger one, with overlap, showing recovery is non-monotonic.The sweep solves 845 questions at at least one budget, versus 690 at K=6.Source: §1 · p. 2, Figure 1 · p. 2
  • With λK = 0.03, Dynamic K-Head reaches 54.44% accuracy at 12.03 average tokens, versus CoDI’s 52.31% at 15.21 tokens, a 2.12-point gain and 20.9% lower token cost.Source: §1 · p. 3
  • End-to-end cost accounting includes latent updates, implementation-specific special tokens, and all decoded text because shorter latent rollouts can shift reasoning into longer explicit continuations.The accounting is applied uniformly across experimental settings.Source: §1 · p. 3
  • The pre-rollout K-head selects executed latent states from the question representation and is jointly trained under variable budgets with difficulty supervision, without direct budget labels or a monotonic mapping.A contrastive objective groups budget representations for questions with similar difficulty; training transitions from fixed-budget warm-up to K-head selection.Source: §1 · pp. 2–3

2 Related Work

Prior work covers explicit and continuous latent reasoning, adaptive computation, and difficulty-based training signals. Dynamic K-Head differs by selecting a question-specific latent budget before rollout and treating budgets as distinct choices.

2.1 Explicit and Latent Reasoning

  • Continuous latent methods compress intermediate reasoning into hidden computation, but the cited approaches do not learn a pre-rollout, question-conditioned budget selector for recurrent latent computation.Explicit CoT incurs autoregressive rationale cost; related approaches include Coconut and CoDI [5] [6].Source: §2.1 · pp. 3–4

2.2 Adaptive Computation

  • Unlike adaptive halting methods, Dynamic K-Head jointly trains the selector and reasoner under executed budgets because non-monotonic latent updates make continuing potentially drift from a useful state.It treats budgets as distinct choices without optimal-K labels or a monotonic mapping [21] [22].Source: §2.2 · p. 4

2.3 Difficulty Supervision

  • Difficulty supervision uses relative neighborhoods in question-level scores rather than regressing budgets, since difficulty need not vary monotonically with the best latent budget.The scores may reflect teacher priors rather than the student’s computation; the contrastive approach uses neither augmentation identity nor discrete class or budget labels.Source: §2.3 · pp. 4–5

3 Background and Problem Formulation

Latent CoT performs intermediate computation in recurrent hidden states, then generates the answer as ordinary tokens. The formulation contrasts a shared fixed update count with selecting the count before rollout.

3 Background and Problem Formulation

  • At each latent update, the model feeds the current final-position hidden state back as a continuous input representation, without selecting a vocabulary token.The resulting hidden state is used for the next recurrent update.Source: §3 · p. 5
  • K denotes the number of latent updates; after exactly K updates, the model decodes the answer autoregressively, conditioned on the question and accumulated latent states.A fixed-budget reasoner uses the same K for every question.Source: §3 · p. 5
  • The selected budget determines how many latent updates run before the model leaves the rollout and generates the answer.The selection occurs before latent reasoning.Source: §3 · p. 5

4 Dynamic K-Head

Dynamic K-Head selects a question-conditioned latent budget before rollout, trains budget representations using relative difficulty, and accounts for latent updates, special tokens, and decoded text as end-to-end cost. Inference executes each selected budget and repacks variable-length prefixes for batched answer decoding.

4.1 Pre-Rollout Budget Selection

  • The K-head maps pre-rollout question representations to budget logits, and training uses an anchor-budget warm-up before drawing per-example budgets from an anchor-mixture prior.The temperature and annealing schedules control the budget distribution during post-warm-up training.Source: §4.1 · pp. 5–6

4.2 Difficulty Supervision

  • The contrastive objective encourages questions with nearby teacher-assigned difficulty scores to have similar budget distributions, without target budgets or a monotonic difficulty-to-budget mapping.The method learns one prototype per candidate budget and forms positive and negative sets from nearest and farthest difficulty scores within a batch.Source: §4.2 · pp. 6–7

4.3 Token-Cost Accounting

  • End-to-end cost counts latent updates, implementation-specific special tokens, and every decoded text token, since shorter rollouts can shift reasoning into explicit continuations.One example uses one latent update, two special tokens, and 28 decoded text tokens, totaling 31 tokens.Source: §4.3 · p. 7

4.4 Deterministic Per-Example Inference

  • At inference, the K-head removes prior mixing and stochastic sampling, then the reasoner executes exactly the selected number of latent updates before answer decoding.Source: §4.4 · p. 7

4.5 Active-Then-Repack Execution

  • Active-then-repack execution drops examples from the active cache after their selected budgets, then repacks variable-length prefixes with explicit masks and position IDs for batched decoding.This avoids further latent updates for examples that have reached their budgets while allowing different budgets in one batch.Source: §4.5 · pp. 7–8

5 Experimental Setup

The experiments evaluate Dynamic K-Head on aligned GSM8K-Aug training versions and three arithmetic datasets for out-of-distribution evaluation. They use question-only teacher difficulty scores, instantiate the method with CoDI, and specify model, training, and hardware settings.

Data

  • The two aligned GSM8K-Aug versions each provide 385,620 training, 500 validation, and 1,319 test examples, with the test split following official GSM8K.The best checkpoint is selected by validation exact-match accuracy and evaluated on the test split.Source: §5 · p. 8
  • Out-of-distribution evaluation uses SVAMP, GSM-Hard, and MultiArith, with training and test data combined for SVAMP and test examples used for the other datasets.The evaluation uses 1,319 GSM-Hard examples and 180 MultiArith test examples.Source: §5 · p. 8

Difficulty scores

  • Main experiments use DeepSeek question-only difficulty scores from 1 to 5, while ablations also test continuation negative log-likelihood as an alternative signal.The scores range from nearly one-step arithmetic to unusually tricky questions.Source: §5 · p. 8

Instantiation

  • Dynamic K-Head is instantiated with CoDI, which trains explicit-rationale and recurrent latent paths through self-distillation, using GPT-2 and Llama 3.2 1B Instruct as base models.Experiments use LoRA fine-tuning, and the K-head uses budgets {1, 2, 3, 4, 6, 8} with a 10% warm-up fraction.Source: §5 · p. 8

Implementation and hardware configuration

  • Training and evaluation use PyTorch and Hugging Face Transformers on NVIDIA H100 80GB GPUs, with two GPUs allocated to each run.Source: §5 · p. 9

6 Results

Dynamic K-Head improves the fixed-budget reasoner on in-domain and out-of-distribution arithmetic results, while its accuracy–cost balance depends on the supervision signal and loss weight. The experiments also show that budget selection and its evaluation require attention to model capacity, difficulty-signal behavior, and how costs are measured.

6.2 Out-of-Distribution Evaluation

  • The selector averages 2.15–2.74 latent updates and never selects K = 8; the authors attribute remaining errors to shifted arithmetic coverage rather than budget exhaustion.The figure compares accuracy and token cost across three datasets, with dataset-specific vertical scales.Source: §6.2 · pp. 9–10, Figure 4 · p. 11

6.1 Overall Accuracy–Cost Results

  • On Llama, Dynamic K-Head raises GSM8K accuracy from 52.31% to 54.44% while cost falls from 15.21 to 12.03 tokens; GSM8K-NL accuracy rises from 46.70% to 50.34% at 1.95 additional tokens.Against CoT SFT on GSM8K, it is 7.88 accuracy points lower while reducing average token cost from 33.96 to 12.03, a 2.82× compression.Source: §6.1 · p. 9, Table 1 · p. 10

6.2 Out-of-Distribution Evaluation

  • Across three OOD datasets, Dynamic K-Head cuts cost by 19.3–25.7% and raises accuracy by 0.30–2.22 points versus fixed-K CoDI.Compared with CoT SFT, it reduces cost by 41.7–69.7%; CoT SFT is more accurate on SVAMP and GSM-Hard, while Dynamic K-Head is 0.56 points more accurate on MultiArith at half the cost.Source: §6.2 · p. 9, Figure 4 · p. 11

6.3 Online vs. External Difficulty Supervision

  • On Llama, DeepSeek difficulty supervision reaches 54.44% at 12.03 tokens and improves both metrics over CoDI and continuation NLL; on GPT-2, it lowers cost to 11.56 tokens but accuracy is 37.68%.The authors suggest GPT-2’s smaller capacity may limit both budget-selection representations and accuracy under shorter computation.Source: §6.3 · p. 12
  • Continuation NLL is recomputed each epoch and tracks student uncertainty, whereas fixed DeepSeek scores provide a stable question-level ordering; neither defines a universal compute target.The value of either signal depends on the student model’s capacity and representations.Source: §6.3 · p. 12

6.4 Difficulty Supervision Controls Budget Selection

  • At λK = 0.03, accuracy is 54.44% at 12.03 tokens, with predictive entropy 0.097 and the highest mutual information, 0.443; λK = 0.01 is most accurate at 54.51%.The 0.03 setting nearly matches the 0.01 accuracy at lower cost and yields the strongest observed balance.Source: §6.4 · p. 12, Table 3 · p. 13
  • At λK = 0.05, predictive entropy is 0.006 and marginal entropy is 0.122; accuracy is 52.62% at the lowest cost, 10.27 tokens, indicating cheaper but less adaptive selection.At λK = 0.10, mutual information is 0.424 while accuracy falls to 49.05%; the low-score train–evaluation shift may contribute to high-weight concentration.Source: §6.4 · p. 12, Table 3 · p. 13

7 Limitations and Broader Considerations

The evaluation is limited to two model families, teacher-assigned difficulty scores, and weak question-level supervision. The authors note that broader model coverage and stronger outcome-based supervision remain open considerations.

Small models

  • Experiments cover only GPT-2 and Llama, so generalization to other model families and larger scales remains unknown.Source: §7 · p. 13

External difficulty scores

  • Teacher-assigned difficulty scores may contain systematic biases and may not reflect the task difficulty experienced by the student model.The authors call for comparison across multiple scoring models.Source: §7 · p. 13

Weak supervision

  • Question-level contrastive supervision gives no explicit per-example credit assignment or counterfactual evidence for unexecuted budgets; collecting outcomes across budgets would require additional rollouts.The learned budget representations may align with difficulty-score neighborhoods without matching budget-specific success patterns.Source: §7 · p. 13

8 Conclusion

Dynamic K-Head treats latent budgets as non-monotonic, instance-conditioned choices and improves CoDI’s GSM8K accuracy while reducing token cost. Budget-selection quality depends on difficulty coverage and supervision, motivating finer-grained data and more direct outcome-based supervision.

8 Conclusion

  • 54.44% GSM8K accuracy at 12.03 average tokens improves fixed-K=6 CoDI by 2.12 points while reducing token cost by 20.9%.The approach also retains joint accuracy–cost gains on three out-of-distribution datasets.Source: §8 · p. 14
  • Budget-selection quality depends on the granularity and balance of difficulty coverage and how directly supervision reflects budget utility.The authors propose finer-grained, better-balanced training data and more direct outcome-based supervision as future work.Source: §8 · p. 14
  • The authors say code, difficulty-annotated datasets, and model weights will be publicly released once prepared.The resources are still being organized for release.Source: Code, Data, and Model Availability · p. 14

A Dataset Formats and Examples

The appendix describes two GSM8K-Aug formats and the input and answer fields used for three out-of-distribution arithmetic datasets. It provides representative records for the formats and evaluation sets.

A.1 GSM8K-Aug Training Data

  • Both GSM8K-Aug formats store a question, a list of steps, and an answer; one uses calculator-style expressions and the other natural-language sentences.The questions and answers align across formats.Source: §A.1 · p. 17
  • The aligned GSM8K-Aug example asks how many of 600 employees received neither a promotion nor a bonus.The example record is presented in both formats.Source: §A.1 · p. 17
  • In the example, 180 promotions plus 60 bonuses yield 240 employees receiving either, leaving 360 receiving neither.The record gives the calculation steps and answer.Source: §A.1 · p. 17

A.2 OOD Evaluation Data

  • SVAMP combines 700 training and 300 test examples, using question_concat as input and Answer as the gold answer.A test example asks for supermarket spending and gives 31 as the answer.Source: §A.2 · p. 17
  • GSM-Hard uses all 1,319 examples, with instruction as input and response as the answer.The representative record gives 3,431,580.0 as the answer.Source: §A.2 · pp. 17–18
  • MultiArith uses its 180-example test split, with question as input and final_ans as the answer.The example asks about songs remaining after deletion and addition, with answer 10.Source: §A.2 · pp. 17–18

B A Small Latent Budget Can Still Produce a Long Continuation

A correctly answered GSM8K example shows that one selected latent update can be followed by substantial explicit reasoning. Counting only latent updates would miss much of its reasoning/output cost.

B A Small Latent Budget Can Still Produce a Long Continuation

  • The example is from an official GSM8K test evaluation of a GPT-2 Dynamic K-Head model, which selected one latent update.The computation then continued in explicit text.Source: B A Small Latent Budget Can Still… · p. 18
  • After that update, the model generated 28 text tokens containing three calculation steps and the final answer.The answer, 91, matches the gold answer.Source: B A Small Latent Budget Can Still… · p. 18
  • Including BOT and EOT tokens, this one-update example costs 31 tokens, showing K alone is not a reliable end-to-end cost proxy.The accounting is 1 latent update + 2 special tokens + 28 generated tokens.Source: B A Small Latent Budget Can Still… · p. 18

C Difficulty-Score Distributions

DeepSeek question-only difficulty scores are more concentrated in augmented training data than in validation and test. The accompanying information quantities characterize per-question budget uncertainty and how soft budget distributions vary across questions.

C Difficulty-Score Distributions

D Information Quantities for Budget Selection

Related papers

Loading 2610.17679v1…