Continuous latent reasoning enables language models to perform intermediate computation without generating explicit reasoning traces. Advances in stochastic latent reasoning allow multiple paths to be sampled in continuous space, but stochasticity alone does not ensure effective exploration. A key challenge is guiding stochastic transitions between intermediate latent states toward useful reasoning paths. We introduce SubLAT, which couples a relative answer potential with Subtrajectory Balance to learn stochastic latent exploration. The relative potential evaluates the answer support of sampled states against replacement states, with a correctness signal at the terminal state. Subtrajectory Balance relates potential differences to forward and backward transition densities across subtrajectories of different lengths. The model is jointly trained with answer supervision, CoT supervision, and representation distillation. Experiments on arithmetic and deductive reasoning show increasing solution coverage with larger sampling budgets across in-domain and out-of-domain benchmarks, as well as improvements on several structural OOD benchmarks. Trajectory analyses reveal progressive divergence among sampled latent paths and higher mean relative potentials for correct trajectories across arithmetic benchmarks.
Figures & tables
Figure 1: Illustration of SubLAT for learning stochastic latent exploration through relative answer potentials and Subtrajectory Balance. The explicit CoT teacher provides a reference representation at the reasoning boundary, while the student samples latent trajectories through a stochastic forward policy PFθ . The relative answer potential Φt provides intermediate credit for sampled latent states. Subtrajectory Balance couples potential differences with forward and backward transition flows for multi-scale credit assignment along latent trajectories.
In-domain
Out-of-domain
Method
GSM8K test
GSM8K-Hard
SVAMP
MultiArith
ASDiv
OOD Avg.
Acc.
#L
Acc.
Acc.
Acc.
Acc.
Acc.
#L
GPT-2 Small
CoT-SFT
41.70±0.82
27.33
8.46±0.14
39.18±0.60
84.33±0.99
49.60±0.44
45.39
18.84
COCONUT
30.14±0.43
6.00
6.14±0.26
30.48±0.65
72.22±1.04
38.71±0.42
36.89
6.00
CODI
29.53±0.81
6.00
6.53±0.09
31.09±0.21
78.13±0.49
40.16±0.37
38.98
6.00
Table 1: Arithmetic reasoning results on GSM8K and four OOD benchmarks. Accuracy (Acc. %) reports the mean and sample standard deviation over five inference runs. #L denotes reasoning length, and † indicates publicly released checkpoints.
Figure 2: Solution coverage (%) versus sampling budget N on GSM8K (top) and ProsQA (bottom), each with four OOD benchmarks. Orange curves denote SubLAT , and blue curves denote CoLaR. Horizontal lines indicate Cov@1 for CoT-SFT, CODI, and COCONUT.
Figure 3: Latent trajectory diversity of SubLAT across reasoning steps on GSM8K (left) and ProsQA (right), including their respective OOD benchmarks. Error bars indicate 95% confidence intervals.
Figure 4: Within-question relative answer potential gap G(t) on GSM8K and four arithmetic OOD benchmarks. Gaps are averaged over questions with both outcomes among 32 sampled trajectories per question. Shaded regions indicate 95% confidence intervals.
Figure 5: Coverage versus latent diversity on GSM8K with GPT-2 at sampling budgets N . MC-dropout (red) and AGN (blue) perturb COCONUT (left) and CODI (right), while policy scaling (green) varies the learned transition standard deviation of SubLAT .
Cov@ N (%)
Method
Greedy (%)
1
8
32
SubLAT
29.71
28.47
40.56
53.15
w/o relative potential
28.32
27.03
39.20
49.76
w/o SubTB
24.94
25.32
36.38
44.03
Table 3: Ablations on GSM8K with GPT-2. Greedy uses the mean latent path, Cov@ N samples N latent trajectories.
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
Dataset
Training
Validation
In-domain test
Transfer benchmarks
GSM8K-Aug
385,620
500
GSM8K (1,319)
GSM-Hard, SVAMP, MultiArith, ASDiv
ProsQA
17,886
300
ProsQA (500)
7-hop, 8-hop, DAG, Dense
Appendix
Table 4: Training and evaluation benchmarks.
GSM8K
ProsQA
Hyperparameter
GPT-2
GPT-2
LLaMA
DeepSeek
Latent Tokens Per Stage
2
1
1
1
Stage 0 Epochs
0
5
5
5
Epochs Per Stage
3
5
5
5
Max Latent Stage
3
6
6
6
Total Epochs
37
50
20
20
Appendix
Table 5: Training hyperparameters for Coconut models.
GSM8K
ProsQA
Hyperparameter
GPT-2
LLaMA
LLaMA
DeepSeek
Latent Loss Weight ( α )
1
1
1
1
CoT Loss Weight ( β )
1
1
1
1
Distillation Loss Weight ( δ )
1
20
1
1
Num Latent Tokens
6
6
6
6
Total Epochs
40
10
20
20
Appendix
Table 6: Training hyperparameters for CODI models.
GSM8K-Aug
ProsQA
Hyperparameter
GPT-2
LLaMA
GPT-2
LLaMA
DeepSeek
Answer Loss Weight
1
1
1
1
1
CoT Loss Weight ( λCoT )
1
1
1
1
1
Distillation Loss Weight ( λKD )
1
1
1
1
1
SubTB Loss Weight ( λflow )
1
1
1
1
1
Drift Loss Weight ( λdrift )
0.05
0.05
0.05
0.05
0.05
Appendix
Table 7: Training hyperparameters for SubLAT .
Configuration
Value
Configuration
Value
Latent horizon
6
Effective batch size
128
Trajectories per question
1
Precision
BF16
Optimizer
AdamW
SubTB length weight
0.9
Training seeds
12 (GSM8K-Aug) / 11, 12, 13 (ProsQA)
Appendix
Table 8: SubLAT training configuration.
Method
Hardware
GPU-h/Epoch
Epochs
CoT-SFT
1 × GH200
0.35
40
CODI
1 × A100
0.78
40
COCONUT
2 × H100
3.71
37
SubLAT
2 × H100
6.91
40
Appendix
Table 9: Training cost on GSM8K-Aug. GPU-hours per epoch are computed from the number of allocated GPUs and elapsed training time.
Figure 7: Validation accuracy over training on GSM8K-Aug and ProsQA.
Token rank
1
2
3
4
5
6
Answer
COCONUT
1
###
⊔ 30
720
⊔ 30
⊔ 280
⊔ 30
840
2
<|start-latent|>
30
⊔ packs
⊔ 50
###
30
3
⊔ with
⊔ 3000
⊔ get
⊔ 60
⊔ packs
⊔ 50
CODI
1
+
⊔ 30
s
⊔ 30
s
⊔ 30
840
Appendix
Table 10: Vocabulary projections of latent states on a GSM8K example. Rows 1–3 show the top three tokens at each latent step using each model’s vocabulary head, with ⊔ indicating a leading space. Teacher shows CODI’s greedy CoT output, and Golden gives the reference solution.
Large Language Models (LLMs) increasingly rely on intermediate reasoning, yet explicit Chain-of-Thought (CoT) suffers from a linguistic space bottleneck: each thought must be decoded into tokens, causing high inference overhead. Latent reasoning moves deliberation into continuous space, but existing methods mostly learn deterministic or reward-maximizing paths, lacking a principled way to allocate probability across trajectories with different correctness and costs. We propose Latent Thought Flow (LTF), which models reasoning as variable-length continuous trajectories and trains a sampler to match a reward-induced posterior over answer quality and computation cost. We instantiate this with a continuous GFlowNet using stochastic latent transitions. To handle sparse answer supervision, we introduce an Entropy-Weighted Subtrajectory Balance objective for intermediate rewards and a reference-prior regularizer to anchor exploration. Experiments under finetuning and transfer learning settings show that LTF outperforms explicit CoT and latent reasoning baselines, improving accuracy by 9.5% while reducing reasoning length by 27.2% on average compared with strong latent reasoning baselines.
Recently, latent reasoning has been introduced into large language models (LLMs) to leverage rich information within a continuous space. However, without stochastic sampling, these methods inevitably collapse to deterministic inference, failing to discover diverse reasoning paths. To bridge the gap, we inject controllable stochasticity into latent reasoning via Gumbel-Softmax, restoring LLMs' exploratory capacity and enhancing their compatibility with Reinforcement Learning (RL). Building on this, we propose Latent Reasoning Policy Optimization~(LEPO), a novel framework that applies RL directly to continuous latent representations. Specifically, in rollout stage, LEPO maintains stochasticity to enable diverse trajectory sampling, while in optimization stage, LEPO constructs a unified gradient estimation for both latent representations and discrete tokens. Extensive experiments show that LEPO significantly outperforms existing RL methods for discrete and latent reasoning.
Yuyan Zhou, Jiarui Yu, Hande Dong +4
Tencent · Zhejiang University · Shanghai Jiao Tong University
Language models relying on verbalized Chain-of-Thought (CoT) are constrained to token-by-token generation in the language space, where many tokens serve textual coherence rather than core reasoning. This entangles genuine differences in logical trajectories with superficial variations in verbalization. Moving reasoning into a compact, continuous latent space offers a promising alternative for efficiently representing and exploring diverse reasoning trajectories. However, current latent models remain largely deterministic, capturing only a single reasoning trajectory rather than modeling a distribution over multiple valid alternatives. While stochastic approaches parameterize distributions over latent reasoning steps, they frequently suffer from unstable optimization and limited empirical performance. Furthermore, their reliance on autoregressive latent prediction during training precludes effective parallelization, leading to prohibitively long training times. To overcome these bottlenecks, we propose Probabilistic Latent Reasoning Models (PLRM), a parallelizable variational framework for probabilistic latent reasoning. PLRM adopts a JEPA-style objective in which the posterior infers latent distributions in parallel from observed reasoning trajectories, while the prior learns to predict the subsequent latent distribution via teacher forcing. To model heterogeneous reasoning processes, we introduce a discrete strategy variable together with a mixture-based prior, allowing the model to preserve the Gaussian-mixture structure induced by multiple distinct reasoning trajectories rather than collapsing them into a single Gaussian. Across diverse benchmarks, PLRM achieves performance competitive with deterministic latent reasoning models on single-trajectory datasets while enabling efficient training and consistently outperforming competing stochastic baselines on multi-trajectory datasets.