Technical Summary: Latent Thought Flow (LTF)

Problem Statement

Large Language Models (LLMs) increasingly rely on intermediate reasoning processes, such as Chain-of-Thought (CoT), to solve complex tasks. However, explicit CoT suffers from a linguistic space bottleneck: every intermediate thought must be decoded into discrete tokens, incurring substantial inference overhead and constraining internal computation to potentially redundant textual forms. This creates an inherent trade-off between accuracy and efficiency.

While latent reasoning (performing computation in continuous hidden-state space) offers a solution to this bottleneck, existing methods primarily learn deterministic paths or optimize for reward maximization under fixed budgets. These approaches lack a principled mechanism to allocate probability mass across reasoning trajectories that vary in both correctness and computational cost. Consequently, they often fail to optimize the latent reasoning manifold, either inheriting the verbosity of training rationales or collapsing into a few high-reward modes, failing to balance accuracy with efficiency.

Methodology: Latent Thought Flow (LTF)

The authors propose Latent Thought Flow (LTF), a framework that models reasoning as variable-length continuous trajectories in the LLM's latent space. LTF trains a sampler to match a reward-induced posterior that balances answer quality against computation cost.

Core Components

1. Variable-Length Latent Trajectories:
   Instead of deterministic tokens, LTF samples a trajectory $\tau = (z_{1:T}, \perp)$, where $z_t$ are continuous latent thought states and $\perp$ is an adaptive stopping decision. The model learns "how much" to reason internally before generating the final answer. The trajectory is sampled from a Gaussian policy $q_\phi$ parameterized by a latent reasoning head.

2. Reward-Proportional Target Distribution:
   LTF defines a utility function $R_{x,y}(\tau)$ combining answer quality ($V_{x,y}$) and computation cost ($C(\tau)$, typically trajectory length). The goal is to learn a sampler where the probability of a trajectory is proportional to its reward:
   $$p^*(\tau | x, y) \propto R_{x,y}(\tau)$$
   This encourages the model to concentrate probability mass on concise, accurate paths while suppressing redundant or erroneous ones.

3. Continuous GFlowNet Objective:
   To learn this distribution, LTF instantiates a continuous GFlowNet. Unlike Reinforcement Learning (RL) which maximizes expected reward (often leading to mode collapse), GFlowNets learn a stochastic policy where terminal samples are proportional to the reward, preserving diverse high-utility solutions.

   • Challenge: Latent thoughts are continuous, making discrete token-level balance equations inapplicable.
   • Solution: LTF adapts Subtrajectory Balance to continuous spaces. It computes flow analytically for every prefix state by allowing the model to stop at any intermediate step, decode an answer, and receive a reward.

4. Entropy-Weighted Subtrajectory Balance (EW-SubTB):
   To handle sparse answer supervision and guide exploration, LTF introduces an entropy-aware weighting mechanism. It reweights the Subtrajectory Balance residuals based on the differential entropy of latent subtrajectories. This ensures stronger supervision is applied to high-entropy regions (where information spread is richer) without altering the reward-proportional target distribution.

5. Reference-Prior Regularizer:
   To prevent random exploration in the continuous latent space from drifting into semantically meaningless states, LTF uses a reference prior derived from teacher-generated reasoning traces (when available). This anchors early exploration to meaningful latent regions, with the GFlowNet reward subsequently reshaping the sampler toward accuracy and efficiency.

Key Contributions

• Formulation of Latent Reasoning as a Flow: The paper reframes latent reasoning as a distribution-matching problem over variable-length continuous trajectories, rather than a deterministic path optimization.
• Continuous GFlowNet for LLMs: It extends GFlowNet theory to the high-dimensional continuous latent space of LLMs, enabling the preservation of diverse reasoning paths that are both accurate and efficient.
• Novel Training Objectives: The introduction of Entropy-Weighted Subtrajectory Balance and a Reference-Prior Regularizer addresses the specific challenges of sparse supervision and unstable exploration in continuous latent spaces.
• Adaptive Computation: The framework enables the model to dynamically allocate its computational budget, assigning shorter trajectories to easier problems and longer ones to harder problems.

Experimental Results

Experiments were conducted on multiple LLM backbones (LLaMA-3.2 1B/3B, LLaMA-3.1 8B, DeepSeek-R1 Distill-Qwen 1.5B) across math reasoning (GSM8K-Aug, ASDiv-Aug, DU, MATH) and transfer learning tasks.

• Performance vs. Explicit CoT: LTF significantly reduces reasoning length (decoding overhead) while maintaining or improving accuracy compared to explicit CoT methods.
• Performance vs. Latent Baselines: Compared to strong latent reasoning baselines like CoLaR and ReGuLaR:
  • Fine-tuning: LTF improved average accuracy by 12.9% while reducing reasoning length by 34.5%.
  • Transfer Learning: LTF improved average accuracy by 6.0% while reducing reasoning length by 19.9%.
  • Specific Gains: On GSM8K-Aug with LLaMA-3.2 1B, LTF improved accuracy from 34.58% (ReGuLaR) to 37.09% while reducing reasoning steps from 3.69 to 3.34.
• Extreme Compression: Under extreme compression settings (e.g., MATH dataset), LTF consistently outperformed baselines, demonstrating superior ability to preserve task-relevant semantic information.
• Ablation Studies: Results confirmed that Entropy-Weighted Subtrajectory Balance improves accuracy with minimal length increase, and that GFlowNet-based objectives provide better accuracy-efficiency trade-offs than standard RL objectives (like GRPO).

Significance and Claims

The paper claims that LTF establishes a new accuracy–efficiency frontier for reasoning in LLMs. By moving deliberation into a continuous latent space and using a reward-proportional sampling framework, LTF overcomes the linguistic bottleneck of explicit reasoning.

The authors argue that the missing distributional view in existing latent reasoning methods is central to their inefficiency. LTF provides a principled way to allocate probability across trajectories, allowing models to dynamically balance correctness and cost. The work demonstrates that meaningful intermediate computation need not be expressed in linguistic words, offering a scalable path for resource-constrained reasoning scenarios.

Limitations: The authors note that current experiments focus on textual tasks (words and symbols), and extending the framework to other modalities (vision, speech) remains future work. Theoretical generalization is also noted as an area for further exploration.