Una red de enrutamiento, también conocida como función de compuerta o función de ponderación, es un componente en los modelos de aprendizaje automático de mezcla de expertos (MoE). Toma una entrada y produce un conjunto de pesos que determinan cuánto contribuye cada red experta a la salida. En los modelos de lenguaje grandes modernos, la red de enrutamiento se utiliza a menudo para seleccionar solo un pequeño subconjunto de expertos para cada token de entrada, una técnica llamada activación dispersa, que reduce el costo computacional mientras mantiene la capacidad del modelo.
El concepto de enrutamiento en MoE se remonta a trabajos tempranos sobre máquinas de comité y mezclas adaptativas de expertos locales. En una arquitectura MoE típica, hay múltiples redes expertas, cada una potencialmente especializada en diferentes regiones del espacio de entrada, y una red de enrutamiento que aprende a asignar entradas a los expertos apropiados. La salida del enrutador puede ser una distribución de probabilidad sobre los expertos, o un conjunto de pesos no negativos que se combinan con las salidas de los expertos mediante una suma ponderada.
Arquitectura Básica
En una configuración MoE estándar, la red de enrutamiento es una función \( w(x) \) que mapea una entrada \( x \) a un vector de pesos \( (w(x)_1, ..., w(x)_n) \), donde \( n \) es el número de expertos. La salida final se calcula como \( f(x) = \sum_{i=1}^n w(x)_i f_i(x) \), donde \( f_i \) es el \( i \)-ésimo experto. Tanto el enrutador como los expertos se entrenan conjuntamente, típicamente mediante descenso de gradiente en una función de pérdida como el error cuadrático medio o la entropía cruzada.
La red de enrutamiento en sí misma es a menudo una red neuronal pequeña, como una capa lineal seguida de una activación softmax. En los modelos MoE dispersos, el enrutador puede usar un mecanismo de selección top-k, donde solo se activan los expertos con los pesos más altos, y el resto se ignora. Esta dispersión es clave para reducir el costo computacional, ya que permite que el modelo use solo una fracción de sus parámetros totales para cada entrada.
Desarrollo Histórico
La idea de usar una red de compuerta para combinar múltiples expertos se exploró en la década de 1990. Un ejemplo temprano es la red meta-pi, reportada por Hampshire y Waibel, que usaba una suma ponderada de salidas de expertos y se entrenaba en una tarea de clasificación de fonemas. En sus experimentos, entrenaron seis expertos, cada uno una red neuronal con retardo temporal, para clasificar el habla de seis hablantes japoneses. Observaron que la red de enrutamiento aprendió a dedicar cinco expertos a cinco hablantes individuales, mientras que la voz del sexto hablante fue manejada por una combinación de expertos para otros hablantes masculinos.
Otro modelo temprano influyente fue el de mezclas adaptativas de expertos locales, que usaba un modelo de mezcla gaussiana. En este enfoque, cada experto predecía una distribución gaussiana con una covarianza fija, y el enrutador era una función lineal-softmax que asignaba pesos basados en la entrada. Este modelo demostró que el enrutamiento podía usarse para particionar el espacio de entrada en regiones homogéneas, con cada experto especializándose en una región local.
Papel en los Modelos de Lenguaje Grandes Modernos
En la década de 2020, las redes de enrutamiento se convirtieron en un componente central en los modelos de lenguaje grandes (LLM) que usan capas MoE. Modelos como los desarrollados por empresas como Google DeepMind y OpenAI han adoptado arquitecturas MoE dispersas para escalar el número de parámetros sin aumentar proporcionalmente el costo de inferencia. En estos modelos, cada capa de transformador puede contener múltiples expertos, y una red de enrutamiento selecciona algunos expertos por token.
Por ejemplo, en un LLM basado en transformadores, las incrustaciones de tokens de entrada se pasan a través de una red de enrutamiento que calcula una distribución de probabilidad sobre los expertos. El enrutador luego selecciona los expertos top-k (por ejemplo, k=2 o k=4) y calcula la salida como una suma ponderada de las salidas de esos expertos. Esto permite que el modelo tenga miles de millones de parámetros mientras solo activa un pequeño subconjunto para cada token, haciendo que el entrenamiento y la inferencia sean más eficientes.
Activación Dispersa y Equilibrio de Carga
La activación dispersa es un beneficio clave de las redes de enrutamiento en los modelos MoE. Al activar solo unos pocos expertos por entrada, el modelo puede tener un gran número total de parámetros pero un costo computacional efectivo mucho menor. Sin embargo, esto introduce desafíos en el equilibrio de carga: si el enrutador selecciona consistentemente los mismos pocos expertos, esos expertos se sobrecargan mientras que otros permanecen subutilizados. Para abordar esto, los modelos MoE modernos a menudo incluyen funciones de pérdida auxiliares que fomentan que el enrutador distribuya los tokens de manera más uniforme entre los expertos.
Se han propuesto varias técnicas para mejorar el rendimiento del enrutador, como agregar ruido a los logits del enrutador durante el entrenamiento para promover la exploración, o usar selección top-k diferenciable. Algunos modelos también usan un esquema de enrutamiento jerárquico, donde un enrutador de primer nivel selecciona un grupo de expertos, y un enrutador de segundo nivel selecciona dentro de ese grupo.
Relación con Otras Técnicas
Las redes de enrutamiento están estrechamente relacionadas con otros conceptos en el aprendizaje automático, como atención de múltiples cabezas y atención cruzada, que también implican ponderar diferentes componentes basados en la entrada. Sin embargo, mientras que los mecanismos de atención ponderan diferentes partes de la secuencia de entrada, una red de enrutamiento pondera diferentes redes expertas, que son típicamente aproximadores de funciones independientes.
Las redes de enrutamiento también comparten similitudes con poda de modelos en que ambas buscan reducir el costo computacional, pero la poda elimina parámetros permanentemente, mientras que el enrutamiento selecciona dinámicamente qué parámetros usar para cada entrada. Además, las redes de enrutamiento pueden verse como una forma de aprendizaje por conjuntos, ya que combinan las salidas de múltiples modelos, pero a diferencia de los conjuntos tradicionales, los expertos se entrenan conjuntamente con el enrutador.
Implementación y Entrenamiento
En la práctica, las redes de enrutamiento se implementan como una capa lineal que toma la representación de entrada y produce logits para cada experto. Los logits luego se pasan a través de una función softmax para producir pesos. Durante el entrenamiento, el enrutador y los expertos se actualizan conjuntamente usando retropropagación. La función de pérdida típicamente incluye tanto la pérdida de la tarea (por ejemplo, entropía cruzada para modelado de lenguaje) como una pérdida auxiliar de equilibrio de carga.
Un desafío en el entrenamiento de redes de enrutamiento es que la selección discreta de expertos (por ejemplo, top-k) no es diferenciable. Para superar esto, muchas implementaciones usan los pesos softmax directamente durante el entrenamiento, o usan un estimador de paso directo para la selección top-k. Algunos modelos también usan una red de enrutamiento separada para cada capa, permitiendo que diferentes capas se especialicen en diferentes tipos de decisiones de enrutamiento.
Aplicaciones Más Allá de los Modelos de Lenguaje
Aunque las redes de enrutamiento son más prominentes en los LLM, también se usan en otros dominios. Por ejemplo, en visión por computadora, las capas MoE con enrutadores se han aplicado a tareas de clasificación y generación de imágenes. En reconocimiento de voz, los primeros modelos MoE usaban enrutadores para particionar el espacio acústico. Las redes de enrutamiento también se usan en aprendizaje por refuerzo, donde diferentes expertos pueden representar diferentes políticas, y el enrutador selecciona la política apropiada basada en el estado.
En el contexto de la investigación en inteligencia artificial, las redes de enrutamiento son un área activa de estudio, con trabajo continuo en mejorar la eficiencia del enrutamiento, la interpretabilidad y la escalabilidad. Investigadores en instituciones como Stanford AI Lab y Berkeley AI Research han contribuido a comprender el comportamiento de los enrutadores en modelos a gran escala.
Desafíos y Direcciones Futuras
A pesar de su éxito, las redes de enrutamiento enfrentan varios desafíos. Un problema es que el enrutador puede convertirse en un cuello de botella, ya que debe procesar cada entrada y tomar decisiones de enrutamiento rápidamente. Otro desafío es que el enrutador puede aprender a enrutar entradas de maneras que no son óptimas para la tarea general, lo que lleva a un rendimiento degradado. Además, la pérdida auxiliar de equilibrio de carga puede interferir con la pérdida principal de la tarea, requiriendo un ajuste cuidadoso.
La investigación futura puede explorar mecanismos de enrutamiento más sofisticados, como políticas de enrutamiento aprendidas que consideren el contenido de la entrada de manera más matizada, o enrutamiento jerárquico que reduzca el número de expertos considerados en cada nivel. También hay interés en hacer que los enrutadores sean más interpretables, para que sea posible comprender por qué una entrada particular se enruta a un cierto experto.
En resumen, la red de enrutamiento es un componente fundamental en los modelos MoE, permitiendo un escalado eficiente y especialización. Su desarrollo desde las primeras funciones de compuerta hasta los enrutadores dispersos modernos en los LLM ilustra la evolución de las técnicas de aprendizaje automático a lo largo de varias décadas.