Rationale Distribution Distillation via Flow Matching for Distributional Continuous Latent Reasoning

Mingyeong Kim; Gyusang Cho; Juho Lee; Giung Nam

cs.CLOct 14, 2026

Rationale Distribution Distillation via Flow Matching for Distributional Continuous Latent Reasoning

Authors: Mingyeong Kim, Gyusang Cho, Juho Lee, Giung Nam

Organizations: Korea Advanced Institute of Science and Technology

Abstract

Sampling multiple chain-of-thought rationales from a language model can uncover correct solutions missed by a single attempt. However, these rationales are typically lengthy, making repeated autoregressive generation costly. Continuous latent reasoning offers a compact alternative, but point-valued methods condition repeated answer sampling on a single rationale, which can limit exploration of diverse paths and answer coverage. We thus propose rationale distribution distillation (RDD), training a conditional flow over fixed-length latent sequences by distilling from multiple textual reasoning paths. Furthermore, we expand the distillation targets through within-query mixup, interpolating between complete latent rationales for the same query. The flow nonautoregressively generates the entire latent sequence, and a frozen language model autoregressively decodes answers from sampled rationale contexts. Empirically, the resulting distributional reasoner achieves higher answer coverage than the evaluated point-valued latent reasoning baselines, particularly on advanced mathematical tasks.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison of reasoning representations. (a) Textual reasoning samples multiple rationales and generates an answer for each. (b) Point-valued reasoning samples multiple answers conditioned on one answer context. (c) Our flow induces a distribution over answer contexts, enabling diverse answer generation.

Panel of Figure 2 — Figure 2: Rationale compression and within-query mixup. (a) Position-weighted embedding pooling maps each text chunk to one $D$ -dimensional vector ( $S=3$ shown). (b) A shared $\lambda$ mixes corresponding vectors of the two complete latent sequences. Original and mixed sequences serve as targets for flow matching. Cell colors depict schematic latent values.

Panel of Figure 3 — Figure 3: Rationale sampling, contextualization, and answer decoding. The sampler produces $\hat{\bm{z}}_{1}$ with an ODE solver; the contextualizer forms KV context $\hat{\bm{u}}$ at $t=1$ . Both share $\phi$ and $\chi$ , while the answer decoder generates $\hat{a}$ with the reasoning adapter disabled. ALM weights and vocabulary projections $W$ are frozen; teal parameter blocks are trainable. Dashed arrows indicate training supervision for rationale generation, local content, and answer prediction.

		Arithmetic benchmarks			Advanced benchmarks
Method	# Token	GSM8K	SVAMP	GSMHard	TAL	MATH500	Gaokao23
CoT	27.4 †	89.5 / 47.5	76.9 / 50.8	23.2 / 10.7	46.0 / 8.1	43.5 / 2.8	30.6 / 1.8
CoT ${}_{\texttt{ST}}$	27.2 †	48.5 / 48.4	53.3 / 52.4	11.4 / 11.4	10.4 / 9.9	4.2 / 3.0	3.3 / 2.2
Coconut ‡	6	59.7 / 32.5	54.3 / 35.2	14.2 / 7.5	20.4 / 5.1	13.2 / 0.8	6.5 / 0.2
CODI	6	76.3 / 37.9	61.9 / 40.3	18.6 / 8.9	33.5 / 6.9	30.7 / 2.5	21.1 / 2.1
PCCoT	6	78.2 / 42.8	64.0 / 46.3	18.8 / 9.9	38.9 / 6.6	35.6 / 1.7	17.5 / 0.9

Table 1: Main results for GPT-2-0.1B on GSM8K-Aug. # Token denotes rationale length. Each result cell reports pass@512 / pass@1 . Excluding CoT , bold-underlined and underlined values respectively denote the best and second-best results for each metric; ties share the same marking. † Averaged across all benchmarks. ‡ Adopted public checkpoint ( Dilgren and Wiegreffe, 2026 , connordilgren/gpt2-gsm8k-coconut ; MIT license;) .

		Arithmetic benchmarks			Advanced benchmarks
