El auge de la mezcla de expertos para modelos de lenguaje grande eficientes

En el mundo del procesamiento de lenguaje natural (NLP), la búsqueda de construir modelos de lenguaje más grandes y capaces ha sido una fuerza impulsora detrás de muchos avances recientes. Sin embargo, a medida que estos modelos crecen en tamaño, los requisitos computacionales para el entrenamiento y la inferencia se vuelven cada vez más exigentes, empujando contra los límites de los recursos de hardware disponibles.

Entonces, entra en juego la Mezcla de Expertos (MoE), una técnica que promete aliviar esta carga computacional al mismo tiempo que permite el entrenamiento de modelos de lenguaje más grandes y poderosos. A continuación, discutiremos MoE, exploraremos sus orígenes, su funcionamiento interno y sus aplicaciones en modelos de lenguaje basados en transformadores.

Los orígenes de la Mezcla de Expertos

El concepto de Mezcla de Expertos (MoE) se puede remontar a los primeros años 90 cuando los investigadores exploraron la idea de computación condicional, donde partes de una red neuronal se activan selectivamente en función de los datos de entrada. Uno de los trabajos pioneros en este campo fue el artículo “Mezcla adaptativa de expertos locales” de Jacobs et al. en 1991, que propuso un marco de aprendizaje supervisado para un conjunto de redes neuronales, cada una especializada en una región diferente del espacio de entrada.

La idea central detrás de MoE es tener múltiples redes neuronales “expertas”, cada una responsable de procesar un subconjunto de los datos de entrada. Un mecanismo de puerta, generalmente una red neuronal en sí misma, determina qué experto(s) debe procesar una entrada determinada. Este enfoque permite al modelo asignar sus recursos computacionales de manera más eficiente al activar solo los expertos relevantes para cada entrada, en lugar de emplear la capacidad completa del modelo para cada entrada.

A lo largo de los años, varios investigadores han explorado y extendido la idea de computación condicional, lo que ha llevado a desarrollos como MoE jerárquicos, aproximaciones de bajo rango para la computación condicional y técnicas para estimar gradientes a través de neuronas estocásticas y funciones de activación de umbral duro.

Mezcla de Expertos en Transformadores

Mezcla de Expertos

Si bien la idea de MoE ha estado alrededor durante décadas, su aplicación a modelos de lenguaje basados en transformadores es relativamente reciente. Los transformadores, que se han convertido en el estándar de facto para modelos de lenguaje de última generación, están compuestos por múltiples capas, cada una con un mecanismo de autoatención y una red neuronal feed-forward (FFN).

La innovación clave en la aplicación de MoE a los transformadores es reemplazar las capas FFN densas con capas MoE dispersas, cada una consistente en múltiples expertos FFN y un mecanismo de puerta. El mecanismo de puerta determina qué experto(s) debe procesar cada token de entrada, permitiendo al modelo activar selectivamente solo un subconjunto de expertos para una secuencia de entrada determinada.

Uno de los primeros trabajos que demostró el potencial de MoE en transformadores fue el artículo “Redes neuronales enormemente grandes: la capa de Mezcla de Expertos con puerta dispersa” de Shazeer et al. en 2017. Este trabajo introdujo el concepto de una capa MoE con puerta dispersa, que empleaba un mecanismo de puerta que agregaba dispersión y ruido al proceso de selección de expertos, asegurando que solo un subconjunto de expertos se activara para cada entrada.

Desde entonces, varios otros trabajos han avanzado aún más la aplicación de MoE a los transformadores, abordando desafíos como la inestabilidad del entrenamiento, el equilibrio de carga y la inferencia eficiente. Ejemplos notables incluyen el Switch Transformer (Fedus et al., 2021), ST-MoE (Zoph et al., 2022) y GLaM (Du et al., 2022).

Beneficios de la Mezcla de Expertos para Modelos de Lenguaje

El beneficio principal de emplear MoE en modelos de lenguaje es la capacidad de escalar el tamaño del modelo mientras se mantiene un costo computacional relativamente constante durante la inferencia. Al activar selectivamente solo un subconjunto de expertos para cada token de entrada, los modelos MoE pueden lograr el poder expresivo de modelos densos mucho más grandes mientras requieren significativamente menos cálculo.

Por ejemplo, consideremos un modelo de lenguaje con una capa FFN densa de 7 mil millones de parámetros. Si reemplazamos esta capa con una capa MoE que consiste en ocho expertos, cada uno con 7 mil millones de parámetros, el número total de parámetros aumenta a 56 mil millones. Sin embargo, durante la inferencia, si solo activamos dos expertos por token, el costo computacional es equivalente al de un modelo denso de 14 mil millones de parámetros, ya que se calculan dos multiplicaciones de matrices de 7 mil millones de parámetros.

