Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance

El artículo presenta ProMoE, un marco de Mixture-of-Experts con una guía de enrutamiento explícita de dos pasos que supera a los métodos actuales en la escalabilidad de Transformadores de Difusión para visión artificial al abordar la redundancia espacial y la heterogeneidad funcional de los tokens visuales.

Lo esencial

¡Claro que sí! Imagina que estás construyendo una gran cocina de restaurante (un modelo de Inteligencia Artificial) para crear pinturas digitales increíbles.

Hasta ahora, los chefs más famosos (los modelos de lenguaje como los que usan en ChatGPT) han tenido un truco genial: en lugar de tener un solo chef que lo haga todo, tienen un equipo de expertos. Si llega un pedido de sushi, el "chef experto en pescado" lo hace; si es pizza, el "chef experto en masa" toma el relevo. A esto se le llama Mezcla de Expertos (MoE).

El problema es que cuando intentaron usar este mismo sistema en la cocina de las imágenes (modelos de difusión como los que crean fotos), no funcionó bien. ¿Por qué?

El Problema: La diferencia entre palabras y píxeles

El paper explica que las palabras y las imágenes son muy diferentes:

1. Las palabras son como fichas de dominó: cada una tiene un significado muy claro y distinto. Si dices "perro", es muy diferente a "gato". Es fácil para un experto especializarse en "animales".
2. Las imágenes son como un lienzo lleno de pintura: si miras una foto de un bosque, hay 1000 trozos de "hojas verdes" que son casi idénticos. Además, en la cocina de imágenes, a veces el chef necesita trabajar con instrucciones (ej: "dibuja un gato") y a veces sin ellas (ej: "dibuja algo al azar").

Los intentos anteriores fallaron porque trataban a todos los trozos de imagen igual, como si fueran palabras. Los expertos se confundían, aprendían cosas repetidas y no se especializaban.

La Solución: ProMoE (El Jefe de Cocina Inteligente)

Los autores crearon ProMoE, un nuevo sistema con un "Jefe de Cocina" (un enrutador) que tiene dos pasos muy claros para organizar a los expertos:

Paso 1: La División por "Tipo de Pedido" (Enrutamiento Condicional)

Imagina que el Jefe de Cocina mira el pedido y dice:

• "¡Oye, este pedido es sin instrucciones (un boceto al azar)!" -> Lo manda inmediatamente al Chef Especialista en Bocetos.
• "¡Este pedido tiene instrucciones (dibuja un gato azul)!" -> Lo manda a la siguiente fase para que los expertos lo analicen mejor.

Esto evita que los expertos de bocetos se mezclen con los de instrucciones específicas, permitiéndoles especializarse mejor.

Paso 2: La Búsqueda por "Semilla" (Enrutamiento Prototípico)

Ahora, para los pedidos con instrucciones, el Jefe no adivina quién lo hace. Tiene una lista de plantillas maestras (prototipos) en su cabeza.

• Si el pedido es "un gato", el Jefe compara la imagen mental con sus plantillas y ve: "¡Ah! Esto se parece mucho a la plantilla del Chef de Animales".
• Si es "un coche", lo manda al Chef de Vehículos.

Además, el paper introduce un truco genial: una regla de oro (una función de pérdida de contraste). Imagina que el Jefe castiga a los expertos si se ponen de acuerdo en hacer cosas aburridas y repetidas. Les dice: "¡Hey, tú y tú, no hagan lo mismo! Ustedes dos deben especializarse en cosas diferentes para que el equipo sea más variado y creativo".

¿Qué logró esto?

Al aplicar esta lógica de "dividir por función" y "guiar por significado":

1. Los expertos se volvieron verdaderos especialistas: Cada uno sabe hacer una cosa muy bien en lugar de intentar hacer todo mediocremente.
2. Ahorro de energía: El modelo crea imágenes de altísima calidad usando menos "chef trabajando" a la vez que los modelos antiguos.
3. Resultados increíbles: En pruebas para crear imágenes (como las de ImageNet), ProMoE superó a los mejores modelos actuales, incluso siendo más pequeño y eficiente.

