¿Qué es la mezcla de expertos?

Las arquitecturas de MoE permiten que los modelos a gran escala, incluso aquellos que comprenden muchos miles de millones de parámetros, reduzcan en gran medida los costos informáticos durante el entrenamiento previo y logren un rendimiento más rápido durante el tiempo de inferencia. En términos generales, logra esta eficiencia activando selectivamente solo a 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 configuraciones de MoE se desarrolló durante (aproximadamente) la última década, la premisa central detrás de los modelos MoE se origina en el documento de 1991 “Adaptive Mixture of Local Experts”. El documento proponía entrenar un sistema de IA compuesto por redes separadas, cada una especializada en un subconjunto diferente de casos de entrenamiento. Esto implicó entrenar tanto a las "redes de expertos" como a una red de acceso que determina qué experto debe usarse para cada subtarea. En comparación con un modelo convencional análogo, los autores descubrieron que su modelo experimental era significativamente más rápido de entrenar: alcanzó 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 aprendizaje profundo utilizados para la IA generativa se volvieron cada vez más grandes y exigentes desde el punto de vista computacional, la MoE 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 exploró sobre todo en el campo del procesamiento de lenguaje natural (PLN): algunos de los principales modelos de lenguaje grandes (LLM), como Mixtral 8x7B de Mistral y (según algunos informes) GPT-4 de OpenAI,² emplearon la arquitectura MoE.

Los modelos modernos de aprendizaje profundo se construyen a partir de redes neuronales artificiales, que comprenden múltiples capas de nodos interconectados (o “neuronas”). Cada neurona tiene una función de activación: una operación matemática realizada en los datos recibidos de la capa anterior, cuya salida informa la entrada que alimenta a la siguiente capa. Las redes neuronales de feed-forward (FFN) clásicas procesan la información pasando progresivamente los datos de entrada de las neuronas en una capa a las neuronas en la siguiente capa hasta que llega a una capa externa donde ocurren las predicciones finales. Algunas arquitecturas de redes neuronales 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 diferentes capas y neuronas están mediadas por parámetros de modelos que se pueden aprender: ponderaciones variables (y sesgos) que amplifican o disminuyen la influencia que una parte determinada de la salida de la red tiene en otras partes de la red. Un modelo de aprendizaje profundo "aprende" ajustando estos parámetros, utilizando algoritmos de optimización, como gradiente descendiente, de manera que aumente la precisión de sus predicciones.

Si bien un mayor número de parámetros aumenta la capacidad del modelo (su habilidad para asimilar información y patrones), también aumenta los recursos computacionales necesarios para entrenar y operar el modelo. En un modelo típico de aprendizaje profundo (lo que en este contexto se denomina modelo denso), toda la red se ejecuta para procesar todas las entradas. Esto crea un equilibrio entre la capacidad y la practicidad del modelo.

A diferencia de los modelos densos convencionales, la MoE emplea la computación condicional para imponer la escasez: en lugar de emplear toda la red para cada entrada, los modelos MoE aprenden una función de mapeo computacionalmente barata que determina qué partes de la red (en otras palabras, qué expertos) son más efectivas para procesar una entrada determinada, como un token individual empleado para representar una palabra o fragmento de palabra en tareas de PLN.

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

Los modelos MoE procesan datos mediante la designación de una serie de “expertos”, donde cada uno tiene su propia subred dentro de una red neuronal más grande, y entrena una red de conexión (o enrutador) para activar solo los expertos específicos que mejor se adapten a una entrada determinada.

El principal beneficio del enfoque 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 costos computacionales.

A nivel arquitectónico, esto se logra reemplazando las capas tradicionales de red de feedforward densa (FFN) con capas (o bloques) de MoE dispersas. En la jerga de las redes neuronales, "bloque" se refiere a un elemento estructural recurrente que realiza una función específica. En un modelo MoE disperso (SMoE), estos bloques expertos pueden ser de una sola capa, FFN autónomas o incluso MoE anidadas en sí mismas.

Por ejemplo, en el modelo de lenguaje Mixtral 8x7B de Mistral, cada capa está compuesta por 8 bloques de feedforward, es decir, expertos, cada uno de los cuales tiene 7 mil millones de parámetros. Para cada token, en cada capa, una red de routers selecciona dos de esos ocho expertos para procesar los datos. Luego combina los resultados 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 bloques dispersos y 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 emplearán realmente para procesar un token individual (ya que pasa por algunos bloques expertos y evita otros) se denomina recuento de parámetros activos, y puede entenderse como una medida de los costos computacionales del modelo. Aunque cada ficha introducida en Mixtral tiene acceso a 46.700 millones de parámetros, solo se emplean 12.900 millones de parámetros activos para procesar un ejemplo determinado.

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

Sin embargo, vale la pena señalar que el recuento general de parámetros de una MoE dispersa no es totalmente irrelevante para los requerimientos computacionales. A pesar de emplear solo un subconjunto de parámetros durante la inferencia, la mayoría de los parámetros del modelo se cargan en la memoria, lo que significa que la eficiencia computacional que disfrutan las SMoE en la mayoría de los aspectos no se aplica a sus requisitos de RAM/VRAM.

