Parallel Test-Time Scaling for Latent Reasoning Models

AI summary

Overview

• Research area: Natural Language Processing / large language model reasoning — specifically test-time scaling (TTS) and latent (continuous-space) reasoning.
• Technical level: Advanced. The paper assumes familiarity with chain-of-thought prompting, parallel decoding strategies (majority voting, best-of-N, beam search), uncertainty estimation, and reward-model training.
• One-sentence scope: The paper shows that parallel test-time scaling, previously limited to token-based chain-of-thought, can be transferred to latent reasoning models by adding two stochastic sampling mechanisms (Monte Carlo Dropout and Additive Gaussian Noise) plus a step-wise contrastively trained Latent Reward Model (LatentRM) for trajectory aggregation.

What This Paper Is About

Latent reasoning models replace verbalized chain-of-thought steps with continuous hidden-state vectors, which is more compact and efficient than writing out tokens. However, parallel test-time scaling — sampling many reasoning paths in parallel and picking or voting among them — has been impossible for these models, because there is no token probability distribution to sample from and no likelihood signal to score or rank trajectories. This paper builds both missing pieces: two ways to inject controlled randomness into latent generation, and a learned scorer that evaluates latent thoughts so that best-of-N and beam search become usable in continuous space.

Key Contributions

1. Introduces parallel test-time scaling to latent reasoning models, a scaling capability the authors describe as previously exclusive to token-based reasoning paradigms.
2. Proposes two complementary stochastic sampling strategies for continuous latent space: Monte Carlo Dropout (MC-dropout), which keeps dropout active at inference to sample different weight configurations and capture epistemic uncertainty, and Additive Gaussian Noise (AGN), which perturbs latent thoughts with isotropic Gaussian noise to simulate aleatoric uncertainty.
3. Designs the Latent Reward Model (LatentRM), a dedicated scoring head attached to the latent reasoning backbone that maps a prompt plus latent trajectory prefix to a scalar score, trained with a step-wise contrastive objective that compares candidates at each reasoning step via softmax rather than isolated binary labels.
4. Provides extensive empirical and visualization analysis of how the two sampling methods scale with compute, how they differ in exploration geometry, and how LatentRM enables best-of-N and beam search aggregation across three benchmarks and multiple backbones.

Main Findings

