Resumen Técnico: Enrutamiento de Sinkhorn Selectivo para un SMoE Disperso Mejorado

Planteamiento del Problema

Los modelos de Mezcla de Expertos Dispersos (SMoE, por sus siglas en inglés) ofrecen un enfoque escalable para aumentar la capacidad del modelo manteniendo la eficiencia computacional al activar solo un subconjunto de expertos por token. Sin embargo, los modelos SMoE convencionales que dependen del enrutamiento basado en Softmax suelen sufrir de "colapso de enrutamiento", donde un pequeño subconjunto de expertos es sobreutilizado mientras otros permanecen subutilizados.

Las estrategias de mitigación existentes dependen típicamente de:

1. Objetivos Auxiliares: Pérdidas de equilibrio de carga o pérdidas z que penalizan el uso desigual de los expertos. Estas pueden introducir desalineación de objetivos, inestabilidad en el entrenamiento y mayor complejidad.
2. Ruido de Puerta (Noisy Gating): Inyección de ruido para fomentar la exploración, lo que añade parámetros entrenables o hiperparámetros.
3. Enrutamiento basado en Sinkhorn: Los métodos previos que utilizan el algoritmo de Sinkhorn para el transporte óptimo (OT) a menudo usan el mapa de transporte únicamente para la selección de expertos (top-k) en lugar de la asignación de pesos. Además, pueden reducir la flexibilidad del enrutamiento porque la matriz de puerta no se optimiza directamente mediante el aprendizaje basado en gradientes de la misma manera que los métodos Softmax, o incurren en una sobrecarga sustancial de entrenamiento debido a la naturaleza iterativa de las actualizaciones de Sinkhorn.

Metodología

El artículo propone el Enrutamiento de Sinkhorn Selectivo (SSR), un mecanismo ligero que reformula la asignación de tokens a expertos como un problema de transporte óptimo con regularización de entropía.

Formulación Central

Los autores modelan la asignación de $m$ tokens a $n$ expertos como un problema de transporte óptimo de costo máximo con regularización de entropía. El objetivo es encontrar un plan de transporte $\hat{\Pi}$ que maximice la compatibilidad mientras satisface restricciones específicas:

• Positividad: $\Pi > 0$.
• Suma de Filas: Cada token enruta toda su masa a los expertos ($\Pi \mathbf{1}_n = \mathbf{1}_m$).
• Suma de Columnas: Cada experto recibe una carga total esperada igual ($\Pi^\top \mathbf{1}_m = \frac{m}{n} \mathbf{1}_n$).

La matriz de costo $C$ se deriva de las puntuaciones de puerta $S$. Los autores proponen dos variantes:

• Costo Lineal: $C = S$.
• Costo Softmax: $C_{i,:} = \text{softmax}(S_{i,:})$, utilizado para prevenir el desbordamiento numérico durante las iteraciones de Sinkhorn.

A diferencia de los métodos previos basados en OT que usan el plan de transporte solo para la selección, el SSR utiliza directamente los valores en $\hat{\Pi}$ para calcular los pesos de enrutamiento. Para cada token, se seleccionan los $k$ mejores expertos basados en las entradas más altas de la fila correspondiente de $\hat{\Pi}$, y los pesos se normalizan a partir de estas entradas.

Estrategia Selectiva e Inyección de Ruido

Un desafío clave identificado es que el enrutamiento de Sinkhorn no actualiza la matriz de pesos de puerta $W_g$ porque $W_g$ está desacoplada del objetivo de OT. Para abordar esto, el SSR emplea una estrategia de entrenamiento híbrida:

• Aplicación Selectiva: Durante el entrenamiento, cada bloque MoE utiliza el enrutamiento de Sinkhorn con una pequeña probabilidad $p$ (por ejemplo, 0.001) y el enrutamiento Softmax estándar en caso contrario. Esto permite que el enrutamiento Softmax actualice $W_g$ de manera efectiva mientras que el enrutamiento de Sinkhorn proporciona señales de utilización de expertos equilibradas.
• Inyección de Ruido: Para prevenir el colapso de los expertos y fomentar la exploración, se añade ruido gaussiano a la matriz de costo ($\tilde{C} = C + \alpha_{noise} \cdot \epsilon$). El análisis teórico (Proposición 4.3) muestra que esto asegura que cada experto tenga una probabilidad no nula de selección.

Comportamiento en la Inferencia

