Reinforcement Learning for Latent-Space Thinking in LLMs
This work addresses inefficiencies in reasoning for LLM users, but it is incremental as it builds on existing latent-space methods without achieving superior performance.
The paper tackled the problem of inefficient discrete language-space reasoning in LLMs by exploring reinforcement learning for latent-space thinking, but found that RL-trained models still underperform traditional language-space CoT models in mathematical reasoning tasks.
Chain-of-Thought (CoT) reasoning typically utilizes the discrete language space for thinking, which is inherently inefficient, as many generated tokens only enforce linguistic rules that are not required for reasoning. To bypass this, latent-space thinking allows models to think using the continuous embedding space. While existing methods for training those models show domain-specific gains, they fail to maintain performance in complex tasks, such as mathematical reasoning. We experimentally demonstrate that the Coconut approach, a form of supervised fine-tuning for latent-space thinking, is highly sensitive to design choices and exhibits several inherent limitations. To address these issues, we investigate reinforcement learning (RL) techniques -- an underexplored direction in latent-space thinking -- including GRPO and design a novel Latent RL method for directly optimizing the latent thinking steps. Our experimental results reveal that these RL-trained models still lag behind traditional language-space CoT models in the mathematical reasoning domain. We make our codebase publicly available.