Esta eficiencia computacional durante la inferencia es particularmente valiosa en escenarios de implementación donde los recursos son limitados, como dispositivos móviles o entornos de computación de borde. Además, los requisitos computacionales reducidos durante el entrenamiento pueden conducir a ahorros de energía sustanciales y una huella de carbono más baja, lo que se alinea con el creciente énfasis en prácticas de IA sostenibles.

Desafíos y Consideraciones

Si bien los modelos MoE ofrecen beneficios convincentes, su adopción y despliegue también conllevan varios desafíos y consideraciones:

1. Inestabilidad del Entrenamiento: Los modelos MoE son conocidos por ser más propensos a inestabilidades de entrenamiento en comparación con sus contrapartes densas. Este problema surge de la naturaleza dispersa y condicional de las activaciones de los expertos, lo que puede generar desafíos en la propagación de gradientes y convergencia. Se han propuesto técnicas como la pérdida z del enrutador (Zoph et al., 2022) para mitigar estas inestabilidades, pero se necesita más investigación.
2. Ajuste y Sobreajuste: Los modelos MoE tienden a sobreajustarse más fácilmente durante el ajuste, especialmente cuando la tarea de ajuste tiene un conjunto de datos relativamente pequeño. Este comportamiento se atribuye a la capacidad y dispersión aumentadas de los modelos MoE, lo que puede conducir a una especialización excesiva en los datos de entrenamiento. Se requieren estrategias de regularización y ajuste cuidadosas para mitigar este problema.
3. Requisitos de Memoria: Si bien los modelos MoE pueden reducir los costos computacionales durante la inferencia, a menudo tienen requisitos de memoria más altos en comparación con modelos densos de tamaño similar. Esto se debe a que todos los pesos de los expertos deben cargarse en la memoria, incluso si solo se activa un subconjunto para cada entrada. Las limitaciones de memoria pueden limitar la escalabilidad de los modelos MoE en dispositivos con recursos limitados.
4. Equilibrio de Carga: Para lograr una eficiencia computacional óptima, es crucial equilibrar la carga entre los expertos, asegurando que ningún experto esté sobrecargado mientras otros permanecen subutilizados. Este equilibrio de carga se logra típicamente a través de pérdidas auxiliares durante el entrenamiento y un ajuste cuidadoso del factor de capacidad, que determina la cantidad máxima de tokens que se pueden asignar a cada experto.
5. Sobrecarga de Comunicación: En escenarios de entrenamiento y inferencia distribuidos, los modelos MoE pueden introducir una sobrecarga de comunicación adicional debido a la necesidad de intercambiar información de activación y gradiente entre expertos que residen en diferentes dispositivos o aceleradores. Se necesitan estrategias de comunicación eficientes y un diseño de modelo consciente del hardware para mitigar esta sobrecarga.

A pesar de estos desafíos, el potencial de los modelos MoE para permitir modelos de lenguaje más grandes y capaces ha impulsado esfuerzos de investigación significativos para abordar y mitigar estos problemas.

Ejemplo: Mixtral 8x7B y GLaM

Para ilustrar la aplicación práctica de MoE en modelos de lenguaje, consideremos dos ejemplos notables: Mixtral 8x7B y GLaM.

Mixtral 8x7B es una variante MoE del modelo de lenguaje Mistral, desarrollado por Anthropic. Consiste en ocho expertos, cada uno con 7 mil millones de parámetros, lo que resulta en un total de 56 mil millones de parámetros. Sin embargo, durante la inferencia, solo se activan dos expertos por token, reduciendo efectivamente el costo computacional a uno equivalente a un modelo denso de 14 mil millones de parámetros.

Mixtral 8x7B ha demostrado un rendimiento impresionante, superando al modelo Llama de 70 mil millones de parámetros mientras ofrece tiempos de inferencia mucho más rápidos. Una versión ajustada de Mixtral 8x7B, llamada Mixtral-8x7B-Instruct-v0.1, también se ha lanzado, mejorando aún más sus capacidades para seguir instrucciones de lenguaje natural.

Otro ejemplo destacado es GLaM (Google Language Model), un modelo MoE de gran escala desarrollado por Google. GLaM emplea una arquitectura de transformador solo decodificador y se entrenó en un conjunto de datos masivo de 1,6 billones de tokens. El modelo logra un rendimiento impresionante en evaluaciones de pocos disparos y un disparo, igualando la calidad de GPT-3 mientras utiliza solo un tercio de la energía necesaria para entrenar a GPT-3.

El éxito de GLaM se puede atribuir a su arquitectura MoE eficiente, que permitió el entrenamiento de un modelo con una cantidad vasta de parámetros mientras se mantenían requisitos computacionales razonables. El modelo también demostró el potencial de los modelos MoE para ser más eficientes en términos de energía y sostenibilidad ambiental en comparación con sus contrapartes densas.

La Arquitectura Grok-1

GROK MEZCLA DE EXPERTOS

Grok-1 es un modelo MoE basado en transformadores con una arquitectura única diseñada para maximizar la eficiencia y el rendimiento. Veamos las especificaciones clave:

