Latent Path Exploration via Potential-Guided Subtrajectory Balance

Xiaoling Zeng; Qi Zhang

cs.CLOct 29, 2026

Latent Path Exploration via Potential-Guided Subtrajectory Balance

Authors: Xiaoling Zeng, Qi Zhang

Organizations: Worcester Polytechnic Institute

Abstract

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

P_{F}^{\theta}

. The relative answer potential

\Phi_{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.

GPT-2 Small
	In-domain		Out-of-domain
Method	GSM8K test		GSM8K-Hard	SVAMP	MultiArith	ASDiv	OOD Avg.
	Acc.	#L	Acc.	Acc.	Acc.	Acc.	Acc.	#L
CoT-SFT	$41.70_{\pm 0.82}$	$27.33$	$8.46_{\pm 0.14}$	$39.18_{\pm 0.60}$	$84.33_{\pm 0.99}$	$49.60_{\pm 0.44}$	$45.39$	$18.84$
COCONUT	$30.14_{\pm 0.43}$	$6.00$	$6.14_{\pm 0.26}$	$30.48_{\pm 0.65}$	$72.22_{\pm 1.04}$	$38.71_{\pm 0.42}$	$36.89$	$6.00$
CODI	$29.53_{\pm 0.81}$	$6.00$	$6.53_{\pm 0.09}$	$31.09_{\pm 0.21}$	$78.13_{\pm 0.49}$	$40.16_{\pm 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.

Hyperparameter	GPT-2	GPT-2	LLaMA	DeepSeek
	GSM8K	ProsQA
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.

Hyperparameter	GPT-2	LLaMA	LLaMA	DeepSeek
	GSM8K		ProsQA
Latent Loss Weight ( $\alpha$ )	1	1	1	1
CoT Loss Weight ( $\beta$ )	1	1	1	1
Distillation Loss Weight ( $\delta$ )	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.

Hyperparameter	GPT-2	LLaMA	GPT-2	LLaMA	DeepSeek
	GSM8K-Aug		ProsQA
Answer Loss Weight	1	1	1	1	1
CoT Loss Weight ( $\lambda_{\mathrm{CoT}}$ )	1	1	1	1	1
Distillation Loss Weight ( $\lambda_{\mathrm{KD}}$ )	1	1	1	1	1
SubTB Loss Weight ( $\lambda_{\mathrm{flow}}$ )	1	1	1	1	1
Drift Loss Weight ( $\lambda_{\mathrm{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 $\times$ GH200	0.35	40
CODI	1 $\times$ A100	0.78	40
COCONUT	2 $\times$ H100	3.71	37
SubLAT	2 $\times$ 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	###	$\sqcup$ 30	720	$\sqcup$ 30	$\sqcup$ 280	$\sqcup$ 30	840
2	<\|start-latent\|>	30	$\sqcup$ packs	$\sqcup$ 50	###	30
3	$\sqcup$ with	$\sqcup$ 3000	$\sqcup$ get	$\sqcup$ 60	$\sqcup$ packs	$\sqcup$ 50
CODI
1	+	$\sqcup$ 30	s	$\sqcup$ 30	s	$\sqcup$ 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

\sqcup

indicating a leading space. Teacher shows CODI’s greedy CoT output, and Golden gives the reference solution.

LLM Reasoning

Large language model reasoning studies how models solve multi-step problems, verify intermediate steps, and use computation or tools reliably.

See all 931 papers

Latent Thought Flow: Efficient Latent Reasoning in Large Language Models
Jun 15, 2026Xiandong Zou, Jing Huang, Jianshu Li +1RL for Language Model Reasoning Efficient Language Model Reasoning
Singapore Management University · Ant Group
LEPO: Latent Reasoning Policy Optimization for Large Language Models
Apr 20, 2026Yuyan Zhou, Jiarui Yu, Hande Dong +4Reinforcement Learning RL for Language Model Reasoning
Tencent · Zhejiang University · Shanghai Jiao Tong University
Parallelizable Probabilistic Latent Reasoning for Diverse Reasoning Trajectories
Oct 21, 2026Seokhyun Lee, Seongjun Lee, Changhee LeeContinuous Latent Reasoning Reasoning Diversity
Korea University

Jun 15, 2026J/K move · Enter open · S save

Latent Thought Flow: Efficient Latent Reasoning in Large Language Models

Xiandong Zou, Jing Huang, Jianshu Li, Pan Zhou

Singapore Management University · Ant Group

RL for Language Model Reasoning Efficient Language Model Reasoning Generative Flow Networks LLM Inference Acceleration Continuous Latent Reasoning Latent Reasoning

Open paper AskPDF arXiv

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.

Abstract

LLM Reasoning

Explore similar work

Latent Thought Flow: Efficient Latent Reasoning in Large Language Models

LEPO: Latent Reasoning Policy Optimization for Large Language Models

Parallelizable Probabilistic Latent Reasoning for Diverse Reasoning Trajectories