Resumen Técnico: DOT-MoE: Transporte Óptimo Diferenciable para la MoEfificación

1. Planteamiento del Problema

El rápido escalado de los Modelos de Lenguaje de Gran Escala (LLMs) ha creado un cuello de botella significativo en la eficiencia de la inferencia. Aunque las arquitecturas de Mezcla de Expertos (MoE) desacoplan el tamaño del modelo del coste de inferencia al activar solo un subconjunto de parámetros por token, el entrenamiento de MoEs desde cero es notoriamente inestable, ávido de datos y computacionalmente costoso.

En consecuencia, la MoEfificación —la conversión de modelos densos preentrenados en MoEs dispersos— ha surgido como una alternativa viable. Sin embargo, los métodos existentes para convertir Redes de Alimentación hacia Adelante (FFN) densas en expertos dispersos dependen en gran medida de estrategias heurísticas (por ejemplo, división aleatoria, agrupamiento basado en pesos o agrupamiento basado en activaciones). Estos enfoques sufren tres limitaciones principales:

1. Separación de Preocupaciones: Tratan típicamente la asignación de neuronas (partición de la FFN) y el entrenamiento del enrutador como etapas separadas y secuenciales, lo que impide que ambas se co-adapten.
2. Optimización de Proxies: Optimizan basándose en estadísticas intermedias (pesos o activaciones) en lugar de la contribución real de las neuronas al resultado, lo que a menudo conduce a una reconstrucción subóptima del comportamiento del modelo denso.
3. Falta de Diferenciabilidad: La naturaleza discreta de la asignación de neuronas impide el flujo de gradiente de extremo a extremo, dificultando la optimización conjunta de la descomposición estructural y la política de enrutamiento.

2. Metodología: DOT-MoE

Los autores proponen DOT-MoE, un marco que formula la descomposición de capas densas en expertos como un problema de Transporte Óptimo Diferenciable (DOT). Este enfoque permite el aprendizaje de extremo a extremo tanto de la descomposición de los expertos como del mecanismo de enrutamiento.

2.1. Formulación de Transporte Óptimo

La idea central es ver la asignación de neuronas como un problema de transporte equilibrado donde la "masa" (neuronas) debe ser transportada desde una distribución de origen (las neuronas de la FFN densa) a una distribución de destino (los expertos) bajo estrictas restricciones de capacidad.

• Origen: $dffn$ neuronas, cada una con una masa unitaria.
• Destino: $E$ expertos, cada uno requiriendo exactamente $s$ neuronas ($s \cdot E = dffn$).
• Función de Coste: En lugar de una matriz de coste fija, el método utiliza una matriz de afinidad aprendible $A$, donde $A_{i,e}$ representa la afinidad de asignar la neurona $i$ al experto $e$. El objetivo es maximizar la afinidad total mientras se satisfacen las restricciones marginales.

2.2. Relajación Diferenciable mediante Sinkhorn-Knopp

Para hacer que el problema de asignación discreta sea diferenciable, DOT-MoE emplea Regularización Entrópica.

• El problema de transporte óptimo se relaja utilizando el algoritmo Sinkhorn-Knopp, que normaliza iterativamente las filas y columnas de la matriz de plan de transporte.
• Esto produce una matriz de asignación suave $M_{soft} \in [0, 1]^{dffn \times E}$ que satisface las restricciones de capacidad en expectativa.
• A medida que el parámetro de temperatura $\tau \to 0$, la solución se aproxima al óptimo discreto.

2.3. Optimización Conjunta con Estimadores de Paso Directo (STE)

Un desafío crítico es que el despliegue requiere asignaciones discretas, y el enrutamiento implica una selección top-k discreta, ambos procesos que no son diferenciables. DOT-MoE aborda esto utilizando Estimadores de Paso Directo (STE):

• Paso hacia adelante (Forward Pass): Utiliza decisiones duras (redondeo codicioso de $M_{soft}$ para la asignación; selección top-k para el enrutamiento).
• Paso hacia atrás (Backward Pass): Enruta los gradientes a través de las contrapartes suaves ($M_{soft}$ y los logits del enrutador $P$).
• Esto permite que los pesos del enrutador ($W_r$) y la matriz de afinidad de asignación ($A$) se co-adapten durante el entrenamiento.

2.4. Objetivo de Entrenamiento

