Technical Summary: SPHERE - Mitigating the Loss of Spectral Plasticity in Mixture-of-Experts for Deep Reinforcement Learning

1. Problem Statement

In Deep Reinforcement Learning (DRL), agents trained on streams of experience often suffer from plasticity loss: a gradual decline in their ability to learn new skills from new experiences as training progresses. This phenomenon is particularly detrimental in Continual Reinforcement Learning (CRL) settings, where an agent must sequentially master a series of tasks.

While Mixture-of-Experts (MoE) architectures have been successfully deployed to scale policy capacity in DRL (e.g., for manipulators, humanoids, and quadrupeds), empirical studies indicate that MoE policies also suffer from severe plasticity loss in continual settings. As training proceeds, MoE performance can degrade, suggesting that the architecture itself is susceptible to the same capacity loss mechanisms observed in dense networks.

The core challenge addressed by this paper is the lack of a formal characterization of plasticity loss specifically within MoE policies and the absence of practical methods to mitigate it. Existing mitigation strategies for dense networks (e.g., spectral normalization, Hessian regularization) do not directly translate to the sparse, routed structure of MoE policies.

2. Methodology: SPHERE

The authors propose SPHERE (Spectral Plasticity via Hyperspherical Expert REgularization), a principled regularizer designed to preserve the learning capacity of MoE policies. The methodology proceeds through four key theoretical and practical steps:

2.1 Formalizing Plasticity Loss as Spectral Plasticity

Building on Neural Tangent Kernel (NTK) theory, the authors formalize plasticity loss as a loss of spectral plasticity.

• Mechanism: The update direction of a policy is governed by the empirical NTK (eNTK) matrix, $K = JJ^\top$, where $J$ is the Jacobian of the network outputs with respect to parameters.
• Definition: Plasticity loss corresponds to a collapse in the spectral-entropy effective rank ($r_e$) of the eNTK. A high effective rank implies an isotropic spectrum, allowing updates in diverse directions. A low effective rank (spectral collapse) restricts updates to a few dominant directions, hindering the ability to learn new tasks.
• Observation: In CRL, the eNTK spectrum of MoE policies collapses, reducing the effective rank and confining functional updates.

2.2 Deriving a Tractable Proxy

Directly computing the eNTK matrix $K$ is intractable for large networks due to memory constraints ($O(N^2P)$). To address this, the authors derive a tractable proxy based on the Gauss-Newton (GN) matrix in parameter space ($G_{GN} = \frac{1}{N}J^\top J$).

• Decomposition: They decompose the GN matrix into block-diagonal components corresponding to gate and expert parameters.
• Layerwise Approximation: Under a block-diagonal approximation, the effective rank of the full eNTK is expressed as a mixture of the effective ranks of individual layer blocks.
• Expert Block Surrogate: For the expert layers, the authors utilize a Kronecker proxy ($\tilde{G}_{GN} = A \otimes G$), where $A$ is the feature Gram matrix and $G$ is the gradient Gram matrix. They prove that the effective rank of the expert block is lower-bounded by a function of the condition number of this Kronecker proxy.

2.3 The SPHERE Regularizer

To maximize the lower bound on spectral plasticity, the authors aim to minimize the condition number of the feature Gram matrix $A$.

• Target: They focus on the weighted expert feature Gram ($A_{last}^{exp}$) at the last hidden expert layer. This matrix is constructed by concatenating expert features weighted by their Top-K gating probabilities.
• Penalty: SPHERE applies a Parseval penalty (quadratic penalty) to the anisotropic component of $A_{last}^{exp}$. The loss function is defined as:
  $$L_{SPHERE}(A_{last}^{exp}) = \left\| A_{last}^{exp} - \frac{\text{Tr}(A_{last}^{exp})}{m} I_m \right\|_F^2$$
  This penalty encourages the eigenvalues of the feature Gram to be equal (isotropic), thereby increasing the effective rank and preserving spectral plasticity.
• Theoretical Guarantee: The authors prove that optimizing this penalty induces a spectral contraction, monotonically decreasing the condition number of the proxy and increasing the lower bound on the eNTK effective rank.

2.4 Implementation

SPHERE is integrated into the PPO training loop. At each gradient update, the weighted expert feature Gram is computed from the forward pass, and the SPHERE penalty is added to the actor's loss function. The regularization coefficient is adaptively scaled based on the ratio of the actor's gradient norm to the SPHERE gradient norm.

3. Key Contributions

1. Formalization: The paper formally defines plasticity loss in MoE policies as a loss of spectral plasticity, quantified by the spectral-entropy effective rank of the empirical NTK.
2. Tractable Proxy: It derives a practical, layerwise surrogate for spectral plasticity that relies only on individual expert feature matrices, avoiding the construction of the full eNTK.
3. Algorithm (SPHERE): It introduces SPHERE, a principled regularizer that enforces isotropy on weighted expert features to counter spectral collapse.
4. Empirical Validation: The method is evaluated on MetaWorld and HumanoidBench, demonstrating consistent improvements over unregularized MoE baselines and prior plasticity mitigation methods.

4. Experimental Results

The authors evaluated SPHERE on two robotics benchmarks:

• MetaWorld (CW10): A suite of 10 manipulation tasks.
• HumanoidBench: A suite of 5 locomotion tasks for the H1 humanoid.

Key Findings:

• Performance Gains:
  • On MetaWorld under CRL, SPHERE improved the average success rate by 133% over an unregularized Top-K MoE baseline.
  • On HumanoidBench under CRL, SPHERE improved average success by 50% over the baseline.
  • SPHERE also showed gains in standard RL settings (single-task training), suggesting spectral plasticity can decay even within a single task due to distributional drift.
• Spectral Plasticity Maintenance: Throughout training, SPHERE maintained a significantly higher effective rank ($r_e(K)$) compared to baselines, which exhibited rapid spectral collapse.
• Ablation Studies:
  • Regularizing only the last hidden expert layer was found to be most effective; regularizing all layers or only the gate yielded lower performance.
  • Cross-expert concatenation (regularizing the joint feature Gram rather than individual experts) was crucial for gains, indicating the importance of managing cross-expert correlations.
  • The method is robust across different activation functions (ReLU, Tanh), expert counts, and MoE variants (Dense-MoE, DS-MoE).
• Qualitative Analysis: t-SNE visualizations showed that without SPHERE, state representations collapsed onto a single dominant feature direction. With SPHERE, representations remained diverse and isotropic.

5. Significance and Claims

The paper claims that SPHERE provides a principled solution to the problem of plasticity loss in MoE-based DRL agents. By framing plasticity loss through the lens of spectral theory (eNTK effective rank), the authors bridge the gap between theoretical understanding and practical regularization.

The significance of the work lies in:

• Enabling Scalable Continual Learning: As MoE architectures become standard for scaling RL agents, SPHERE ensures these models retain the ability to learn new skills over long sequences of tasks, preventing systematic failures on later tasks.
• Theoretical Insight: It offers a new perspective on MoE dynamics, linking routing mechanisms and expert specialization to the spectral properties of the learning landscape.
• Practical Efficacy: The method is computationally efficient (adding only ~3-6% overhead even with 1000 experts) and integrates seamlessly into existing PPO pipelines without requiring architectural changes.

The authors acknowledge limitations, noting that their derivation relies on theoretical approximations (block-diagonal and within-layer independence) and that the current scope is limited to continuous-control RL. They identify extending SPHERE to large-scale foundation models as a direction for future work.