Towards Inference-time Scaling for Continuous Space Reasoning

AI summary

Overview

• Research area: Inference-time scaling for large language model reasoning, applied to continuous (latent-space) reasoning rather than text-based chain-of-thought.
• Technical level: Intermediate. The paper is readable without deep math background, but familiarity with LLM sampling, best-of-N reranking, and reward models helps.
• Scope: A single-backbone (COCONUT, GPT-2 scale) empirical study on GSM8k asking whether sampling plus reward-model reranking transfers from discrete text reasoning to continuous thought space.

What This Paper Is About

Inference-time scaling works well for text reasoning: generate many candidate chains of thought, then use a Process Reward Model (PRM) or Outcome Reward Model (ORM) to rerank them. This paper asks whether the same recipe works when a model reasons in continuous latent space, using COCONUT as the backbone. The authors show that diverse continuous reasoning paths can indeed be sampled, and that a large accuracy ceiling exists, but they find that adapting existing reward-model training recipes yields only marginal gains and then diagnose why.

Key Contributions

1. First exploration of inference-time scaling for continuous space reasoning (to the authors' knowledge). They introduce a dropout-based sampling method that injects controlled stochasticity only into the iterative hidden-state generation phase, while keeping dropout disabled during final text answer generation, thereby producing diverse continuous reasoning trajectories from an otherwise deterministic model.
2. Adaptation of the MATH-Shepherd annotation pipeline to continuous thoughts. They treat each continuous thought vector as a reasoning step, generate Monte Carlo completions from partial trajectories, and produce hard and soft step-level labels, training PRM and ORM variants on top of the COCONUT backbone itself (since latent thoughts are model-specific and cannot be interpreted by a different architecture). This yielded 238k PRM training samples and 324k ORM training samples.
3. A multi-angle diagnostic analysis of why reward models fail in continuous space, covering classification performance (accuracy, precision, recall, F1, specificity, confusion matrices), geometric properties of thought vectors (IsoScore* isotropy, Hoyer sparsity, t-SNE visualization), trajectory dynamics (compactness, curvature, local smoothness, straightness with t-tests and Cohen's d), and behavior under Gaussian noise perturbation.
4. A diagnosis and proposed direction. They argue that current limitations stem from the absence of key inductive biases in continuous thought representations, because training supervises only text tokens and optimizes final answer accuracy without promoting structural differentiation in latent space, and they propose incorporating geometric inductive biases.

Main Findings

• Sampling diversity is achievable and the ceiling is high. Deterministic COCONUT reasoning shows nearly a 10-point accuracy gap versus text CoT, consistent with the original COCONUT paper. Dropout-based sampling gives a Pass@1 slightly lower than deterministic COCONUT, but Pass@N rapidly surpasses deterministic COCONUT as sample size grows and exceeds the text CoT baseline at N=32. Pass@N rises from 31.08 at N=1 to 44.43 at N=32.
• Unique answers grow logarithmically, not linearly. The average number of unique answers goes from 1.00 at N=1 to 3.17 at N=32, which the authors interpret as evidence that reranking could be computationally efficient.
• Reward models deliver only marginal gains. PRM-HE is the most consistent reranker, reaching 33.36% at N=16 versus the 31.08% baseline, but the improvement is less than 2.3 points despite a theoretical upper bound of 42.61% at N=16. The authors note these gains pale in comparison to text-based RM verification with similar methodology.
• Model-intrinsic reranking fails. Confidence-based reranking provides no improvement (e.g., 30.48 at N=4, 29.87 at N=8), indicating poor calibration in COCONUT. Self-consistency shows only marginal gains (32.15 at N=16 and N=32). Across all N values, correct answers consistently fall below the count of dominant incorrect answers (at N=32: 9.84 correct answers versus 15.88 major incorrect answers).
• Score aggregation strategy does not matter. Testing last-step, min, max, and mean aggregation for the PRM shows negligible variation across methods, suggesting the core limitation is deeper than scoring mechanics.
• Reward models classify poorly. PRM achieves accuracy 62.98, precision 41.60, recall 77.28, F1 54.09, specificity 57.36. ORM achieves accuracy 73.72, precision 39.11, recall 75.76, F1 51.59, specificity 73.26. PRM's confusion matrix shows 5,535 false positives (30.6%) against 3,943 true positives (21.8%).
• Continuous thoughts are geometrically homogeneous. Correct and incorrect thoughts show negligible differences in isotropy (IsoScore* around 0.0134 versus 0.013) and sparsity (Hoyer around 0.21 versus 0.22) across the entire set, PRM+ and PRM- groups. t-SNE shows correct and incorrect thoughts completely intermixed; the two visible clusters arise from the two vectors per reasoning step (c=2), not from correctness.
• Trajectory dynamics show no substantial correctness signal. Most metrics differ minimally. Significant differences appear only for compactness on the entire set (p=0.023, Cohen's d=0.17) and in the PRM+ subset (p=0.022, d=1.14), plus local smoothness in the PRM+ subset (p=0.049, d=−0.97).
• COCONUT is robust to noise, which raises questions about latent thoughts. Pass@5 is 39.20 at noise ratio 0.0 and 38.67 at 0.2, with 76.35% of majority answers unchanged at 0.2. Even at complete noise corruption (ratio=1.0), Pass@5 remains non-zero at 12.59%, suggesting the model can sometimes answer correctly without relying on continuous reasoning.
• Evaluation dataset characteristics. The diagnostic test set contains 3,014 samples with 18.48% correct answers and 28.21% correct reasoning steps, deliberately left unbalanced to reflect the natural distribution.

Methodology in Plain English

The authors start from COCONUT, which reasons by feeding its own hidden states back as inputs for a fixed number of steps instead of writing out text. Because that process is deterministic, all samples would be identical. To fix this, they turn dropout back on only during the continuous reasoning phase, so each sampled trajectory follows a slightly different path, while answer generation stays unchanged.

They reproduce COCONUT on GSM8k using GPT-2, training for 6 epochs in the initial stage and 3 epochs in each of the remaining 3 stages, and running inference with T = 3 × c = 6 continuous thought steps (c=2). They then evaluate with sample sizes of 1, 2, 4, 8, 16, and 32, comparing against text CoT and deterministic COCONUT, and reporting Pass@N alongside the number of unique answers.

To build rerankers, they copy MATH-Shepherd's automatic annotation idea but apply it to latent vectors. They generate M=5 reasoning paths per GSM8k training problem (1.32 samples per problem after deduplicating by final answer) and N=10 Monte Carlo completions per step, producing hard (binary: did any completion reach the right answer?) and soft (fraction of correct completions) step labels, plus an outcome label for the final answer. Positive and negative samples are balanced at a 1:1 ratio. PRM and ORM share the COCONUT backbone with small prediction heads (two linear layers with ReLU and a sigmoid), trained for 10 epochs with peak learning rate 1e-4, 500 warmup steps, batch size 128, on a single A100 GPU. Evaluation uses Best-of-N with N from 1 to 32 on the GSM8k test set with no deduplication, compared against confidence reranking and self-consistency.

Finally, they probe why reranking underperforms: they inspect reward-model classification metrics, measure the geometry of thought vectors, measure properties of the reasoning trajectory, and inject Gaussian noise at varying ratios (ratio × noise + (1 − ratio) × thought) to see how much the continuous thoughts actually matter.

Why This Matters

Impact on research. The paper maps out a specific failure mode in continuous reasoning models: even when a high accuracy ceiling exists (Pass@N of 44.43 at N=32), the representations lack the structure that reward models need to reach it. It argues that optimizing for answer accuracy alone is insufficient, and that training objectives must explicitly build in inductive biases that support discrimination at inference time. It also establishes that reward-model training data is not transferable across architectures in continuous space, unlike in text space.

Real-world applications:

• Efficient reasoning assistants that spend less computation per reasoning step while still allowing quality to be scaled up at inference time.
• Latent reasoning models deployed under latency or cost constraints, where generating many samples and reranking them is cheaper than generating long text chains.
• Reward-model and verifier pipelines for any system operating in non-text representation spaces, such as speech or multimodal latent reasoning.
• Robustness testing of models that reason in latent space, since the paper shows answers can remain correct even when latent thoughts are fully corrupted.

Industry relevance. Inference-time scaling is the practical lever most deployments use to trade compute for accuracy. This paper gives a concrete warning: the reranking half of that recipe does not automatically port to latent reasoning, so organizations building or adopting continuous reasoning models need to budget for training-time changes, not just more sampling. The finding that 12.59% Pass@5 survives complete noise corruption is also directly relevant to anyone relying on latent thoughts for interpretability or monitoring.

Future Directions

• Adding geometric inductive biases during training, such as encouraging higher isotropy in thought representations, so latent thoughts occupy the space more evenly and become more separable.
• Promoting trajectory diversity by explicitly varying geometric patterns across reasoning paths, rather than relying on dropout alone.
• Contrastive learning objectives that teach the model to distinguish correct from incorrect reasoning patterns, rather than supervising only final text tokens.
• Reassessing what continuous thoughts contribute, given that the model still solves some problems at 12.59% Pass@5 under full noise corruption; the paper leaves open whether latent vectors carry genuine semantic content or partly act as positional placeholders, and whether scaling to larger base models changes these findings (the study uses GPT-2, and no larger-scale results are reported).

Target Audience

Researchers working on inference-time scaling, reward modeling, and verifier design; practitioners building or evaluating latent/continuous reasoning architectures; and anyone interested in the trade-off between interpretability and computational efficiency in LLM reasoning. The paper is also useful for readers studying representational geometry and diagnostics of hidden states, since much of its value lies in the analytical methodology rather than in a headline accuracy improvement.

Authors’ abstract

Inference-time scaling through multiple sample generation in combination with Process- or Outcome-Reward Model (PRM or ORM) re-ranking has proven effective for text-based reasoning in large language models. This paper investigates whether such established techniques can be successfully adapted to reasoning in the continuous space, using COCONUT (Hao et al. 2024) continuous space reasoning LM as the backbone. We demonstrate the feasibility of generating diverse reasoning paths through dropout-based sampling. Our Pass@N analysis on the generated samples reveals the potential that could enable a significant gain in performance akin to observed gain in the discrete space. However, we highlight unique challenges faced for materializing this gain in the continuous thought space. In particular, working recipes for data generation and training PRM and ORM models in the discrete space unlocks only marginal improvements in the continuous space. Through probing various aspects including geometric properties and trajectory dynamics we identify the underlying reasons that prevent effective discrimination between correct and incorrect reasoning (essential for the functioning of PRM and ORM). Our findings reveal that current limitations stem from the absence of key inductive biases in continuous thought representations. We argue that the training frameworks for continuous reasoning LMs require not only to optimize for accuracy but also to explicitly incorporate inductive biases that could be utilized during inference-time for discrimination of correct and incorrect thoughts.\footnote{Our code and data will be publicly available.}