En resumen

El papel nos dice que no puedes tratar a las imágenes como si fueran texto. Para que la Inteligencia Artificial pinte mejor, necesitas un sistema que sepa diferenciar entre "borrar un lienzo en blanco" y "pintar un paisaje específico", y luego asignar a los mejores artistas para cada tarea, asegurándose de que todos tengan algo único que aportar. ProMoE es ese nuevo sistema de gestión de talento que hace que la cocina de la IA sea más eficiente y creativa.

Resumen técnico

Aquí tienes un resumen técnico detallado del artículo "Routing Matters in MoE: Scaling Diffusion Transformers with Explicit Routing Guidance" (ProMoE), presentado en ICLR 2026.

1. El Problema: La Brecha entre LLMs y Modelos de Difusión en MoE

Aunque la arquitectura de Mezcla de Expertos (MoE) ha permitido escalar masivamente los Modelos de Lenguaje Grande (LLM) manteniendo la eficiencia computacional, su aplicación en Transformers de Difusión (DiT) para generación de imágenes ha tenido resultados limitados. Métodos anteriores como DiT-MoE, EC-DiT y DiffMoE no han logrado las mejoras significativas observadas en el lenguaje.

Los autores identifican dos diferencias fundamentales entre los tokens de lenguaje y los visuales que impiden la especialización efectiva de los expertos en visión:

1. Alta Redundancia Espacial: A diferencia de los tokens de texto, que son semánticamente densos y discretos, los tokens visuales (parches de imagen) son continuos, espacialmente acoplados y altamente redundantes. Esto hace que los expertos tiendan a aprender características homogéneas en lugar de especializarse.
2. Heterogeneidad Funcional: Los modelos de difusión utilizan Classifier-Free Guidance (CFG), lo que introduce dos tipos de tokens funcionalmente distintos: condicionales (con etiquetas o texto) y no condicionales (sin información). Los enfoques de MoE existentes tratan ambos tipos de manera uniforme, ignorando sus roles distintos, lo que obstaculiza la diversidad entre expertos.

2. Metodología: ProMoE

Para abordar estos desafíos, los autores proponen ProMoE, un marco MoE que introduce guía de enrutamiento explícita mediante un enrutador de dos pasos. El objetivo es fomentar la coherencia intra-experto (que un experto procese patrones similares) y la diversidad inter-experto (que diferentes expertos se especialicen en tareas distintas).

A. Enrutador de Dos Pasos

El proceso de enrutamiento se divide en dos etapas secuenciales:

1. Enrutamiento Condicional (Paso 1 - Funcionalidad):

   • Separa explícitamente los tokens de imagen en dos conjuntos basándose en su rol funcional: tokens no condicionales (derivados de entradas sin etiqueta) y tokens condicionales (derivados de entradas con etiqueta).
   • Los tokens no condicionales se asignan determinísticamente a un conjunto dedicado de expertos no condicionales.
   • Los tokens condicionales se envían a la segunda etapa para un enrutamiento más fino.
   • Objetivo: Forzar la segregación funcional y permitir que los expertos se especialicen en la tarea de "guía" vs. "generación".

2. Enrutamiento Prototípico (Paso 2 - Semántica):

   • Para los tokens condicionales, el enrutamiento no se basa en una capa lineal estándar, sino en la similitud coseno entre las incrustaciones de los tokens y un conjunto de prototipos aprendibles ($P$), donde cada prototipo corresponde a un experto específico.
   • Esto permite asignar tokens semánticamente similares al mismo experto en el espacio latente.
   • Se utiliza una función de activación de identidad (en lugar de softmax) para preservar el orden relativo y mejorar la estabilidad del entrenamiento.

B. Pérdida de Contraste de Enrutamiento (Routing Contrastive Loss - RCL)

Para superar las limitaciones de las guías de enrutamiento implícitas (como el clustering de K-means) o las que requieren etiquetas manuales costosas, se propone la RCL.

