PaperlayerView exact revision on arXiv

Source-linked AI summary

Rationale Distribution Distillation via Flow Matching for Distributional Continuous Latent Reasoning

Mingyeong Kim, Gyusang Cho, Juho Lee, Giung Nam

arXiv:2610.19871v1cs.CLcs.AI

TL;DR

Lengthy textual rationales make repeated autoregressive reasoning costly, while point-valued latent methods reuse a single rationale context and may limit exploration of alternatives. RDD distills multiple textual paths into a query-conditioned flow over latent rationales, and its answer coverage exceeds evaluated point-valued latent reasoning baselines across arithmetic and advanced mathematical tasks.

  • Problem

    Textual rationales are often lengthy and costly to generate repeatedly, while point-valued latent methods can struggle to exploit multiple reasoning paths.

    Source: §1 · p. 1

  • Method

    RDD distills multiple answer-correct textual rationales into a query-conditioned latent distribution, augmented with within-query mixup, and samples rationale contexts for a frozen language model to decode answers.

    Source: §1 · p. 2, §3.3 · p. 4

  • Results

    RDD improves answer coverage over evaluated point-valued latent reasoning baselines across arithmetic and advanced mathematical tasks, with particularly pronounced improvements on advanced mathematics.

    Source: §4.1 · p. 6, §5 · p. 8

  • Takeaways & Limitations

    Sampling a distribution of latent rationale contexts can improve answer coverage compared with repeatedly decoding from a single fixed context.

    Source: §4.1 · p. 6

    Experiments use backbones up to 1.7B parameters and shorter, less elaborate rationales than modern large language models’ extended reasoning traces, leaving broader applicability open.

    Source: §6 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

Textual rationales can improve reasoning but their length makes repeated generation costly, motivating more compact continuous latent representations. Because multiple paths can reveal correct answers missed by individual attempts, RDD distills multiple answer-correct textual rationales into a query-conditioned latent distribution.

1. Introduction

  • Lengthy CoT traces create inference latency and memory overhead, motivating efforts to shorten reasoning while preserving accuracy.Continuous latent reasoning carries intermediate reasoning in continuous states rather than requiring each step to be verbalized as discrete tokens.Source: §1 · p. 1
  • The figure contrasts textual rollouts with point-valued reasoning that reuses one answer context and a flow that induces a distribution over contexts.The flow representation supports diverse answer generation from sampled contexts.Source: Figure 1 · p. 1
  • Existing latent methods can struggle to exploit multiple reasoning paths, although exploring alternatives can uncover correct answers missed by individual attempts.Prior evidence for latent superposition largely concerns synthetic settings.Source: §1 · pp. 1–2
  • RDD distills multiple answer-correct textual CoTs into a query-conditioned latent distribution rather than one fixed latent rationale.It samples a complete latent rationale for each answer rollout, enabling exploration of alternative contexts.Source: §1 · p. 2

2. Related Work

Related work spans point-valued continuous reasoning and methods that induce distributions over continuous rationales. RDD differs by distilling each query’s empirical distribution of complete textual rationales into a conditional distribution over continuous rationales.

Point-valued latent reasoning

  • Point-valued continuous reasoning methods produce a single continuous rationale per query, including Coconut, CODI, SIM-CoT, MUX, PCCoT, and LOTUS.PCCoT and LOTUS use parallel refinement approaches to reduce sequential computation.Source: §2 · p. 2

Distributional latent reasoning

  • Distributional approaches induce continuous-rationale variation through stochasticity, embedding mixtures, variational autoencoding, or flow matching.Examples include GTS, Soft Thinking, Stochastic Soft Thinking, Multiplex Thinking, CoLaR, and LaDiR.Source: §2 · p. 2
  • RDD explicitly distills each query’s empirical distribution of complete textual CoT rationales into a conditional distribution over continuous rationales.It generates full rationales non-autoregressively and freezes the answer model.Source: §2 · p. 2

3. Method

The method evaluates answer coverage by sampling complete rationale–answer rollouts, then replaces textual rationale generation with a query-conditioned flow over fixed-length latent rationales. Mixup augments its targets, while answer-aware contextualization trains sampled rationales to support answer decoding by a frozen language model.

3.1. Textual CoT and answer coverage

  • Textual CoT samples a rationale and then an answer for each rollout, while pass@k measures whether at least one of k rollouts yields a correct answer.The continuous counterpart updates all rationale positions jointly but leaves answer decoding autoregressive.Source: §3.1 · p. 3

3.2. Continuous CoT and superposition

  • Soft Thinking mixes token embeddings, but analysis finds its later predictions can be dominated by the highest-probability token, resembling greedy decoding.Controlled multi-chain experiments found that Coconut and CODI did not benefit from additional valid textual CoTs in the evaluated settings.Source: §3.2 · p. 3

