---
title: Latent-Variable Inference for Chain-of-Thought
url: https://www.emergentmind.com/topics/training-chain-of-thought-via-latent-variable-inference
type: topic
---

# Latent-Variable Inference for Chain-of-Thought

Training chain-of-thought (CoT) via latent-variable inference refers to the probabilistic modeling paradigm in which the intermediate reasoning processes of large language models (LLMs) are formalized as latent variables. This approach enables marginalization over unobserved rationales, principled amortized inference, and a range of EM, variational, and RL-based fine-tuning strategies for both language and vision-language models. Below, key theoretical and methodological facets are detailed, with representative models and empirical results.

## 1. Probabilistic Foundations: CoT as Latent-Variable Modeling

The central premise is to model the answer $y$ to an input $x$ as being mediated by a stochastic reasoning chain or rationale $z$. Formally, the generative model factorizes as
\[
p_\theta(y, z \mid x) = p_\theta(z \mid x) \, p_\theta(y \mid x, z)
\]
with $z$ ranging over all possible textual or structural rationales consistent with $x$. In practice, the log-marginal likelihood objective,
\[
\mathcal{L}(\theta) = \log p_\theta(y \mid x) = \log \sum_z p_\theta(z \mid x) p_\theta(y \mid x, z),
\]
is intractable, prompting the use of approximate inference—typically expectation-maximization (EM), variational inference, or their amortized, RL-augmented, or multi-sample variants [2312.02179][2503.19618][2312.04684][2510.23925][2505.02391][2507.08182][2509.25020].

## 2. Training Algorithms: EM, Variational, and RL Approaches

Latent-variable CoT training bifurcates into several algorithmic strategies:

### EM-based Algorithms

Classic EM-style solutions, as in TRICE [2312.02179], maintain per-example latent rationale chains and alternate between:

- **E-step**: Approximate $p_\theta(z \mid x, y)$ via MCMC (e.g., Metropolis–Hastings using $p_\theta(z \mid x)$ as proposal),
- **M-step**: Update $\theta$ using expected gradients under the imputed $z$.

Variance reduction is achieved by control-variate techniques, with the variance of the gradient estimator provably vanishing as the model improves [2312.02179].

### Variational Inference

Variational techniques introduce a tractable inference distribution $q_\phi(z \mid x, y)$ to lower-bound $\log p_\theta(y \mid x)$:
\[
\mathcal{L}_{\rm ELBO} = \mathbb{E}_{q_\phi(z \mid x, y)} \left[\log p_\theta(z, y \mid x) - \log q_\phi(z \mid x, y)\right].
\]
JEPO [2503.19618] simplifies this to Jensen's lower bound by setting $q_\phi(z \mid x, y) = p_\theta(z \mid x)$, yielding a pragmatic, scalable estimator without auxiliary inference nets. GFlowNet-based amortized inference for sequential or multimodal CoT is also adopted [2510.23925].

### Distributional and Reinforcement-Learning Extensions

CoT is further cast as a Markov Decision Process (MDP) with latent states (CTRLS) [2507.08182], or as a Markov chain of continuous thoughts (MARCOS) [2509.25020]. On-policy RL and distributional RL (e.g., entropy-regularized policy gradient, epsilon-greedy exploration) are used to refine the latent transition dynamics, supporting diverse trajectory exploration and robust reasoning.

## 3. Model Architectures and Inference

### Amortized Encoders and Latent-Skill Modeling

In LaRS [2312.04684], rationales are embedded into a latent skill space $\mathbf{z} \in \mathbb{R}^d$ via a conditional variational autoencoder (CVAE):
- **Encoder** $q_\omega(\mathbf{z} \mid Q, R)$ maps question-rationale pairs to latent skill vectors.
- **Prior/Policy** $\pi_\phi(\mathbf{z} \mid Q)$ predicts the skill distribution conditioned on the new input.
- **Decoder** $p_\psi(R \mid \mathbf{z}, Q)$ reconstructs rationales from (Q, z).

At inference, latent skills inferred from the test question steer example selection by matching latent skills to the example bank (cosine similarity), leading to highly efficient and theory-grounded in-context learning pipelines.

### Latent Markov and Continuous-State Models

MARCOS [2509.25020] employs continuous latent thoughts $\mathbf{z}_{1:T}$ transitioned via a “Thinker” transformer, with each emission of observable CoT steps mediated by a “Speaker” model. The approximate posterior is parameterized as $q_\phi(R_t \mid x_t)$, with training consisting of both local pretraining (omitting prior KL) and global ELBO maximization.

