Beyond Imitation: Reinforcement Learning for Active Latent Planning

TL;DR

This work tackles the problem of suboptimal latent reasoning policies arising from imitation-based latent CoT methods that map diverse, equivalent CoTs to a single latent trajectory. It introduces ATP-Latent, a two-stage framework that first uses a variational auto-encoder to create a smooth latent token space with an optional stop mechanism, and then applies reinforcement learning guided by a coherence-based reward derived from VAE-decoded contents to actively plan latent reasoning. The approach yields superior accuracy with substantially fewer latent tokens across multiple math benchmarks on LLaMA-3.2-1B-Instruct, demonstrating that active latent planning can close training-testing gaps and promote diverse, robust reasoning strategies. The results suggest practical gains in reasoning efficiency and generalizability, with the coherence reward providing a soft, unsupervised signal to constrain latent planning.

Abstract

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.

Figures (8)

• Figure 1: Equivalent Language CoTs may lead to different latent reasoning policies. Existing latent reasoning methods (b) imitate one of them, leading to suboptimal policies. The proposed ATP-Latent method (c) actively optimizes the latent reasoning policies in a well-defined space, employing both the verifiable accuracy of answers and the coherence of decoded latent CoTs as rewards.
• Figure 2: Equivalent Language CoTs represent different reasoning policies. Existing latent reasoning methods (b) imitate one out of them, leading to suboptimal latent policies. The proposed ATP-Latent method (c) actively optimizes latent policies, employing both the accuracy of answers and the coherence of latent CoTs as reward.
• Figure 3: GRPO validation curve over Coconut and SIM-CoT and noises in different scales.
• Figure 4: GRPO validation curve of SIM-CoT finetuned after Coconut training.
• Figure 5: Ablation study and case studies on ATP-Latent. (a) exhibits the validation accuracy curve during RL on GSM8K. (b) shows the loss curve during the SFT stage on the validation set of GSM8K-Aug (500 instances). (c) presents the relationship between the correctness of questions on the training dataset (RL split) and the coherence reward before the RL training in ATP-Latent.