PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Imitation: Reinforcement Learning for Active Latent Planning

Zhi Zheng, Wee Sun Lee

arXiv:2601.21598v1cs.AI

TL;DR

Existing latent-reasoning methods imitate arbitrary chain-of-thought labels, despite multiple equivalent solutions, motivating ATP-Latent’s smoother latent space and active reinforcement-learning planning. On LLaMA-1B, it achieves 4.1% more accuracy and 3.3% fewer tokens than advanced baselines on numerical reasoning benchmarks.

  • Problem

    Existing latent-reasoning methods imitate arbitrary labeled chain-of-thought trajectories, which can produce overfitted representations and biased policies despite multiple equivalent correct solutions.

  • Method

    ATP-Latent combines VAE-based latent-token training and stop generation with reinforcement learning using answer accuracy and VAE-decoded coherence rewards.

  • Results

    4.1% more accuracy and 3.3% fewer tokens on numerical reasoning benchmarks compared to advanced baselines.

  • Takeaways & Limitations

    ATP-Latent provides a smoother latent space and soft RL state constraints for pursuing more generalizable latent reasoning policies.

  • Takeaways & Limitations

    The approach assumes latent tokens occupy a small range of mutable values, motivating a very small β = 1e−3.

Abstract

from arXiv · show

Aiming at efficient and dense chain-of-thought (CoT) reasoning, latent reasoning methods fine-tune Large Language Models (LLMs) to substitute discrete language tokens with continuous latent tokens. These methods consume fewer tokens compared to the conventional language CoT reasoning and have the potential to plan in a dense latent space. However, current latent tokens are generally supervised based on imitating language labels. Considering that there can be multiple equivalent but diverse CoT labels for a question, passively imitating an arbitrary one may lead to inferior latent token representations and latent reasoning policies, undermining the potential planning ability and resulting in clear gaps between training and testing. In this work, we emphasize the importance of active planning over the representation space of latent tokens in achieving the optimal latent reasoning policy. So, we propose the \underline{A}c\underline{t}ive Latent \underline{P}lanning method (ATP-Latent), which models the supervision process of latent tokens as a conditional variational auto-encoder (VAE) to obtain a smoother latent space. Moreover, to facilitate the most reasonable latent reasoning policy, ATP-Latent conducts reinforcement learning (RL) with an auxiliary coherence reward, which is calculated based on the consistency between VAE-decoded contents of latent tokens, enabling a guided RL process. In experiments on LLaMA-1B, ATP-Latent demonstrates +4.1\% accuracy and -3.3\% tokens on four benchmarks compared to advanced baselines. Codes are available on https://github.com/zz1358m/ATP-Latent-master.

1. Introduction

The introduction motivates latent reasoning as a lower-latency alternative to lengthy language CoT, then proposes ATP-Latent to actively plan smoother latent representations and policies. Its VAE-guided RL uses coherence signals, while experiments report improved accuracy with fewer tokens.

  • Long CoT reasoning paths can hinder LLM applications requiring real-time responses.
  • Latent reasoning replaces sampled discrete language tokens with high-dimensional latent tokens to represent and reason in continuous latent space.This approach aims to reduce reasoning latency and avoid redundant natural-language tokens.
  • Existing methods imitate language CoT policies despite multiple equivalent correct CoTs for each question, potentially limiting generalizable latent reasoning.Their imitation strategies include curriculum learning, self-distillation, and compression, while still reducing token cost.
  • ATP-Latent introduces active latent-token planning through a two-stage method that trains a VAE and stop-head for interpretability and smoother representations.The VAE extends imitation-based SFT to construct a smoother latent token representation space.
  • ATP-Latent uses VAE-decoded-content consistency as an auxiliary RL reward, providing unsupervised training signals and soft constraints for latent planning.The supervised VAE decoder decodes latent tokens during RL to guide reasoning-path optimization.
  • 4.1% more accuracy and 3.3% fewer tokens are achieved by ATP-Latent with LLaMA-1B on numerical reasoning benchmarks compared to advanced baselines.

2. Preliminaries

