¿Qué es la mezcla de expertos?

Las arquitecturas de mezcla de expertos permiten que los modelos a gran escala, incluso los que comprenden muchos miles de millones de parámetros, reduzcan en gran medida los costes de cálculo durante el preentrenamiento y logren un rendimiento más rápido durante el tiempo de inferencia. En términos generales, logra esta eficiencia mediante la activación selectiva solo de los expertos específicos necesarios para una tarea determinada, en lugar de activar toda la red neuronal para cada tarea.

Aunque gran parte de la implementación moderna de las configuraciones de mezcla de expertos se desarrolló durante (aproximadamente) la última década, la premisa central que subyace a los modelos MoE tiene su origen en el artículo de 1991 "Adaptive Mixture of Local Experts". El documento proponía entrenar un sistema de IA compuesto por redes separadas, cada una de las cuales se especializaba en un subconjunto diferente de casos de entrenamiento. Esto implicó entrenar tanto a las propias "redes de expertos" como a una red de control que determinase qué experto debía utilizarse en cada subtarea. En comparación con un modelo convencional análogo, los autores observaron que su modelo experimental se entrenaba mucho más rápido: alcanzaba el umbral de precisión objetivo en la mitad de épocas de entrenamiento que el modelo convencional.¹

En los últimos años, a medida que los principales modelos de deep learning utilizados para la IA generativa se han vuelto cada vez más grandes y computacionalmente exigentes, la combinación de expertos ofrece un medio para abordar el equilibrio entre la mayor capacidad de los modelos más grandes y la mayor eficiencia de los modelos más pequeños. Esto se ha explorado sobre todo en el campo del procesamiento del lenguaje natural (PNL): algunos de los principales modelos de lenguaje de gran tamaño (LLM), como el Mixtral 8x7B de Mistral y (según algunos informes) el GPT-4,² de OpenAI, han empleado la arquitectura MoE.

Los modelos modernos de deep learning se construyen a partir de redes neuronales artificiales, que comprenden varias capas de nodos interconectados (o "neuronas"). Cada neurona tiene una función de activación: una operación matemática realizada sobre los datos recibidos de la capa anterior, cuyo resultado informa la entrada alimentada a la capa siguiente. Las redes neuronales de retroalimentación (FFN) clásicas procesan la información pasando progresivamente los datos de entrada de las neuronas de una capa a las neuronas de la siguiente capa hasta que llegan a una capa exterior donde se producen las predicciones finales. Algunas arquitecturas de red neuronal incorporan elementos adicionales, como los mecanismos de autoatención de los modelos transformadores, que capturan patrones y dependencias adicionales en los datos de entrada.

Las conexiones entre las diferentes capas y neuronas están mediadas por parámetros del modelo que se pueden aprender: pesos variables (y sesgos) que amplifican o disminuyen la influencia que una parte determinada de la producción de la red tiene en otras partes de la red. Un modelo de deep learning "aprende" mediante el ajuste de estos parámetros, el uso de algoritmos de optimización como el descenso de gradientes, de manera que aumenta la precisión de sus predicciones.

Aunque un mayor número de parámetros aumenta la capacidad del modelo, su capacidad para absorber la información y los patrones que contiene, también aumentan los recursos informáticos necesarios para entrenar y hacer funcionar el modelo. En un modelo típico de deep learning, lo que en este contexto se conoce como un modelo denso, toda la red se ejecuta para procesar todas y cada una de las entradas. Esto crea un equilibrio entre la capacidad del modelo y la practicidad.

enA diferencia de los modelos densos convencionales, la mezcla de expertos utiliza la computación condicional para hacer cumplir la escasez: en lugar de utilizar toda la red para cada entrada, los modelos MoE aprenden una función de mapeo computacionalmente barata que determina qué partes de la red, es decir, qué expertos, son más efectivas para procesar una entrada determinada, como un token individual utilizado para representar una palabra o un fragmento de palabra en tareas de PLN.

