Next Thoughts Are Distributions: Generative Autoregressive Reasoning in the Latent Space

Yang Li; Yi Wang; Shiyuan Huang; Yang Liu; Hao Wang; Chengzhi Mao

cs.AISep 27, 2026

Next Thoughts Are Distributions: Generative Autoregressive Reasoning in the Latent Space

Authors: Yang Li, Yi Wang, Shiyuan Huang, Yang Liu, Hao Wang, Chengzhi Mao

Organizations: Rutgers University · Amazon · University of Illinois at Urbana-Champaign

Abstract

Reasoning problems often admit multiple valid ways to proceed. Continuous reasoning promises to move computation beyond language tokens into a more compact latent space, but representing several plausible ways to think next remains difficult. We introduce Autoregressive Thought Flow (ATF), which models the next continuous thought as a multimodal distribution. A causal autoregressive model performs the reasoning computation, while a lightweight diffusion head generates a plausible next thought from the resulting condition. The sampled thought is fed back into the model, allowing continuous reasoning to unfold for a variable number of steps while preserving the pretrained backbone. Across mathematical reasoning tasks, ATF improves accuracy with compact latent traces and benefits from reinforcement learning and additional test-time thinking. Multi-sample evaluation shows broader solution coverage, indicating that its multimodal predictions capture useful diversity among reasoning paths. Our results suggest that continuous reasoning is more effective when multiple possible next thoughts remain available rather than being collapsed into a single prediction.

Figures & tables

Figure 1: Our ATF enables multiple latent continuations. Given the same math prompt, a point-estimate continuous-thought model collapses the next-thought distribution into a single averaged trajectory, producing similar decoded continuations and an incorrect final answer. ATF instead samples next thoughts from a multi-mode latent distribution and feeds each realized thought back autoregressively, allowing distinct reasoning paths to be explored. In this example, one sampled ATF trajectory follows the correct taxicab-counting path and reaches the correct answer.

Figure 2: ATF generates each next thought conditioned on the realized history. A causal backbone produces the condition for a lightweight flow head, which samples a continuous thought and feeds it back into the reasoning history.

Method	Budget	GSM8K			MATH-500			OlympiadBench			AMC 12
		Acc.	$N_{\rm think}$	$N_{\rm ans}$	Acc.	$N_{\rm think}$	$N_{\rm ans}$	Acc.	$N_{\rm think}$	$N_{\rm ans}$	Acc.	$N_{\rm think}$	$N_{\rm ans}$
Qwen3-1.7B
Discrete CoT	$\infty$	79.15	1252	242	64.40	2582	324	27.37	4132	166	33.73	4127	209
Discrete CoT	$\leq 1000$	67.17	599	535	60.60	879	627	28.06	986	469	26.51	984	711
Discrete CoT	$\leq 500$	76.65	438	337	63.20	494	556	22.20	501	472	25.30	501	646
MSE regression	$\leq 32$	82.41	4.0	300	62.40	4.1	550	27.88	4.6	604	37.35	4.6	669

Table 1: ATF improves compact reasoning over discrete and regression baselines. Accuracy (%) and generation lengths on Qwen3 backbones, with saved answers rescored using a common grader. Budget limits intermediate text tokens for discrete CoT and latent steps for continuous methods.

N_{\rm think}

and

N_{\rm ans}

report the corresponding realized thinking length and final-answer tokens. ATF-RL applies answer-policy refinement. Bold and underline mark best and second-best accuracy within each model size, excluding unbounded discrete CoT.

Family	Method	MATH-500		GSM8K		OlympiadBench
Family	Method	Pass@1	Pass@100	Pass@1	Pass@100	Pass@1	Pass@100
Hidden states	iCoT ( Deng et al., 2023 )	$35.2$	$37.9$	$61.8$	$63.9$	$4.3$	$7.1$
	Coconut ( Hao et al., 2024 )	$37.3$	$39.3$	$68.3$	$74.3$	$5.8$	$6.3$
	CODI ( Shen et al., 2025 )	$38.5$	$45.1$	$76.3$	$81.7$	$7.6$	$14.8$
	TaH+ ( Fu et al., 2025 )	$46.1$	$49.4$	$\mathbf{85.9}$	$89.7$	$12.2$	$15.2$