The paper contrasts discrete language-based reasoning with latent CoT reasoning, where continuous vectors can reduce token cost and enable denser reasoning. It also identifies imitation of arbitrary CoT labels as a source of inferior latent policies and representations.

  • Language reasoning: Language reasoning generates discrete CoT tokens R followed by answer tokens A under an LLM next-token policy πθ.The question, reasoning, and answer sequences all belong to the discrete language-token domain.
  • Language reasoning: SFT maximizes accuracy across annotated CoT and answer positions, but its effectiveness depends heavily on difficult-to-obtain, high-quality CoT labels.RL methods instead sample multiple candidate CoTs and evaluate them using answer-quality rewards, such as relative advantages in GRPO.
  • Latent CoT reasoning: Latent CoT methods replace the original reasoning process with |L| continuous latent tokens, each represented as a d-dimensional real vector.These methods target efficient reasoning by moving from the pretrained language domain to a continuous latent space.
  • Latent CoT reasoning: Most latent CoT approaches fine-tune LLMs to imitate language labels, mapping latent tokens to the role of one or more language tokens.Prior work progressively substitutes language reasoning tokens with latent tokens using a curriculum-based language-to-latent transformation.
  • Latent CoT reasoning: Imitating an arbitrary correct CoT can produce inferior latent token representations and reasoning policies because one question may admit several correct CoTs, creating a training–testing gap.Latent reasoning can also remove redundant syntax or explanatory tokens, supporting dense CoT reasoning and reducing token cost.

3. Studies on Exploring Latent CoTs

Latent-CoT exploration methods apply Gaussian noise and GRPO to deterministic latent tokens, but unconstrained random exploration and nonsmooth latent spaces hinder reliable policy improvement. Experiments further show that latent-space supervision and initialization affect whether reinforcement learning activates useful exploration.

  • Gaussian-noise exploration: Current methods add Gaussian noise to deterministic latent tokens and optimize them with GRPO using Gaussian reparameterization.The perturbed token is defined as l̂_t = l_t + N(0, σ^2I_d).
  • Limitations: These methods show no clear primary-setting accuracy gains, likely because random unconstrained exploration fails to preserve a smooth latent state space.Both limitations make it difficult for reinforcement learning to discover a good latent reasoning policy.
  • Empirical comparison: Standard GRPO can significantly improve Coconut with appropriate noise, whereas SIM-CoT fails to show stable reinforcement-learning benefits despite stronger latent supervision and benchmark performance.The contrast indicates that better latent-space supervision alone does not guarantee stable gains from reinforcement learning.
  • Motivation for active planning: Equivalent language CoTs can encode different reasoning policies, so ATP-Latent actively optimizes latent policies using answer accuracy and latent-CoT coherence rewards.This contrasts with existing methods that imitate one of multiple equivalent language CoTs.
  • Empirical comparison: Running SIM-CoT on a pretrained Coconut before reinforcement learning activates its exploration capabilities, suggesting initialization influences latent-policy exploration.This result is presented as an explanation for the different reinforcement-learning behavior between Coconut and SIM-CoT.

4. Method: Active Latent Planning

ATP-Latent combines VAE-like supervised fine-tuning with stop-controlled latent generation and reinforcement learning to smooth latent representations and actively explore better reasoning policies. Its coherence reward uses consistency between decoded latent contents to guide answer-correct latent reasoning and support unsupervised RL.

  • SFT stage: ATP-Latent uses an SFT stage followed by RL, with the LLM reasoning pipeline as VAE encoder and an additional LLM as decoder.The decoder reveals latent tokens in language, while the method also enables a stop mechanism for smooth representations.
  • SFT stage: The encoder predicts Gaussian latent-token distributions instead of deterministic tokens, reducing overfitting to arbitrary language-label reasoning policies.An MLP latent head predicts each token’s mean µt and standard deviation σt from the question and preceding latent tokens.
  • SFT stage: ATP-Latent adds a stop head that learns whether each sampled latent token should continue or stop, improving information-density consistency across latent tokens.The stop policy is trained over the CoT label’s K steps and supports variable-length latent reasoning.
  • SFT stage: The decoder reconstructs the corresponding language-label spans with cross-entropy, outputs an empty string after the final stage, and jointly optimizes KL loss.The SFT objective uses β-VAE training, with β = 1e−3 to prioritize latent reasoning over decoding.
  • RL stage: ATP-Latent’s RL stage explores VAE-trained latent-token representations and uses decoded-content coherence as a soft reward for correct latent reasoning.For math reasoning, RCoh(L) measures whether equation results recur in later left-hand sides or the final answer; the reward also helps avoid question-to-answer shortcuts.
  • RL stage: ATP-Latent also supports unsupervised RL using only latent-CoT coherence and format rewards.This setting removes reliance on answer correctness while retaining coherence-based guidance over latent trajectories.

5. Experiment

