dynamic-gradient-gating-rlvr

# Dynamic Gradient Gating for Sample-Efficient RLVR Methodology for safely reusing rollout batches in Reinforcement Learning with Verifiable Rewards (RLVR) by monitoring the lm_head gradient norm as a real-time signal of catastrophic policy shift. ## Core Problem RLVR has become the dominant paradigm for advanced reasoning in LLMs, but rollout samples are expensive. Reusing rollout batches for multiple gradient updates (standard in classical RL) amplifies policy shift in RLVR, causing severe performance degradation. Detecting the onset of degradation early enough to stop reuse is challenging. ## Key Discovery: Disproportionate Weight Divergence (DWD) - Performance degradation is **synchronized** with a sharp surge in `lm_head` weight change - Intermediate layers remain stable during this degradation - DWD emerges consistently across diverse LLMs and tasks ### Theoretical Foundations 1. Harmful gradients concentrate at `lm_head` while intermediate layers are structurally attenuated 2. `lm_head` gradient norm **lower-bounds** the policy divergence 3. This establishes `lm_head` gradient norm as a principled, real-time signal of catastrophic policy shift ## Dynamic Gradient Gating (DGG) Methodology ### Core Mechanism - Monitor `lm_head` gradient norm in real-time during batch reuse - When the norm exceeds a threshold, intercept harmful gradients before they corrupt the optimizer - Lightweight intervention: no additional forward passes needed ### Implementation Pattern ```python class DynamicGradientGating: def __init__(self, model, threshold_scale=1.5): self.lm_head = model.get_output_embeddings() self.baseline_norm = None self.threshold_scale = threshold_scale def should_gate(self, grad_norm): if self.baseline_norm is None: self.baseline_norm = grad_norm return False return grad_norm > self.threshold_scale * self.baseline_norm def step(self, batch, optimizer): loss = model(batch) loss.backward() # Monitor lm_head gradient norm lm_head_grad = self.lm_head.weight.grad current_norm = lm_head_grad.norm().item() if self.should_gate(current_norm): # Intercept harmful gradients lm_head_grad.zero_() optimizer.step() optimizer.zero_grad() ``` ### Key Design Choices - **Threshold**: Scale factor above baseline `lm_head` gradient norm - **Gating action**: Zero out lm_head gradients (or clip to threshold) - **Baseline**: Computed from first-use gradients, updated periodically ## Performance - Matches or exceeds standard single-use baseline - Up to **2.93x sample efficiency** improvement - Up to **2.14x wall-clock speedup** - Validated across: math reasoning, ALFWorld, WebShop, search-augmented QA ## When to Use - RLVR training with expensive rollout samples - When you want to reuse batches but need protection against policy degradation - Any RL setup where rollout generation is the bottleneck - Especially useful for reasoning tasks with verifiable rewards ## Activation DGG, dynamic gradient gating, RLVR sample efficiency, lm_head monitoring, batch reuse RL, policy shift detection, gradient clipping RL, verifiable reward RL ## Reference arXiv: 2605.19425v1 - "When to Stop Reusing: Dynamic Gradient Gating for Sample-Efficient RLVR"