El enrutamiento MoE disperso es un mecanismo en las redes neuronales de Mixture of experts que determina qué subredes expertas procesan cada token de entrada. En una capa de mezcla de expertos dispersa, un enrutador o función de compuerta evalúa la entrada y selecciona solo un pequeño número de expertos, típicamente uno o dos, de un grupo más grande. Esto contrasta con la mezcla de expertos densa, donde todos los expertos contribuyen a cada salida. La decisión de enrutamiento se toma por token, permitiendo que diferentes partes de una secuencia de entrada sean procesadas por diferentes expertos. Este diseño aumenta el número total de parámetros de un modelo mientras mantiene el costo computacional por token aproximadamente constante, porque solo los expertos seleccionados se activan durante los pases hacia adelante y hacia atrás.
El concepto surgió de trabajos anteriores sobre mezcla de expertos en la década de 1990, pero el enrutamiento disperso se volvió prácticamente importante en las décadas de 2010 y 2020 como una forma de escalar modelos de lenguaje grandes más allá de los límites de las arquitecturas densas de Transformer (architecture). Al enrutar tokens a expertos especializados, los modelos pueden capturar patrones diversos en los datos sin requerir que todos los parámetros estén activos para cada entrada. El enrutamiento MoE disperso es ahora un componente central de varios modelos de producción prominentes, incluidos los desarrollados por Google DeepMind, OpenAI y otras organizaciones de investigación.
Orígenes Históricos
La idea fundamental de la mezcla de expertos se remonta a principios de la década de 1990. Investigadores como Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan y Geoffrey E. Hinton introdujeron la arquitectura de mezclas adaptativas de expertos locales en 1991, donde una red de compuerta ponderaba las salidas de varias redes expertas. Alrededor del mismo tiempo, la red meta-pi de John Hampshire y Alex Waibel aplicó un esquema de ponderación similar a la clasificación de fonemas del habla, entrenando seis redes neuronales con retardo temporal en datos de seis hablantes japoneses. Estos sistemas tempranos usaban ponderación densa, lo que significa que todos los expertos contribuían a cada salida, pero establecieron los componentes centrales: funciones expertas, una función de compuerta y una combinación ponderada.
El enrutamiento disperso, donde solo se activa un subconjunto de expertos, se exploró a finales de la década de 1990 y principios de la de 2000, pero ganó tracción con el auge del aprendizaje profundo. En 2017, investigadores de Google DeepMind publicaron la capa de mezcla de expertos con compuerta dispersa, que introdujo una red de compuerta entrenable que seleccionaba los top-k expertos para cada entrada. Este trabajo demostró que la activación dispersa podía aumentar drásticamente la capacidad del modelo sin un aumento proporcional en la computación, allanando el camino para aplicaciones a gran escala posteriores.
Mecanismo de Enrutamiento
En una capa MoE dispersa, el mecanismo de enrutamiento consiste en una función de compuerta, típicamente una capa lineal seguida de un softmax sobre la dimensión de expertos. Para una representación de token de entrada \(x\), el enrutador calcula una puntuación para cada uno de los \(n\) expertos, a menudo como \(w(x)_i = \text{softmax}(W_g x)_i\), donde \(W_g\) es una matriz de pesos aprendible. El enrutador luego selecciona los top-k expertos con las puntuaciones más altas, donde k es un hiperparámetro, usualmente establecido en 1 o 2. La salida de la capa es una suma ponderada de las salidas de los expertos seleccionados, con pesos normalizados entre los expertos seleccionados.
Esta selección por token permite que el modelo se especialice: diferentes expertos pueden aprender a manejar diferentes tipos de entradas, como diferentes idiomas, dominios o patrones sintácticos. La decisión de enrutamiento es determinista durante la inferencia, pero durante el entrenamiento, el enrutador debe entrenarse conjuntamente con los expertos. Debido a que la selección top-k no es diferenciable, los investigadores usan técnicas como el estimador de paso recto o la compuerta top-k ruidosa, que añade ruido aprendible a los logits durante el entrenamiento para fomentar la exploración y el equilibrio de carga.
Equilibrio de Carga y Pérdidas Auxiliares
Un desafío importante en el enrutamiento MoE disperso es asegurar que los expertos se usen de manera aproximadamente igual. Sin intervención, el enrutador puede converger a un estado donde unos pocos expertos reciben la mayoría de los tokens, mientras que otros permanecen subutilizados, reduciendo la capacidad efectiva. Para abordar esto, los modelos incorporan pérdidas auxiliares de equilibrio de carga que penalizan la distribución desequilibrada de tokens entre expertos. Un enfoque común, introducido en Switch Transformer (2021) por investigadores de Google (AI), añade un término de pérdida proporcional a la fracción de tokens enrutados a cada experto multiplicada por la probabilidad promedio del enrutador para ese experto. Esto fomenta que el enrutador distribuya los tokens de manera más uniforme.
Otra técnica es la capacidad de experto, que limita el número de tokens que cada experto puede procesar en un lote dado. Si un experto alcanza su capacidad, los tokens excedentes se descartan o se enrutan a una conexión residual. Esto previene que cualquier experto se convierta en un cuello de botella y asegura una computación predecible. Métodos más recientes, como el enrutamiento sin pérdida auxiliar usado en DeepSeek-V3 (2024), utilizan ajustes de sesgo dinámicos para equilibrar la carga sin pérdidas auxiliares explícitas.
Escalado y Eficiencia
El enrutamiento MoE disperso permite escalar los parámetros del modelo mucho más allá de lo que permiten los modelos densos. Por ejemplo, Switch Transformer introdujo modelos con hasta 1.6 billones de parámetros, aunque cada token activaba solo una pequeña fracción de ellos. Similarmente, Mixtral 8x7B, lanzado por Mistral AI en 2023, usa ocho expertos por capa con dos activos por token, logrando un rendimiento comparable a modelos densos más grandes mientras usa menos computación por token. GShard de Google (2020) aplicó MoE disperso a la traducción automática, y modelos posteriores como GLaM (2021) y PaLM (2022) incorporaron capas MoE para eficiencia.
Las ganancias de eficiencia provienen del hecho de que el costo computacional por token escala con el número de expertos activos, no con el número total de parámetros. Esto permite que los modelos tengan más parámetros, lo que puede mejorar la precisión, mientras mantiene los costos de inferencia y entrenamiento manejables. Sin embargo, los modelos MoE dispersos requieren más memoria para almacenar todos los parámetros de los expertos, e introducen sobrecarga de comunicación en el entrenamiento distribuido, ya que los tokens deben enrutarse a expertos que pueden residir en diferentes dispositivos.
Aplicaciones en Modelos de Lenguaje Grandes
El enrutamiento MoE disperso se ha convertido en una técnica estándar en los modelos de lenguaje grandes. Muchos modelos de peso abierto y propietarios usan capas MoE para equilibrar calidad y costo. Por ejemplo, Mixtral 8x7B y Mixtral 8x22B de Mistral AI, Qwen1.5-MoE de Alibaba, y DeepSeek-V2 y V3 de DeepSeek emplean enrutamiento disperso. En el espacio propietario, se informa ampliamente que GPT-4 de OpenAI usa una arquitectura de mezcla de expertos, aunque la compañía no ha confirmado detalles. Los modelos Claude de Anthropic y los modelos Gemini de Google probablemente también incorporan capas MoE, basándose en declaraciones públicas y patentes.
Estos modelos usan enrutamiento para especializar expertos en dominios, idiomas o tareas de razonamiento. Por ejemplo, algunos expertos pueden manejar razonamiento matemático, mientras que otros se centran en generación de código o texto multilingüe. El enrutador aprende a dirigir tokens a expertos apropiados basándose en el contexto de entrada, mejorando el rendimiento general sin requerir que todos los expertos procesen cada token.
Desafíos de Entrenamiento
Entrenar modelos MoE dispersos presenta desafíos únicos más allá del equilibrio de carga. El enrutador puede sufrir inestabilidad, donde pequeños cambios en los parámetros conducen a grandes cambios en las decisiones de enrutamiento, causando divergencia en el entrenamiento. Técnicas como la compuerta top-k ruidosa y el dropout en el enrutador ayudan a estabilizar el entrenamiento. Además, los expertos pueden especializarse demasiado temprano, llevando a una forma de colapso modal donde el modelo falla en generalizar. Los investigadores han propuesto métodos como dropout de expertos y regularización de enrutamiento para mitigar esto.
Otro desafío es la sobrecarga de comunicación en el entrenamiento distribuido. En una configuración típica, cada experto se coloca en un dispositivo diferente, y los tokens deben enviarse desde su dispositivo original al dispositivo que aloja el experto seleccionado. Esta comunicación de todos a todos puede convertirse en un cuello de botella, especialmente para tamaños de lote grandes. Implementaciones eficientes, como las de las bibliotecas Megatron y DeepSpeed, usan horarios de comunicación optimizados para reducir esta sobrecarga.
Desarrollos Recientes
La investigación reciente se ha centrado en mejorar la calidad del enrutamiento y reducir el costo de los modelos MoE. El enfoque de Mezcla de Profundidades (2024) extiende la idea de dispersión a la dimensión de profundidad, enrutando tokens no solo a diferentes expertos sino también omitiendo capas por completo. Otra dirección es la segmentación fina de expertos, donde los expertos son más pequeños y más numerosos, permitiendo una especialización más fina. DeepSeek-V3, por ejemplo, usa 256 expertos por capa con solo 8 activos, logrando un rendimiento de vanguardia a un costo de entrenamiento reducido.
El co-diseño de hardware y software también ha avanzado. Compañías como NVIDIA y AMD han optimizado sus aceleradores para inferencia MoE, y marcos como vLLM y TensorRT-LLM soportan el servicio eficiente de MoE. El creciente interés en el despliegue en el borde ha llevado a investigación sobre comprimir modelos MoE, como poda de expertos y cuantización, para caber en dispositivos con memoria limitada.
Direcciones Futuras
El enrutamiento MoE disperso sigue siendo un área activa de investigación. Preguntas abiertas incluyen cómo determinar automáticamente el número óptimo de expertos y la granularidad de enrutamiento, cómo hacer el enrutamiento más interpretable, y cómo combinar MoE con otras técnicas de eficiencia como Model Pruning y cuantización. A medida que los modelos continúan escalando, el enrutamiento disperso probablemente jugará un papel central en equilibrar capacidad y costo computacional, permitiendo sistemas de IA más potentes sin demandas de recursos prohibitivas.