PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent Reasoning with Normalizing Flows

Guancheng Tu, Xiangjun Fu, Suhao Yu, Yao Tang, Haoqiang Kang, Lianhui Qin, Yizhe Zhang, Jiatao Gu

arXiv:2606.06447v1cs.CLcs.LG

TL;DR

Existing latent-reasoning methods can sacrifice language modeling’s native sampling, scoring, and decoding interface. NF-CoT uses normalizing flows to model continuous thoughts within an LLM’s causal stream, improving code-generation pass rates while reducing intermediate-reasoning cost.

  • Problem

    Existing latent CoT methods can lack native left-to-right likelihood modeling, sampling, scoring, and decoding compatibility with current LLMs.

  • Method

    NF-CoT places an autoregressive normalizing flow inside a shared LLM causal stream to jointly model continuous thoughts and text with exact likelihoods.

  • Results

    68.8 pass@1, improving Qwen3-8B-Base from 55.8 by +13.0% across five code-generation benchmarks.

  • Takeaways & Limitations

    Likelihood-based latent reasoning provides a practical interface for sampling, scoring, and refining continuous thoughts in LLMs.

  • Takeaways & Limitations

    Validation focuses on code-generation benchmarks, while fixed-length trajectories, latent-budget assumptions, and verifier-dependent rewards limit demonstrated generality.

Abstract

from arXiv · show

Large language models often improve reasoning by generating explicit chain-of-thought (CoT), demonstrating the importance of intermediate computation. However, textual CoT forces this computation through a discrete, serial, and communication-oriented token stream: each reasoning step must be verbalized before the model can proceed, even when the underlying update is semantic, uncertain, or only partially formed. Latent reasoning offers a higher-bandwidth alternative by performing intermediate computation in compact continuous states before committing to text. Yet existing latent-reasoning methods often sacrifice key advantages that make CoT effective in autoregressive language models, including native left-to-right generation, probabilistic sampling, compatibility with KV-cache decoding, and tractable likelihood estimation. We propose NF-CoT, a latent reasoning framework that preserves these advantages by modeling continuous thoughts with normalizing flows. NF-CoT instantiates a TARFlow-style normalizing flow inside the LLM backbone, defining a tractable probability model over compact continuous thoughts distilled from explicit CoT. Continuous-thought positions are generated by an NF head, while text positions are generated by the standard LM head within the same causal stream. This design provides exact likelihoods for latent thoughts, enables probabilistic left-to-right decoding with the original KV cache, and supports direct policy-gradient optimization in the latent reasoning space. On code-generation benchmarks, NF-CoT improves pass rates over explicit-CoT and prior latent-reasoning baselines while substantially reducing intermediate-reasoning cost.

1 Introduction

NF-CoT addresses the limitations of latent chain-of-thought methods by modeling continuous thoughts with a normalizing flow inside a shared causal LLM stream. It preserves likelihood-based sampling, scoring, and decoding while supporting joint latent-text modeling and policy-gradient optimization.

  • Motivation: CoT improves reasoning by inserting sampled intermediate variables between prompts and answers, but textual reasoning forces computation through discrete token sequences.CoT changes direct input-output prediction into conditioned prediction through a sampled reasoning path.
  • Motivation: Latent CoT methods replace textual traces with continuous states, but existing approaches are either local and token-space-bound or require iterative diffusion denoising.These limitations prevent native left-to-right likelihood modeling and challenge the sampling, scoring, and decoding interface of current LLMs.
  • NF-CoT: NF-CoT places a scalable normalizing flow inside the LLM’s causal stream, using one autoregressive backbone for continuous thoughts and text answers.During training, explicit CoT supervision is converted into continuous thoughts and optimized with an end-to-end likelihood objective.
  • Contributions: NF-CoT enables supervised likelihood training and policy-gradient refinement in continuous reasoning space, improving code-generation pass rates while reducing intermediate-reasoning cost.Evaluations cover MBPP, MBPP+, HumanEval, HumanEval+, and LiveCodeBench v6, with comparisons against explicit-CoT and prior latent-reasoning baselines.
  • Contributions: NF-CoT preserves explicit CoT’s sampling, scoring, and decoding interface while jointly modeling latent thoughts and text answers with NF and LM heads.The architecture uses separate NF and LM heads over the shared LLM backbone.

