CTRLS: Chain-of-Thought Reasoning via Latent State-Transition

Chain-of-thought (CoT) reasoning enables large language models (LLMs) to break down complex problems into interpretable intermediate steps, significantly enhancing model transparency and performance in reasoning tasks. However, conventional CoT methods rely on heuristic sampling without structured modeling of reasoning transitions, constraining their ability to systematically explore and discover diverse and effective reasoning trajectories. In this work, we introduce CTRLS, a framework that formulates CoT reasoning as a Markov decision process (MDP) with latent state transitions, enabling principled and state-aware exploration via distributional reinforcement learning. By modelling reasoning actions as explicit probability distributions in latent space, our approach explicitly models epistemic uncertainty, facilitating robust exploration of the reasoning space. As part of our framework, we introduce an on-policy reinforcement learning strategy incorporating epsilon-greedy exploration and entropy-based regularization to iteratively refine latent state transitions without requiring additional fine-tuning of the underlying LLM. Theoretical analyses provide evidence lower bounds (ELBO), theoretically grounding our transition-aware modeling of latent reasoning dynamics. Further experiments demonstrate improvements in reasoning accuracy, diversity, and exploration efficiency across benchmark reasoning tasks.

💡 Research Summary

The paper introduces CTRLS, a novel framework that reconceptualizes chain‑of‑thought (CoT) reasoning as a Markov decision process (MDP) operating in a latent semantic space. Instead of generating each reasoning step token‑by‑token in an autoregressive fashion, CTRLS first abstracts the current prompt and all previously generated steps into a continuous latent state sₜ. A learned transition model Pθ(sₜ₊₁ | sₜ, aₜ) governs how the reasoning process moves from one latent state to the next, while the “action” aₜ is not a discrete token but a probability distribution over possible next reasoning fragments. The authors model this action distribution with a Dirichlet (or similar) distribution, thereby explicitly representing epistemic uncertainty about which logical step to take next.

The framework consists of three tightly coupled components: (1) a stochastic encoder ρϕ that maps the prompt and intermediate steps into latent states via a variational posterior Qϕ(z₁:T | x₁:T); (2) a state‑conditioned UNet adapter that injects the latent representation into the token embeddings of a frozen large language model (LLM) and generates the next textual fragment conditioned on the sampled action; (3) a policy network πθ that, given the current latent state, outputs the parameters of the Dirichlet action distribution. Training proceeds in two phases. First, an off‑line variational pre‑training optimizes a unified evidence lower bound (ELBO) that jointly learns the encoder, transition dynamics, and policy. Second, an on‑policy reinforcement‑learning stage uses the sparse terminal reward R(τ) = 1{y = y*} (correct answer) to fine‑tune πθ via policy‑gradient methods. To encourage systematic exploration, the authors combine ε‑greedy sampling with an entropy regularizer, preventing premature convergence to a single reasoning trajectory.

By treating actions as distributions, CTRLS leverages distributional reinforcement learning (DRL). The distributional Bellman operator T_Z propagates not only expected returns but full return distributions, allowing the agent to reason about the variability of future rewards. This yields richer exploratory behavior: the policy can maintain multiple plausible reasoning paths early on and gradually shift probability mass toward higher‑reward trajectories as learning progresses.

Empirical evaluation on standard reasoning benchmarks (GSM8K, MathQA, MATH) demonstrates that CTRLS outperforms strong baselines such as zero‑shot CoT, Self‑Consistency, and recent CoT‑Self‑Verification methods. Accuracy improvements range from 2 to 5 percentage points, with the most pronounced gains on harder problems. Moreover, the framework generates a larger set of diverse reasoning trajectories, as measured by n‑gram diversity and the number of distinct solution paths per prompt. Human evaluation confirms that the additional trajectories provide meaningful alternative explanations, enhancing interpretability.

A key practical advantage is that the underlying LLM remains frozen; only the latent encoder, transition model, and policy are updated. This means that existing large models can be retrofitted with structured, controllable reasoning capabilities without costly full‑model fine‑tuning. The paper also discusses limitations, such as the reliance on a first‑order Markov assumption for latent states and the need for more sophisticated reward shaping to handle multi‑step feedback.

In summary, CTRLS offers a principled, state‑aware formulation of CoT reasoning that integrates variational latent modeling with distributional reinforcement learning. By explicitly modeling uncertainty in reasoning transitions and enabling systematic exploration, it advances both the performance and explainability of LLM‑based problem solving. Future work may extend the approach to multi‑modal reasoning, tool‑use integration, and human‑in‑the‑loop reward design.