Measure gradients, not activations! Enhancing neuronal activity in deep reinforcement learning

Deep reinforcement learning (RL) agents frequently suffer from neuronal activity loss, which impairs their ability to adapt to new data and learn continually. A common method to quantify and address this issue is the tau-dormant neuron ratio, which uses activation statistics to measure the expressive ability of neurons. While effective for simple MLP-based agents, this approach loses statistical power in more complex architectures. To address this, we argue that in advanced RL agents, maintaining a neuron’s learning capacity, its ability to adapt via gradient updates, is more critical than preserving its expressive ability. Based on this insight, we shift the statistical objective from activations to gradients, and introduce GraMa (Gradient Magnitude Neural Activity Metric), a lightweight, architecture-agnostic metric for quantifying neuron-level learning capacity. We show that GraMa effectively reveals persistent neuron inactivity across diverse architectures, including residual networks, diffusion models, and agents with varied activation functions. Moreover, resetting neurons guided by GraMa (ReGraMa) consistently improves learning performance across multiple deep RL algorithms and benchmarks, such as MuJoCo and the DeepMind Control Suite.

💡 Research Summary

The paper addresses a pervasive problem in deep reinforcement learning (RL): the gradual loss of neuronal activity, often termed “dormant neurons,” which reduces a network’s capacity to adapt to new data and hampers continual learning. The dominant metric for detecting such neurons, the τ‑dormant neuron ratio, relies solely on activation statistics. While effective for simple multilayer perceptrons (MLPs) with ReLU activations, the authors demonstrate that this metric loses statistical power in modern RL architectures that incorporate residual connections, multi‑branch pathways, mixture‑of‑experts, diffusion‑based policies, and various normalization layers.

Through a series of diagnostic experiments, the authors show that activation‑based methods (e.g., ReDo) misidentify neurons in complex networks: they fail to flag neurons that have lost learning capacity (low gradient magnitude) and mistakenly reset neurons that are still learning effectively but happen to have low instantaneous activation. Three architectural factors are highlighted as culprits: (i) branch fusion obscuring individual neuron contributions, (ii) non‑ReLU activations (Leaky ReLU, GELU, etc.) breaking the clear “dead vs. alive” dichotomy, and (iii) normalization layers (LayerNorm, BatchNorm) that rescale activations and mask true inactivity.

To overcome these limitations, the authors propose a paradigm shift: evaluate neuronal health via gradient magnitude rather than activation magnitude. They introduce GraMa (Gradient Magnitude Neural Activity Metric), which computes the L2 norm of each neuron’s gradient during the backward pass, normalizes it across the layer, and compares it against a threshold τ. Because gradients are already computed for parameter updates, GraMa incurs negligible additional computational or memory overhead.

Building on GraMa, the paper presents ReGraMa, a targeted neuron‑reset mechanism. At regular intervals, ReGraMa identifies the bottom‑percentile of neurons with the smallest gradient magnitudes (indicating low learning capacity) and reinitializes their parameters. The reset proportion is kept modest (typically the top 10‑25% of low‑gradient neurons) to avoid catastrophic forgetting.

Extensive empirical evaluation spans several RL algorithms (SAC, TD3, PPO) and benchmarks (MuJoCo, DeepMind Control Suite). The authors test a variety of architectures: standard SAC‑SA‑C, residual BR‑O‑net, diffusion‑based policy DACER, and mixture‑of‑experts models. Across all settings, ReGraMa consistently outperforms activation‑based ReDo, delivering 8‑12 % higher cumulative rewards on average and reducing learning curve variance. In visual‑input tasks where activation‑based metrics essentially fail, GraMa continues to detect dormant neurons and guide effective resets, preserving performance.

Theoretical analysis links gradient magnitude to the expected parameter update size, establishing that a neuron with near‑zero gradient contributes little to learning regardless of its activation level, while a neuron with substantial gradient can still be highly plastic even if its activation is temporarily low. The authors also employ moving‑average smoothing of gradients to mitigate stochastic noise and retain compatibility with the existing τ‑dormant threshold framework.

In summary, the paper makes four key contributions: (1) it empirically shows the inadequacy of activation‑based neuronal health metrics in complex RL architectures, (2) it introduces GraMa, a lightweight, architecture‑agnostic gradient‑based metric that accurately reflects learning capacity, (3) it develops ReGraMa, an efficient neuron‑reset strategy guided by GraMa, and (4) it validates the approach with comprehensive experiments demonstrating improved performance and stability across diverse tasks and models. The work opens avenues for applying gradient‑based neuronal health monitoring to continual learning, multi‑agent systems, and large‑scale vision/language models, where maintaining plasticity without sacrificing stability is increasingly critical.