2 Preliminaries

The preliminaries frame explicit CoT as a probabilistic intermediate variable but note its verbose token cost, then introduce continuous CoT and normalizing flows as alternatives for stochastic latent reasoning. Existing approaches either produce deterministic thoughts or separate latent generation from the autoregressive likelihood interface.

  • Explicit chain-of-thought: Explicit CoT is an intermediate text variable that can be generated, sampled, and scored autoregressively, but it imposes verbose natural-language reasoning at high token cost.This formulation applies when an explicit CoT trace is available.
  • Continuous chain-of-thought: Continuous CoT replaces text rationales with K continuous codes of dimension D that serve as the intermediate reasoning variable.The codes are denoted e1:K ∈ R^K×D.
  • Prior continuous-reasoning approaches: Coconut feeds distilled hidden-state thoughts back autoregressively, but its deterministic representations do not define a distribution over reasoning paths.This limitation concerns stochastic reasoning-path modeling.
  • Prior continuous-reasoning approaches: VAE-based methods learn stochastic continuous codes, then train a separate generative model for pθ(e1:K | q), separating latent-thought generation from the native autoregressive likelihood interface.The VAE objective combines reconstruction loss with a KL regularizer before separate latent generation pretraining.
  • Normalizing flows: Normalizing flows provide invertible mappings with exact likelihood evaluation and direct sampling, while autoregressive variants use prior-position context and triangular Jacobians for causal sequence modeling.TarFlow and STARFlow scale this family with Transformer blocks for high-dimensional structured latents.

3 Method

NF-CoT uses invertible normalizing flows to transform continuous CoT targets into an LLM-facing thought space that supports autoregressive sampling and exact likelihoods. It integrates flow-based thought generation with standard answer-token generation, then enables supervised and policy-gradient training in the same causal model.

  • Flow-based thought space: NF-CoT maps continuous CoT targets e1:K to information-equivalent thoughts u1:K through an invertible stack of shallow autoregressive flow blocks.The transformation preserves information rather than performing lossy compression.
  • Flow-based thought space: The reparameterized likelihood combines the autoregressive likelihood of u1:K with the flow Jacobian log-determinant, enabling exact end-to-end training.The objective is log pθ(e1:K | q) = log pθ(u1:K | q) + log |det JFθ(e1:K; q)|.
  • Architecture: The LLM backbone processes continuous thoughts and answer tokens in one causal stream, using an NF head for thought densities and the standard LM head for text.Continuous thoughts are projected into the token embedding dimension before processing.
  • Training: NF-CoT is trained on encoder-derived continuous CoT targets with a unified objective that scores both latent reasoning and final-answer tokens.The flow term assigns exact likelihood to the original continuous target, while the text term is the standard autoregressive likelihood conditioned on continuous thoughts.
  • Inference: At inference, NF-CoT samples thoughts directly left-to-right in u-space, then switches to the LM head for answer generation without running the shallow flow blocks.This avoids diffusion-style iterative latent generation and uses stochastic sampling analogous to standard token sampling.
  • Reinforcement learning: After supervised training, NF-CoT applies policy gradients directly to the Gaussian thought likelihood and token likelihood while freezing the shallow flow blocks.The policy likelihood decomposes into continuous-thought and answer-token terms, making both jointly optimizable.

4 Experiments

