Resumen Técnico: SPHERE - Mitigación de la Pérdida de Plasticidad Espectral en Mezclas de Expertos para Aprendizaje por Refuerzo Profundo

1. Declaración del Problema

En el Aprendizaje por Refuerzo Profundo (DRL), los agentes entrenados en flujos de experiencia a menudo sufren pérdida de plasticidad: un declive gradual en su capacidad para aprender nuevas habilidades a partir de nuevas experiencias a medida que avanza el entrenamiento. Este fenómeno es particularmente perjudicial en entornos de Aprendizaje por Refuerzo Continuo (CRL), donde un agente debe dominar secuencialmente una serie de tareas.

Si bien las arquitecturas de Mezcla de Expertos (MoE) se han implementado con éxito para escalar la capacidad de la política en DRL (por ejemplo, para manipuladores, humanoides y cuadrúpedos), estudios empíricos indican que las políticas MoE también sufren una pérdida de plasticidad severa en entornos continuos. A medida que avanza el entrenamiento, el rendimiento de MoE puede degradarse, lo que sugiere que la propia arquitectura es susceptible a los mismos mecanismos de pérdida de capacidad observados en las redes densas.

El desafío central abordado por este artículo es la falta de una caracterización formal de la pérdida de plasticidad específicamente dentro de las políticas MoE y la ausencia de métodos prácticos para mitigarla. Las estrategias de mitigación existentes para redes densas (por ejemplo, normalización espectral, regularización de Hessiano) no se traducen directamente a la estructura dispersa y enrutada de las políticas MoE.

2. Metodología: SPHERE

Los autores proponen SPHERE (Plasticidad Espectral mediante REgularización Hiperesférica de Expertos), un regularizador principista diseñado para preservar la capacidad de aprendizaje de las políticas MoE. La metodología procede a través de cuatro pasos teóricos y prácticos clave:

2.1 Formalización de la Pérdida de Plasticidad como Plasticidad Espectral

Basándose en la teoría del Kernel Tangente Neural (NTK), los autores formalizan la pérdida de plasticidad como una pérdida de plasticidad espectral.

• Mecanismo: La dirección de actualización de una política está gobernada por la matriz empírica del NTK (eNTK), $K = JJ^\top$, donde $J$ es el Jacobiano de las salidas de la red con respecto a los parámetros.
• Definición: La pérdida de plasticidad corresponde a un colapso en la rango efectivo de entropía espectral ($r_e$) del eNTK. Un rango efectivo alto implica un espectro isotrópico, permitiendo actualizaciones en diversas direcciones. Un rango efectivo bajo (colapso espectral) restringe las actualizaciones a pocas direcciones dominantes, obstaculizando la capacidad de aprender nuevas tareas.
• Observación: En CRL, el espectro eNTK de las políticas MoE colapsa, reduciendo el rango efectivo y confinando las actualizaciones funcionales.

2.2 Derivación de un Proxy Tractable

El cálculo directo de la matriz eNTK $K$ es intratable para redes grandes debido a las restricciones de memoria ($O(N^2P)$). Para abordar esto, los autores derivan un proxy tratable basado en la matriz Gauss-Newton (GN) en el espacio de parámetros ($G_{GN} = \frac{1}{N}J^\top J$).

• Descomposición: Descomponen la matriz GN en componentes diagonales por bloques correspondientes a los parámetros de puerta y experto.
• Aproximación por Capas: Bajo una aproximación diagonal por bloques, el rango efectivo del eNTK completo se expresa como una mezcla de los rangos efectivos de los bloques de capas individuales.
• Surrogado de Bloque de Experto: Para las capas de expertos, los autores utilizan un proxy de Kronecker ($\tilde{G}_{GN} = A \otimes G$), donde $A$ es la matriz Gram de características y $G$ es la matriz Gram de gradientes. Demuestran que el rango efectivo del bloque de experto está limitado inferiormente por una función del número de condición de este proxy de Kronecker.

2.3 El Regularizador SPHERE

Para maximizar la cota inferior de la plasticidad espectral, los autores buscan minimizar el número de condición de la matriz Gram de características $A$.

• Objetivo: Se centran en la Gram de características de experto ponderada ($A_{last}^{exp}$) en la última capa oculta de experto. Esta matriz se construye concatenando características de expertos ponderadas por sus probabilidades de puerta Top-K.
• Penalización: SPHERE aplica una penalización de Parseval (penalización cuadrática) al componente anisotrópico de $A_{last}^{exp}$. La función de pérdida se define como:
  $$L_{SPHERE}(A_{last}^{exp}) = \left\| A_{last}^{exp} - \frac{\text{Tr}(A_{last}^{exp})}{m} I_m \right\|_F^2$$
  Esta penalización fomenta que los valores propios de la Gram de características sean iguales (isotrópicos), aumentando así el rango efectivo y preservando la plasticidad espectral.