El artículo argumenta que imponer restricciones de equilibrio (como las sumas de columnas iguales) durante la inferencia en una sola entrada es contraproducente, ya que fuerza un enrutamiento uniforme independientemente de las puntuaciones de compatibilidad específicas de la entrada. Por lo tanto, tanto el enrutamiento de Sinkhorn como la inyección de ruido se desactivan durante la inferencia, volviendo al enrutamiento Softmax determinista para asegurar predicciones consistentes y dependientes de la entrada.

Contribuciones Clave

1. Nuevo Marco de Enrutamiento: Un método que integra el transporte óptimo con regularización de entropía con la inyección de ruido estocástico para promover un uso equilibrado de los expertos sin pérdidas de equilibrio auxiliares.
2. Perspectivas Teóricas: Demostraciones que muestran cómo el enrutamiento basado en Sinkhorn y la inyección de ruido ayudan en el entrenamiento al fomentar la exploración y el equilibrio, pero deben desactivarse en la inferencia para evitar asignaciones distorsionadas.
3. Eficiencia y Rendimiento: Evaluaciones extensas que muestran que el SSR mejora la eficiencia del entrenamiento, la precisión y la robustez en comparación con los modelos de referencia de vanguardia.

Resultados Experimentales

Los autores evaluaron el SSR en modelado de lenguaje (WikiText-103, Enwik-8) y clasificación de imágenes (ImageNet-1K, ImageNet-A, ImageNet-O, ImageNet-R).

• Modelado de Lenguaje:
  • En WikiText-103, las variantes de SSR superaron consistentemente al Vanilla SMoE y a los modelos de referencia que utilizan pérdida de equilibrio de carga, pérdida z o ruido.
  • SSR-L con ruido logró una perplejidad (PPL) de prueba de 34.367, una reducción de 1.183 en comparación con Vanilla SMO de 35.550, superando al siguiente mejor modelo de referencia por un margen significativo.
  • Crucialmente, el SSR incurrió en una sobrecarga de entrenamiento mínima (0.33% – 0.65%), mientras que un SMoE completo basado en Sinkhorn incurrió en un 72.47% de sobrecarga.
  • Bajo configuraciones de Momento (Momentum), el SSR preservó la evolución característica de la norma de Vanilla SMoE mientras mejoraba el rendimiento, mientras que otros métodos de equilibrio degradaron el rendimiento.
• Modelado a Nivel de Byte (Enwik-8):
  • SSR-S con ruido logró el mejor Bits-por-Carácter (BPC) de prueba de 1.128, mejorando al Vanilla SMoE en 0.010 BPC.
  • Fue 2.28x más rápido que el SMoE completo basado en Sinkhorn.
• Tareas de Visión:
  • En ImageNet-1K, SSR-L con ruido logró una precisión Top-1 de 77.420, una ganancia de +2.368 sobre Vanilla SMoE.
  • También mostró una fuerte robustez en ImageNet-O e ImageNet-R.
• Estudios de Ablación:
  • La probabilidad $p$ de aplicar el enrutamiento de Sinkhorn es crítica; valores muy pequeños (por ejemplo, $10^{-3}$ o $10^{-4}$) fueron suficientes para proporcionar señales de equilibrio sin restringir demasiado el enrutador.
  • La inyección de ruido mejoró consistentemente el rendimiento en todas las configuraciones.
  • Desactivar los mecanismos de equilibrio en la inferencia produjo los mejores resultados, confirmando las afirmaciones teóricas.

Significado y Reivindicaciones

El artículo afirma que el Enrutamiento de Sinkhorn Selectivo ofrece una solución práctica y efectiva para el diseño de SMoE. Al reemplazar las complejas pérdidas auxiliares con una optimización de Sinkhorn intermitente y eficiente, el SSR logra:

• Mejor Estabilidad de Entrenamiento: Mediante un equilibrio de expertos integrado sin desalineación de objetivos.
• Mayor Precisión: Demostrada en tareas de texto y visión.
• Robustez: Rendimiento mejorado frente a la corrupción de la entrada y ejemplos adversarios.
• Eficiencia: Sobrecarga computacional mínima en comparación tanto con los métodos de pérdida auxiliar como con el enrutamiento completo de Sinkhorn.

Los autores enfatizan que el método es ligero, no requiere parámetros adicionales entrenables (en la variante de ruido) y mantiene la flexibilidad de la selección de expertos asegurando al mismo tiempo un uso equilibrado durante la fase de entrenamiento.