Mixtral of Experts

Albert Q. Jiang; Alexandre Sablayrolles; Antoine Roux; Arthur Mensch; Blanche Savary; Chris Bamford; Devendra Singh Chaplot; Diego de las Casas; Emma Bou Hanna; Florian Bressand; Gianna Lengyel; Guillaume Bour; Guillaume Lample; Lélio Renard Lavaud; Lucile Saulnier; Marie-Anne Lachaux; Pierre Stock; Sandeep Subramanian; Sophia Yang; Szymon Antoniak; Teven Le Scao; Théophile Gervet; Thibaut Lavril; Thomas Wang; Timothée Lacroix; William El Sayed

Artículos · arXiv 2401.04088

Mixtral of Experts

Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Emma Bou Hanna, Florian Bressand, Gianna Lengyel, Guillaume Bour et al.

Publicado 2024-01-08cs.LGcs.CL arXiv PDF

Leer con traducción e IA

Resumen en tres líneas

🎯 Escalar la capacidad y el rendimiento de los transformers sin incurrir en una latencia de inferencia y costes de cómputo prohibitivos.
🔧 Una arquitectura tipo decoder-only Sparse Mixture of Experts (SMoE) que sustituye los bloques feedforward estándar por 8 expertos SwiGLU, seleccionando 2 expertos por token mediante un enrutador aprendido.
📈 Accede a un total de 47B de parámetros con solo 13B activos por token, superando a Llama 2 70B en MMLU (70.6% vs. 69.9%) y GSM8K (74.4% vs. 69.6%).

Resumen

Presentamos Mixtral 8x7B, un modelo de lenguaje Sparse Mixture of Experts (SMoE). Mixtral comparte la misma arquitectura que Mistral 7B, con la diferencia de que cada capa está compuesta por 8 bloques feedforward (es decir, expertos). Para cada token y en cada capa, una red de enrutamiento selecciona dos expertos para procesar el estado actual y combinar sus salidas. Aunque cada token solo interactúa con dos expertos, los expertos seleccionados pueden variar en cada paso temporal. Como resultado, cada token tiene acceso a 47B de parámetros, pero solo utiliza 13B de parámetros activos durante la inferencia. Mixtral fue entrenado con un tamaño de contexto de 32k tokens y supera o iguala a Llama 2 70B y a GPT-3.5 en todos los benchmarks evaluados. En particular, Mixtral supera ampliamente a Llama 2 70B en matemáticas, generación de código y benchmarks multilingües. También proporcionamos un modelo ajustado para seguir instrucciones, Mixtral 8x7B - Instruct, que sobrepasa a GPT-3.5 Turbo, Claude-2.1, Gemini Pro y al modelo Llama 2 70B - chat en evaluaciones humanas. Tanto el modelo base como el instruct se publican bajo la licencia Apache 2.0.

Resumen original (inglés)

We introduce Mixtral 8x7B, a Sparse Mixture of Experts (SMoE) language model. Mixtral has the same architecture as Mistral 7B, with the difference that each layer is composed of 8 feedforward blocks (i.e. experts). For every token, at each layer, a router network selects two experts to process the current state and combine their outputs. Even though each token only sees two experts, the selected experts can be different at each timestep. As a result, each token has access to 47B parameters, but only uses 13B active parameters during inference. Mixtral was trained with a context size of 32k tokens and it outperforms or matches Llama 2 70B and GPT-3.5 across all evaluated benchmarks. In particular, Mixtral vastly outperforms Llama 2 70B on mathematics, code generation, and multilingual benchmarks. We also provide a model fine-tuned to follow instructions, Mixtral 8x7B - Instruct, that surpasses GPT-3.5 Turbo, Claude-2.1, Gemini Pro, and Llama 2 70B - chat model on human benchmarks. Both the base and instruct models are released under the Apache 2.0 license.

Guía de lectura

Problema y motivación

Los modelos de lenguaje autorregresivos densos incurren en costes computacionales directamente proporcionales a su recuento total de parámetros para cada token procesado. A medida que los modelos escalan a 70B de parámetros o más, desplegarlos con baja latencia y alto rendimiento (throughput) resulta computacionalmente muy costoso. Los autores buscan escalar la capacidad del modelo manteniendo el coste de inferencia por token comparable al de modelos mucho más pequeños, creando un modelo de pesos abiertos (open-weights) capaz de igualar o superar a Llama 2 70B y a líneas base propietarias como GPT-3.5.