Experiments evaluate NF-CoT for Python code generation against autoregressive, diffusion-based, and continuous latent-reasoning baselines using pass@1 and pass@k. NF-CoT improves benchmark performance, scales with sampling, and supports latent-space reinforcement learning without pass@k collapse.

  • Baselines: Under the Qwen3-8B-Base backbone, comparisons include autoregressive, diffusion, looped, and continuous-reasoning baselines alongside NF-CoT Dual-Path and Unified variants.The same setup reports pass@1 and pass@k across multiple benchmarks.
  • Datasets and evaluation: Experiments use Ling-Coder’s 1.4M Python samples and five Python benchmarks, reporting average pass@1 from 16 sampled solutions per problem.The benchmarks are HumanEval, MBPP, HumanEval+, MBPP+, and LiveCodeBench v6.
  • Pass@1 improvements on code generation: 68.8 average pass@1: NF-CoT Unified improves Qwen3-8B-Base from 55.8 to 68.8 (+13.0%), while Dual-Path improves the base by +9.4%.On four shared non-LiveCodeBench benchmarks, Unified reaches 80.0 versus Standard SFT’s 67.5 from the same base-model starting point of 65.3.
  • Pass@k scaling: 87.5 at k = 128: NF-CoT’s MBPP+ pass@1 of 72.1 matches the base model’s pass@128 of 72.0, while HumanEval+ rises from 78.3 to 97.5 (+19.2).NF-CoT dominates the base model and LaDiR across the full k range, with LaDiR rising from 73.2 to 90.2 (+17.0%) and flattening at larger budgets.
  • Execution-guided latent RL: 70.1 average pass@1: GRPO raises supervised NF-CoT Unified from 68.8 to 70.1 across all five benchmarks, while latent-space RL preserves upward pass@k scaling.Unlike token-space GRPO, NF-CoT + RL remains above the supervised checkpoint across both MBPP+ and HumanEval+ without collapsing the latent trajectory distribution.

5 Analysis

NF-CoT’s analysis shows that autoregressive flow generation improves efficiency while preserving strong code-generation performance, and that its training curriculum materially benefits benchmark results. Perturbation experiments further indicate that continuous thoughts control solution trajectories more than functional correctness.

  • Efficiency: 64 latent tokens encode approximately 385 text tokens, yielding an effective compression rate of approximately 6.0× versus explicit-CoT.The comparison uses matched HumanEval settings with 16 candidates per problem, a Qwen3-8B backbone, and vLLM decoding.
  • Efficiency: 2.85× sample and 5.71× token throughput, with 6.66× fewer total FLOPs, are achieved by NF-CoT Unified over LaDiR during training.NF-CoT Dual-Path also reaches 1.90× higher sample throughput, while Unified replaces iterative denoising with autoregressive flow generation.
  • Training ablations: Removing the frozen-backbone warm-up lowers pass@1 from 84.4 to 81.5 on HumanEval, 78.7 to 75.5 on HumanEval+, and 23.1 to 21.4 on LiveCodeBench v6.The default curriculum trains flow components with the backbone frozen before joint optimization; the ablation skips this warm-up.
  • Latent-space analysis: As perturbation noise increases from σ=0 to 3.0, cosine similarity falls from 1.000 to 0.116 while pass@1 decreases only from 86.0 to 83.6.The overlapping confidence intervals indicate local smoothness in the learned continuous-CoT space.
  • Latent-space analysis: Perturbations mainly switch implementation style rather than correctness, showing that NF-CoT latent variables control solution trajectories rather than whether the model solves the problem.Canonical-solution log-probability rises under moderate noise but degrades under large noise.

6 Related Work

Prior work contrasts token-space chain-of-thought with latent reasoning based on continuous thoughts, while normalizing flows provide invertible mappings and exact, tractable likelihoods. Related flow research spans image generation and continuous latent representations for NLP.

  • CoT and latent reasoning: CoT generates intermediate reasoning steps in token space before producing an answer.Supervised or RL-based training and test-time sampling and scoring can further improve CoT.
  • CoT and latent reasoning: Latent reasoning feeds continuous thoughts back into the model instead of reasoning entirely in token space.
  • Normalizing flows: Normalizing flows transform a simple base distribution through invertible networks while providing exact, tractable likelihoods.Applications include natural image generation and NLP models that lift discrete text into continuous latent spaces.

7 Conclusion

