Technical Summary: Dropout-GRPO for Continuous Latent Reasoning

1. Problem Statement

The paper addresses a fundamental structural incompatibility between Group Relative Policy Optimization (GRPO) and continuous latent reasoning models (specifically the COCONUT architecture).

• The GRPO Requirement: GRPO relies on the diversity of $K$ rollouts within a group to compute a non-zero group-mean advantage ($A^{(k)} = r^{(k)} - \mu_r$). If all rollouts produce identical trajectories, the advantage collapses to zero, stalling optimization.
• The Latent Reasoning Challenge: Unlike standard Chain-of-Thought (CoT) models that derive diversity from stochastic token sampling, continuous latent reasoning models (like COCONUT) operate via deterministic hidden state recurrence. Given a prompt and fixed parameters, multiple rollouts yield identical latent trajectories and answers.
• The Consequence: Applying standard GRPO to these models results in zero variance ($\sigma_r = 0$), causing the policy gradient to vanish and preventing any learning signal from being generated.

2. Methodology: Dropout-GRPO

The authors propose Dropout-GRPO, a method that injects necessary stochasticity into the latent reasoning process through structured dropout while maintaining the theoretical guarantees of variational inference.

Core Mechanism: Shared-Mask Variational Dropout

Instead of standard token-level sampling, the method introduces stochasticity at the parameter level:

1. Mask Generation: For each rollout $k$, a single Bernoulli mask $\xi^{(k)}$ is drawn.
2. Consistency: This mask is held constant across all $T$ latent recurrence steps within that specific rollout.
3. Perturbation: The mask induces a structured perturbation $\tilde{\theta}(\xi)$ to the network weights. The latent recurrence becomes $h_t = f_{\tilde{\theta}(\xi)}(h_{t-1})$.
4. Mask Replay: Crucially, the RNG seed for the mask is stored during the rollout phase. During the policy update, the same mask is regenerated and applied. This ensures that the trajectory evaluated for the reward is bit-identical to the trajectory used for gradient computation when $\theta = \theta_{old}$.

Theoretical Interpretation

• Bayesian Model Average: Following Gal and Ghahramani, the shared mask treats each rollout as a single posterior sample from a structured distribution over parameters. The marginal policy $\bar{\pi}_\theta$ is interpreted as a Bayesian model average.
• Gradient Unbiasedness: The authors prove that the surrogate gradient using the mean-only advantage ($A^{(k)} = r^{(k)} - \mu_r$) is an unbiased estimator of the gradient of the expected reward of the marginal policy, scaled by $(1 - 1/K)$.
• Variance Reduction: By reusing the same mask (Common Random Numbers) during the update as was used during the rollout, the method reduces variance compared to resampling a fresh mask at update time.

Implementation Refinements

To stabilize training, four specific refinements are introduced:

1. Huberized $k_3$ KL: Replaces standard clamping in the KL divergence estimator with a Huber-style extension to prevent gradients from vanishing for tokens with large log-ratios, avoiding "escaped" tokens that drift without regularization.
2. Trust-Region Annealing: The KL coefficient $\beta$ is annealed proportionally with the learning rate $\eta_t$ to prevent the policy from being pinned at an equilibrium where KL pull cancels the policy gradient.
3. Group-Level Accuracy Filter: Groups with empirical accuracy $\mu_r$ outside a configured window (e.g., $[0.20, 0.95]$) are dropped. This acts as an emergent curriculum, focusing compute on prompts with non-trivial variance.
4. DDP-Coordinated Skip: Ensures distributed training consistency when dropping groups.

3. Key Contributions

1. New Paradigm for Latent RL: The paper identifies the deterministic nature of latent recurrence as the cause of GRPO failure and proposes shared-mask variational dropout as a minimal, theoretically clean solution. It is the first method to successfully apply group-relative RL to continuous latent-reasoning models.
2. Theoretical Justification:
   • Proves the well-definedness of the latent Jacobian $\partial h_T / \partial \theta$ under standard Transformer regularity.
   • Proves the unbiasedness of the surrogate gradient for the marginal policy.
   • Demonstrates that mask replay provides common-random-numbers variance reduction.
3. Empirical Validation: Successfully applies GRPO to the COCONUT model on the GSM8K benchmark, a regime where deterministic GRPO fails completely.
4. Reference Implementation: Releases code that realizes bit-identical mask replay across the entire latent recurrence.

4. Experimental Results

Experiments were conducted using Qwen2.5-1.5B as the base model with a latent depth of $T=6$.

• Performance Improvement: Dropout-GRPO improved the GSM8K pass@1 accuracy from a COCONUT baseline of 27.29% to 29.01%.
• Recovery of Capabilities: The method reversed a 1.33-point performance regression on the SVAMP dataset caused by curriculum training, recovering performance to 44.00% (matching the base model's variance).
• Ablation Studies:
  • Deterministic GRPO: Failed completely, yielding a flat training trajectory identical to the SFT initialization due to zero advantage.
  • REINFORCE + Dropout (EMA): Failed due to high variance; the single-rollout gradient was too noisy to stabilize the SFT initialization.
  • Dropout-GRPO: The only configuration that produced a viable learning signal, confirming that group-relative baselines are essential for variance reduction in this setting.

5. Significance and Claims

The paper positions this work as a practical and theoretically grounded approach for post-training latent-reasoning Large Language Models (LLMs).

• Feasibility: It demonstrates that RL post-training is feasible for models that reason directly in hidden states, a capability previously blocked by the lack of rollout diversity.
• Mechanism: It re-frames dropout not as a regularizer (which is often disabled in modern LLMs) but as a source of structured exogenous stochasticity required for variational policy optimization.
• Limitations: The authors are modest about the magnitude of improvement (a 1.72 point gain on GSM8K), attributing this to the small scale of the backbone (1.5B) and the scarcity of suitable concise-step reasoning datasets. They note that the deployed policy (inference without dropout) differs from the optimized marginal policy, introducing a bias dependent on Lipschitz constants.

In summary, Dropout-GRPO bridges the gap between group-relative reinforcement learning and continuous latent reasoning, enabling the optimization of models that previously could not utilize GRPO due to their deterministic internal dynamics.