ReLaX: Reasoning with Latent Exploration for Large Reasoning Models

1. Problem Statement

The paper addresses a critical bottleneck in training Large Reasoning Models (LRMs) using Reinforcement Learning with Verifiable Rewards (RLVR).

• Entropy Collapse & Over-Determinism: While RLVR improves reasoning, it often drives the policy distribution toward over-determinism. The policy prematurely converges to a narrow subspace, causing "entropy collapse." This results in ineffective exploration and suboptimal performance.
• Limitations of Token-Level Methods: Existing solutions attempt to mitigate this by increasing token-level entropy (e.g., reshaping rewards, entropy regularization, or anchoring high-entropy tokens). However, the authors argue these methods are superficial:
  • They conflict with the natural tendency of RL to minimize entropy.
  • They fail to capture the rich computational structure of multimodal models (MLLMs), where internal cross-modal computations are misaligned with unimodal text outputs.
  • They do not address the root cause: the loss of flexibility in the model's latent dynamics (the hidden state evolution).

2. Methodology: ReLaX Framework

The authors propose ReLaX (Reasoning with Latent eXploration), a framework that shifts the focus from token statistics to the latent dynamics of the model's hidden states.

A. Theoretical Foundation: Koopman Operator Theory

To make the nonlinear, high-dimensional latent dynamics of LLMs analytically tractable, ReLaX leverages Koopman operator theory.

• Linearization: The theory allows representing the nonlinear evolution of hidden states ($x_{t+1} = f(x_t)$) as a linear evolution in an infinite-dimensional space of observables.
• Neural Koopman Dictionary: The authors use ResKoopNet to learn a neural dictionary ($g(x) = \sigma(Wx)$) that maps hidden states into a linear representation space. This dictionary is learned once on initial policy data and then frozen.

B. Key Metric: Dynamic Spectral Dispersion (DSD)

Using the linearized representation, the authors introduce Dynamic Spectral Dispersion (DSD) as a metric to quantify policy exploration.

• Definition: DSD is defined as the variance of the magnitudes of the Koopman operator's eigenvalues ($\Lambda$) for a given trajectory:
  $$\text{DSD}(x) = \text{Var}(|\Lambda|)$$
• Interpretation:
  • Low DSD: Indicates a concentrated spectrum, implying degenerate, repetitive, and rigid internal computations (entropy collapse).
  • High DSD: Indicates a dispersed spectrum, signifying rich, heterogeneous, and flexible internal dynamics capable of supporting diverse reasoning trajectories.

C. Optimization Objective

ReLaX integrates DSD into the Group Relative Policy Optimization (GRPO) objective via a regularization term ($\mathcal{L}_{xp}$).

• Advantage-Shaped Regularization: To ensure exploration is beneficial, the DSD penalty is weighted by the positive part of the advantage ($\hat{A}_i$). This encourages flexibility only along trajectories that yield positive rewards.
  $$\tilde{\mathcal{L}}_{xp} = \log \left( \frac{1}{R} \sum_{i=1}^R \exp(-\text{clip}(\hat{A}_i, 0) \cdot \text{DSD}(x_i)) \right)$$
• Adaptive KL Regularization: To prevent instability from excessive dispersion, an adaptive KL divergence penalty is applied only to trajectories where DSD exceeds a threshold $\xi$.
• Final Objective:
  $$J(\theta) = J_{\text{GRPO}}(\theta) + \alpha \tilde{\mathcal{L}}_{xp} + \beta \sum_{i \in I} D_{\text{KL}}(\pi_\theta \parallel \pi_{\text{ref}})$$

3. Key Contributions

1. DSD Metric: A novel metric that captures the heterogeneity of latent reasoning dynamics, providing a deeper characterization of exploration than surface-level token entropy.
2. ReLaX Framework: A principled approach to regulate the exploration-exploitation trade-off by explicitly optimizing for flexible latent dynamics rather than token diversity.
3. State-of-the-Art Performance: Demonstrated significant improvements over existing token-level methods across diverse benchmarks.

4. Experimental Results

The authors evaluated ReLaX on 3B and 7B parameter models (both LLMs and MLLMs) across multiple benchmarks.

• Multimodal Reasoning (VLMs):

  • Tested on 7 benchmarks (e.g., MMMU, MathVista, EMMA).
  • ReLaX-VL-7B achieved a new SOTA average score of 53.2, surpassing the previous best (VL-Rethinker-7B at 52.5).
  • ReLaX-VL-3B outperformed several 7B-level models, demonstrating strong scaling efficiency.
  • Training dynamics showed ReLaX maintained higher entropy and DSD compared to vanilla GRPO, preventing the rapid collapse seen in baselines.

• Text-Only Mathematical Reasoning (LLMs):

  • Tested on 6 benchmarks (e.g., MATH500, AIME, AMC).
  • ReLaX-7B surpassed the previous SOTA (FR3E) by 4.3 points on Qwen2.5-7B-Base and 6.3 points on Qwen2.5-7B-Math.
  • On the AIME 2024 benchmark, ReLaX achieved a score of 36.9, significantly outperforming competitors.

• Generalization:

  • The method generalized well to different model families (Llama3.2, Qwen3) and scales.
  • Qualitative analysis showed ReLaX models performed more robust self-verification (using mathematical principles) compared to token-level methods that often hallucinated code execution or failed on visual variants.

5. Significance and Conclusion

• Paradigm Shift: ReLaX moves beyond "token-level" interventions, arguing that the true bottleneck in LRM scaling is the rigidity of internal latent computations. By optimizing the latent space directly, it achieves a more effective exploration-exploitation balance.
• Multimodal Superiority: The method is particularly effective for MLLMs, where token-level entropy fails to capture the complexity of cross-modal reasoning.
• Efficiency: The computational overhead is minimal. The Koopman dictionary is learned once (negligible cost), and the DSD calculation adds only ~10% to the per-step runtime.
• Future Direction: The paper establishes the latent space as a powerful frontier for advancing reasoning capabilities, suggesting that steering exploration via dynamic spectral properties is a more principled solution than heuristic token manipulation.