La clave del concepto (y la eficiencia) de las MoE es que solo algunos de los expertos (y, por lo tanto, parámetros) en 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 desacoplar 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 documento de 2017 de Shazeer et al. “Outrageously Large Neural Networks: The Sparsely-Gate 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 donde los patrones y dependencias son a menudo complejos y no lineales: por ejemplo, en las tareas de PLN que requieren un modelo para procesar una larga secuencia de texto, cada palabra generalmente solo está relacionada con un pequeño subconjunto de otras palabras en esa secuencia. Esto convierte a las SMoE en un área de tremendo potencial en el campo de los LLM, donde los modelos MoE bien calibrados pueden disfrutar de los beneficios de la escasez sin sacrificar el rendimiento. Los modelos MoE escasamente activados también se han aplicado con éxito a tareas de visión artificial,^{5 6} y siguen siendo un área de estudio activo en ese campo. 

Esta escasez se logra mediante el cálculo condicional: la activación dinámica de parámetros específicos en respuesta a entradas específicas. El diseño eficaz de la red de acceso (o "enrutador"), que aplica esa computación condicional, es esencial para el éxito de los modelos MoE.

Se pueden usar varios mecanismos de activación para seleccionar qué expertos se emplean en una situación determinada. La función de activación correcta es fundamental para el rendimiento del modelo, ya que una estrategia de enrutamiento deficiente puede provocar que algunos expertos no estén capacitados o estén demasiado especializados y reducir la eficacia de toda la red.

Un mecanismo de activación típico en una configuración tradicional de MoE, presentado en el artículo seminal de Shazeer, emplea la función softmax: para cada uno de los expertos, por ejemplo, el enrutador predice un valor de probabilidad (basado en las ponderaciones de las conexiones de ese experto al parámetro actual) de ese experto que produce el mejor resultado para una entrada determinada; en lugar de calcular la salida de todos los expertos, el enrutador calcula solo la salida de (lo que predice que serán) los expertos top k para ese ejemplo. Como se describió anteriormente, Mixtral utiliza esta estrategia clásica de enrutamiento top-k: específicamente, utiliza enrutamiento top-2, es decir, k = 2, seleccionando a los 2 mejores expertos (de su 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 escalar el modelo hasta un billón de parámetros, esta configuración mejoró la velocidad de preentrenamiento en un 400 %.⁶

A pesar de sus muchos beneficios, las MoE agregan una complejidad significativa al proceso de entrenamiento. Una desventaja importante de la estrategia "vainilla" de enrutamiento top-k es la posibilidad de que la red de activación converja en activar solo a unos pocos expertos. Se trata de un problema que se refuerza a sí mismo: si al principio se selecciona a unos cuantos expertos de forma desproporcionada, esos expertos se entrenarán más rápido y seguirán siendo seleccionados en mayor medida, ya que ahora realizan predicciones más confiables que los demás expertos menos entrenados. 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 una activación de ruido de top-k: se agrega algo de ruido gaussiano a los valores de probabilidad predichos para cada experto, introduciendo cierta aleatoriedad que fomenta mejor la activación distribuida uniformemente de los expertos. También agregaron dos términos de regularización entrenables a la selección de expertos: minimizar la pérdida de equilibrio de carga penaliza la dependencia excesiva de cualquier experto, mientras que minimizar la pérdida de diversidad de expertos recompensa la utilización equitativa de todos los expertos.

El documento 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 experto "superior" en su configuración de los 2 principales se selecciona mediante la función softmax estándar, el segundo experto se elige de forma semialeatoria (con la probabilidad de que cualquier experto sea elegido proporcional a la ponderación 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 lo sea.

• Capacidad experta: los autores establecen un umbral que define la cantidad máxima de tokens que puede procesar cualquier experto. Si alguno de los 2 mejores expertos elegidos está al máximo de su capacidad, el token se considera "desbordado" y se salta a la siguiente capa de la red.⁷

Como se mencionó anteriormente, las ventajas de los modelos dispersos a veces se ven mitigadas por su complejidad añadida. Los desafíos de implementar MoE son particularmente evidentes en el proceso de ajuste. Los modelos dispersos son más propensos al sobreajuste que los modelos densos tradicionales, y la presencia de capas de MoE dispersas y capas FFN densas complica un enfoque único para todos.

Se propusieron varias observaciones y enfoques para reducir la inestabilidad al afinar los ME. Los autores del artículo de Switch Transforlmers observaron que las variantes con menos expertos tuvieron más éxito en el ajuste fino, lo que sugiere que los beneficios de un mayor número de expertos en el preentrenamiento pueden verse contrarrestados 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 preciso de todos los parámetros ("All"), solo de los parámetros no MoE ("Non MoE"), solo de los parámetros MoE ("MoE"), solo de los parámetros de autoatención y atención codificador-decodificador ("Attention") y solo de los parámetros FFN no MoE ("FFN").

• Casi no se encontraron diferencias entre All y No MoE

• El ajuste preciso de los parámetros de atención solo provocó una disminución menor del rendimiento.

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

• FFN fue el único enfoque que mejoró el rendimiento en relación con la línea de referencia All. 

Los autores plantearon la hipótesis de que, debido a que las capas expertas representaban solo una cuarta parte de las capas totales de su modelo, y un token verá como máximo solo dos expertos por capa, aislar los parámetros de la MoE da como resultado actualizaciones de ponderación menos completas (y, por lo tanto, 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 mediante equivalentes de los LLM T5 y Flan-T5 de Google, una versión de T5 ajustada con el protocolo Flan de Google,  como referencia. Su experimento comparó cuatro configuraciones: el ajuste de un modelo T5 denso, el ajuste de un modelo Flan-T5 denso, el ajuste de un modelo MoE y el ajuste de un modelo Flan-MoE ajustado por instrucciones.

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

Es alentador que esto sugiera que, a pesar de sus dificultades con el ajuste preciso estándar, los modelos MoE en realidad se benefician más del ajuste de instrucciones que sus contrapartes densas. 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.