Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

Resumen Técnico: Tres Fases del Enrutamiento de Expertos

1. El Problema

Las arquitecturas de Mezcla de Expertos (MoE) escalan la capacidad de los modelos de lenguaje distribuyendo cada token a un subconjunto de redes especializadas ("expertos"). El desafío de ingeniería central es el balance de carga: sin intervención, los tokens tienden a concentrarse en unos pocos expertos de alta calidad, dejando a los demás inactivos.
Actualmente, la solución estándar es una pérdida de balance auxiliar (auxiliary balance loss) con un coeficiente ajustable ($\alpha$). Sin embargo, la literatura existente no explica cómo evoluciona la compensación (trade-off) entre la calidad del experto y el balance de carga durante el proceso de entrenamiento. La mayoría de los análisis se realizan a posteriori en modelos ya convergidos, ignorando la dinámica temporal.

2. Metodología

El autor propone un marco teórico que modela el enrutamiento de tokens en MoE como un juego de congestión (congestion game) bajo la teoría de Juegos de Campo Medio (Mean-Field Games, MFG).

• Formulación del Juego:
  • Agentes: Los tokens.
  • Recursos: Los expertos ($M$).
  • Costo: Un token elige un experto basándose en la calidad del experto ($q_i$) menos un costo de congestión proporcional a la carga del experto ($\gamma \mu_i$).
  • Parámetro Clave: El coeficiente de congestión efectivo ($\gamma_{eff}$), que cuantifica la fuerza del trade-off entre calidad y balance.
• Equilibrio Teórico: Se demuestra que, para modelos bien balanceados, el equilibrio del juego de campo medio se reduce matemáticamente a un softmax escalado por temperatura, donde la temperatura corresponde al coeficiente de congestión.
• Análisis Empírico:
  • Se rastreó $\gamma_{eff}$ a lo largo de 20 puntos de control (checkpoints) durante el preentrenamiento de OLMoE-1B-7B (con muestreo denso en la región de "auge") y 6 checkpoints de OpenMoE-8B.
  • Se estimaron las cualidades de los expertos a partir de los logits de la puerta (router) y se ajustó $\gamma_{eff}$ minimizando la distancia $L_1$ entre la distribución de carga observada y la predicha por el modelo.
  • Se utilizaron intervalos de confianza de remuestreo (bootstrap) sobre 50 lotes de texto independientes.

3. Contribuciones Clave

1. Descubrimiento de la Trayectoria de Tres Fases:
   El hallazgo central es que $\gamma_{eff}$ sigue una trayectoria no monótona en forma de "U invertida" durante el entrenamiento, invisible en el análisis de modelos finales:

   • Fase 1: Auge (Surge, pasos 5K–50K): El router aprende a balancear la carga. $\gamma_{eff}$ aumenta drásticamente (de ~14 a un pico de 36–39). La entropía de enrutamiento aumenta, indicando una distribución más uniforme.
   • Fase 2: Estabilización (pasos 100K–400K): El coeficiente se estabiliza (~24–28). Bajo este régimen fijo, los expertos comienzan a especializarse (la dispersión de calidad $B_0$ disminuye). El router ha encontrado su punto de operación.
   • Fase 3: Relajación (pasos 400K–1.2M): A medida que los roles de los expertos se solidifican, el router relaja la exigencia de balance para priorizar la calidad. $\gamma_{eff}$ cae significativamente (de ~27 a ~8.5). El modelo sacrifica un poco de balance para maximizar la calidad de la respuesta.

2. Descomposición de la Congestión Efectiva:
   Se introduce la fórmula $\gamma_{eff} = \gamma_{explicit} + \gamma_{implicit}$.

   • $\gamma_{explicit}$ proviene de la pérdida auxiliar ($\alpha M$).
   • $\gamma_{implicit}$ es la carga de balance internalizada por el optimizador a través de la dinámica de gradiente.
   • Hallazgo: En la convergencia, $\gamma_{eff}$ es ~13 veces mayor que $\gamma_{explicit}$. Esto demuestra que el balance no es solo un efecto de la pérdida auxiliar, sino que el optimizador internaliza el balance de manera masiva.

3. Extensión Multi-Tipo (Multi-Type MFG):
   Se propone una extensión que modela la heterogeneidad de los tokens agrupándolos en $K$ tipos. Esto mejora la predicción de la carga en todas las capas (mejora media del 30%) al capturar la estructura de la población, superando al softmax simple por tipo.

4. Diagnostics de Alcance (Scope Diagnostics):
   Se definen métricas para determinar cuándo el modelo de campo medio es aplicable:

   • Umbral Anti-concentración ($\gamma_c$): Si $\gamma_{eff} < \gamma_c$, los expertos podrían colapsar. El modelo se mantiene en una zona segura durante todo el entrenamiento.
   • Relación $K/M$: El modelo funciona bien cuando se seleccionan múltiples expertos ($K > 1$), pero falla en enrutamiento Top-1.

4. Resultados Principales

• OLMoE-1B-7B:
  • Pico de $\gamma_{eff}$: ~36–39 (alrededor del paso 30K-40K).
  • Final de entrenamiento: $\gamma_{eff} \approx 8.5$.
  • Relación Pico/Final: $\ge 4.2\times$.
  • La dispersión de calidad ($B_0$) disminuye monótonamente, mientras que $\gamma_{eff}$ sube y luego baja, demostrando que son dinámicas desacopladas.
• OpenMoE-8B:
  • Replicó el patrón de tres fases, aunque con una fase inicial "dormida" (donde $\gamma_{eff}=0$) debido a su arquitectura más dispersa (solo 4 capas MoE).
• Validación Estática:
  • En un modelo convergido, el modelo de juego de congestión (MFG) es estadísticamente indistinguible del softmax escalado por temperatura (Error $L_1$ de 0.199 vs 0.200). Esto confirma que el valor del enfoque no es la predicción estática, sino la dinámica y la descomposición.
• Análisis por Capas:
  • 12 de 16 capas muestran el patrón de "auge" y 10 muestran "relajación". Las capas tardías (12-15) a veces no desarrollan estructura de congestión, lo que indica que el supuesto de campo medio se rompe cuando las representaciones de los tokens están muy diferenciadas.

5. Significado e Implicaciones

• Cambio de Paradigma en la Optimización: El entrenamiento de MoE sigue un principio general: primero se reduce la varianza (priorizando el balance en la fase de auge) y luego se reduce el sesgo (priorizando la calidad en la fase de relajación).
• Origen del Balance: El balance de carga no es solo un producto de la pérdida auxiliar explícita; es una propiedad emergente que el optimizador internaliza fuertemente (13x a 60x más que la señal explícita).
• Herramienta de Diagnóstico: El parámetro $\gamma_{eff}$ sirve como un monitor de entrenamiento. Una caída prematura de $\gamma_{eff}$ o una aproximación al umbral crítico $\gamma_c$ podría señalar un colapso de expertos antes de que sea visible en otras métricas.
• Limitaciones: El estudio se basa en dos modelos abiertos. La teoría asume congestión lineal y el muestreo de tipos de tokens es heurístico (k-means).

Conclusión:
El artículo revela que el enrutamiento en MoE es un proceso dinámico y evolutivo, no estático. La "tensión" fundamental en la optimización de MoE es el cambio de prioridad del balance hacia la calidad a medida que avanza el entrenamiento, un fenómeno cuantificable únicamente a través del coeficiente de congestión efectivo $\gamma_{eff}$.