ERMoE: Eigen-Reparameterized Mixture-of-Experts for Stable Routing and Interpretable Specialization

Resumen Técnico: ERMoE

1. El Problema

Los modelos de Mezcla de Expertos (MoE) han permitido escalar la capacidad de los modelos de IA activando solo un subconjunto de parámetros por entrada. Sin embargo, las arquitecturas MoE actuales enfrentan dos desafíos estructurales críticos:

• Desalineación entre el Enrutador y los Expertos: Los logits de enrutamiento aprendidos a menudo no coinciden con la estructura real de los expertos. Esto provoca una asignación inestable de tokens, subutilización de expertos y una especialización deficiente. Evidencias (como en R2-T2) muestran que re-enrutar tokens en tiempo de prueba sin actualizar los expertos mejora la precisión, lo que indica que el enrutador inicial elige frecuentemente el experto incorrecto.
• Desequilibrio de Carga y Colapso de Representación: Los mecanismos de enrutamiento basados en "elección de token" tienden a crear dinámicas de "el rico se hace más rico", donde unos pocos expertos absorben la mayor parte del tráfico, creando cuellos de botella (stragglers) y dejando a otros expertos inactivos.
• Limitaciones de las Pérdidas Auxiliares: Para mitigar el desequilibrio, la práctica estándar añade una pérdida auxiliar de balanceo de carga (LBL). Sin embargo, esto introduce gradientes de interferencia que pueden debilitar la especialización de los expertos y reducir la precisión en tareas aguas abajo.

2. Metodología: ERMoE

Los autores proponen ERMoE (Mezcla de Expertos Reparametrizada por Autovalores), una arquitectura que redefine fundamentalmente cómo se calculan los expertos y cómo se realiza el enrutamiento.

• Reparametrización en Base de Autovalores (Eigenbasis):
  En lugar de usar pesos libres, cada experto $e$ reparametriza su matriz de pesos $W^{(e)}$ en una base ortonormal aprendida ($U^{(e)}, V^{(e)}$):
  $$W^{(e)} = U^{(e)} \text{diag}(s^{(e)}) V^{(e)\top}$$
  Esto fuerza a los expertos a operar en subespacios representacionales decorrelacionados, reduciendo la redundancia de características.

• Enrutamiento por Puntuación de Autovalores (Eigenbasis Score):
  Se elimina el enrutador basado en logits aprendidos. En su lugar, el enrutador calcula una puntuación basada en la alineación geométrica:

  1. Se proyecta el token de entrada y su contexto ponderado por atención en la base de autovalores de cada experto.
  2. La puntuación de enrutamiento es la similitud del coseno entre estas dos proyecciones.
     $$\text{Score}_e(i) = \cos(u_i^{(e)}, v_i^{(e)})$$
     Esta puntuación mide qué tan bien el token y su contexto se alinean con el subespacio intrínseco del experto.

• Enrutamiento Umbralizado Top-k:
  Se introduce un umbral de confianza $T$. Solo los expertos con una puntuación $\ge T$ son candidatos. De estos, se seleccionan los top-$k$.

  • Ventaja: Esto elimina la necesidad de pérdidas auxiliares de balanceo de carga (LBL), ya que la selección se basa en la compatibilidad de contenido (geometría) en lugar de forzar un equilibrio artificial.
  • Mecanismo de respaldo: Si ningún experto supera el umbral, se activa un mecanismo de respaldo (fallback) para garantizar el progreso.

• Adaptación 3D (ERMoE-ba):
  Para neuroimagen, se adapta a una arquitectura ViT 3D. Se introducen dos tipos de expertos:

  • Expertos Regionales: Entrenados específicamente en tejidos cerebrales (Materia Blanca, Materia Gris, Líquido Cefalorraquídeo).
  • Expertos Libres: Que procesan el cerebro completo.

3. Contribuciones Clave

1. Nueva Arquitectura de Enrutamiento: Introduce un mecanismo de enrutamiento basado en la alineación de características con subespacios de autovalores, eliminando la dependencia de logits de enrutamiento libres y pérdidas de balanceo auxiliares.
2. Estabilidad y Especialización: Al vincular la asignación de tokens a los subespacios representacionales de los expertos, se logra una especialización estable y se mitiga el colapso de representaciones sin sacrificar la dispersidad (sparsity).
3. Interpretabilidad Anatómica: En el dominio médico, demuestra que los expertos aprenden a especializarse en tejidos biológicos específicos (WM, GM, CSF) de manera interpretable, algo difícil de lograr con MoE estándar.
4. Eficiencia: Elimina los gradientes de interferencia de las pérdidas de balanceo, mejorando la convergencia y la precisión final.

4. Resultados Experimentales

• Clasificación de Imágenes (Natural):

  • En ImageNet-1K, ERMoE alcanza un 88.03% de precisión Top-1, superando a V-MoE (87.41%) y otros baselines, a pesar de no utilizar pérdidas de balanceo de carga.
  • En tareas de pocos disparos (few-shot) en CIFAR y Tiny-ImageNet, ERMoE muestra una mejora significativa (hasta un 7% en Tiny-ImageNet), indicando que las características aprendidas son más linealmente separables y menos redundantes.

• Recuperación de Imagen-Texto:

  • Integrado en CLIP (CLIP-ERMoE), logra resultados de vanguardia en recuperación de imagen a texto (I2T) en COCO y Flickr30K, superando a modelos densos y otros MoE, demostrando que las representaciones diversas de los expertos se generalizan bien a tareas multimodales.

• Neuroimagen (Predicción de Edad Cerebral):

  • En el conjunto de datos ADNI (MRI 3D), la variante ERMoE-ba logra un Error Absoluto Medio (MAE) de 2.31 años, superando significativamente a CNNs 3D, ViT 3D y Swin Transformer 3D.
  • Muestra una mejor calibración y una relación más estrecha entre la edad cronológica y la edad cerebral predicha.
  • Interpretabilidad: Los expertos regionales (WM, GM, CSF) muestran una alta confianza y especialización específica para sus tejidos correspondientes tras el entrenamiento, validando la capacidad del modelo para aprender subespacios anatómicos significativos.

• Análisis de Carga:

  • ERMoE presenta una distribución de carga de expertos mucho más plana (menos "cola larga") que V-MoE, acercándose al equilibrio de métodos "balanceados por construcción" (como BASE) pero sin sus limitaciones de contenido.

5. Significado e Impacto

El trabajo de ERMoE representa un cambio de paradigma en el diseño de MoE:

• Resuelve el compromiso Precisión-Estabilidad: Demuestra que es posible lograr un enrutamiento estable y un equilibrio de carga natural sin recurrir a pérdidas auxiliares que degradan el rendimiento de la tarea principal.
• Interpretabilidad en IA Médica: Proporciona un marco donde la "caja negra" de los MoE se vuelve interpretable, permitiendo que los expertos se alineen con estructuras biológicas reales (tejidos cerebrales), lo cual es crucial para aplicaciones clínicas.
• Eficiencia Computacional: Al eliminar la necesidad de hiperparámetros complejos de balanceo y gradientes de interferencia, simplifica el entrenamiento y mejora la eficiencia en la inferencia y el uso de memoria.

En conclusión, ERMoE ofrece una solución robusta a los problemas de inestabilidad y desequilibrio en los MoE, logrando resultados de vanguardia en visión por computadora y neuroimagen mediante una reparametrización geométrica inteligente de los expertos.