Ideas clave

Sparse Mixture of Experts (SMoE): Sustituir los bloques de redes feedforward (FFN) densas en las 32 capas del transformer por 8 redes de expertos diferenciadas, enrutando cada token únicamente hacia los 2 mejores expertos (top-2).
Recuentos de parámetros desacoplados: La arquitectura alberga 47B de parámetros en total a través de sus capas, pero solo ejecuta 13B de parámetros activos por pasada hacia adelante (forward pass) de cada token, igualando la huella computacional de un modelo denso de tamaño moderado.
Ventana de contexto completa de 32k: Soporte nativo para secuencias de 32,768 tokens de longitud con una precisión de recuperación del 100% en benchmarks de clave de paso (passkey).
Enrutamiento sintáctico en lugar de temático: La inspección empírica del enrutamiento de tokens revela que la selección de expertos está impulsada principalmente por la función sintáctica y la localidad temporal, más que por dominios temáticos de alto nivel.

Método (con las ecuaciones clave)

Mixtral modifica la arquitectura de Mistral 7B ( $32$ capas, dimensión oculta de $4096$ , $32$ cabezales de atención, $8$ cabezales de clave-valor) sustituyendo cada subbloque FFN por un módulo MoE compuesto por $n = 8$ expertos ( $E_{0}, \dots, E_{n - 1}$ ), implementado cada uno como una red SwiGLU con dimensión intermedia de $14336$ .

