Redes neuronales descomunalmente grandes: la capa Mixture-of-Experts con compuertas dispersas

Guía de lectura

Problema y motivación

Los modelos de aprendizaje profundo se benefician enormemente de la capacidad de parámetros cuando se entrenan en grandes conjuntos de datos. Sin embargo, las arquitecturas estándar activan toda la red para cada token o ejemplo de entrada. Esto conduce a un aumento inviable en el costo computacional a medida que escalan tanto el tamaño de la red como el volumen de datos.

Aunque el cómputo condicional (activar solo subconjuntos de un modelo por ejemplo) se propuso teóricamente, las implementaciones anteriores tuvieron dificultades en los aceleradores modernos debido a:

• Las GPU favorecen la aritmética uniforme sobre la bifurcación condicional.
• El "problema de reducción del batch" (shrinking batch problem), donde enrutar subconjuntos de ejemplos deja a las subredes individuales con tamaños de batch diminutos e ineficientes.
• Cuellos de botella en el ancho de banda de red entre dispositivos.
• Enrutamiento desbalanceado, donde unas pocas subredes dominan mientras que otras quedan desatendidas.

Ideas clave

• Capa MoE con compuertas dispersas: Las subredes feed-forward ("expertos") se colocan entre capas recurrentes (como LSTMs apiladas). Solo se evalúa una selección dispersa top-$k$ de expertos por posición.
• Noisy Top-K Gating: Se agrega ruido gaussiano ajustable a los logits de compuerta antes de la selección top-$k$ y la función softmax para proporcionar gradientes continuos y retropropagables, y fomentar una exploración equilibrada entre los expertos.
• Paralelismo combinado de datos y modelos: Las capas estándar y las compuertas se ejecutan con paralelismo de datos entre dispositivos, mientras que los expertos se particionan con paralelismo de modelo entre GPU. Los lotes de todos los dispositivos se combinan para que los expertos reciban minilotes agregados lo suficientemente grandes.
• Restricciones de balanceo: Pérdidas auxiliares explícitas sobre la importancia y la carga de los expertos evitan que la red de compuertas colapse en soluciones degeneradas que sobreutilizan un puñado de expertos.

Método (con las ecuaciones clave)

Una capa MoE consta de $n$ redes expertas feed-forward $E_1,\dots ,E_n$ y una red de compuertas $G(x)$. La salida de la capa es: $$y=\sum _{i=1}^{n}G(x{)}_i{E}_i(x)$$ El cómputo se evita siempre que $G(x{)}_i=0$.

Noisy Top-K Gating

Para combinar la dispersión con compuertas diferenciables, se inyecta ruido en las proyecciones lineales: $$H(x{)}_i=(x\cdot W_g{)}_i+\text{StandardNormal}()\cdot \text{Softplus}((x\cdot W_{\text{noise}}{)}_i)$$ $$\text{KeepTopK}(v,k{)}_i=\{\begin{array}{ll}{v}_i& \text{si }{v}_i\text{ estaˊ en los mejores }k\text{ elementos de }v\\ -\infty & \text{en caso contrario}\end{array}$$ $$G(x)=\text{Softmax}(\text{KeepTopK}(H(x),k))$$

Pérdidas de balanceo de carga

Para imponer pesos de compuerta totales iguales, se define una pérdida de importancia sobre un lote $X$ utilizando el coeficiente de variación ($CV$): $$\text{Importance}(X)=\sum _{x\in X}G(x)$$ $$L_{\text{importance}}(X)=w_{\text{importance}}\cdot CV(\text{Importance}(X){)}^2$$ Para asegurar una cantidad equitativa de ejemplos por experto, el estimador suave de carga utiliza la CDF normal $\Phi$: $$P(x,i)=\Phi \left(\frac{(x\cdot W_g{)}_i-\text{kth\_excluding}(H(x),k,i)}{\text{Softplus}((x\cdot W_{\text{noise}}{)}_i)}\right)$$ $$\text{Load}(X{)}_i=\sum _{x\in X}P(x,i)$$ $$L_{\text{load}}(X)=w_{\text{load}}\cdot CV(\text{Load}(X){)}^2$$

Resultados (con los números más importantes)

• 1 Billion Word Benchmark: Un modelo MoE-4096-h computacionalmente equivalente logró una perplejidad de prueba de 34.1 (una reducción del 24% en comparación con los LSTM de referencia equiparables en cómputo). Un modelo MoE de alto presupuesto alcanzó una perplejidad de prueba de 28.0 en 47 horas sobre 32 GPU K40, superando el resultado de última generación previo de 34.7 (que requirió 59 horas sobre 32 GPU K40).
• Corpus de 100 mil millones de palabras: Escalar hasta 65,536 expertos (68 mil millones de parámetros) redujo la perplejidad de prueba en un 39% en comparación con la referencia densa (28.9 frente a 47.0), manteniendo 0.72 TFLOPS/GPU.
• Traducción automática: GNMT aumentada con MoE logró 40.56 BLEU en WMT'14 En$\to$Fr (+1.34 BLEU sobre GNMT sin RL) y 26.03 BLEU en WMT'14 En$\to$De (+1.12 BLEU) requiriendo a la vez menos de la mitad de las operaciones por paso de tiempo.
• Traducción multilingüe: Un único modelo multilingüe MoE superó a una referencia multilingüe de GNMT en 11 de 12 pares de idiomas (hasta +5.84 BLEU en coreano$\to$inglés) y redujo la perplejidad de validación en un 19%.

Limitaciones y preguntas abiertas

• La dispersión extrema puede conducir a rendimientos decrecientes o regresiones en el rendimiento; la perplejidad empeoró al escalar de 65,536 a 131,072 expertos en el conjunto de datos de 100B palabras.
• La sobrecarga de comunicación distribuida sigue dependiendo de la relación entre dimensiones ocultas y de entrada/salida, lo que hace que capacidades de expertos pequeñas sean menos eficientes en el uso del ancho de banda.
• Aplicar MoE de forma recurrente dentro de los pasos de recurrencia de una RNN rompe la optimización de desenrollado de lotes a lo largo de las posiciones de la secuencia, requiriendo esquemas de recomputación de activaciones.