Token-based	Disc. Latent ( Su et al., 2025 )	$43.2$	$47.3$	$83.9$	$88.6$	$\mathbf{13.3}$	$17.8$

Table 2: Multimodal next-thought prediction unlocks strong test-time scaling. ATF is competitive with prior methods at Pass@1, but the gap widens sharply as more latent trajectories are sampled. At Pass@100, ATF outperforms every baseline, reaching

83.2

on MATH-500,

97.6

on GSM8K, and

52.4

on OlympiadBench—a

34.6

point gain over the strongest prior result on OlympiadBench. The result suggests that ATF does not merely improve a single reasoning trajectory; its multimodal next-thought distribution reaches diverse solution modes that become increasingly valuable under test-time sampling. With ATF (Greedy), the final answer is decoded deterministically, isolating continuous thinking as the source of this diversity. All methods use LLaMA-3.1-8B; baseline scores are reported by LaDiR ( Kang et al., 2025 ) . Bold marks the best result for each metric and benchmark.

Distribution	Method	GSM8K		MATH
Distribution	Method	Acc. (%) $\uparrow$	Latent steps $\downarrow$	Acc. (%) $\uparrow$	Latent steps $\downarrow$
Single Gaussian	CoLaR	$26.8\pm 0.17$	$\mathbf{5.57}\pm 0.02$	$7.08\pm 0.07$	$16.1\pm 0.14$
Flow-based	ATF	$\mathbf{30.17}\pm 1.26$	$6.43\pm 0.01$	$\mathbf{8.59}\pm 0.05$	$\mathbf{7.17}\pm 0.03$

Table 3: Comparison with CoLaR. CoLaR ( Tan et al., 2025 ) and ATF use LLaMA-3.2-1B-Instruct under the GSM8K-Aug/MATH recipe; mean

\pm

standard deviation over five runs. We directly reuse CoLaR’s number and setting, which is the cloest related work that use a Gaussian distribution (single mode distribution) to model the next prediction. For the ATF here we use exactly CoLaR’s setting for fair comparison. Bold marks the better mean within each dataset.

Figure 3: ATF benefits from more latent steps. MATH-500 accuracy at 4B for stopping margins $m\in\{0,4,8\}$ , with $\gamma=1$ and greedy answer decoding. The horizontal axis reports realized mean latent length. ATF achieves a stronger accuracy–latent-length trade-off than the matched regression baseline in this sweep.

$\gamma$	AIME 24			AIME 25
$\gamma$	Pass@1	Pass@16	Diversity	Pass@1	Pass@16	Diversity
1.0	8.54	30.00	0.7503	6.25	23.33	0.7518
2.5	9.38	26.67	0.6357	7.08	20.00	0.6407
5.0	9.58	26.67	0.4147	6.88	20.00	0.4499

Table 4: CFG provides a built-in control over latent mode exploration. For the 4B ATF on AIME, the guidance scale

\gamma

trades off single-sample accuracy against sampled solution coverage. Larger

\gamma

concentrates generation around high-confidence continuations, favoring exploitation, whereas smaller

\gamma

preserves more diverse latent modes for exploration. Bold marks the best result for each metric.

Method	Trace	Decoded excerpt
MSE regression	A	“Okay, so I need to find all natural numbers $N$ such that $N(N-101)$ is a perfect square …First, I remember that $N(N-101)$ needs to be a square …”
MSE regression	B	“If $N$ is a natural number, then $N(N-101)$ must equal the square …So, $N(N-101)$ must be a perfect square …”
ATF	A	“ Expressed in the form $N(N-101)=x^{2}$ … $N=x^{2}$ is a perfect square …”
ATF	B	“ $N$ is the perimeter of a square …identify which numbers from $1\times 1$ to $69$ are squares.”

Table 5: Decoded latent thoughts reveal regression collapse versus generative next-thought diversity. MSE regression produces similar premise-restatement fragments, while ATF realizes distinct continuations from the same autoregressive state. Blue highlights shared premise restatements; orange highlights contrasting continuations.