3.3. Rationale distribution distillation

  • RDD models a query-conditioned distribution over complete continuous rationales, converts each sampled rationale into a context, and uses the frozen ALM to decode an answer.It trains on multiple answer-correct teacher traces compressed into fixed-length latent sequences, with within-query mixup adding interpolated targets between same-query rationales.Source: §3.3 · pp. 3–4

3.4. Answer-aware contextualization

  • An ODE sampler generates complete rationales from Gaussian noise with all latent positions updated jointly; a contextualizer maps them to KV contexts for autoregressive answer decoding.Answer and local-content supervision train the sampler and contextualizer jointly, while pretrained ALM weights and the vocabulary projection remain fixed.Source: §3.4 · pp. 4–5

4. Experiments

Across arithmetic and advanced mathematics benchmarks, RDD improves answer coverage over evaluated point-valued latent reasoning baselines. Additional experiments examine multi-path supervision, answer-distribution fidelity, a larger backbone, and design choices in RDD.

4.1. Main results

  • RDD exceeds every evaluated point-valued baseline in pass@512 across all six benchmarks, including Coconut, CODI, PCCoT, and MUX at the same rationale length; it approaches textual CoT on TAL and MATH500 and surpasses it on Gaokao23.Table 1 reports the main results for GPT-2-0.1B on GSM8K-Aug.Source: §4.1 · p. 6, Table 1 · p. 6
  • Among distributional methods, RDD has the highest coverage on all three advanced mathematics benchmarks and outperforms CODIMCD on all six at S=6.CoTSST leads distributional methods on arithmetic, while Monte Carlo dropout reduces coverage relative to CODI throughout.Source: §4.1 · p. 6, Table 2 · p. 6

4.2. Additional results

  • With multi-path supervision, CODI's mean pass@512 rises from 28.4% to 37.0% on advanced mathematics, and RDD reaches 54.7% on arithmetic and 39.8% on advanced mathematics.CODI's arithmetic mean decreases slightly, from 52.3% to 52.0%.Source: §4.2 · p. 7, Table 3 · p. 7
  • RDD achieves the lowest average JSD and highest average TSC against textual CoT, especially on advanced tasks; it also has the lowest mean JSD and TV on all three advanced benchmarks and highest TSC on all six.Figure 4 pools query-level JSD and TSC by benchmark group; Figure 5 provides per-benchmark results.Source: §4.2 · p. 7, Figure 4 · p. 7
  • On SLM-2-1.7B trained with GSM8K-Aug-NL, RDD still outperforms evaluated latent baselines at S=6 and retains its answer-coverage advantage despite stronger compression.The longer natural-language rationales are compressed into six latent positions, which may contribute to the generally larger gap from textual CoT and CoTSST.Source: §4.2 · p. 7, Table 4 · p. 8
  • Midpoint interpolation of aligned, answer-correct rationale embeddings preserves nearly all CoT accuracy: 99.99%, versus 100% with the original rationale and 44.33% with a rationale from another query.The probe uses 10,000 triples in an aligned, uncompressed setting.Source: §4.2 · p. 7

4.3. Ablation study

  • Sampling a fresh context for every rollout improves pass@512 in both benchmark groups; using three Euler steps yields comparable coverage while halving flow evaluations.At S=6, this also reduces sequential model forward passes relative to autoregressive baselines.Source: §4.3 · pp. 7–8
  • Stochastic mixup improves coverage over no mixing in both groups, whereas midpoint mixup improves only advanced coverage; causal attention performs better than bidirectional attention in these experiments.The mixup results support a regularization interpretation.Source: §4.3 · p. 8

5. Conclusion

RDD distills multiple answer-correct textual reasoning paths into a query-conditioned distribution of continuous latent rationales, augmented by within-query mixup. The paper reports improved answer coverage over evaluated point-valued latent baselines across arithmetic and advanced mathematics, and identifies compression and sampling strategies as future research directions.

5. Conclusion

  • RDD explicitly learns a query-conditioned distribution of latent rationales from multiple answer-correct textual CoT traces, augmented through within-query mixup.Parallel refinement of fixed-length latent sequences replaces token-by-token rationale generation.Source: §5 · p. 8
  • The experiments show improved answer coverage over evaluated point-valued latent reasoning baselines across arithmetic and advanced mathematical tasks, nearly matching textual CoT on advanced benchmarks.The paper presents distributional modeling as a promising direction and points to more effective rationale compression and sampling strategies for larger scales.Source: §5 · pp. 8–9

6. Limitations

