El enrutamiento en Mixture of Experts (MoE) es un mecanismo central en las arquitecturas modernas de aprendizaje automático que asigna dinámicamente cada token de entrada a un subconjunto de redes expertas especializadas. En lugar de activar todos los parámetros para cada entrada, un enrutador, o red de compuerta, aprende a dirigir los tokens a los expertos más relevantes, equilibrando la capacidad del modelo con la eficiencia computacional. Este enfoque permite que los modelos de lenguaje de gran escala escalen a billones de parámetros mientras mantienen los costos de inferencia manejables, ya que solo una fracción de la red está activa por token.
El concepto se origina en los trabajos de principios de la década de 1990 sobre mezclas adaptativas de expertos locales, pero ganó prominencia en la década de 2010 con la introducción de capas MoE de activación dispersa en el aprendizaje profundo. Las implementaciones modernas, como las utilizadas en modelos de Google DeepMind y OpenAI, dependen de funciones de enrutamiento aprendidas que generan distribuciones de probabilidad sobre los expertos, a menudo incorporando pérdidas auxiliares para garantizar un uso equilibrado. El enrutamiento se ha convertido en una elección de diseño crítica en los sistemas generativos de IA de última generación, influyendo tanto en la eficiencia del entrenamiento como en la calidad final del modelo.
Desarrollo Histórico
La idea de enrutar entradas a componentes especializados precede al aprendizaje profundo, con raíces en métodos de conjunto y redes neuronales modulares. En 1991, Michael I. Jordan y Robert Jacobs publicaron un trabajo fundamental sobre mezclas jerárquicas de expertos, donde una red de compuerta ponderaba las salidas de varias redes feedforward. Este marco temprano estableció el principio central: aprender a particionar el espacio de entrada para que diferentes expertos manejen diferentes regiones.
En 2017, investigadores de Google DeepMind (entonces Google Brain) introdujeron la capa MoE de activación dispersa en un artículo liderado por Noam Shazeer. Este diseño utilizaba una función de compuerta softmax para seleccionar los top-k expertos para cada token, con k típicamente establecido en 1 o 2. Los autores demostraron que las capas MoE podían escalar a cientos de expertos, logrando resultados de última generación en benchmarks de modelado de lenguaje y traducción automática mientras reducían el costo computacional por ejemplo. Este trabajo influyó directamente en arquitecturas posteriores como el Switch Transformer (2021), que simplificó el enrutamiento para seleccionar solo un experto por token, y el marco GShard para entrenamiento masivamente paralelo.
Mecanismos de Enrutamiento
Las funciones de enrutamiento se pueden categorizar en varios tipos. El más común es el enrutamiento por elección de token, donde cada token selecciona independientemente sus top-k expertos basándose en una puntuación aprendida. Esto se implementa como una proyección lineal del estado oculto del token seguida de un softmax sobre los índices de los expertos. Los expertos seleccionados procesan entonces el token, y sus salidas se ponderan por las probabilidades de enrutamiento y se suman.
Una alternativa es el enrutamiento por elección de experto, donde cada experto selecciona los top-k tokens de un lote, asegurando que todos los expertos reciban una carga mínima. Este enfoque, popularizado en el Switch Transformer y refinado posteriormente en modelos como Mixtral, aborda problemas de desequilibrio de carga pero requiere un manejo cuidadoso de las asignaciones token-a-experto.
Otras variantes incluyen el enrutamiento basado en hash, que utiliza funciones hash deterministas para asignar tokens sin parámetros aprendidos, y el enrutamiento jerárquico, donde un sistema de compuerta de dos niveles primero selecciona un grupo de expertos y luego un experto específico dentro de ese grupo. Cada método equilibra flexibilidad, sobrecarga computacional y estabilidad del entrenamiento.
Equilibrio de Carga y Pérdidas Auxiliares
Un desafío importante en el enrutamiento MoE es el desequilibrio de carga: el enrutador puede aprender a seleccionar siempre algunos expertos populares, dejando a otros subentrenados. Para mitigar esto, la mayoría de las implementaciones añaden una pérdida auxiliar de equilibrio de carga que penaliza la distribución desigual de tokens entre expertos. La formulación estándar, introducida en el artículo de Shazeer de 2017, calcula la fracción de tokens enrutados a cada experto y fomenta la uniformidad mediante un término de entropía cruzada.
Enfoques más recientes, como la pérdida de equilibrio de carga del Switch Transformer, utilizan un coeficiente más simple que multiplica el producto punto entre la probabilidad promedio del enrutador y el recuento real de tokens por experto. Algunos sistemas, como GShard, emplean factores de capacidad que limitan cuántos tokens puede procesar cada experto, forzando al enrutador a distribuir el trabajo. Estas técnicas son esenciales para un entrenamiento estable y para prevenir el colapso de expertos, donde algunos expertos se convierten en pesos muertos.
Integración Arquitectónica
El enrutamiento se aplica típicamente dentro de los bloques del transformer, reemplazando la red feedforward (FFN) con una capa MoE. En un Transformer (architecture) estándar, cada token pasa a través de una subcapa de atención multi-cabeza seguida de una FFN por posición. En un transformer MoE, la FFN se reemplaza por un conjunto de FFN expertas, cada una con sus propios parámetros, y un enrutador que selecciona qué expertos activar.
Esta integración permite que los modelos tengan un gran recuento total de parámetros mientras mantienen constante el recuento de parámetros activos por token. Por ejemplo, un modelo con 64 expertos, cada uno con 1 billón de parámetros, tiene 64 billones de parámetros totales pero solo activa 2 billones por token si k=2. Esta propiedad es crucial para escalar a modelos de billones de parámetros, como se ve en sistemas como los modelos de Mixture of Experts de Anthropic y otros laboratorios.
El enrutador en sí es una pequeña red neuronal, a menudo una sola capa lineal seguida de un softmax, con sus propios parámetros entrenados conjuntamente con los expertos. Algunas arquitecturas utilizan enrutadores separados para diferentes capas o cabezas de atención, y trabajos recientes exploran políticas de enrutamiento aprendidas que se adaptan según el tipo de token o la posición.
Dinámicas de Entrenamiento
Entrenar modelos MoE con enrutamiento introduce desafíos únicos. Las decisiones discretas del enrutador (seleccionar top-k expertos) no son diferenciables, por lo que los gradientes fluyen solo a través de las salidas de los expertos seleccionados ponderadas por las probabilidades de enrutamiento. Esto crea un problema de objetivo móvil: a medida que los expertos mejoran, las preferencias del enrutador cambian, lo que potencialmente causa inestabilidad.
Para abordar esto, los investigadores utilizan técnicas como el gating top-k ruidoso, donde se añade ruido gaussiano a los logits de enrutamiento durante el entrenamiento para fomentar la exploración. Otro enfoque es usar un estimador de paso recto para la decisión de enrutamiento, tratando la selección como una asignación dura pero pasando gradientes a través de los pesos de probabilidad. Además, algunos métodos ajustan la temperatura del softmax durante el entrenamiento para afinar gradualmente las decisiones de enrutamiento.
Las pérdidas de equilibrio de carga se ponderan típicamente con un coeficiente pequeño (por ejemplo, 0.01) para evitar dominar la pérdida de la tarea principal. En la práctica, entrenar modelos MoE requiere un ajuste cuidadoso de hiperparámetros, y muchos sistemas emplean pérdidas auxiliares tanto para el equilibrio de carga como para la confianza del enrutador.
Aplicaciones en Modelos de Lenguaje de Gran Escala
El enrutamiento en MoE se ha convertido en una técnica estándar en los modelos de lenguaje de gran escala. Ejemplos notables incluyen:
- El Switch Transformer (2021) de Google, que escaló a 1.6 billones de parámetros con una arquitectura MoE dispersa.
- GShard (2020), que demostró un entrenamiento eficiente de modelos MoE en miles de núcleos TPU.
- Mixtral 8x7B (2023) de Mistral AI, que utiliza 8 expertos con enrutamiento top-2, logrando un rendimiento comparable a modelos densos más grandes.
- DeepSeek-V3 (2024), que emplea un MoE de grano fino con 256 expertos y enrutamiento top-8.
- Se informa que varios modelos de OpenAI y Anthropic utilizan capas MoE, aunque los detalles exactos suelen ser propietarios.
Estos modelos demuestran que el enrutamiento permite un escalado rentable, ya que el cómputo por token permanece constante incluso cuando crecen los parámetros totales. Esto ha hecho que MoE sea una opción preferida para servir modelos grandes en entornos de producción, incluidas plataformas en la nube como Amazon Web Services, Microsoft Azure y Google Cloud.
Eficiencia y Consideraciones de Hardware
El enrutamiento introduce sobrecarga de comunicación en el entrenamiento e inferencia distribuidos, ya que los tokens deben enviarse al experto apropiado, que puede residir en un dispositivo diferente. Este patrón de comunicación de todos-a-todos puede convertirse en un cuello de botella, especialmente con muchos expertos. Los proveedores de hardware como NVIDIA (aunque no listado, implícito) y AMD han desarrollado kernels optimizados para operaciones MoE, y los aceleradores de IA especializados de Groq y SambaNova están diseñados para manejar patrones de activación dispersa de manera eficiente.
Para reducir la comunicación, algunos sistemas utilizan paralelismo de expertos, donde los expertos se replican entre dispositivos, o enfoques híbridos que combinan paralelismo de datos y de expertos. La elección de la granularidad del enrutamiento (a nivel de token vs. a nivel de bloque) también afecta la eficiencia. El enrutamiento a nivel de token ofrece un control más fino pero más comunicación, mientras que el enrutamiento a nivel de bloque agrupa tokens para reducir la sobrecarga.
Desafíos y Direcciones Futuras
A pesar de su éxito, el enrutamiento en MoE enfrenta varios problemas abiertos. Uno es el equilibrio entre especialización de expertos y generalización: los expertos demasiado especializados pueden no transferirse bien a nuevas tareas. Otro es la dificultad de ajustar finamente los modelos MoE, ya que el enrutador puede necesitar adaptarse a nuevas distribuciones de datos. La investigación sobre enrutamiento adaptativo, donde el número de expertos activos varía por token, y sobre políticas de enrutamiento aprendidas que consideran contexto más allá del token actual, está en curso.
Las direcciones futuras incluyen combinar el enrutamiento con otras técnicas de eficiencia como poda de modelos y cuantización, y desarrollar enrutadores que puedan manejar entradas multimodales. A medida que los modelos continúan escalando, el enrutamiento seguirá siendo un mecanismo clave para equilibrar capacidad y cómputo, con aplicaciones potenciales más allá del lenguaje, en visión y aprendizaje por refuerzo.
Véase También
- Transformer (architecture)
- Large language model
- Multi-Head Attention
- Model Pruning
- Google DeepMind
- OpenAI
- Anthropic
Referencias
- Shazeer, N., et al. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.
- Fedus, W., et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.
- Lepikhin, D., et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
- Jiang, A. Q., et al. (2023). Mixtral of Experts.
- DeepSeek-AI (2024). DeepSeek-V3 Technical Report.