Safety

Dropout-GRPO: Variational Stochasticity for Continuous Latent Reasoning

arXiv:2606.10184v1 Announce Type: cross Abstract: Group Relative Policy Optimization (GRPO) relies on the diversity of K rollouts within each group; otherwise, the group-mean advantage A^{(k)} = r^{(k

DGX agentpaper
safetyarxiv-cs-ai

arXiv:2606.10184v1 Announce Type: cross Abstract: Group Relative Policy Optimization (GRPO) relies on the diversity of K rollouts within each group; otherwise, the group-mean advantage A^{(k)} = r^{(k)} - mu_r collapses to zero. This presents a structural challenge for latent-reasoning models like Coconut, which feed continuous hidden states recurrently in place of discrete chain-of-thought tokens. Because the latent phase is inherently deterministic given the parameters and prompt, multiple rollouts produce identical trajectories, stalling GRPO's progress. Consequently, applying group-relative reinforcement learning to continuous latent reasoning has proven difficult. To address this, we propose sourcing the necessary stochasticity through structured dropout. By applying a single Bernoulli mask held constant across all latent recurrence steps for a given rollout, we generate essential trajectory variance. This shared mask effectively treats each rollout as a posterior sample from a variational distribution over parameters, allowing GRPO to optimize the expected reward of a Bayesian model-average policy. We provide both theoretical justification for this method -- including unbiasedness, variance reduction, and the well-definedness of the latent gradient -- and empirical validation. On GSM8K, dropout-GRPO improves a Coconut baseline from 27.29% to 29.01% pass@1, demonstrating the viability of GRPO learning for latent-reasoning models. Our work positions this as a practical, theoretically grounded approach for post-training latent-reasoning LLMs.

Source: arXiv cs.AI | 2026-06-10

Loading related sources…