• Mecanismo: Esta pérdida de contraste actúa directamente sobre los prototipos.
  • Atracción: Empuja cada prototipo hacia el centroide de los tokens asignados a él (fomentando la coherencia intra-experto).
  • Repulsión: Empuja los prototipos lejos de los centroides de los tokens asignados a otros expertos (fomentando la diversidad inter-experto).
• Ventaja: No requiere etiquetas manuales, es robusto y actúa como un regularizador de balanceo de carga basado en semántica, superando a las pérdidas de balanceo tradicionales.

3. Contribuciones Clave

1. Análisis de Diferencias Token-Lenguaje vs. Token-Visión: Demostración empírica de que la redundancia espacial y la heterogeneidad funcional en visión requieren un enfoque de MoE diferente al de los LLMs.
2. Arquitectura ProMoE: Diseño de un enrutador de dos pasos que combina la segregación funcional (condicional/no condicional) con la asignación semántica basada en prototipos.
3. Pérdida de Contraste de Enrutamiento (RCL): Una nueva función de pérdida que mejora explícitamente el proceso de enrutamiento prototípico, asegurando una especialización robusta de los expertos sin necesidad de etiquetas adicionales.
4. Validación Exhaustiva: Demostración de que ProMoE supera a los modelos densos y a los métodos MoE de última generación (SOTA) en múltiples configuraciones y objetivos de entrenamiento.

4. Resultados Experimentales

Los experimentos se realizaron principalmente en el benchmark ImageNet (256x256) bajo objetivos de Rectified Flow (RF) y DDPM.

• Rendimiento Superior: ProMoE supera consistentemente a los modelos DiT densos y a los SOTA de MoE (como DiffMoE y EC-DiT).
  • En Rectified Flow, ProMoE-L reduce el FID en un 24.8% y aumenta el IS en un 19.7% en comparación con el modelo denso equivalente, a pesar de usar menos parámetros activados.
  • ProMoE-L (1.063B parámetros totales) supera al DiffMoE-L más grande (1.846B parámetros totales), demostrando una mayor eficiencia de parámetros.
• Escalabilidad: El modelo muestra mejoras monotónicas al escalar desde tamaños Small (S) hasta XL, y al aumentar el número de expertos.
• Eficiencia Computacional: ProMoE logra tiempos de inferencia y GFLOPs comparables o inferiores a los métodos SOTA, confirmando que las ganancias de rendimiento provienen del diseño metodológico y no de un mayor costo computacional.
• Generalización: En tareas de texto-a-imagen (benchmark GenEval), ProMoE supera significativamente a las bases densas y a los MoE de tipo "Token-Choice", demostrando su capacidad para manejar instrucciones complejas.
• Análisis de Especialización: Las visualizaciones t-SNE y el análisis de uso de expertos confirman que ProMoE logra una especialización clara (diferentes expertos se activan para diferentes clases de imágenes), a diferencia de los métodos basales donde la asignación es homogénea.

5. Significado e Impacto

El trabajo ProMoE es significativo porque resuelve la brecha fundamental en la aplicación de MoE a la generación de imágenes. Al reconocer que los tokens visuales no son semánticamente densos como los de texto, el artículo propone un cambio de paradigma: pasar de un enrutamiento implícito a uno explícito y guiado.

• Eficiencia: Permite escalar modelos de difusión a capacidades masivas (miles de millones de parámetros) sin el costo computacional lineal de los modelos densos.
• Calidad: Mejora sustancialmente la fidelidad y la diversidad de las imágenes generadas, superando los límites actuales de los DiT.
• Dirección Futura: Establece que la guía semántica explícita es crucial para el éxito de MoE en visión, abriendo nuevas vías para el diseño de arquitecturas eficientes en modelos generativos multimodales.

En resumen, ProMoE demuestra que "el enrutamiento importa" y que, con la guía adecuada, los modelos de difusión pueden beneficiarse plenamente de la escalabilidad de la Mezcla de Expertos.