Figure 4: PCA feature visualizations of ATF’s latent reasoning modes. We sample 500 times to estimate the distribution for our diffusion model, and plot their density map. Compared with MSE regression, ATF preserves more diverse next-thought continuations (top). Our diffusion head design allows us to apply the classifier-free guidance then adjusts the concentration of this learned distribution, trading broader exploration for more focused exploitation (bottom). The rows illustrate different questions; Pass@1 refers to the bottom-row example. Tables 5 and 6 give decoded excerpts.

CFG $\gamma$	Trace	Decoded excerpt
$1$	A	“Okay, so the problem is about finding the units digit …If $n=7$ , then row 7 is $(1,5,6,5,1)$ …”
$1$	B	“Okay, let’s see …triangular array with 15th row …additional entries are constructed by reflecting a row with TicTacToe …”
$5$	A	“Alright, let’s try to figure out this problem …The first rows are $1$ , $1\ 1$ , $1\ 3\ 1$ , $1\ 5\ 5\ 1$ , $1\ 7\ 11\ 7\ 1$ …”
$5$	B	“Alright, let’s try to figure out this problem …The rows are $1$ , $1\ 1$ , $1\ 3\ 1$ , $1\ 5\ 5\ 1$ , $1\ 7\ 11\ 7\ 1$ …”

Table 6: Stronger CFG concentrates decoded continuations around the reference reasoning pattern. Low guidance produces more varied, sometimes off-mode continuations, while high guidance favors similar row-pattern reconstructions. Orange marks contrasting low-guidance fragments; blue marks the repeated high-guidance row sequence.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Method	MATH-500	GSM8K	OlympiadBench
iCoT	$35.2\,/\,37.9$	$61.8\,/\,63.9$	$4.3\,/\,7.1$
Coconut	$37.3\,/\,39.3$	$68.3\,/\,74.3$	$5.8\,/\,6.3$
CODI	$38.5\,/\,45.1$	$76.3\,/\,81.7$	$7.6\,/\,14.8$
Discrete Latent	$43.2\,/\,47.3$	$83.9\,/\,88.6$	$\mathbf{13.3}\,/\,17.8$
Soft Think	$44.3\,/\,46.7$	$83.7\,/\,86.6$	$10.4\,/\,13.1$
TaH+	$46.1\,/\,49.4$	$\mathbf{85.9}\,/\,89.7$	$12.2\,/\,15.2$

Appendix

Table 7: Complete single-sample and multi-sample comparison. All methods use the LLaMA-3.1-8B backbone; ATF is trained on DART-MATH-derived data and RL-refined. Each cell reports Pass@1 / Pass@100 (%). Baseline scores are reported by LaDiR. Bold marks the best value for each metric.

Figure 5: Causal-Register VAE architecture. Register tokens inserted after reasoning blocks provide the encoder states used to parameterize Gaussian posteriors. Sampled latent registers condition the transformer decoder, which reconstructs the input text using mask placeholders. Causal encoding preserves the ordering of the reasoning prefix; the resulting latent registers serve as continuous thought targets for ATF.

VAE Comp. Ratio	AR Recon Acc.	GSM8K	Olympiad
16:1	99.1	76.80	20.31
64:1	75.2	78.09	22.03
128:1	25.7	80.52	21.17

Appendix

Table 8: More compact CR-VAE latents can improve downstream reasoning despite lower token reconstruction at 1.7B. Using the Qwen3-1.7B-Base backbone, AR reconstruction accuracy measures in-distribution reconstruction of discrete thoughts, while GSM8K and OlympiadBench report CFG = 1 reasoning accuracy. Higher VAE compression discards token-level surface form, yet downstream accuracy remains strong, indicating that compact latents preserve the reasoning structure needed by ATF.

Benchmark	Method	Think (s)	Answer (s)	Total (s)	Speedup
GSM8K	Text CoT	10.53	3.49	14.02	$1.00\times$
	ATF	0.22	3.80	4.02	$3.49\times$
