Latent Reasoning with Normalizing Flows

Abstract: LLMs 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 LLMs, 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.

• The paper introduces a novel framework, integrating autoregressive normalizing flows into LLMs to enable continuous and tractable chain-of-thought reasoning.
• It achieves state-of-the-art code generation performance with significant improvements in pass@1 metrics and reduced computational costs compared to diffusion-based methods.
• The unified training approach ensures consistency between latent and output distributions, facilitating effective policy-gradient reinforcement learning directly on the latent space.

Latent Reasoning with Normalizing Flows: A Technical Exposition

Overview

The paper "Latent Reasoning with Normalizing Flows" (2606.06447) introduces a novel latent reasoning framework, denoted here as \methodname{}, that integrates autoregressive normalizing flows (NFs) within LLMs to facilitate continuous and tractable chain-of-thought (CoT) reasoning. This approach addresses the inefficiencies of explicit text-based CoT by compactly representing intermediate reasoning steps in a continuous latent space, while crucially preserving the core computational interfaces of language modeling: left-to-right generation, KV-cache compatibility, likelihood scoring, and probabilistic sampling.

Figure 1: Training and inference pipeline of \methodname{}. Latent states are distilled from explicit chain-of-thought traces and integrated within a single causal LLM stream via a normalizing flow head.

Motivation and Background

Standard explicit CoT methods encode intermediate reasoning as serial, discrete token sequences, which are verbose and inflexible for non-linguistic or partial computations. Latent reasoning, using continuous or embedding-based representations, promises higher semantic bandwidth and efficiency. However, extant latent CoT techniques, including deterministic state recycling and diffusion models, face critical limitations: they either collapse to a single solution mode, lack stochasticity, require iterative (and slow) denoising, or decouple latent reasoning from the autoregressive interfaces intrinsic to LLMs. Notably, latent diffusion approaches such as LaDiR (Kang et al., 6 Oct 2025) offer stochastic latent generation but incur heavy computational overhead and make likelihood estimation intractable.

Methodology

Latent CoT with Normalizing Flows

\methodname{} connects explicit CoT traces (text) and continuous latent thought trajectories through a frozen VAE encoder, which produces information-rich continuous codes. These codes are mapped to an LLM-facing latent space via shallow, invertible autoregressive flows. The LLM backbone then models the distribution over continuous thought vectors with a deep autoregressive Gaussian flow, parameterized to allow left-to-right, tractable sampling (see Eq. 3.1 in the paper):

$$p_\theta(u_{1:K} \mid q) = \prod_{i=1}^K \mathcal{N}\left(u_i; \mu_\theta(q, u_{<i}), \operatorname{diag}(\sigma_\theta^2(q, u_{<i}))\right)$$

At each position, either an NF head (for latent slots) or the standard LM head (for answer tokens) reads from the same causal stream, enabling efficient hybrid continuous/discrete generation.

Unified Training and End-to-End Likelihood

A two-stage curriculum first aligns the latent reparameterization with the frozen LLM backbone, then unfreezes all parameters for end-to-end optimization. The objective is the sum of the normalizing flow negative log-likelihood (for the latent thoughts) and the cross-entropy loss (for answer tokens), computed in a unified stream:

$$\mathcal{L}_{\mathrm{sup}} = \lambda_{\mathrm{flow}}\, \mathcal{L}_{\mathrm{flow}} + \lambda_{\mathrm{text}}\, \mathcal{L}_{\mathrm{text}}$$

This setup ensures that continuous reasoning and answer decoding are trained on the exact trajectories used at inference, avoiding interface mismatch.

Inference and Policy Optimization

During inference, \methodname{} samples continuous latent trajectories in a single forward pass, then decodes answers conditionally with the same KV cache. Iterative denoising is eliminated, substantially reducing computational cost relative to diffusion-based latent methods. Furthermore, the explicit density over latent trajectories enables effective policy-gradient reinforcement learning (RL) directly on the latent space, allowing for reward-driven refinement of reasoning strategies.

Figure 2: Pass@k scaling on MBPP+ and HumanEval+. \methodname{} consistently outperforms both the base and diffusion-based baselines as k increases, indicating more effective coverage over solution space.

Empirical Results

Code Generation Performance

On standard Python code-generation benchmarks (MBPP, MBPP+, HumanEval, HumanEval+, LiveCodeBench v6), \methodname{} delivers strong empirical gains:

• Pass@1 jump: On Qwen3-8B-Base, \methodname{} (Unified) improves pass@1 from 55.8 to 68.8 (+13.0).
• Outperforms latent and explicit baselines: Surpassing LaDiR (diffusion latent) and Soft Thinking (embedding mixture) by 7–13 points on average. It even slightly outperforms top open-sourced AR models such as OlympicCoder on average pass@1.
• Efficiency: Autoregressive latent flow sampling is up to 2.7x faster and 2.48x less computationally expensive per-sample than diffusion models (Table 3 in the paper), with higher throughput during training.

Sample Efficiency and Solution Diversity

The flow-based latent reasoning enables substantial gains in sample efficiency. On MBPP+ and HumanEval+, the pass@1 of \methodname{} at $k=1$ matches or exceeds the base model's performance at $k=128$, with continued improvement as k increases—demonstrating that the sampled latent trajectories yield genuinely diverse, functional solutions rather than superficial textual paraphrases.

Figure 3: Pass@k diversity before and after RL on MBPP+ and HumanEval+. Whereas token-space RL saturates with larger k, latent-space RL via NF-CoT preserves sample diversity and upward scaling.

Moreover, structural similarity metrics (AST-based) across multiple sampled correct solutions confirm that \methodname{} increases true solution diversity, exploring distinct algorithmic strategies rather than only rewording a single correct plan.

Figure 4: Pairwise output similarity among passing HumanEval programs. Lower values under \methodname{} indicate higher structural diversity within the solution set.

Robustness in Latent Space

Perturbation experiments reveal that moderate noise in the continuous thought trajectory almost never changes overall pass@1 but dramatically alters surface program form—demonstrating the model's robustness and the smoothness of the learnt latent reasoning manifold.

Figure 5: As perturbation strength increases, solutions remain functionally correct despite lower textual overlap, highlighting latent space smoothness.

Reinforcement Learning Compatibility

Policy-gradient RL directly in the continuous latent space further raises average pass@1 (from 68.8 to 70.1) without sacrificing the diversity expected for optimal pass@k scaling. Unlike token-space RL, which risks mode collapse, RL in the NF-CoT latent space refines reasoning strategies while preserving a broad support over solution modes.

Architectural and Training Analysis

Comprehensive ablations confirm the necessity of the two-stage curriculum—removing the warm-up stage degrades benchmark performance due to misaligned early gradients. Unified causal modeling (single stream for NF and LM objectives) outperforms dual-path variants by maintaining a coherent interface for both training and inference, eliminating mismatches between latent and answer-conditioning spaces.

Theoretical and Practical Implications

By placing the normalizing flow inside the causal LLM architecture, \methodname{} offers a compelling bridge between the expressivity and efficiency of continuous latent reasoning and the essential probabilistic, left-to-right, and likelihood-based characteristics of modern Transformer LMs. This interface enables both tractable probabilistic inference and seamless integration of policy optimization, eliminating classic barriers (like intractable likelihoods, slow iterative sampling, KV-cache incompatibility) present in other latent reasoning frameworks.

Practically, the reduction in intermediate reasoning cost paired with enhanced solution diversity makes \methodname{} well-suited for deployments where computational budget and inference latency are constraints, such as code generation assistants, test-time ensemble methods, and execution-guided RL loops.

Theoretically, the exact likelihood characterization over latent reasoning trajectories opens analytic pathways for measuring, regularizing, and controlling reasoning strategies within LLMs, laying the groundwork for more interpretable, robust, and verifiable latent reasoning systems.

Future Directions

The current study focuses on code reasoning tasks; generalizing \methodname{} to other complex, multi-step reasoning domains (such as mathematical proof, multimodal tasks, or open-domain dialogue) remains a promising line of investigation. Moreover, adaptive or dynamically allocated latent budgets, as well as unsupervised or semi-supervised latent reasoning objective variants, could further improve flexibility and coverage. Closing the gap between “probed” (decoded) latent thoughts and human-interpretable explanations is also an open challenge.

Conclusion

\methodname{} establishes a tractable and efficient approach for latent reasoning in LLMs, leveraging normalizing flows to jointly model and optimize continuous chains-of-thought within the standard language modeling pipeline. The framework achieves state-of-the-art latent CoT performance, maintains efficiency, greatly enhances solution diversity, and supports direct policy-gradient reinforcement learning—highlighting the utility of likelihood-based latent reasoning as a practical and theoretically attractive paradigm for next-generation LLMs.