El marco se entrena mediante una Fase de Alineación donde los pesos del modelo denso se mantienen congelados. El modelo minimiza una combinación de pérdidas:

1. Reconstrucción de Salida: Medida mediante la divergencia KL entre la salida del profesor denso y la del estudiante MoE disperso.
2. Modelado de Lenguaje: Pérdida de entropía cruzada sobre el objetivo de modelado de lenguaje.
3. Auxiliares de Estabilidad: Pérdida $z$ del enrutador (para prevenir la explosión de logits) y pérdida de equilibrio de carga (para prevenir el colapso de expertos).

3. Contribuciones Clave

1. Marco de Transporte Óptimo: El artículo introduce una formulación novedosa de la conversión de denso a MoE como un problema de transporte equilibrado, resuelto mediante iteraciones de Sinkhorn diferenciables para imponer estrictas restricciones de capacidad de expertos.
2. Co-Adaptación de Doble Nivel: A diferencia de los métodos anteriores que fijan las asignaciones antes de entrenar los enrutadores, DOT-MoE optimiza conjuntamente la descomposición de neurona-a-experto y el enrutamiento de token-a-experto utilizando STEs complementarios.
3. Construcción Consciente de la Salida: Al modelar directamente la contribución de la neurona al resultado final (en lugar de optimizar proxies como los pesos o las activaciones intermedias), el método logra una mayor fidelidad de reconstrucción.

4. Resultados Experimentales

Los autores evaluaron DOT-MoE a través de tres familias de modelos (LLaMA-2, LLaMA-3, Qwen2.5) y seis benchmarks (ARC-Challenge, Winogrande, HellaSwag, PIQA, SciQ, BoolQ).

• Rendimiento vs. Poda (Pruning): En LLaMA-2 7B con un presupuesto de parámetros del 50%, DOT-MoE logró una perplejidad de 7.99 en WikiText-2, superando significativamente al método de poda estructurada de vanguardia DISP-LLM (9.84).
• Rendimiento vs. MoEfificación Existente: En configuraciones zero-shot, DOT-MoE superó sustancialmente a métodos como CMoE y LLaMA-MoE-v2. Por ejemplo, en LLaMA-2 7B, DOT-MoE alcanzó un 61.5% de precisión promedio comparado con el 44.5% de CMoE.
• Eficiencia: DOT-MoE retiene aproximadamente el 90% del rendimiento del modelo denso original mientras reduce los parámetros activos en un 50%.
• Ajuste Fino (Fine-Tuning): Con solo 1.2B de tokens de ajuste fino continuo, los modelos DOT-MoE recuperaron un rendimiento cercano a sus contrapartes densos (por ejemplo, 66.6% frente a 76.6% en LLaMA-2 7B), superando a los baselines que requerían presupuestos de datos similares o mayores.
• Estudios de Ablación:
  • Granularidad: El rendimiento mejora con un aumento de la granularidad de los expertos hasta un punto de saturación, tras el cual las ganancias disminuyen.
  • Velocidad de Inferencia: A pesar de un mayor número de expertos, el rendimiento de la inferencia se mantiene estable cuando se utilizan kernels de MoE fusionados (por ejemplo, vLLM), ya que la dimensión intermedia total fusionada permanece constante.
  • Dispersión (Sparsity): Los modelos entrenados con niveles de dispersión más altos (75%) se generalizan de forma robusta a regímenes de dispersión aún más extremos en el momento de la inferencia.

5. Significado y Reivindicaciones

El artículo afirma que DOT-MoE representa un avance significativo en la eficiencia de la inferencia de LLMs al proporcionar un método diferenciable y con principios para convertir modelos densos en MoEs dispersos.

• Cambio de Paradigma: Se aleja de la partición heurística y estática hacia una descomposición dinámica y aprendible que respeta la coherencia funcional de los expertos.
• Practicidad: El método permite la esparcidad post-hoc eficiente de modelos preentrenados sin necesidad de un reentrenamiento completo desde cero, reteniendo una alta capacidad mientras reduce drásticamente los parámetros activos.
• Robustez: El marco demuestra que preservar la fidelidad de la FFN requiere modelar explícitamente la contribución de la neurona al resultado, un principio que los métodos basados en proxies existentes no logran capturar.

Los autores concluyen que DOT-MoE ofrece una inicialización superior para modelos dispersos, cerrando la brecha de rendimiento entre arquitecturas densas y dispersas con un mínimo de datos de entrenamiento adicionales.