Para una representación de token de entrada $x$ , una red de enrutamiento calcula una distribución dispersa de activación (gating) entre los expertos utilizando los $K$ logits superiores ( $K = 2$ ): $G (x) := Softmax (TopK (x \cdot W_{g}))$ donde el operador Top-K asigna $- \infty$ a todas las coordenadas fuera de los $K$ valores más grandes de los logits $ℓ = x \cdot W_{g} \in R^{n}$ : $(TopK (ℓ))_{i} := {ℓ_{i} - \infty si ℓ_{i} est \overset{a}{ˊ} en las K coordenadas superiores de ℓ, en caso contrario.$

La salida general de la capa $y$ se calcula como la suma ponderada por el mecanismo de gating sobre las salidas de los expertos: $y = i = 0 \sum n - 1 G (x)_{i} \cdot SwiGLU_{i} (x)$ La ejecución eficiente se logra mediante kernels como MegaBlocks, que representan el cómputo de MoE como multiplicaciones de matrices dispersas, junto con Expert Parallelism a través de múltiples GPUs.

Resultados (con las cifras más importantes)

Benchmarks generales y de razonamiento: En MMLU (5-shot), Mixtral 8x7B logra un 70.6% frente al 69.9% de Llama 2 70B y al 70.0% de GPT-3.5. En GSM8K (8-shot, maj@8), alcanza el 74.4% frente al 69.6% de Llama 2 70B. En el benchmark MATH (4-shot, maj@4), Mixtral consigue un 28.4% frente al 13.8% de Llama 2 70B.
Generación de código: Mixtral supera sustancialmente a Llama 2 70B en HumanEval (40.2% vs. 29.3%) y en MBPP (60.7% vs. 49.8%).
Tareas multilingües: El sobremuestreo (upsampling) de datos multilingües permite a Mixtral superar a Llama 2 70B en francés, alemán, español e italiano en ARC-Challenge, HellaSwag y MMLU (por ejemplo, MMLU en francés: 70.9% vs. 64.3%).
Ajuste fino de instrucciones y Elo en Arena: Mixtral 8x7B – Instruct (entrenado con SFT y Direct Preference Optimization) obtiene 8.30 en MT-Bench y alcanzó un Elo de 1121 en LMSys Chatbot Arena, superando a Claude-2.1 (1117), GPT-3.5-Turbo (1117) y Llama-2-70b-chat (1077).
Recuperación de contexto: Logra una precisión de recuperación del 100% en la tarea sintética de passkey retrieval a lo largo de toda la ventana de contexto de 32k.

Limitaciones y preguntas abiertas

Huella de memoria vs. cómputo: Aunque el coste computacional escala con los 13B de parámetros activos, el despliegue requiere suficiente VRAM para almacenar los 47B de parámetros completos, lo que exige configuraciones multi-GPU o técnicas de descarga (offloading).
Cuellos de botella en Expert Parallelism: El artículo identifica una alta localidad temporal: los tokens consecutivos comparten las mismas asignaciones de expertos con una frecuencia mucho mayor que la esperada por azar (hasta un 28.4% de repetición de la primera opción en la capa 15 vs. 12.5% aleatorio). Esto puede generar desequilibrios de carga y sobreutilización de GPUs específicas bajo Expert Parallelism.
Falta de especialización temática: La selección del enrutador no se divide de forma natural por dominios conceptuales (p. ej., biología vs. filosofía), lo que deja abiertos interrogantes sobre cómo estructurar el gating si se busca explícitamente una especialización de expertos por dominios.

Términos clave

Sparse Mixture of Experts (SMoE): Arquitectura de red neuronal donde el cómputo condicional enruta cada token de entrada hacia un subconjunto disperso de subredes feedforward (expertos) en lugar de activar todos los parámetros.
Active Parameter Count: Número de parámetros realmente cargados y computados durante la inferencia para un solo token, que es de 13B en Mixtral frente a sus 47B de parámetros dispersos totales.
SwiGLU: Variante de activación de unidad lineal con compuerta (gated linear unit) que combina Swish y multiplicación elemento a elemento, utilizada como la arquitectura interna de cada bloque feedforward de los expertos.
TopK Gating: Mecanismo de enrutamiento que aplica una proyección lineal a la representación del token, retiene los K valores más altos mientras enmascara el resto hacia infinito negativo, y los normaliza mediante softmax.
Expert Parallelism (EP): Estrategia de particionamiento distribuido que mapea diferentes expertos en distintas GPUs y enruta las activaciones de los tokens a través de los dispositivos hacia su experto asignado.
Direct Preference Optimization (DPO): Algoritmo de alineamiento que optimiza los parámetros del modelo de lenguaje directamente sobre pares de preferencias sin necesidad de entrenar un modelo explícito de recompensa por aprendizaje por refuerzo.
Passkey Retrieval: Prueba de evaluación sintética que mide la capacidad de un modelo para localizar y devolver una clave corta posicionada aleatoriamente dentro de secuencias de contexto largo de hasta 32k tokens.
Positional Locality: Tendencia observada por la cual tokens consecutivos en una secuencia son asignados repetidamente al mismo experto con una probabilidad muy superior a la aleatoria uniforme.

Preguntas

¿Por qué Mixtral se llama 8x7B si solo tiene 47B de parámetros en lugar de 56B?

Aunque hay 8 expertos por capa y la arquitectura se deriva de Mistral 7B, los pesos de autoatención y las capas de embeddings se comparten entre todos los expertos. Solo los subbloques feedforward se replican 8 veces, lo que resulta en un total de 47B de parámetros en lugar de 8 × 7B = 56B.

¿Cuántos expertos se activan para cada token durante la inferencia?

El enrutador selecciona exactamente 2 expertos ( $K = 2$ ) de los 8 disponibles en cada capa para cada token, lo que resulta en 13B de parámetros activos por pasada hacia adelante (forward pass).

¿Se especializan los diferentes expertos en distintas materias académicas o dominios temáticos?

No se observó una especialización temática evidente. El análisis sobre The Pile mostró distribuciones de asignación de expertos prácticamente idénticas en áreas como ArXiv, PubMed y PhilPapers; en su lugar, las decisiones de enrutamiento se correlacionan en gran medida con la sintaxis, la identidad del token y la localidad posicional.

¿Bajo qué licencia de software se publican los modelos Mixtral 8x7B?

Tanto el modelo base Mixtral 8x7B como Mixtral 8x7B – Instruct se publican bajo la permisiva licencia Apache 2.0 para uso académico y comercial.

La explicación de esta página la escribió una IA a partir del artículo y puede contener errores. Consulte el artículo para los detalles.

Leer con traducción e IA