LEPO: \underline{L}atent R\underline{e}asoning \underline{P}olicy \underline{O}ptimization for Large Language~Models

2026-04-20

AI Summary

Existing latent reasoning approaches suffer from insufficient controllable stochasticity, leading to deterministic reasoning paths that hinder effective integration with reinforcement learning. This work proposes a Gumbel-Softmax–based latent reasoning framework that preserves diverse sampling during rollout while unifying gradient estimation for continuous latent representations and discrete tokens during optimization. For the first time, this enables direct reinforcement learning in continuous latent spaces. By restoring the exploration capability of large language models and establishing a unified gradient estimation mechanism, the proposed method significantly outperforms current discrete and latent reasoning reinforcement learning approaches across multiple experimental benchmarks.

Technology Category

• Machine Learning: Reinforcement Learning
• Natural Language Processing: (Large) Language Models
• Reasoning under Uncertainty: Sequential Decision Making

Application Category

• Search and Retrieval-Augmented AI: Large language models for search
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Recently, latent reasoning has been introduced into large language models (LLMs) to leverage rich information within a continuous space. However, without stochastic sampling, these methods inevitably collapse to deterministic inference, failing to discover diverse reasoning paths. To bridge the gap, we inject controllable stochasticity into latent reasoning via Gumbel-Softmax, restoring LLMs' exploratory capacity and enhancing their compatibility with Reinforcement Learning (RL). Building on this, we propose \textbf{\underline{L}}atent R\textbf{\underline{e}}asoning \textbf{\underline{P}}olicy \textbf{\underline{O}}ptimization~(\textbf{LEPO}), a novel framework that applies RL directly to continuous latent representations. Specifically, in rollout stage, LEPO maintains stochasticity to enable diverse trajectory sampling, while in optimization stage, LEPO constructs a unified gradient estimation for both latent representations and discrete tokens. Extensive experiments show that LEPO significantly outperforms existing RL methods for discrete and latent reasoning.

Problem

• latent reasoning
• stochasticity
• large language models
• reinforcement learning
• diverse reasoning paths

Innovation

• Latent Reasoning
• Reinforcement Learning
• Gumbel-Softmax
• Policy Optimization
• Large Language Models