NF-CoT gives continuous chain-of-thought the same modeling status as language tokens by embedding an autoregressive normalizing flow within the LLM’s causal stream. It provides exact likelihoods for reasoning trajectories, enabling supervised likelihood training and policy-gradient refinement, while improving code-generation accuracy and running faster than LaDiR.

  • Conclusion: NF-CoT gives continuous CoT the same modeling status as language tokens through an autoregressive normalizing flow inside the LLM’s causal stream.This integrates latent reasoning directly into the model’s causal computation.
  • Conclusion: Exact likelihoods over reasoning trajectories support both supervised likelihood training and policy-gradient refinement in continuous reasoning space.NF-CoT models an explicit distribution over reasoning trajectories.
  • Conclusion: NF-CoT improves accuracy over SFT and prior latent-reasoning baselines on Qwen3-8B-Base across code-generation benchmarks.The comparison is reported across code-generation benchmarks using Qwen3-8B-Base.
  • Conclusion: NF-CoT runs faster than LaDiR in both training and inference.The conclusion reports speed advantages over LaDiR for both stages.

Limitations

NF-CoT is validated only on code-generation benchmarks and inherits limitations from fixed-length VAE trajectories, fixed latent budgets, and domain-specific reward supervision. Extending it to other reasoning tasks and settings without verifiers remains future work.

  • Validation focuses on code-generation benchmarks, leaving extension to other reasoning tasks for future work.
  • Training on fixed-length VAE-encoded explicit-CoT trajectories causes the learned density to inherit trace biases, while fixed latent budgets may not fit every problem.
  • Latent likelihood is not a correctness oracle, and RL refinement depends on unit-test rewards available only for code.
  • Generalization to domains without verifiers remains open.

A Qualitative Analysis of Decoded Latent CoTs … A.3 HumanEval/89: encrypt

NF-CoT’s decoded latent CoTs are imperfect textual reconstructions but serve as qualitative probes of the algorithmic routes underlying generated programs. Across HumanEval examples, sampled latent trajectories support multiple plausible implementations of the same task, including distinct dynamic-programming strategies and alternative cipher or rotation mechanisms.

  • A Qualitative Analysis of Decoded Latent CoTs: Decoded latent CoTs are reconstructed from sampled latent vectors and should be treated as qualitative probes, not faithful natural-language explanations.The text is unedited for correctness or fluency and reconstructed through a frozen VAE decoder.
  • A Qualitative Analysis of Decoded Latent CoTs: The showcase uses NF-CoT’s unified checkpoint with 64 latent tokens, guidance scale 1.0, NF temperature 0.9, and answer temperature 0.6 on held-out HumanEval tasks.Latent CoTs are converted back to VAE-latent space and greedily decoded with the VAE decoder.
  • A.1 HumanEval/63: fibfib: The fibfib task yields three passing answers using rolling-state dynamic programming, explicit tabulation, and recursive memoization.The decoded latent CoTs correspond to fixed-state storage, list-based tabulation, and cached recursion.
  • A.2 HumanEval/154: cycpattern_check: For cycpattern_check, different latent samples instantiate the same rotation-substring idea with different surface forms.Generated programs either loop over rotations with early return or materialize rotations and apply any.
  • A.3 HumanEval/89: encrypt: The encrypt task shows latent-controlled diversity between alphabet-index lookup and ASCII/modulo arithmetic for the same Caesar-style cipher.One solution uses an explicit alphabet string and index(), while another uses ord/chr modulo arithmetic.
  • A.3 HumanEval/89: encrypt: Across examples, decoded latent CoTs consistently reflect the generated program’s coarse algorithmic route despite token-level imperfections.The fibfib variants particularly support a distribution-over-solutions view rather than storage of one hidden solution trace.

B Experimental Setup … C Additional Analysis

