Latent Reasoning with Normalizing Flows

2026-06-04

AI Summary

This work addresses the limitations of existing language models that rely on discrete, sequential chain-of-thought (CoT) reasoning, which constrains both inference efficiency and representational flexibility, while current implicit reasoning approaches struggle to simultaneously support autoregressive generation, probabilistic sampling, KV cache compatibility, and tractable likelihood computation. To overcome these challenges, the authors propose NF-CoT, a novel framework that integrates invertible normalizing flows into the backbone of large language models. This enables compact, continuous-state modeling of intermediate reasoning steps while preserving the native autoregressive structure. NF-CoT supports exact likelihood estimation, permits probabilistic left-to-right decoding compatible with KV caching, and allows direct policy gradient optimization in the latent space. Experiments demonstrate that NF-CoT significantly outperforms both explicit CoT and existing implicit methods on code generation tasks, achieving higher pass rates and substantially reduced inference overhead.

Abstract

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.

Problem

• latent reasoning
• chain-of-thought
• normalizing flows
• large language models
• intermediate computation

Innovation

• latent reasoning
• normalizing flows
• chain-of-thought
• KV-cache decoding
• probabilistic decoding