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

TL;DR

This paper introduces GTS, a learnable Gaussian sampler that improves inference-time scaling in latent reasoning models by replacing heuristic perturbations with explicit, optimized probabilistic sampling, leading to more reliable reasoning.

Contribution

The paper proposes GTS, a novel module that learns conditional distributions for latent exploration, replacing heuristic methods with an explicit, trainable sampling policy.

Findings

GTS outperforms heuristic perturbations across multiple benchmarks.

Explicit probabilistic sampling improves the reliability of inference-time scaling.

GTS demonstrates better control and optimization of latent exploration.

Abstract

Inference-time scaling (ITS) in latent reasoning models typically relies on heuristic perturbations, such as dropout or fixed Gaussian noise, to generate diverse candidate trajectories. However, we show that stronger perturbations do not necessarily yield better sampling quality: they often induce larger distribution shifts without producing more useful reasoning paths or better final decisions. A key limitation is that these perturbations inject stochasticity without defining an explicit conditional sampling distribution, making latent exploration difficult to control or optimize. To address this, we propose the Gaussian Thought Sampler (GTS), a lightweight module that reformulates latent exploration as sampling from a learned conditional distribution over continuous reasoning states. GTS predicts context-dependent perturbation distributions and is trained with GRPO-style policy optimization while keeping the backbone frozen, turning heuristic perturbation into an explicit probabilistic sampling policy. Experiments across multiple benchmarks and two latent reasoning architectures show that GTS yields more reliable inference-time scaling than heuristic baselines, suggesting that effective latent ITS requires better-controlled and optimizable sampling rather than simply amplifying stochasticity.