ProbMoE: Differentiable Probabilistic Routing for Mixture-of-Experts

Resumen Técnico: ProbMoE: Enrutamiento Probabilístico Diferenciable para Mezcla de Expertos

Planteamiento del Problema
Los modelos de Mezcla de Expertos (MoE) escalan de manera eficiente activando solo un pequeño subconjunto de expertos por token. Sin embargo, el entrenamiento de estos modelos sigue siendo un desafío porque el enrutamiento estándar depende de un operador top-$k$ discreto y no diferenciable. Los procedimientos de entrenamiento convencionales evitan esto ignorando la dependencia de la selección top-$k$ de los logits del enrutador, propagando gradientes solo a través de las probabilidades softmax de los expertos seleccionados. Esta aproximación limita la señal de aprendizaje para subconjuntos de expertos alternativos, lo que a menudo conduce a distribuciones de enrutamiento concentradas, una utilización deficiente de los expertos y dinámicas de entrenamiento inestables. Los intentos existentes para mitigar esto, como la retropropagación densa o los estimadores de gradiente heurísticos, a menudo fallan al no modelar explícitamente la distribución sobre los subconjuntos de expertos factibles o al proporcionar mecanismos limitados para explorar combinaciones alternativas durante el entrenamiento.

Metodología
Los autores proponen ProbMoE, un marco que reformula el enrutamiento de expertos como un problema de inferencia probabilística sobre subconjuntos de expertos con restricción de cardinalidad. En lugar de tratar el enrutamiento como un operador determinista, ProbMoE modela la selección de expertos como una distribución sobre subconjuntos discretos.

1. Formulación Probabilística:

   • Enrutamiento Exact-$k$: El marco define una distribución sobre subconjuntos $S$ donde $|S|=k$. Condiciona las probabilidades de selección Bernoulli independientes (derivadas de los logits del enrutador mediante sigmoide) a la restricción de que se seleccionen exactamente $k$ expertos. La constante de normalización para esta distribución, $Z_k$, se computa exactamente mediante programación dinámica, lo que hace que la formulación sea tratable.
   • Enrutamiento Dinámico-$k$: El marco se generaliza a un entorno de rango restringido donde el tamaño del subconjunto $|S|$ se encuentra dentro de $[k_{min}, k_{max}]$. La distribución está condicionada a este rango, lo que permite al modelo inferir conjuntamente las identidades de los expertos y la cardinalidad del enrutamiento para cada token.

2. Pase hacia Adelante (Forward Pass):

   • Durante el entrenamiento, ProbMoE muestrea un subconjunto de expertos $S$ de la distribución definida (Exact-$k$ o de rango restringido).
   • La capa MoE computa la salida utilizando solo los expertos muestreados, preservando el patrón de ejecución dispersa de los modelos MoE estándar.
   • En la inferencia, se selecciona el conjunto de Máxima Posterior (MAP), asegurando que los costos de inferencia coincidan con el enrutamiento MoE convencional.

3. Pase hacia Atrás (Backward Pass):

   • Para permitir la optimización basada en gradientes, ProbMoE emplea el estimador de gradiente SIMPLE (Ahmed et al., 2023).
   • En lugar de diferenciar a través de la muestra discreta, el pase hacia atrás utiliza las probabilidades marginales de cada experto de ser seleccionado bajo la distribución restringida.
   • Los pesos de enrutamiento se construyen utilizando un estimador de paso directo (straight-through estimator): $w = (\text{stopgrad}(z - m) + m) \odot \pi$, donde $z$ es la máscara muestreada, $m$ es el vector de probabilidades marginales exactas y $\pi$ son las probabilidades de enrutamiento softmax.
   • Esto permite que los gradientes fluyan a través de las probabilidades marginales, proporcionando actualizaciones informativas al enrutador que reflejan cómo los cambios en las probabilidades de selección afectan la pérdida esperada, mientras mantiene la ejecución dispersa de los expertos.

Contribuciones Clave

1. Formulación de Inferencia Probabilística: El artículo plantea el enrutamiento de MoE como una inferencia probabilística sobre subconjuntos de expertos con restricción de cardinalidad, yendo más allá de las aproximaciones heurísticas de gradiente.
2. Optimización Basada en Marginales Tratable: Al aprovechar el estimador SIMPLE, los autores combinan el muestreo de subconjuntos estocásticos con gradientes de enrutador exactos basados en marginales, permitiendo un entrenamiento diferenciable sin computación densa del lado del experto.
3. Marco Unificado Dinámico-$k$: El método se extiende naturalmente al enrutamiento dinámico, donde el modelo infiere tanto las identidades de los expertos como el número de expertos activos por token dentro de un rango predefinido, sin depender de reglas de gating heurísticas.
4. Mejoras Empíricas:
   • Exact-$k$: ProbMoE logra un sólido rendimiento en benchmarks (GSM8K, Law, Translation, MBPP, MMLU) en backbones OLMoE y Qwen, igualando o superando a menudo a baselines competitivos como DenseMixer y el Top-$k$ Convencional.
   • Dinámico-$k$: ProbMoE Dynamic-$k$ logra un rendimiento comparable al de Exact-$k$ mientras activa menos expertos en promedio, demostrando una asignación adaptativa efectiva.
5. Diversidad de Enrutamiento: El análisis muestra que ProbMoE induce distribuciones de enrutamiento más amplias y una utilización de expertos más equilibrada en comparación con los baselines, mitigando el colapso de expertos.

Resultos y Significado
El artículo demuestra que ProbMoE mejora la utilización de los expertos y la diversidad del enrutamiento. En estudios de ablación, se muestra que la combinación de muestreo estocástico y optimización basada en marginales es el principal motor de las ganancias de rendimiento, superando configuraciones que utilizan solo uno de estos componentes.

Los autores afirman que ProbMoE proporciona un método con principios para la selección de expertos que mejora la estabilidad del entrenamiento y la eficiencia computacional. Al modelar el enrutamiento de forma probabilística, el marco permite la exploración de combinaciones de expertos alternativas durante el entrenamiento, lo que conduce a distribuciones de enrutamiento mejor calibradas. La variante Dynamic-$k$ muestra además que la cardinalidad adaptativa puede aprenderse para ajustarse a la complejidad del nivel de token, reduciendo los costos efectivos de enrutamiento sin sacrificar el rendimiento de la tarea. El trabajo sugiere que las direcciones futuras incluyen la extensión de ProbMoE al pre-entrenamiento y optimizaciones a nivel de sistema para explotar plenamente la dispersión en entornos dinámicos.