1. Parámetros: Con un total de 314 mil millones de parámetros, Grok-1 es el modelo LLM abierto más grande hasta la fecha. Sin embargo, gracias a la arquitectura MoE, solo se activan el 25% de los pesos (aproximadamente 86 mil millones de parámetros) en cualquier momento, mejorando las capacidades de procesamiento.
2. Arquitectura: Grok-1 emplea una arquitectura de Mezcla de 8 Expertos, donde cada token se procesa con dos expertos durante la inferencia.
3. Capas: El modelo consiste en 64 capas de transformadores, cada una con mecanismos de autoatención y bloques densos.
4. Tokenización: Grok-1 utiliza un tokenizador SentencePiece con un tamaño de vocabulario de 131.072 tokens.
5. Incrustaciones y Codificación Posicional: El modelo cuenta con incrustaciones de 6.144 dimensiones y emplea incrustaciones de posición rotativas, lo que permite una interpretación más dinámica de los datos en comparación con las codificaciones de posición fijas tradicionales.
6. Atención: Grok-1 utiliza 48 cabezas de atención para consultas y 8 cabezas de atención para claves y valores, cada una con un tamaño de 128.
7. Longitud de Contexto: El modelo puede procesar secuencias de hasta 8.192 tokens de longitud, utilizando precisión bfloat16 para cálculos eficientes.

Rendimiento y Detalles de Implementación

Grok-1 ha demostrado un rendimiento impresionante, superando al modelo LLaMa 2 70B y Mixtral 8x7B con una puntuación MMLU del 73%, mostrando su eficiencia y precisión en diversas pruebas.

Sin embargo, es importante destacar que Grok-1 requiere recursos GPU significativos debido a su tamaño. La implementación actual en la versión de código abierto se centra en validar la corrección del modelo y emplea una implementación de capa MoE ineficiente para evitar la necesidad de kernels personalizados.

No obstante, el modelo admite particionamiento de activación y cuantización de 8 bits, lo que puede optimizar el rendimiento y reducir los requisitos de memoria.

En un movimiento notable, xAI ha lanzado Grok-1 bajo la licencia Apache 2.0, haciendo que sus pesos y arquitectura estén disponibles para la comunidad global para su uso y contribución.

El lanzamiento de código abierto incluye un repositorio de código de ejemplo JAX que demuestra cómo cargar y ejecutar el modelo Grok-1. Los usuarios pueden descargar los pesos de los puntos de control utilizando un cliente de torrent o directamente a través del HuggingFace Hub, facilitando el acceso a este modelo innovador.

El Futuro de la Mezcla de Expertos en Modelos de Lenguaje

A medida que la demanda de modelos de lenguaje más grandes y capaces continúa creciendo, se espera que la adopción de técnicas MoE gane aún más impulso. Los esfuerzos de investigación en curso se centran en abordar los desafíos restantes, como mejorar la estabilidad del entrenamiento, mitigar el sobreajuste durante el ajuste y optimizar los requisitos de memoria y comunicación.

Una dirección prometedora es la exploración de arquitecturas MoE jerárquicas, donde cada experto está compuesto por múltiples subexpertos. Este enfoque podría potencialmente permitir una mayor escalabilidad y eficiencia computacional mientras se mantiene el poder expresivo de los modelos grandes.

Además, el desarrollo de sistemas de hardware y software optimizados para modelos MoE es un área de investigación activa. Aceleradores especializados y marcos de entrenamiento distribuidos diseñados para manejar eficientemente los patrones de cálculo dispersos y condicionales de los modelos MoE podrían mejorar aún más su rendimiento y escalabilidad.

Por otra parte, la integración de técnicas MoE con otros avances en modelado de lenguaje, como mecanismos de atención dispersa, estrategias de tokenización eficientes y representaciones multimodales, podría conducir a modelos de lenguaje aún más poderosos y versátiles capaces de abordar una amplia gama de tareas.

Conclusión

La técnica de Mezcla de Expertos ha surgido como una herramienta poderosa en la búsqueda de modelos de lenguaje más grandes y capaces. Al activar selectivamente expertos en función de los datos de entrada, los modelos MoE ofrecen una solución prometedora a los desafíos computacionales asociados con el escalado de modelos densos. Aunque aún hay desafíos que superar, como la inestabilidad del entrenamiento, el sobreajuste y los requisitos de memoria, los beneficios potenciales de los modelos MoE en términos de eficiencia computacional, escalabilidad y sostenibilidad ambiental los convierten en un área emocionante de investigación y desarrollo.

A medida que el campo del procesamiento de lenguaje natural continúa empujando los límites de lo que es posible, la adopción de técnicas MoE probablemente desempeñará un papel crucial en la habilitación de la próxima generación de modelos de lenguaje. Al combinar MoE con otros avances en arquitectura de modelos, técnicas de entrenamiento y optimización de hardware, podemos esperar modelos de lenguaje aún más poderosos y versátiles que puedan comprender y comunicarse con los humanos de manera natural y fluida.