The evaluation is limited to backbones up to 1.7B parameters and rationales shorter and less elaborate than modern extended reasoning traces. RDD improves answer coverage, but its single-rollout accuracy generally trails the strongest baselines.

6. Limitations

  • Experiments use backbones up to 1.7B parameters and short symbolic or natural-language rationales with verbal explanations.Whether the benefits extend to larger models and longer, more verbal traces remains open.Source: §6 · p. 9
  • The evaluated rationales are substantially shorter and less elaborate than extended reasoning traces produced by modern large language models.The generality of the reported benefits to those longer, more verbal traces remains unresolved.Source: §6 · p. 9
  • RDD improves answer coverage (pass@k), while its single-rollout accuracy (pass@1) generally trails the strongest baselines.The experiments focus on supervised distillation; whether reward-based optimization can narrow this gap remains open.Source: §6 · p. 9

7. Ethical Concerns

Continuous latent reasoning raises transparency and oversight concerns because intermediate states are not directly interpretable in natural language. Textual reasoning traces can also be unfaithful, but can support chain-of-thought monitoring.

7. Ethical Concerns

  • Latent intermediate states are not directly interpretable in natural language, raising concerns about transparency and oversight.Source: §7 · p. 9
  • Textual reasoning traces can be unfaithful to model behavior, yet can still support chain-of-thought monitoring.Source: §7 · p. 9
  • The reduced access to monitoring signals motivates oversight methods for latent reasoning.Source: §7 · p. 9

A. Additional Related Work

Prior work studies distribution distillation, multiple textual rationales, and latent reasoning, but the paper identifies no prior method explicitly distilling per-query empirical distributions of complete textual CoTs into conditional continuous rationale distributions. Its mixup analysis characterizes reference-target geometry while cautioning that this geometry does not determine generated-latent dispersion or answer coverage.

A. Additional Related Work

  • Prior distribution-distillation work models predictive outputs, while prior language- and latent-reasoning work does not explicitly distill per-query empirical distributions of complete textual CoTs into conditional continuous rationale distributions.Source: A. Additional Related Work · p. 15
  • Unlike related flow and latent-augmentation approaches that ultimately generate verbalized text or condition textual generation, this approach keeps sampled rationales continuous and decodes answers from them.Source: A. Additional Related Work · p. 15

B. Analysis of Within-Query Mixup

  • Within-query mixup contracts reference-target covariance, retains supervision on original targets, and gives exact interpolation of linear-path flow-matching regression examples.The analysis does not establish a causal link between latent variance and answer coverage.Source: B. Analysis of Within-Query Mixup · p. 15, §B.1 · p. 16

B.1. Reference geometry

  • For fixed teacher targets, covariance contraction strengthens in the order no mixing, Beta(2, 2), Beta(5, 5), then midpoint mixing, but this ordering does not determine generated dispersion or answer coverage.The listed covariance factors are 0, 1/5, 5/22, and 1/4, respectively.Source: §B.1 · p. 16

B.2. Effect on the training objective

Within-query mixup interpolates the flow-matching regression inputs and velocity targets while keeping the query and time fixed. It changes the expected flow-matching objective but preserves the expected auxiliary objective at fixed parameters.

B.2. Effect on the training objective

  • Mixup interpolates both sides of the flow-matching regression problem while keeping the query and time fixed.A shared source sample is used for the analytical coupling; training may instead draw fresh noise independently for each example.Source: §B.2 · pp. 16–17
  • With finite teacher samples and limited model capacity, intermediate regression examples may discourage over-specialization to observed embeddings.This regularization interpretation does not require the learned velocity field to be linear; the mixing-weight ablation provides empirical support.Source: §B.2 · p. 17
  • At fixed parameters, mixup changes the expected flow-matching objective but preserves the expected auxiliary objective for the same teacher set and query distribution.The auxiliary losses depend on generated rationales, which are independent of the sampled teacher indices and mixing weight.Source: §B.2 · p. 17
  • The generated rationale distribution need not match reference targets or inherit their covariance ordering, because auxiliary losses, model capacity, optimization, and numerical integration also shape it.The expected auxiliary objective remains unchanged by mixing at fixed parameters, but this does not guarantee matching generated behavior.Source: §B.2 · p. 17

B.3. Empirical evidence

The empirical analysis finds that stochastic target mixing improves answer coverage across arithmetic and advanced-mathematics benchmark groups, while generated-latent dispersion and trajectory compatibility provide separate evidence about the behavior and scope of interpolation. The experiments use two backbones, multiple answer-correct training rationales per query, and fixed flow-solver settings.

