GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

1. Problem Statement

Inference-Time Scaling (ITS) is a technique to improve Large Language Model (LLM) reasoning by generating multiple candidate reasoning trajectories at test time and selecting the best one (e.g., via Best-of-N or Self-Consistency).

• Discrete Models: In standard token-based LLMs, ITS relies on explicit conditional probability distributions over the vocabulary, allowing for principled sampling strategies (e.g., temperature scaling, nucleus sampling).
• Latent Reasoning Models: Recent models perform reasoning directly in continuous hidden state space without generating intermediate tokens. These models typically lack explicit probability distributions over latent states.
• The Limitation: Current ITS methods for latent reasoning rely on heuristic perturbations (e.g., Dropout or fixed Gaussian noise) to induce diversity. The authors identify two critical flaws:
  1. Lack of Control: Heuristic perturbations do not define an explicit conditional sampling distribution. They inject stochasticity without a principled policy, making exploration difficult to optimize.
  2. Inefficiency: Stronger perturbations do not guarantee better results. They often cause large distribution shifts that move the model away from correct reasoning paths (over-exploration) or fail to move them enough (under-exploration), leading to a mismatch between diversity and decision quality.

2. Methodology: Gaussian Thought Sampler (GTS)

The authors propose GTS, a lightweight module that reformulates latent exploration as sampling from a learned conditional distribution over continuous reasoning states.

Core Architecture

• Frozen Backbone: The underlying latent reasoning model (e.g., COCONUT, CODI) remains frozen.
• Learnable Sampler: GTS is a lightweight MLP that predicts a context-dependent Gaussian distribution for perturbations at each reasoning step.
  • Input: The deterministic hidden state $h_t^{det}$ at step $t$.
  • Output: Parameters for a diagonal Gaussian distribution: mean $\mu_\phi(h_t^{det})$ and log-standard-deviation $\log \sigma_\phi(h_t^{det})$.
  • Perturbation: The actual latent state is updated as $\tilde{h}_t = h_t^{det} + z_t$, where $z_t \sim \mathcal{N}(\mu, \text{diag}(\sigma^2))$.
• Reparameterization: The sampling process uses the reparameterization trick ($z_t = \mu + \sigma \odot \epsilon$) to allow gradient flow through the sampling process.

Training Objective (Policy Optimization)

Since the final answer correctness is non-differentiable with respect to the sampler, GTS is trained using Reinforcement Learning, specifically a GRPO (Group Relative Policy Optimization) style approach:

1. Rollout: For a given prompt, sample $N$ perturbation trajectories.
2. Reward Design:
   • Correctness: Binary reward based on whether the decoded answer matches the ground truth.
   • Shaping: A confidence-based term derived from the normalized log-probability of the answer. This encourages high-confidence correct answers and penalizes high-confidence wrong answers.
3. Optimization:
   • Compute the density ratio between the current policy and a reference policy (Exponential Moving Average of the current policy).
   • Minimize a clipped surrogate loss (similar to PPO/GRPO) combined with a KL-divergence penalty to prevent the policy from collapsing to a deterministic point or drifting too far from the reference.

3. Key Contributions

1. Diagnostic Insight: The paper provides a rigorous analysis showing that in latent reasoning, distribution shift does not correlate with sampling quality. Heuristic methods often increase divergence (JS Divergence) while decreasing the probability of finding the correct answer (Sampling Gain).
2. Formulation Shift: It reformulates latent exploration from a heuristic perturbation problem to an explicit probabilistic sampling policy problem, bridging the gap between discrete token sampling and continuous latent reasoning.
3. GTS Module: Introduces a lightweight, trainable Gaussian sampler that learns where and how much to perturb based on the context, rather than using fixed noise scales.
4. Empirical Superiority: Demonstrates that learned, context-conditioned sampling consistently outperforms fixed heuristic baselines (Dropout, Standard Gaussian) across multiple architectures and budgets.

4. Experimental Results

The authors evaluated GTS on two latent reasoning architectures (COCONUT and CODI) using the GSM8K dataset and three out-of-distribution benchmarks (MultiArith, SVAMP, GSM8K-Hard).

• Scaling Performance (Pass@N):
  • GTS achieves the strongest scaling behavior as the sampling budget ($N$) increases.
  • While heuristic baselines (especially strong Dropout $p=0.5$) often degrade performance or show diminishing returns, GTS maintains consistent improvement.
  • On COCONUT, GTS significantly outperforms baselines; on CODI, it shows more moderate but consistent gains.
• Sampling Quality Analysis:
  • Sampling Gain (SG): GTS achieves the highest SG, meaning the sampled trajectories significantly increase the model's confidence in the correct answer compared to the deterministic baseline.
  • Distribution Shift: GTS achieves high SG with controlled distribution shift (lower Jensen-Shannon divergence than strong heuristics). This confirms GTS explores useful regions of the latent space rather than just random noise.
• Failure Modes of Baselines:
  • Under-exploration: Weak heuristics (e.g., Dropout $p=0.1$) fail to move the trajectory enough to find better paths.
  • Over-exploration: Strong heuristics (e.g., Dropout $p=0.5$) disrupt the reasoning dynamics, leading to negative SG.
• Step-wise Analysis (SNR):
  • GTS adapts its perturbation strategy dynamically. For COCONUT, it shows a "saw-tooth" pattern of deterministic vs. stochastic influence, aligning with the model's internal curriculum. For CODI, it applies stronger deterministic steering early and more exploration later.

5. Significance and Conclusion

• Paradigm Shift: The paper argues that effective Inference-Time Scaling in continuous spaces requires better-controlled sampling, not just more stochasticity.
• Optimizable Policy: By treating latent exploration as a learnable policy, GTS allows the system to adapt to the specific geometry of the reasoning space and the task requirements, which fixed heuristics cannot do.
• Generalizability: The method works across different backbone architectures (GPT-2 based, LLaMA-3 based) and generalizes to out-of-distribution tasks without retraining the sampler.
• Future Direction: This work suggests that future latent reasoning models should incorporate learnable, context-aware exploration mechanisms rather than relying on static noise injection, potentially unlocking the full potential of continuous reasoning for complex problem-solving.