Reinforcement Learning for Latent-Space Thinking in LLMs

TL;DR

This paper explores reinforcement learning methods to improve latent-space thinking in large language models, aiming to enhance reasoning capabilities beyond traditional language-space approaches.

Contribution

It introduces a novel Latent RL method for optimizing latent thinking steps and evaluates its effectiveness compared to existing approaches.

Findings

Latent RL models still underperform compared to language-space CoT models in math reasoning.

Coconut fine-tuning is highly sensitive to design choices and has inherent limitations.

Reinforcement learning offers a promising but currently limited avenue for latent-space reasoning improvement.

Abstract

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.