NF-CoT is evaluated through dual-path and unified-path normalizing-flow variants sharing the same backbone, latent geometry, training schedule, and decoding pipeline. The setup also specifies causal latent inference, execution-guided reinforcement learning, and backbone-drift analysis.

  • B Experimental Setup: The dual-path and unified-path variants differ in whether the backbone is invoked twice or once per training step, while sharing core model and data configurations.The dual-path model separately computes flow-density and answer cross-entropy losses; the unified model computes both in a single causal forward.
  • B.1 Dual-Path Model / B.2 Unified-Path Model: The dual-path model uses separate flow and CE projectors with distinct boundary tokens, whereas the unified model fuses both paths into one causal sequence.The unified path computes NF NLL on latent positions and CE on answer positions from one backbone forward.
  • B.3 Training Hyperparameters: Training uses a frozen-backbone to joint-finetuning curriculum with a combined NF-NLL and teacher-forcing cross-entropy objective.The loss weights are λNF = λCE = 1, and Stage 1 trains shallow blocks before Stage 2 unfreezes the backbone and shallow blocks.
  • B.4 Inference Configuration: Inference generates N = 64 latent thoughts with an autoregressive NF reverse pass using KV-cache reuse, then decodes answers from a soft latent prefix.The two phases are latent generation followed by answer generation; standard evaluation draws 16 independent samples and reports pass@1.
  • B.5 Execution-Guided RL Details: Execution-guided RL updates only the shared backbone, treating the continuous latent trajectory and answer tokens as actions in a joint policy.The latent objective uses a sequence-level PPO surrogate, while answer tokens use a token-level PPO objective; rollouts sample n = 8 trajectories per prompt.
  • B.5 Execution-Guided RL Details: RL training runs for 150 steps with learning rate 3 × 10−6, βKL = 10−3, PPO clip range 0.2, and 8 rollouts per prompt.Rollouts use answer temperature 0.9 and latent temperature 1.0, while post-RL evaluation uses 16 samples per problem.
  • C Additional Analysis: The additional backbone-drift analysis finds a small overall drift difference, with substantially larger final-layer drift in the stage-2-only run.This comparison is made relative to the original Qwen3-8B-Base checkpoint.

C.1 Backbone Drift Diagnostic

The frozen-backbone warm-up produces a better-calibrated flow before joint training, reducing early NF-driven gradients and concentrating less additional drift in the final transformer layer. Without warm-up, global drift changes modestly, but the final layer moves substantially more.

  • Early training dynamics: The default warm-up starts stage 2 with LNF ≈−0.42 and log-determinant approximately −0.92, indicating a non-identity flow with a nontrivial latent-density model.The flow-side components are trained while the backbone remains frozen.
  • Early training dynamics: The stage-2-only run starts with LNF ≈0.47, near-zero log-determinant, and a larger initial gradient norm of 1.96 versus 0.96.These values indicate an almost-identity, poorly calibrated flow exposing the pretrained backbone to larger early gradients.
  • Final backbone drift: 1.6% larger global relative L2 distance and 3.3% greater global cosine drift show that stage-2-only training changes overall backbone drift only modestly.The comparison is between the stage-2-only model and the default curriculum.
  • Final backbone drift: 18% larger relative L2 drift occurs in the final transformer layer, while embedding, middle-layer, and embedding/head/normalization parameters remain nearly unchanged.This layer directly feeds both the LM and NF heads, linking its concentrated movement to early gradients from the uncalibrated NF branch.

C.2 Pairwise Output Diversity

NF-CoT yields more structurally diverse passing HumanEval solutions than Qwen3-8B-Base under matched sampling, as measured by a structure-aware AST metric. Its latent sampling changes algorithmic strategy before answer generation rather than merely varying wording or formatting.

  • Metric: Similarity uses parent-child AST-node-type 2-grams plus canonical-AST hashing after α-renaming variables and stripping docstrings/comments.This structure-aware metric captures control-flow and syntactic patterns while identifying truly identical canonical programs.
  • Results: 0.548 to 0.469: mean intra-prompt similarity falls from Qwen3-8B-Base to NF-CoT, a relative reduction of approximately 14%.The decrease is consistent across all four HumanEval tasks, with the largest difference on HumanEval/124.
  • Mechanism: Sampling latent z ∼ N(0, τz^2 I) steers decoding toward different algorithmic regions before answer tokens, changing solution strategy rather than surface realization.The diversity comparison uses the main evaluation configuration rather than analysis-specific tuning.
Loading 2606.06447v1…