Method	# Token	GSM8K	SVAMP	GSMHard	TAL	MATH500	Gaokao23
CoT	27.4 †	89.5 / 47.5	76.9 / 50.8	23.2 / 10.7	46.0 / 8.1	43.5 / 2.8	30.6 / 1.8
CoT ${}_{\texttt{SST}}$	27.6 †	81.6 / 48.4	68.3 / 51.6	21.2 / 11.1	37.1 / 8.7	34.6 / 2.9	27.2 / 1.9
CODI ${}_{\texttt{MCD}}$	6	65.0 / 37.7	56.3 / 40.0	15.6 / 8.8	27.9 / 7.4	20.0 / 2.8	16.3 / 2.3
RDD (ours)	6	78.6 / 40.5	66.0 / 41.2	19.4 / 9.2	44.4 / 5.6	43.4 / 1.5	31.5 / 1.1

Table 2: Distributional rationale results for GPT-2-0.1B on GSM8K-Aug. # Token denotes rationale length. Each result cell reports pass@512 / pass@1 . † Averaged across all benchmarks.

Method	MPS	Arithmetic	Advanced
CoT		63.2 / 36.3	40.0 / 4.2
CODI		52.3 / 29.0	28.4 / 3.8
CODI	$\checkmark$	52.0 / 27.4	37.0 / 3.5
RDD (ours)	$\checkmark$	54.7 / 30.3	39.8 / 2.7

Table 3: Multi-path supervision results for GPT-2-0.1B on GSM8K-Aug. MPS denotes our multi-path supervision protocol. Each cell reports pass@512 / pass@1 , averaged over the three benchmarks in each group.

		Arithmetic benchmarks			Advanced benchmarks
Method	# Token	GSM8K	SVAMP	GSMHard	TAL	MATH500	Gaokao23
CoT	87.9 †	96.7 / 55.3	94.0 / 61.9	46.1 / 16.2	77.0 / 22.4	69.2 / 7.1	52.2 / 5.0
CoT ${}_{\texttt{ST}}$	77.9 †	57.5 / 54.6	61.7 / 61.4	17.4 / 16.2	28.7 / 22.6	18.2 / 7.1	10.9 / 4.7
CoT ${}_{\texttt{SST}}$	78.8 †	94.2 / 54.2	91.7 / 62.0	42.8 / 16.0	74.6 / 22.0	66.4 / 6.9	52.2 / 5.1
CODI	6	84.8 / 27.4	77.3 / 38.9	22.7 / 6.6	61.0 / 12.0	52.0 / 3.6	33.7 / 3.9
MUX	6	87.6 / 27.6	76.0 / 39.1	24.3 / 6.9	62.0 / 13.1	55.0 / 5.0	32.6 / 4.8

Table 4: Additional results for SLM-2-1.7B on GSM8K-Aug-NL. # Token denotes rationale length. Each result cell reports pass@512 / pass@1 . Excluding CoT , bold-underlined and underlined values respectively denote the best and second-best results for each metric; ties share the same marking. † Averaged across all benchmarks.

Figure 4: Answer-distribution fidelity to textual CoT. Query-level JSD (

\downarrow

) and TSC (

\uparrow

) pooled by group; see Figure 5 for per-benchmark results.

Ablation	Arithmetic	Advanced
Per-rollout rationale sampling
Fixed	54.0 / 30.5	38.5 / 2.7
Sampled	54.7 / 30.3	39.8 / 2.7
Number of function evaluations
NFE=3	54.6 / 22.7	40.7 / 3.0
NFE=6	54.7 / 30.3	39.8 / 2.7

