LEPO: 대규모 언어 모델을 위한 잠재 추론 정책 최적화

LEPO: Latent Reasoning Policy Optimization for Large Language Models

최근, 대규모 언어 모델(LLM)에 잠재 추론이 도입되어 연속적인 공간 내의 풍부한 정보를 활용합니다. 그러나 확률적 샘플링이 없으면 이러한 방법은 필연적으로 결정론적 추론으로 수렴하여 다양한 추론 경로를 발견하는 데 실패합니다. 이러한 격차를 해소하기 위해, 우리는 Gumbel-Softmax를 통해 잠재 추론에 제어 가능한 확률성을 주입하여 LLM의 탐색 능력을 복원하고 강화 학습(RL)과의 호환성을 향상시킵니다. 이를 바탕으로, 우리는 RL을 직접적으로 연속적인 잠재 표현에 적용하는 새로운 프레임워크인 **LEPO (L**atent **E**asoning **P**olicy **O**ptimization)를 제안합니다. 구체적으로, LEPO는 롤아웃 단계에서 다양한 경로 샘플링을 가능하게 하기 위해 확률성을 유지하고, 최적화 단계에서는 잠재 표현과 이산적인 토큰 모두에 대한 통일된 그래디언트 추정 방법을 구축합니다. 광범위한 실험 결과, LEPO는 이산적 및 잠재적 추론을 위한 기존의 강화 학습 방법보다 훨씬 뛰어난 성능을 보이는 것을 확인했습니다.

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.