Esto permite aumentar la capacidad del modelo (ampliando el número total de parámetros) sin un aumento correspondiente de la carga computacional necesaria para entrenarlo y ejecutarlo (porque no todos esos parámetros se utilizarán necesariamente en un momento dado).

Los modelos MoE procesan los datos al designar a varios "expertos", cada uno su propia subred dentro de una red neuronal más grande, y entrenar una red de cierre (o enrutador) para que active solo a los expertos específicos que mejor se adapten a una entrada determinada.

La principal ventaja del enfoque del MoE es que al aplicar la escasez, en lugar de activar toda la red neuronal para cada token de entrada, se puede aumentar la capacidad del modelo y, al mismo tiempo, mantener constantes los costes computacionales.

A nivel arquitectónico, esto se logra reemplazando las capas tradicionales de red de retroalimentación densa (FFN) con capas (o bloques) de MoE dispersas. En el lenguaje de las redes neuronales, "bloquear" se refiere a un elemento estructural recurrente que desempeña una función específica. En un modelo MoE disperso (SMoE), estos bloques expertos pueden ser de una sola capa, FFN autónomos o incluso MoE anidados en sí mismos.

Por ejemplo, en el modelo lingüístico Mixtral 8x7B de Mistral, cada capa se compone de 8 bloques de retroalimentación (es decir, expertos), cada uno de los cuales tiene 7000 millones de parámetros. Para cada token, en cada capa, una red de enrutador selecciona dos de esos ocho expertos para procesar los datos. A continuación, combina las salidas de esos dos expertos y pasa el resultado a la siguiente capa. Los expertos específicos seleccionados por el enrutador en una capa determinada pueden ser expertos diferentes de los seleccionados en la capa anterior o siguiente.³

Un SMoE puede estar compuesto en su totalidad por capas MoE dispersas, pero muchas arquitecturas de modelos MoE implican tanto bloques dispersos como densos. Lo mismo ocurre con Mixtral, en el que los bloques responsables del mecanismo de autoatención del modelo se comparten entre los 8 expertos. En la práctica, esto hace que designaciones como "8x7B" sean potencialmente engañosas: dado que muchos de los parámetros del modelo son compartidos por cada subred experta de 7 mil millones de parámetros, Mixtral tiene un total de alrededor de 47 mil millones de parámetros, no 56 mil millones, como se podría suponer a través de una simple multiplicación.

Este recuento global de parámetros suele denominarse recuento de parámetros dispersos y, en general, puede entenderse como una medida de la capacidad del modelo. El número de parámetros que se utilizarán realmente para procesar un token individual (a medida que transita por algunos bloques expertos y elude otros) se denomina recuento de parámetros activos, y puede entenderse como una medida de los costes computacionales del modelo. Aunque cada token introducido en Mixtral tiene acceso a 46 700 millones de parámetros, sólo se utilizan 12 900 millones de parámetros activos para procesar un ejemplo dado.

Comprender esta utilización óptima de los recuentos de parámetros es clave para entender las ventajas de los modelos ME. Por ejemplo, Mixtral supera a la variante de 70 000 millones de parámetros de Llama 2 de Meta en la mayoría de las pruebas de referencia, con una velocidad mucho mayor, a pesar de tener un tercio menos de parámetros totales y utilizar menos del 20 % de parámetros activos en el momento de la inferencia.³

Cabe señalar, sin embargo, que el recuento global de parámetros de un ME disperso no es totalmente irrelevante para los requisitos computacionales. A pesar de que sólo se utiliza un subconjunto de parámetros durante la inferencia, la totalidad de los parámetros del modelo deben cargarse en memoria, lo que significa que la eficiencia computacional de la que disfrutan los SMoE en la mayoría de los aspectos no se aplica a sus requisitos de RAM/VRAM.

La clave del concepto (y de la eficacia) de los MoE es que solo algunos de los expertos (y, por lo tanto, los parámetros) de una capa dispersa se activarán en un momento dado, lo que reduce los requisitos computacionales activos.