MATH-500	Text CoT	26.43	5.46	31.89	$1.00\times$
	ATF	0.22	6.97	7.19	$4.44\times$
OlympiadBench	Text CoT	48.69	3.18	51.86	$1.00\times$
	ATF	0.26	7.89	8.15	$6.40\times$

Appendix

Table 9: End-to-end latency estimates for Qwen3-4B. Reasoning, answer, and total times on one RTX PRO 6000, batch size 1. ATF uses 24 flow-head denoising steps per thought.

Method	GFLOPs / thought	Latency / thought (ms)
ATF	25.2	35.8
LaDiR, 10 steps	139.6	41.4
LaDiR, 50 steps	697.9	207.0

Appendix

Table 10: Per-thought compute under matched conditions. LLaMA-3.1-8B backbone,

5.5{:}1

compression, identical hardware, and batch size 1. The step counts denote flow-head evaluations for ATF and full-backbone evaluations for LaDiR.

Method	MATH-500		AMC 12
	Pass@16	Self-consistency Acc.	Pass@16	Self-consistency Acc.
MSE regression	77.60	69.00	72.29	62.65
ATF (Ours)	79.80	73.00	74.70	57.83

Appendix

Table 11: Multi-sample rollouts show broader answer coverage for ATF at 4B. Using the Qwen3-4B-Base backbone, we evaluate

K=16

adaptive-stopping rollouts at temperature

0.1

on MATH-500 and AMC 12 . Pass@16 measures whether any sampled answer is correct, while self-consistency accuracy uses majority voting across the same samples.

Explore similar work

Latent Recurrent Thoughts: Recurrent Refinement of Proposed Latents for Reasoning with Frozen LLMs
Sep 1, 2026Zhaoliang Chen, Jie FuLatent Thoughts Chain-of-Thought Reasoning
Emory University · IQuest Research
Reasoning with Continuous Latent Diffusion
Sep 28, 2026Xiang ChengEfficient Latent Reasoning Latent Flow
Duke University Department of Electrical and Computer Engineering
A Model with No Head and Many Thoughts
Aug 31, 2026Nikita Koriagin, Yaroslav Aksenov, George Bredis +3Latent Thoughts Thoughts
1Yandex · 2T-Tech

Sep 1, 2026, cs.AIJ/K move · Enter open · S save

Latent Recurrent Thoughts: Recurrent Refinement of Proposed Latents for Reasoning with Frozen LLMs

Zhaoliang Chen, Jie Fu

Emory University · IQuest Research

Latent Thoughts Chain-of-Thought Reasoning Frozen Encoder

Open paper AskPDF arXiv

Chain-of-thought reasoning unfolds in discrete token space: each step is committed as text, errors propagate, and eliciting good traces presupposes traces to imitate. Reasoning instead in a model's continuous representation space - where intermediate states are vectors rather than words - sidesteps these constraints, but leaves open how those latent states should be computed. We approach this along two axes. First, we keep a large language model (LLM) frozen and use it for what it is already good at - modeling and decoding sequences - while a small auxiliary network supplies continuous latent thoughts as input. Second, we produce those latents by recurrence: a tiny recurrent reasoner refines them over many steps, decoupling the depth of computation from the size of the model, so that the latents are a product of iterative processing rather than a single forward pass. We instantiate this as Latent Recurrent Thoughts (LRT): a task-dedicated proposer supplies base latents, a recurrent reasoner refines them through bounded residual corrections, and the frozen LLM decodes the answer. On symbolic reasoning with answer supervision but no reasoning traces (Countdown-4, Sudoku) and on natural-language reasoning (HumanEval, MBPP, StrategyQA), LRT substantially outperforms prior frozen-decoder continuous-space reasoning methods under an identical decoder, prompt, data, and training budget, and outperforms non-thinking-mode chain-of-thought prompting on the same backbone at a small fraction of its inference compute.

Abstract

Explore similar work

Latent Recurrent Thoughts: Recurrent Refinement of Proposed Latents for Reasoning with Frozen LLMs

Reasoning with Continuous Latent Diffusion

A Model with No Head and Many Thoughts