Experiments evaluate ATP-Latent on LLaMA-3.2-1B-Instruct across in-domain and out-of-domain math benchmarks, comparing it with latent-reasoning and supervised-CoT baselines. ATP-Latent achieves the strongest reported accuracy-efficiency trade-off, and ablations support the contributions of its VAE, stop head, and RL stage.

  • Training & Testing Settings: Experiments use LLaMA-3.2-1B-Instruct, training on 385k GSM8K-Aug instances split 80% for SFT and 20% for RL.The dataset’s language CoTs are built with equation steps.
  • Training & Testing Settings: Evaluation covers GSM8K in-domain and GSM-Hard, SVAMP, and Multi-Arith out-of-domain benchmarks.Accuracy and generated-token counts jointly measure correctness and efficiency.
  • Baselines: Baselines include Coconut, SIM-CoT, iCoT, CoLaR, CoT-SFT, and Answer-SFT, with experiments run on 8× NVIDIA H200 GPUs.CoLaR uses compression factor 5 for similar token counts.
  • Main Results: 47.7% average accuracy and 8.4 tokens on average mark ATP-Latent’s strongest overall trade-off across four math reasoning benchmarks.Acc is higher-is-better answer accuracy, while #Token is lower-is-better average generated tokens.
  • Main Results: +4.1% accuracy and -3.3% tokens are achieved versus the reimplemented SIM-CoT result.The paper attributes these gains to active planning in a smooth latent token representation space and reports continued superiority against RL-boosted baselines in Appendix D.
  • Ablation Study: Removing the VAE or stop head causes average accuracy drops of -0.5% and -0.8%, respectively, validating both components.The ablations also include removing the RL stage after SFT.

6. Discussion

The discussion attributes ATP-Latent’s gains to coherence-guided reinforcement learning and VAE-based imitation, supported by smoother training, interpretable decoded reasoning, and stronger Pass@K performance. Coherence also correlates with correctness and improves GSM8K early in RL training.

  • Ablations on the Reward Function: Removing the coherence reward causes inferior in-domain performance, highlighting its role in reweighting positive samples toward suitable latent reasoning policies.The coherence reward is part of ATP-Latent’s final reward and guides reinforcement learning.
  • Ablations on the VAE Training: VAE training with a stop head yields smoother SFT validation-loss convergence than the SIM-CoT baseline using the Coconut backbone.This ablation models the imitation process as VAE training with a stop head.
  • The Validity of Coherence in Reward: 800 steps: coherence alone improves GSM8K performance during the early reinforcement-learning process.The study treats coherence as an unsupervised reward model for accuracy.
  • The Validity of Coherence in Reward: 0.26 average group-normalized coherence: correct examples; -0.33: incorrect examples.The coherence values are computed on ATP-Latent latent tokens after the SFT stage.
  • Examples for the Latent Reasoning Policy: RL produces more consistent latent reasoning steps and more accurate results, with improvements attributed to the coherence reward.The VAE decoder provides optional interpretability for latent tokens, while RL favors correct and coherent step-by-step policies.
  • Examples for the Latent Reasoning Policy: K=1 to 64: ATP-Latent achieves significantly higher Pass@K than baselines and ATP-Latent-SFT after RL training.The evaluation samples 64 runs with different L values and uses Gaussian noise with varying variances.

7. Conclusion … Prompt for ATP-Latent.