• Garantía Teórica: Los autores demuestran que optimizar esta penalización induce una contracción espectral, disminuyendo monótonamente el número de condición del proxy y aumentando la cota inferior del rango efectivo del eNTK.

2.4 Implementación

SPHERE se integra en el bucle de entrenamiento de PPO. En cada actualización de gradiente, se calcula la Gram de características de experto ponderada a partir del paso hacia adelante, y se añade la penalización SPHERE a la función de pérdida del actor. El coeficiente de regularización se escala adaptativamente en función de la relación entre la norma del gradiente del actor y la norma del gradiente de SPHERE.

3. Contribuciones Clave

1. Formalización: El artículo define formalmente la pérdida de plasticidad en políticas MoE como una pérdida de plasticidad espectral, cuantificada por el rango efectivo de entropía espectral del NTK empírico.
2. Proxy Tractable: Deriva un sustituto práctico y por capas para la plasticidad espectral que depende únicamente de las matrices de características de expertos individuales, evitando la construcción del eNTK completo.
3. Algoritmo (SPHERE): Introduce SPHERE, un regularizador principista que impone isotropía en las características de expertos ponderadas para contrarrestar el colapso espectral.
4. Validación Empírica: El método se evalúa en MetaWorld y HumanoidBench, demostrando mejoras consistentes sobre las líneas base MoE sin regularizar y métodos anteriores de mitigación de plasticidad.

4. Resultados Experimentales

Los autores evaluaron SPHERE en dos benchmarks de robótica:

• MetaWorld (CW10): Una suite de 10 tareas de manipulación.
• HumanoidBench: Una suite de 5 tareas de locomoción para el humanoide H1.

Hallazgos Clave:

• Ganancias de Rendimiento:
  • En MetaWorld bajo CRL, SPHERE mejoró la tasa de éxito promedio en un 133% sobre una línea base MoE Top-K sin regularizar.
  • En HumanoidBench bajo CRL, SPHERE mejoró el éxito promedio en un 50% sobre la línea base.
  • SPHERE también mostró ganancias en entornos de RL estándar (entrenamiento de tarea única), lo que sugiere que la plasticidad espectral puede decaer incluso dentro de una sola tarea debido a la deriva distribucional.
• Mantenimiento de la Plasticidad Espectral: A lo largo del entrenamiento, SPHERE mantuvo un rango efectivo significativamente más alto ($r_e(K)$) en comparación con las líneas base, las cuales exhibieron un colapso espectral rápido.
• Estudios de Ablación:
  • Se encontró que regularizar solo la última capa oculta de experto era lo más efectivo; regularizar todas las capas o solo la puerta arrojó un rendimiento inferior.
  • La concatenación entre expertos (regularizar la Gram de características conjunta en lugar de expertos individuales) fue crucial para las ganancias, lo que indica la importancia de gestionar las correlaciones entre expertos.
  • El método es robusto frente a diferentes funciones de activación (ReLU, Tanh), cantidades de expertos y variantes MoE (Dense-MoE, DS-MoE).
• Análisis Cualitativo: Las visualizaciones t-SNE mostraron que sin SPHERE, las representaciones de estado colapsaban hacia una única dirección de característica dominante. Con SPHERE, las representaciones permanecieron diversas e isotrópicas.

5. Significado y Afirmaciones

El artículo afirma que SPHERE proporciona una solución principista al problema de la pérdida de plasticidad en agentes de DRL basados en MoE. Al enmarcar la pérdida de plasticidad a través de la lente de la teoría espectral (rango efectivo del eNTK), los autores cierran la brecha entre la comprensión teórica y la regularización práctica.

El significado del trabajo radica en:

• Habilitación del Aprendizaje Continuo Escalable: A medida que las arquitecturas MoE se vuelven estándar para escalar agentes de RL, SPHERE asegura que estos modelos retengan la capacidad de aprender nuevas habilidades a lo largo de largas secuencias de tareas, evitando fallos sistemáticos en tareas posteriores.
• Perspectiva Teórica: Ofrece una nueva perspectiva sobre la dinámica de MoE, vinculando los mecanismos de enrutamiento y la especialización de expertos a las propiedades espectrales del paisaje de aprendizaje.
• Eficacia Práctica: El método es computacionalmente eficiente (añadiendo solo ~3-6% de sobrecarga incluso con 1000 expertos) y se integra perfectamente en las tuberías PPO existentes sin requerir cambios arquitectónicos.

Los autores reconocen limitaciones, señalando que su derivación depende de aproximaciones teóricas (diagonal por bloques e independencia dentro de la capa) y que el alcance actual se limita al RL de control continuo. Identifican extender SPHERE a modelos fundacionales a gran escala como una dirección para trabajos futuros.