Aunque la computación condicional se había propuesto durante mucho tiempo como un medio teórico para desvincular las demandas computacionales del aumento de la capacidad del modelo, los desafíos algorítmicos y de rendimiento para su ejecución exitosa no se superaron hasta el artículo de 2017 de Shazeer et al "Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer".⁴

Las ventajas de las capas dispersas sobre las capas densas son más evidentes cuando se trata de datos de alta dimensión en los que los patrones y las dependencias suelen ser complejos y no lineales: por ejemplo, en las tareas de PLN que requieren que un modelo procese una larga secuencia de texto, cada palabra suele estar relacionada con un pequeño subconjunto de otras palabras en esa secuencia. Esto hace que los SMoE sean un área de enorme potencial en el campo de los LLM, donde los modelos MoE bien calibrados pueden disfrutar de los beneficios de la dispersión sin sacrificar el rendimiento. También se han aplicado con éxito modelos MoE de computación dispersa a tareas de visión por ordenador,^{5 6} y siguen siendo un área de estudio activo en ese campo. 

Esta dispersión se consigue mediante la computación condicional: la activación dinámica de parámetros específicos en respuesta a entradas específicas. Por lo tanto, el diseño eficaz de la red de puertas (o "enrutador"), que imponga ese cálculo condicional, es esencial para el éxito de los modelos de MoE.

Se pueden utilizar varios mecanismos de compuerta para seleccionar qué expertos se utilizan en una situación determinada. La función de compuerta adecuada es crucial para el rendimiento del modelo, ya que una mala estrategia de enrutamiento puede dar lugar a que algunos expertos no estén suficientemente formados o estén excesivamente especializados y reducir la eficacia de toda la red.

Un mecanismo de compuerta típico en una configuración tradicional de MoE, presentado en el artículo seminal de Shazeer, utiliza la función softmax: para cada uno de los expertos, el enrutador predice un valor de probabilidad (basado en los pesos de las conexiones de ese experto con el parámetro actual) de que ese experto produzca el mejor resultado para una entrada dada, en lugar de calcular el resultado de todos los expertos, el enrutador calcula sólo la salida de (lo que predice que son) los mejores k expertos para ese ejemplo. Como se ha descrito anteriormente, Mixtral utiliza esta estrategia clásica de enrutamiento top-k: en concreto, utiliza el enrutamiento de los 2 mejores, es decir, k=2, para seleccionar a los 2 mejores expertos (de un total de 8).

En su influyente artículo de 2021, "Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficiency Sparsity", Fedus et al. llevaron el enrutamiento top-k al extremo: trabajando con el LLM T5 de Google, sustituyeron las capas FFN del modelo por 128 expertos e implementaron k=1, también llamado "enrutamiento duro". Incluso al ampliar el modelo hasta un billón de parámetros, esta configuración mejoró las velocidades de preentrenamiento en un 400 %.⁶

A pesar de sus muchos beneficios, los MoE agregan una complejidad significativa al proceso de capacitación. Un inconveniente importante de la estrategia "vainilla" de enrutamiento de top-k es la posibilidad de que la red de enrutamiento converja a la activación de sólo unos pocos expertos. Se trata de un problema que se refuerza a sí mismo: si al principio se selecciona a un puñado de expertos de forma desproporcionada, esos expertos se formarán más rápidamente y seguirán siendo seleccionados en mayor medida, ya que ahora realizan predicciones más fiables que los demás expertos menos formados. Este desequilibrio de la carga significa que otros expertos acaban siendo, en sentido figurado y literal, un peso muerto.

Para mitigar esto, Shazeer et al introdujeron la compuerta de ruido top-k: se añaden algunos ruidos gaussianos a los valores de probabilidad previstos para cada experto, introduciendo algo de aleatoriedad que fomenta la activación de expertos de manera más uniforme. También añadieron dos términos de regularización entrenables a la selección de expertos: la minimización de la pérdida de equilibrio de carga penaliza la dependencia excesiva de cualquier experto, mientras que la minimización de la pérdida de diversidad de expertos recompensa la utilización equitativa de todos los expertos.

El artículo de Google de 2020, "GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding", introdujo dos medios adicionales de equilibrio de carga:

• Enrutamiento aleatorio: mientras que el "mejor" experto de su configuración de los dos mejores se selecciona mediante la función softmax estándar, el segundo experto se elige de forma semialeatoria (con la probabilidad de que se elija a un experto proporcional al peso de su conexión). Por lo tanto, es más probable que se seleccione al segundo experto mejor clasificado, pero ya no se garantiza que sea seleccionado.

• Capacidad de expertos: los autores establecen un umbral que define el número máximo de tokens que puede procesar cualquier experto. Si alguno de los dos primeros expertos elegidos está al límite de su capacidad, la ficha se considera "desbordada" y se salta a la siguiente capa de la red^.7

Como ya se ha mencionado, las ventajas de los modelos dispersos a veces se ven atenuadas por su complejidad añadida. Los retos de la aplicación de los MoE son especialmente evidentes en el proceso de puesta a punto. Los modelos dispersos son más propensos al sobreajuste que los modelos densos tradicionales, y la presencia tanto de capas MoE dispersas como de capas FFN densas complica un enfoque único.

Se han propuesto varias observaciones y enfoques para reducir la inestabilidad a la hora de ajustar los MoE. Los autores del artículo de Switch Transforlmers observaron que las variantes con menos expertos tuvieron más éxito en el ajuste, lo que sugiere que las ventajas de un mayor número de expertos en el preentrenamiento pueden verse contrarrestadas por su obstáculo a la especialización en tareas posteriores. 

En el artículo de 2022 "ST-MoE: Designing Stable and Transferable Sparse Expert Models", Zoph et al compararon los resultados de 5 enfoques diferentes: ajuste de todos los parámetros ("Todos"), sólo de los parámetros no MoE ("Non MoE"), sólo de los parámetros MoE ("MoE"), sólo de los parámetros de autoatención y atención codificador-decodificador ("Attention") y sólo de los parámetros FFN no MoE ("FFN").

• Prácticamente no se encontraron diferencias entre Todos y No MoE

• El ajuste fino sólo de los parámetros de atención dio lugar a una disminución menor del rendimiento.

• Actualizar solo los parámetros del MoE degradó significativamente el rendimiento del modelo, a pesar de que aproximadamente el 80 % de los parámetros del modelo residían en las escasas capas del MoE.

• El FFN fue el único enfoque que mejoró el rendimiento con respecto a la línea de base de Todos. 

Los autores plantearon la hipótesis de que, dado que las capas de expertos representaban sólo una cuarta parte del total de capas de su modelo, y un token verá como mucho sólo dos expertos por capa, aislar los parámetros de MoE da lugar a actualizaciones de pesos menos exhaustivas (y, por tanto, a un mayor sobreajuste y pérdida de entrenamiento).⁸

Un documento de julio de 2023, “Mixture-of-Experts Meets Instruction Tuning”, exploró el impacto del ajuste de instrucciones en los modelos MoE utilizando equivalentes de T5 y Flan-T5 de Google, una versión de T5 ajustada con el protocolo Flan de Google, LLMs as a línea de base. Su experimento comparó cuatro configuraciones: ajuste fino de un modelo T5 denso, ajuste fino de un modelo Flan-T5 denso, ajuste fino de un modelo MoE y ajuste fino de un modelo Flan-MoE ajustado por instrucciones.

Como era de esperar, el equivalente T5 denso superó al MoE tras el ajuste fino. Pero, a la inversa, el modelo Flan-MoE ajustado superó significativamente al modelo Flan-T5 ajustado. Además, la mejora del Flan-MoE en comparación con el MoE fue incluso mayor que la mejora del Flan-T5 sobre el T5 original.9

Esto sugiere que, a pesar de sus dificultades con el ajuste fino estándar, los modelos ME se benefician más del ajuste de las instrucciones que sus homólogos densos. Este hallazgo se hizo realidad con el aclamado lanzamiento de Mixtral 8x7B Instruct, una variante de Mixtral ajustada a las instrucciones que se ofrece como modelo fundacional en IBM watsonx.ai.