The paper concludes that ATP-Latent combines VAE-based latent-space smoothing with active planning and coherence-guided reinforcement learning, alongside related latent-reasoning approaches and specified prompting conventions. The surrounding sections distinguish latent-reasoning paradigms, RL-based methods, soft-thinking, vision-language applications, and ATP-Latent’s prompt format.

  • 7. Conclusion: ATP-Latent actively plans over a well-defined, smooth latent-token representation space instead of relying heavily on single annotated CoT trajectories.It uses a VAE framework and a stop-generation mechanism to make the latent space smoother and more expressive.
  • A. Related Work: Latent reasoning decouples reasoning from explicit natural language and includes token-wise auto-regressive methods that generate chains with dense latent embeddings.The related-work passage also identifies auxiliary strategies as a second broad category.
  • A.1. Latent Reasoning: ATP-Latent focuses on token-wise auto-regressive latent reasoning and builds step-by-step imitation as a VAE.This design targets the optimal latent reasoning policy through the VAE formulation described in the passage.
  • A.2. Latent Reasoning with Reinforcement Learning: Existing latent-reasoning RL methods use Gaussian noise and reparameterization to address deterministic latent tokens, missing randomness, and RL gradient problems.They search for an optimal latent reasoning policy within a given Gaussian space using batch data.
  • A.3. Soft-Thinking Reasoning: Soft-thinking conveys abstract concepts through probability-weighted sums of token embeddings and aims to improve information density.The passage states that soft-thinking can outperform discrete-token reasoning in many cases.
  • A.4. Latent Reasoning in Vision Language Model: Recent work applies latent reasoning to visual tasks, but none of the acknowledged methods adopts ATP-Latent’s coherence idea.Some of these vision-language methods also adopt the RL idea from Tan et al. (2025).
  • B. Prompt for ATP-Latent: The ATP-Latent prompt contains only the question, followed by an instruction to think step by step and output the final answer after ### <latent>.The format reward RFormat is 1 when ### is successfully output after <latent>, and 0 otherwise.
  • Prompt for ATP-Latent.: The prompt examples define latent, CoT-SFT, and Answer-SFT formats, while generated-token counts exclude <latent> and < / latent>.The latent example places the answer after the closing latent marker, whereas CoT-SFT includes language CoT.

C. Detailed parameters

ATP-Latent uses LLaMA-3.2-1B-Instruct for both its reasoning model and decoder, with GELU-activated two-layer MLPs for the stop and latent heads.

  • C. Detailed parameters: ATP-Latent uses the LLaMA-3.2-1B-Instruct LLM for both the reasoning LLM θ and decoder ϕ.
  • C. Detailed parameters: Both the stop head and latent head are two-layer MLPs using GELU activation.

D. Supplementary Experiments … E.1. Baselines

The supplementary experiments show that ATP-Latent benefits substantially from RL, avoids latent-stage instability, and develops more diverse reasoning paths. The baseline section details the compared supervised and latent-reasoning systems and their reproduction setup.

  • D. Supplementary Experiments: The experiments compare latent reasoning methods and RL fine-tuning using accuracy (Acc) and average reasoning length (#Token).
  • D.1. Comparison with Baselines + RL: After RL tuning, ATP-Latent achieves the highest overall average accuracy, 47.7%, outperforming all baselines.RL provides especially substantial gains for ATP-Latent compared with CoLaR, Coconut, and Coconut + SIM-CoT.
  • D.2. Comparison on the Self-Extending Ability: Coconut’s latent instability prevents models trained at a maximum stage of 3 from reliably extending to larger stages.SIM-CoT can mitigate performance degradation when generalizing to larger maximum numbers of latent tokens.
  • D.2. Comparison on the Self-Extending Ability: ATP-Latent’s stop mechanism prevents performance drops as the maximum number of latent stages increases, including for ATP-Latent-SFT.This behavior is reported on GSM8K across different numbers of latent stages.
  • D.3. Case Study: Exploring Different Ways: RL training enables ATP-Latent-SFT to develop new patterns and diverse reasoning paths, improving Pass@K.Decoded latent tokens reveal wider varieties of valid reasoning-path branches for the same problem after RL training.
  • E.1. Baselines: The baseline suite includes CoT-SFT, Answer-SFT, CoLaR, iCoT, Coconut, and SIM-CoT on LLaMA-3.2-1B-Instruct.
  • E.1. Baselines: The supervised baselines are implemented with the LLaMA-Factory framework, using prompts specified in Appendix B.
  • E.1. Baselines: All latent-reasoning baselines are fine-tuned from the base LLM using their provided code, and reproduction results in Table 1 generally match reported results.This setup avoids reproducing potentially complex stages such as SIM-CoT’s SFT warm-up.

E.2. Datasets · E.3. Licenses

The evaluation uses GSM8K for in-domain testing and GSM-Hard, SVAMP, and Multi-Arith for out-of-domain testing. Licenses for the base LLM, datasets, and frameworks are listed in Table 5.

  • E.2. Datasets: GSM8K is used as the in-domain evaluation test set.
  • E.2. Datasets: GSM-Hard is included in the out-of-domain evaluation.
  • E.2. Datasets: SVAMP is included in the out-of-domain evaluation.
  • E.2. Datasets: Multi-Arith is included in the out-of-domain evaluation.
  • E.2. Datasets: The paper follows Wei et al. (2025) in adopting the GSM8K test set.
  • E.3. Licenses: Licenses for the base LLM, datasets, and frameworks are listed in Table 5.
Loading 2601.21598v1…