Table 5: Ablation results for GPT-2-0.1B. Each cell reports pass@512 / pass@1 , averaged over the three benchmarks in each group.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Resource	Reference	Hugging Face repository	Reported license
Model backbones
GPT-2-0.1B	Radford et al. (2019)	openai-community/gpt2	MIT
SLM-2-1.7B	Ben Allal et al. (2025)	HuggingFaceTB/SmolLM2-1.7B	Apache-2.0
Training data sources
GSM8K-Aug	Deng et al. (2023) ; Wu et al. (2025)	whynlp/gsm8k-aug	Not stated †
GSM8K-Aug-NL	Deng et al. (2023) ; Wu et al. (2025)	whynlp/gsm8k-aug-nl	Not stated †

Appendix

Table 6: Model backbones and dataset sources. References, public Hugging Face releases, and reported licenses are listed for each resource. Licenses are those declared by the linked release; “not stated” means no license declaration was found in its card or repository metadata.

Optimization
Hyperparameter	Value
Global batch size	128
Optimizer	AdamW
$(\beta_{1},\beta_{2})$	$(0.9,0.95)$
Weight decay	0.1
Peak learning rate	$10^{-3}$ ( $A$ ), $4{\times}10^{-3}$ ( $B$ )

Appendix

Table 7: Training hyperparameters. Values shared across methods and backbones.

Figure 5: Answer-distribution fidelity by benchmark. Per-question JSD (top), TV (middle), and TSC (bottom). Lower JSD and TV and higher TSC indicate better fidelity. Boxes show the interquartile range with median lines; whiskers extend to observations within

1.5

times the interquartile range, and dots denote outliers.

Figure 6: Pass@

k

curves for the methods in Table 1 . Results for GPT-2-0.1B on the arithmetic benchmarks (top row) and advanced benchmarks (bottom row). The horizontal axis uses a base-2 logarithmic scale; the vertical axis reports pass@

k

(%). Markers indicate evaluated

k

values, with lines connecting adjacent points.

Figure 7: Pass@

k

curves for the methods in Table 2 . Results for GPT-2-0.1B on the arithmetic benchmarks (top row) and advanced benchmarks (bottom row). The horizontal axis uses a base-2 logarithmic scale; the vertical axis reports pass@

k

(%). Markers indicate evaluated

k

values, with lines connecting adjacent points.

Figure 8: Pass@

k

curves for the methods in Table 4 . Results for SLM-2-1.7B on the arithmetic benchmarks (top row) and advanced benchmarks (bottom row). The horizontal axis uses a base-2 logarithmic scale; the vertical axis reports pass@

k

(%). Markers indicate evaluated

k

values, with lines connecting adjacent points.

LLM Reasoning

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

See all 926 papers

Next Thoughts Are Distributions: Generative Autoregressive Reasoning in the Latent Space
Sep 27, 2026Yang Li, Yi Wang, Shiyuan Huang +3Numerical Reasoning in Language Models Continuous Latent Reasoning
Rutgers University · Amazon · University of Illinois at Urbana-Champaign
Structural Rationale Distillation via Reasoning Space Compression
May 8, 2026Jialin Yang, Jiankun Wang, Jiajun Wu +3CoT Distillation Structured Reasoning
University of Calgary · Calgary, Canada · University of Michigan +1
Latent Recurrent Thoughts: Recurrent Refinement of Proposed Latents for Reasoning with Frozen LLMs
Sep 1, 2026Zhaoliang Chen, Jie FuEfficient Language Model Reasoning Continuous Latent Reasoning
Emory University · IQuest Research

Sep 27, 2026J/K move · Enter open · S save

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

Yang Li, Yi Wang, Shiyuan Huang, Yang Liu+2

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

Numerical Reasoning in Language Models Continuous Latent Reasoning Latent Reasoning Latent CoT Reasoning

Open paper AskPDF arXiv

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.

Abstract

LLM Reasoning

Explore similar work

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

Structural Rationale Distillation via Reasoning Space Compression

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