B.3. Empirical evidence

  • Both stochastic mixing settings improve coverage over no mixing across arithmetic and advanced-mathematics benchmarks, although the strongest setting differs by group.Beta(2, 2) has the highest arithmetic pass@512, Beta(5, 5) the highest advanced-mathematics coverage, and midpoint mixing improves advanced but not arithmetic coverage.Source: §B.3 · p. 17
  • 1.7259×10−2 for no mixing and 1.7257 × 10−2 for midpoint mixing are the mean generated-endpoint variances in the GSM8K comparison.The comparison uses 500 queries and 128 shared Gaussian initial states per query; nearly unchanged variance does not determine the separate answer-coverage result.Source: §B.3 · pp. 17–18
  • 100%, 99.99%, and 44.33% accuracy result from the original rationale, same-query midpoint interpolation, and cross-query replacement, respectively.The controlled probe uses 10,000 triples with answer-correct rationales aligned in reasoning-step counts and token counts at each step.Source: §B.3 · p. 18

C.1. Models and datasets

  • Training data are sanitized and decontaminated, then filtered to retain answer-correct paths among ten sampled responses per question.The resulting datasets contain 378,647 training instances and 500 validation instances; data preparation also filters selected evaluation subsets.Source: §C.1 · p. 18

C.2. Hyperparameters

  • Textual CoT training runs for 300k steps with GPT-2-0.1B and 50k with SLM-2-1.7B.LoRA is applied to linear layers; the B-matrix learning rate is four times the A-matrix rate.Source: §C.2 · p. 18

C.1. Models and datasets

  • The listed model backbones are GPT-2-0.1B and SLM-2-1.7B; training data sources are GSM8K-Aug and GSM8K-Aug-NL; evaluation benchmarks are GSM8K, SVAMP, GSMHard, TAL, MATH500, and Gaokao23.Source: §C.2 · p. 19

C.2. Hyperparameters

  • RDD uses six Euler steps for flow integration and draws mixup weights from Beta(2, 2), or Beta(0.5, 0.5) for SLM-2-1.7B.Each teacher rationale is split into six nearly equal chunks; latent methods are fine-tuned for 100k steps with GPT-2-0.1B and 10k with SLM-2-1.7B.Source: §C.2 · pp. 18–19

C.3. Evaluation

The evaluation measures answer coverage with pass@k and compares sampled answer distributions with textual CoT using JSD, TV, and teacher support coverage. Experiments ran on NVIDIA RTX 3090 and RTX A6000 GPUs.

C.3. Evaluation

  • For each query, pass@k uses the unbiased estimator based on correct answers among 512 rollouts sampled at temperature 0.7, for 1 ≤ k ≤ 512.Source: §C.3 · p. 19
  • JSD and TV compare empirical textual-CoT and method answer distributions, while TSC measures teacher probability mass on answers observed from the evaluated method.Each distribution uses 512 rollouts; equivalent numeric answers share a category, unextractable answers share a category, and incorrect answers remain included.Source: §C.3 · pp. 19–20
  • Lower JSD and TV indicate closer answer distributions, while higher TSC indicates less missing teacher mass.All three fidelity metrics lie in [0, 1].Source: §C.3 · p. 20

C.4. Software and hardware

  • Experiments used torch, transformers, and peft, and ran on NVIDIA RTX 3090 and RTX A6000 GPUs.The reported library versions are torch v2.11.0, transformers v5.14.1, and peft v0.20.0.Source: §C.4 · p. 20

D. Supplementary Results

The supplementary results show that RDD improves answer-distribution fidelity on advanced mathematics and has a pass@k advantage that becomes clearer at larger k, particularly on advanced tasks. The figures provide benchmark-level views of these comparisons.

D. Supplementary Results

  • RDD has the lowest mean JSD and TV on all three advanced-mathematics benchmarks and the highest mean teacher support coverage on all six benchmarks.Figure 5 reports per-question JSD, TV, and TSC by benchmark; lower JSD and TV and higher TSC indicate better fidelity.Source: D. Supplementary Results · p. 20, Figure 5 · p. 21
  • RDD’s pass@k advantage becomes apparent at larger k, particularly on advanced mathematics, where it matches or exceeds textual CoT coverage for GPT-2-0.1B.Figures 6–8 show benchmark-wise curves; CoTST, CoTSST, and LOTUS use more latent tokens than other methods.Source: D. Supplementary Results · p. 20, Figure 6 · p. 22, Figure 7 · p. 23

E. Declaration of AI Usage

  • The authors used language models to assist with manuscript drafting, editing, proofreading, and experimental-code development and debugging, then reviewed and validated the assisted content.The authors state that they take responsibility for the final manuscript and experimental results.Source: E. Declaration of AI Usage · p. 20

Related papers

Loading 2610.19871v1…