Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer/es
| Research Paper | |
|---|---|
| Authors | Noam Shazeer; Azalia Mirhoseini; Krzysztof Maziarz; Andy Davis; Quoc Le; Geoffrey Hinton; Jeff Dean |
| Year | 2017 |
| Topic area | Machine Learning |
| Difficulty | Research |
| arXiv | 1701.06538 |
| Download PDF | |
Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer es un artículo de 2017 de Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton y Jeff Dean que introdujo la capa Sparsely-Gated Mixture-of-Experts (MoE) para redes neuronales profundas. El artículo demostró que una sola capa compuesta por miles de subredes feed-forward, con una red de gating entrenable que selecciona solo unas pocas por ejemplo, puede aumentar la capacidad del modelo en más de 1000× manteniendo aproximadamente constante el cómputo por ejemplo. Aplicados entre capas LSTM apiladas para modelado de lenguaje y traducción automática, los modelos MoE con hasta 137 mil millones de parámetros establecieron nuevos resultados de estado del arte con un costo computacional menor que el de las líneas base densas.
Resumen
La capacidad de una red neuronal — su habilidad para absorber información — está acotada por su número de parámetros, pero el número de parámetros está normalmente ligado al cómputo por ejemplo. La computación condicional había sido propuesta como una forma de romper este acoplamiento: activar solo un subconjunto disperso de la red para cada entrada, de modo que la capacidad crezca sin un aumento proporcional en FLOPs. Trabajos previos habían identificado la idea pero no la habían realizado a gran escala, debido a desafíos algorítmicos y de sistemas (tamaños de batch desiguales entre subredes activas, presión sobre el ancho de banda de red, colapso de expertos).
Este artículo presenta la primera realización práctica de computación condicional a gran escala. La capa Sparsely-Gated Mixture-of-Experts consta de hasta miles de redes expertas feed-forward y una red de gating que, dada una entrada, produce un vector disperso de pesos de expertos. Solo se evalúan los expertos con gates distintos de cero, por lo que el costo por ejemplo permanece acotado incluso cuando el número de expertos (y por lo tanto de parámetros) escala a miles de millones. La capa se inserta de forma convolucional — aplicada de manera idéntica en cada paso temporal — entre dos capas LSTM apiladas en los modelos de los autores.
Contribuciones clave
- Capa Sparsely-Gated MoE. Una sola capa de red neuronal con hasta 131.072 expertos y una red de gating top-k que activa solo un pequeño número constante de expertos por ejemplo, desacoplando la capacidad del cómputo.
- Noisy Top-K Gating. Un mecanismo de gating que añade ruido gaussiano entrenable y ajustable antes de la selección top-k, proporcionando tanto sparsity como balanceo de carga estocástico.
- Soluciones al problema del shrinking-batch. Una mezcla de paralelismo de datos y de modelo, junto con la aplicación convolucional a lo largo del tiempo, que mantiene el batch efectivo de cada experto suficientemente grande para una ejecución eficiente en GPU.
- Pérdidas suaves de balanceo de carga. Pérdidas auxiliares de "importance" y "load" que evitan que unos pocos expertos dominen y aseguran una utilización balanceada en toda la población.
- Resultados empíricos a una escala sin precedentes. Modelos MoE con hasta 137 mil millones de parámetros, logrando 24% menos perplejidad en el 1 Billion Word Benchmark y superando el estado del arte previo en WMT'14 En→Fr y En→De y en una tarea de traducción multilingüe En→{Fr, De, Es, It, Pt, ...} — todo con menor costo de cómputo que las líneas base densas.
Métodos
La capa MoE
Una capa MoE contiene $ n $ redes expertas $ E_1, \dots, E_n $ (típicamente redes feed-forward ReLU de dos capas) y una red de gating $ G $ que produce un vector disperso de dimensión $ n $. La salida de la capa para una entrada $ x $ es
- $ y = \sum_{i=1}^{n} G(x)_i \, E_i(x). $
Donde $ G(x)_i = 0 $, el experto correspondiente se omite. Con un gate top-k, solo se evalúan k expertos por ejemplo, dando un cómputo por ejemplo constante independientemente de $ n $.
Noisy Top-K Gating
La red de gating combina un softmax sobre una proyección lineal aprendida con ruido gaussiano entrenable aditivo, y luego conserva solo los k logits más grandes antes de normalizar:
- $ H(x)_i = (x \cdot W_g)_i + \epsilon \cdot \mathrm{softplus}((x \cdot W_{noise})_i),\quad \epsilon \sim \mathcal{N}(0,1) $
- $ G(x) = \mathrm{Softmax}(\mathrm{KeepTopK}(H(x), k)). $
El ruido tanto disipa las activaciones como ayuda a balancear la carga entre expertos durante el entrenamiento; fijar todas las entradas distintas de las top k a $ -\infty $ garantiza una sparsity exacta en inferencia.
Ingeniería de rendimiento
Cuando se eligen $ k $ expertos de $ n $ para cada uno de $ b $ ejemplos, cada experto ve en promedio $ kb/n $ ejemplos, lo que es demasiado pequeño para multiplicaciones de matrices eficientes en GPU. Los autores abordan esto mediante:
- Mezclar paralelismo de datos y de modelo. La capa MoE se reparte entre $ d $ dispositivos de modo que cada experto vive en exactamente un dispositivo, y todas las $ d $ réplicas data-parallel enrutan sincrónicamente sus ejemplos seleccionados a los expertos que necesitan. El batch efectivo de cada experto se vuelve $ k \cdot b \cdot d / n $, que escala con el tamaño del clúster.
- Aplicación convolucional. La MoE se aplica independientemente en cada paso temporal del LSTM, multiplicando el batch por experto por la longitud de secuencia desplegada.
- MoE jerárquica. Para $ n $ muy grande, una primera gate selecciona un grupo de expertos y una segunda gate selecciona expertos dentro de ese grupo, manteniendo el ruteo manejable.
- Optimización de ancho de banda. Los expertos se hacen pesados en cómputo (capas ocultas grandes) para que su razón de cómputo sobre bytes de entrada absorba el costo de mover ejemplos entre dispositivos.
Pérdidas de balanceo de carga
Se añaden dos pérdidas auxiliares al objetivo de entrenamiento para evitar que la red de gating colapse hacia unos pocos expertos favoritos:
- $ L_{importance}(X) = w_{importance} \cdot \mathrm{CV}\!\left(\sum_{x \in X} G(x)\right)^2, $
donde CV es el coeficiente de variación. Una segunda pérdida $ L_{load} $ iguala el número de ejemplos que recibe cada experto — no solo la importance ponderada por el gate — usando un estimador suave y diferenciable construido a partir de la distribución de gating con ruido. Juntas imponen tanto un peso igual como un conteo de ejemplos igual por experto.
Resultados
Modelado de lenguaje
En el 1 Billion Word Language Modeling Benchmark, los autores entrenaron modelos MoE con 4, 32, 256, 1024, 4096, 32768 y 65536 expertos, todos igualados a aproximadamente 8M ops/timestep. El modelo de 4096 expertos redujo la perplejidad en el conjunto de prueba en un 24% respecto al mejor modelo previamente publicado, mientras que variantes con mayor cómputo de hasta 4 mil millones de parámetros establecieron nuevas perplejidades de estado del arte con un cómputo significativamente menor que los modelos densos previos. La eficiencia computacional en 16–32 GPUs Tesla K40 alcanzó 0,74–1,56 TFLOPS/GPU.
En el 100 Billion Word Google News Corpus — aproximadamente 100× más grande que el 1 Billion Word — entrenaron una MoE con 131.072 expertos, totalizando 137 mil millones de parámetros, y siguieron viendo mejoras de perplejidad a medida que crecía la capacidad, lo que indica que los beneficios de MoE se siguen acumulando cuando hay datos de entrenamiento abundantes.
Traducción automática
En WMT'14 En→Fr y En→De, encoders/decoders LSTM apilados aumentados con MoE superaron las puntuaciones BLEU del estado del arte previo de Google Neural Machine Translation (GNMT) usando menos cómputo de entrenamiento. En un sistema de traducción multilingüe en producción que cubre doce pares de lenguas simultáneamente, un único modelo aumentado con MoE mejoró el BLEU en cada par comparado con un modelo GNMT dedicado entrenado en cada par individualmente.
Impacto
Este artículo es una entrada fundacional en la literatura moderna de MoE. Estableció la receta — gating top-k disperso, balanceo de carga basado en ruido, sharding de expertos, pérdidas suaves de balanceo — sobre la cual el trabajo posterior construyó directamente. Inspiró directamente a GShard (2020), que escaló la misma idea a encoders Transformer para traducción multilingüe, y a Switch Transformer (2021), que simplificó el gate a ruteo top-1 y escaló a más de un billón de parámetros. También sustenta a GLaM, ST-MoE, Mixtral 8×7B y los backbones MoE usados en muchos modelos de lenguaje grandes de frontera. La frase "outrageously large" del título resultó ser modesta: a comienzos de la década de 2020, capas MoE con cientos de miles de millones de parámetros se desplegaban rutinariamente en modelos de lenguaje en producción, y el principio de computación condicional introducido aquí se convirtió en una herramienta central para intercambiar número de parámetros por costo por token.
Véase también
- Multi-gate Mixture-of-Experts — MoE condicionado por tarea para aprendizaje multitarea.
- Attention Is All You Need — la arquitectura Transformer en la que GShard y Switch Transformer insertaron posteriormente capas MoE.
- Dropout: A Simple Way to Prevent Neural Networks from Overfitting — una forma relacionada de computación condicional estocástica.
- Language Models are Few-Shot Learners — modelo de lenguaje denso grande al que el trabajo de escalado de MoE apuntaba a superar a menor costo.
Referencias
- Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. International Conference on Learning Representations (ICLR).
- Bengio, Y., Léonard, N., & Courville, A. (2013). Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation.
- Eigen, D., Ranzato, M. A., & Sutskever, I. (2013). Learning Factored Representations in a Deep Mixture of Experts.
- Jacobs, R. A., Jordan, M. I., Nowlan, S. J., & Hinton, G. E. (1991). Adaptive Mixtures of Local Experts. Neural Computation, 3(1), 79–87.
- Jordan, M. I., & Jacobs, R. A. (1994). Hierarchical Mixtures of Experts and the EM Algorithm. Neural Computation, 6(2), 181–214.
- Lepikhin, D. et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding.
- Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.