Parallel Test-Time Scaling for Latent Reasoning Models

1. Problem Statement

Large Language Models (LLMs) have demonstrated significant performance gains through Test-Time Scaling (TTS), particularly via parallel scaling (e.g., sampling multiple reasoning paths and aggregating them). However, this paradigm is currently limited to token-based reasoning (explicit Chain-of-Thought), where models generate discrete tokens with inherent probability distributions that allow for standard sampling (e.g., top-k, nucleus).

Latent Reasoning Models (or Continuous CoT) represent a promising alternative where intermediate reasoning steps are represented as continuous vectors (hidden states) rather than discrete tokens. While more efficient and compact, these models face two critical barriers to parallel TTS:

1. Lack of Sampling Mechanism: Continuous latent spaces lack an explicit probability distribution over a vocabulary, making it impossible to use standard token-based sampling strategies to generate diverse reasoning trajectories.
2. Lack of Aggregation Signals: Without token-level probabilities or log-likelihoods, there is no inherent metric to rank or aggregate multiple sampled latent trajectories (e.g., via Best-of-N or Beam Search).

2. Methodology

The authors propose a framework to enable parallel TTS for latent reasoning by addressing sampling and aggregation separately.

A. Stochastic Sampling in Latent Space

To generate diverse latent trajectories $\{h^{(n)}\}_{n=1}^N$, the authors introduce two uncertainty-inspired stochastic strategies to inject controlled noise into the generation process:

1. Monte Carlo Dropout (MC-Dropout):

   • Mechanism: Keeps dropout active during inference, sampling binary masks to create different weight configurations for each forward pass.
   • Theoretical Basis: Captures epistemic uncertainty (uncertainty due to limited model knowledge).
   • Effect: Produces structured, directed exploration, often drifting toward unconventional solutions.

2. Additive Gaussian Noise (AGN):

   • Mechanism: Adds isotropic Gaussian noise ($\epsilon \sim \mathcal{N}(0, \sigma^2 I)$) directly to the latent vectors at each reasoning step.
   • Theoretical Basis: Simulates aleatoric uncertainty (noise inherent in the input/data).
   • Effect: Drives broad, isotropic exploration around the deterministic center, enriching diversity without necessarily shifting the trajectory's direction as drastically as dropout.

B. Latent Trajectory Aggregation

Since latent vectors lack natural scores, the authors propose the Latent Reward Model (LatentRM):

• Architecture: A dedicated scorer (a small neural head) appended to the latent reasoning backbone that maps the input prompt and the current latent trajectory prefix to a scalar score $r_t$.
• Training Objective: Trained using a step-wise contrastive objective. Instead of binary labels, the model learns to rank $N$ candidate thoughts at each step $t$ relative to each other using a softmax over their scores.
• Aggregation Strategies:
  • Best-of-N: Selects the trajectory with the highest cumulative sum of latent rewards ($\sum r_t$).
  • Beam Search: Expands the top-$B$ partial trajectories at each step based on cumulative rewards.

3. Key Contributions

1. First Parallel TTS for Latent Reasoning: Successfully extends the parallel scaling paradigm from token-based to continuous latent spaces, unlocking a key capability previously exclusive to explicit CoT.
2. Novel Sampling Strategies: Proposes MC-Dropout and AGN as principled methods to induce stochasticity in continuous spaces, grounded in uncertainty estimation theory.
3. Latent Reward Model (LatentRM): Designs a specialized reward model trained with step-wise contrastive supervision to evaluate and guide latent reasoning, solving the aggregation problem in the absence of token probabilities.
4. Comprehensive Analysis: Provides deep insights into the exploration dynamics of different sampling methods (e.g., MC-Dropout excels on hard tasks; AGN maintains robustness on easy tasks).

4. Experimental Results

The framework was evaluated on benchmarks including GSM8K, GSM8K-Hard, MultiArith, MATH500, GPQA, and AIME using various backbones (GPT-2, Llama-3.2, Qwen3-VL).

• Scaling Effectiveness: Both MC-Dropout and AGN show monotonic improvements in solution coverage as the number of sampled trajectories ($N$) increases.
• Sampling Dynamics:
  • MC-Dropout consistently achieves higher coverage across most $N$, particularly on harder problems, due to its ability to explore structured, directed paths.
  • AGN provides superior diversity and maintains accuracy even under high noise, making it robust for easy tasks where the correct solution is near the deterministic center.
• Aggregation Performance:
  • LatentRM significantly outperforms non-parametric baselines like Majority Voting.
  • Best-of-N aggregation guided by LatentRM yields the most consistent gains, improving accuracy by up to 2-3% on complex math tasks (e.g., RoT-4B on MATH500).
• Generalization: The approach scales effectively to larger models (up to 4B parameters) and diverse architectures.
• Efficiency: Latent reasoning with TTS achieves competitive accuracy to explicit CoT but with significantly lower computational costs (fewer decoding steps) and lower wall-clock latency.

5. Significance and Future Work

This work bridges a critical gap between efficient latent reasoning and scalable inference. By enabling parallel test-time scaling in continuous spaces, it demonstrates that latent models can leverage additional compute to improve reasoning without retraining or parameter updates.

Significance:

• It validates that "thinking in vectors" can be as scalable as "thinking in tokens."
• It provides a practical toolkit (sampling + LatentRM) for deploying latent reasoning models in high-stakes scenarios requiring reliability.
• It offers geometric insights into how stochasticity reshapes the latent reasoning landscape.

Future Directions:

• Integrating sampling and aggregation into a Reinforcement Learning (RL) framework for iterative optimization.
• Adapting ensemble techniques specifically for continuous latent representations.

In conclusion, the paper establishes a new pathway for scalable inference in latent space, proving that with redesigned sampling and aggregation mechanisms, latent reasoning models can fully harness the power of test-time compute.