• Both sampling strategies scale with compute: Coverage increases monotonically as the number of sampled trajectories N grows, for both MC-dropout and AGN.
• Returns diminish with N: Marginal coverage improvement shrinks as N grows, indicating a saturation effect where extra samples contribute less.
• Sampling narrows model gaps: At N = 64, COCONUT and CODI achieve nearly equivalent coverage, even though CODI is clearly superior at N = 1.
• MC-dropout generally wins on coverage: It achieves higher coverage across nearly all values of N, making it the more reliable stochastic sampling approach in the main sampling experiments.
• Diversity has a sweet spot: Across models and methods, coverage peaks at moderate diversity; too much or too little stochasticity hurts performance.
• AGN is more robust at high diversity: At larger diversity levels AGN tends to maintain or even improve coverage, while MC-dropout shows a sharp decline.
• Different step-wise dynamics: MC-dropout maintains similar diversity across reasoning steps for a given dropout rate p (stochasticity follows the model's own uncertainty), whereas AGN fluctuates markedly because a fixed σ perturbs latent vectors of varying scales unevenly.
• Distinct latent geometries: In t-SNE visualizations, dropout produces a dense, contiguous "directional drift" along specific directions, while AGN produces an isotropic radial "firework" dispersion with broader area but lower local density. On easy questions, AGN keeps probability mass near the deterministic center and preserves accuracy, whereas large p makes dropout drift away and degrade; on hard questions, dropout's larger displacement and denser exploration increase the chance of reaching correct solutions.
• Generalization across backbones: The framework improves Latent-SFT (Llama-3.2-1B) on GSM8K from 0.445 deterministic accuracy to 0.585 coverage@8 and 0.649 coverage@16, and on MultiArith from 0.934 to 0.962 and 0.967. RoT-4B (Qwen3-VL) on MATH500 improves from 20.3% to 22.0% at N = 16, and RoT-2B improves on MATH500 from 0.115 to 0.128 (Cov@8) and 0.130 (Cov@16).
• Aggregation scales: Accuracy increases monotonically with N across GSM-Test, GSM-Hard and MultiArith, and both Best-of-N and Beam Search consistently outperform Majority Voting, confirming LatentRM can distinguish promising trajectories.
• Best-of-N is the strongest route: Beam Search performs comparably to Best-of-N on GSM-Test and MultiArith but trails on GSM-Hard, which the authors attribute to early-step score noise causing premature pruning on harder problems. Gains are most pronounced on MultiArith.
• Ablations confirm design choices matter: In the reported table (Best-of-8 rows with LatentRM; the caption labels the setting Best-of-N with N = 32 while the text describes the ablation as N = 8), Best-of-8 with LatentRM scores 35.4 on GSM-Test and 7.8 on GSM-Hard; removing the step-wise contrastive loss (using BCE) drops this to 33.5 and 7.4; removing stochastic rollouts drops it to 30.7 and 6.0; and an untrained random scalar head gives only 28.9 and 5.8, below Majority Voting (33.6 and 6.1).
• Harder benchmarks show limited absolute performance: On GPQA and AIME the method produces consistent gains, but absolute performance remains limited; the paper attributes this to two factors, though the explanation is cut off in the provided content.

Methodology in Plain English

The authors start from the structure of latent reasoning: instead of writing words, the model repeatedly emits the last hidden state of its transformer as the "thought" for the next step, until an end-of-thinking token is produced and the final answer is decoded normally. To create multiple different trajectories, they inject randomness in two ways.

For Monte Carlo Dropout, dropout stays switched on during inference, so each pass through the network uses a different randomly sampled weight mask. This is justified with an approximate Bayesian argument: stochastic forward passes approximate sampling from a posterior over weights, and the resulting variability reflects epistemic uncertainty, i.e., uncertainty caused by the model's limited knowledge.

For Additive Gaussian Noise, they simply add zero-mean isotropic Gaussian noise with standard deviation σ to the latent thought at each step before continuing generation. Because the noise scale is fixed externally and independent of the model parameters, this is framed as simulating aleatoric uncertainty, with variance proportional to the network's local sensitivity (Jacobian) at that point.

To rank the resulting trajectories, they train LatentRM, an extra scoring head on the latent reasoning backbone that outputs a scalar for the prompt plus the trajectory so far. Training labels come from Monte Carlo estimation: for each thought in each of N sampled trajectories, they roll out M stochastic completions, then label the thought with the fraction of those completions that match the ground-truth answer. Rather than the usual binary cross-entropy on each thought in isolation, they compare all N candidates at the same step through a softmax and train with a negative log-likelihood loss, so the model learns relative quality among contemporaneous thoughts.

At inference, trajectories are ranked by the sum of per-step scores. The authors derive that the normalizing term in the cumulative log-probability is identical across trajectories, so cumulative logits alone determine ranking. This supports both Best-of-N selection and a beam search where the beam size B = √N keeps compute comparable to Best-of-N, and N full samples are used for Majority Voting as the non-parametric baseline. Cumulative rewards are pre-normalized by trajectory length to avoid favoring longer chains.

Experiments use GSM8K-Test (1,319 test samples), GSM8K-Hard (1,319 test samples) and MultiArith (600 test samples), with five latent reasoning models: COCONUT, CODI, CoLaR, Latent-SFT (Llama-3.2-1B-Instruct), and Render-of-Thought (RoT) with Qwen3-VL at 2B and 4B scales. COCONUT and CODI are backboned on GPT-2, CoLaR on Llama-3.2-1B. COCONUT and CODI fix T = 6 latent thoughts; CoLaR uses thinking speed 2; CoLaR, Latent-SFT and RoT use a maximum of 64 latent thoughts. Hyperparameters were tuned by binary search to maximize coverage@64, sweeping σ over [0.01, 1.5] and p over [0, 1], and the authors also give heuristic starting ranges: p ∈ [0.1, 0.3] and σ ∈ [0.5, 0.7] for GPT-2-style models, and p ∈ [0.01, 0.03] and σ ∈ [0.01, 0.03] for Llama-3.2-1B-style models.

Why This Matters

Impact on research. The paper removes a structural barrier that kept latent reasoning models out of the parallel test-time scaling literature. It supplies both a sampling primitive for continuous representations and a scoring signal for trajectory comparison, and it characterizes the two sampling methods as having complementary geometric footprints in latent space rather than treating randomness as a single knob. It also connects latent reasoning to uncertainty estimation theory, framing objective sampling design rather than arbitrary noise injection. The authors contrast their setting with Soft Thinking work in token-probability space (including Stochastic Soft Thinking, which uses Gumbel-Softmax), arguing that pure latent vectors are unconstrained by vocabulary structure and reveal different reasoning patterns.

Potential real-world applications (the paper does not enumerate application domains; these follow from the setting it targets):

• Deploying compact reasoning models on edge or on-device hardware, where latent reasoning's compressed trajectories reduce token-generation cost and parallel scaling can be traded against an inference budget.
• Math and quantitative tutoring or homework-checking tools that need reliable step-by-step verification without verbose natural-language rationales.
• Agentic or multi-step planning pipelines where many candidate plans are generated and one must be selected cheaply.
• Scientific or engineering calculation assistants that benefit from parallel exploration of solution paths under a fixed compute ceiling.

Industry relevance. Parallel test-time scaling is one of the main levers for improving model quality without retraining, so extending it to latent reasoning gives a second axis of control for serving systems. The paper's compute-matching conventions (Best-of-N versus Majority Voting with N samples versus Beam Search with beam size √N) and its heuristic hyperparameter ranges per backbone family are directly aimed at practitioners deciding how to allocate inference compute.

Future Directions

1. Integrate sampling and aggregation into a reinforcement learning framework, optimizing latent trajectories through iterative feedback and reward shaping. This is the first direction the paper names explicitly.
2. A second future direction is listed but the text is truncated in the provided content, so its content cannot be reported.
3. Scale to harder benchmarks. The authors report consistent but limited absolute gains on GPQA and AIME and attribute this to at least two factors, the explanation for which is cut off; closing that gap is an open problem.
4. Investigate score noise in early steps. Beam Search trails Best-of-N on GSM-Hard because early-step score noise causes premature pruning, which points to improving early-step scoring or pruning criteria as a concrete next step.
5. Evaluate under variable-length latent reasoning paths. The paper flags a "Revised Performance with variable-length latent reasoning paths on GSM8K-Test" and defers extended analysis of LatentRM under variable-length settings (Appendix F.3), along with latent-versus-explicit comparisons with TTS (Appendix F.1) and wall-clock comparisons (Appendix F.2), to appendices.

Target Audience

Researchers and engineers working on LLM inference, test-time compute scaling, and reasoning efficiency, particularly those interested in chain-of-thought alternatives such as COCONUT, CODI, CoLaR, Latent-SFT and Render-of-Thought. It is also relevant to practitioners building reward models or process reward models, since the step-wise contrastive formulation and the Monte Carlo thought-labeling scheme transfer to other settings, and to readers tracking uncertainty estimation applied to LLM decoding. Readers without background in parallel decoding, reward modeling, or latent/continuous reasoning will find the paper demanding.

Authors’ abstract

Parallel test-time scaling (TTS) is a pivotal approach for enhancing large language models (LLMs), typically by sampling multiple token-based chains-of-thought in parallel and aggregating outcomes through voting or search. Recent advances in latent reasoning, where intermediate reasoning unfolds in continuous vector spaces, offer a more efficient alternative to explicit Chain-of-Thought, yet whether such latent models can similarly benefit from parallel TTS remains open, mainly due to the absence of sampling mechanisms in continuous space, and the lack of probabilistic signals for advanced trajectory aggregation. This work enables parallel TTS for latent reasoning models by addressing the above issues. For sampling, we introduce two uncertainty-inspired stochastic strategies: Monte Carlo Dropout and Additive Gaussian Noise. For aggregation, we design a Latent Reward Model (LatentRM) trained with step-wise contrastive objective to score and guide latent reasoning. Extensive experiments and visualization analyses show that both sampling strategies scale effectively with compute and exhibit distinct exploration dynamics, while LatentRM enables effective trajectory selection. Together, our explorations open a new direction for scalable inference in continuous spaces. Code and checkpoints released at https://github.com/ModalityDance/LatentTTS