Parallel architectures, such as PCCoT [2506.18582], further decouple sequential dependency via Jacobi iteration, yielding order-of-magnitude efficiency gains while maintaining accuracy.

### Manifold and Quality-Based Latent Steering

GeoSteer [2601.10229] maps hidden states of CoT segments into a low-dimensional VAE manifold, learns a differentiable segment-quality function, and applies latent gradient-based steering at each inference step. This approach aligns internal activations with regions of high-quality reasoning, preserving both geometric coherence and quantitative answer accuracy.

## 4. Evidence Lower Bounds, Gradient Estimation, and Variance Reduction

A central technical feature is the use of evidence lower bounds (ELBOs) and their optimization via stochastic gradient estimators. For instance:

- **EM/ELBO Gradient**: The key estimator computes
\[
\nabla_\theta \log p_\theta(y \mid x) = \mathbb{E}_{z \sim p_\theta(z \mid x, y)} \left[ \nabla_\theta \log p_\theta(z \mid x) \right],
\]
with practical algorithms approximating this via MCMC or variational sampling [2312.02179][2503.19618].
- **Variance Reduction**: Control-variate (score function) methods use baseline terms (e.g., acceptance rates or leave-one-out averages) to minimize estimator variance [2312.02179][2505.02391].
- **Dynamic Sampling**: GVM-RAFT [2505.02391] allocates sampling budgets adaptively using gradient-norm and acceptance-rate analysis, provably accelerating convergence for gradient-based EM and RL.

## 5. Inference-Time Strategies and Prompt Selection

In models such as LaRS [2312.04684], the amortized inference network assigns latent skill vectors to queries and to banked examples, enabling similarity-based retrieval of supporting rationales without laborious hand-labeling or repeated LLM calls. Bayesian inference-scaling and best-of-$n$ marginal likelihood approximations further enable cost-effective, robust answer selection in multimodal and autoregressive settings [2510.23925].

GeoSteer [2601.10229] demonstrates that steering model activations in the latent manifold space increases both the coherence of intermediate steps and the exact-match metric on math reasoning benchmarks.

## 6. Empirical Results and Applications

A variety of datasets and models have validated these latent-variable techniques:

| Model         | Dataset(s)            | Relative Gains                     | Notable Features                    |
|---------------|----------------------|------------------------------------|--------------------------------------|
| TRICE         | GSM8K, BBH           | +5–10 pt over baselines            | MCMC-EM, CV variance reduction       |
| LaRS          | TabMWP, GSM8K, Spider, COGS | Up to +6 pt over non-latent retrieval | CVAE skill modeling, fast ICL        |
| MARCOS        | GSM8K, SVAMP, MultiArith   | +4.7 pt over strong SFT, 15× speedup | Continuous Markov chain of thoughts  |
| GeoSteer      | GSM8k                | +2.6 pt EM, +5.3 pairwise win      | VAE manifold steering                |
| GVM-RAFT      | Math500, Minerva     | 2–4× faster, +2–3 pt accuracy      | Dynamic sample allocation            |
| CTRLS         | GSM8K, MATH          | +10% in exploration/sample efficiency | Transition-aware latent RL         |

*All performance claims verbatim from cited sources; actual statistical confidence intervals can be found in the referenced works.*

## 7. Significance and Outlook

Latent-variable-based CoT training provides a rigorous, flexible foundation for reasoning in LLMs, supporting principled marginalization, efficient fine-tuning absent gold rationales, and robust exploration of diverse reasoning trajectories. The framework accommodates a wide spectrum of practical instantiations—EM-MCMC, amortized inference, variational and GFlowNet losses, continuous-state Markov models, manifold steering, and on-policy RL. Recent advances resolve sampling and variance bottlenecks and yield statistically significant improvements in accuracy and efficiency across math, text, SQL, and multimodal reasoning benchmarks [2312.02179][2312.04684][2509.25020][2601.10229][2505.02391][2510.23925][2507.08182][2503.19618][2506.18582].

A plausible implication is the emergence of new algorithmic paradigms unifying planning, inference, and learning across symbolic, continuous, and black-box latent spaces—potentially facilitating further advances in scalable, interpretable, and controllable multi-step reasoning under uncertainty.

Source: https://www.emergentmind.com/topics/training-chain-of-thought-via-latent-variable-inference