dMoE: dLLMs with Learnable Block Experts

Resumen Técnico: dMoE – Modelos de Lenguaje de Gran Escala de Difusión con Expertos de Bloque Aprendibles

Planteamiento del Problema
Los Modelos de Lenguaje de Gran Escala de Difusión (dLLMs) han surgido como una alternativa competitiva a los modelos autorregresivos, ofreciendo la ventaja de la decodificación en paralelo. Para escalar su capacidad, los dLLMs se están integrando cada vez más con arquitecturas de Mezcla de Expertos (MoE). Sin embargo, existe un desajuste fundamental entre la naturaleza de decodificación por bloques de los dLLMs y el paradigma de selección de expertos a nivel de token heredado de los modelos MoE autorregresivos. En los dLLMs, una única pasada hacia adelante procesa un bloque de tokens con dependencias bidireccionales, mientras que las capas MoE convencionales enrutan cada token de forma independiente. Esta independencia provoca un aumento drástico en el número de expertos activados de forma única dentro de una sola pasada hacia adelante. El análisis empírico revela que esta explosión en la activación de expertos únicos convierte el acceso a la memoria en el principal cuello de botella de la inferencia, haciendo que la latencia de MoE domine el tiempo de inferencia de extremo a extremo y aumente linealmente con el número de expertos únicos.

Metodología: dMoE
Los autores proponen dMoE, un marco de trabajo MoE a nivel de bloque diseñado para reducir el número de expertos activados de forma única sin alterar el número de expertos seleccionados por token. El método se basa en dos observaciones empíricas clave:

1. Los puntajes de expertos a nivel de token (pesos del enrutador) sirven como indicadores efectivos de la importancia de los expertos.
2. La concentración de expertos varía significativamente a través de diferentes pasos de eliminación de ruido y bloques.

Basándose en estas observaciones, dMoE emplea una estrategia de enrutamiento de grueso a fino:

• Agregación: En lugar de enrutar tokens individualmente, dMoE primero agrega los puntajes de expertos a nivel de token dentro de un bloque para generar una distribución de expertos a nivel de bloque unificada.
• Selección de Coreset Adaptativa: En lugar de imponer un tamaño de coreset fijo, el método normaliza los puntajes a nivel de bloque y aplica un criterio top-p para seleccionar un coreset de expertos adaptativo. Esto permite que el tamaño del coreset se ajuste dinámicamente según la concentración de la distribución de expertos (un coreset más pequeño para distribuciones concentradas, uno más grande para las dispersas).
• Enrutamiento Restringido: El enrutamiento original a nivel de token se restringe entonces para que ocurra únicamente dentro de este coreset seleccionado.
• Estrategia de Entrenamiento: Durante el entrenamiento, dMoE utiliza un paradigma de autodestilación, donde el modelo genera objetivos de supervisión para sí mismo utilizando el mismo procedimiento de enrutamiento a nivel de bloque, asegurando la alineación entre el entrenamiento y la inferencia.

Contribuciones Clave

• Identificación del Cuello de Botella: El artículo establece que la latencia limitada por la memoria, causada por la excesiva activación de expertos únicos, es la principal barrera de eficiencia en los dLLM con MoE, distinta de los desafíos en los modelos autorregresivos.
• Marco de Enrutamiento a Nivel de Bloque: dMoE introduce un mecanismo novedoso para pasar del enrutamiento a nivel de token al enrutamiento a nivel de bloque, comprimiendo eficazmente el recuento de expertos únicos mientras se mantiene la carga computacional por token.
• Adaptabilidad Dinámica: Al utilizar un criterio top-p, el método se adapta dinámicamente a las variaciones en las características de enrutamiento a través de diferentes pasos de eliminación de ruido y bloques, evitando la rigidez de la selección de tamaño fijo.

Resultados Experimentales
Evaluado en el modelo de código abierto de última generación LLaDA2.0-mini a través de cuatro bancos de pruebas (MATH500, GSM8K, ARC-C y MMLU), dMoE demuestra ganancias significativas de eficiencia con una pérdida de rendimiento mínima:

• Compresión de Expertos: En promedio, dMoE reduce el número de expertos activados de forma única de 69.5 a 14.6 (una reducción de ~4.77×), logrando una reducción del 79.04% en el recuento de expertos únicos.
• Retención de Rendimiento: El método retiene el 99.11% del rendimiento del modelo original en promedio. Incluso bajo una compresión agresiva (reduciendo los expertos cerca del límite de 8 por token), la degradación del rendimiento sigue siendo insignificante.
• Ganancias de Eficiencia:
  • Uso de Memoria: Reducido en un 76.64% a 79.84% en comparación con el modelo original.
  • Latencia: Logra una aceleración de la latencia de extremo a extremo de 1.14× a 1.66×, impulsada principalmente por la reducción de la latencia del kernel de MoE.
• Comparación: dMoE supera a las líneas base existentes (incluyendo la selección Top-4 y variantes de DES) en la relación rendimiento-eficiencia, particularmente en ratios de compresión altos donde otros métodos sufren caídas significativas de precisión.

Significancia
El artículo posiciona a dMoE como un paso fundacional hacia el enrutamiento eficiente a nivel de bloque en los dLLM con MoE. Demuestra que el cuello de botella de la memoria en los modelos de MoE basados en difusión puede mitigarse eficazmente agregando las decisiones de enrutamiento a nivel de bloque. El trabajo establece una base sólida para la investigación futura en el escalado de dLLMs, probando que la compresión agresiva de expertos es alcanzable sin comprometer las capacidades de razonamiento del modelo. El método se presenta como ajustable, permitiendo a los profesionales ajustar el número de expertos activados para adaptarse a restricciones de hardware específicas o requisitos de aplicación.