Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer/paper/es

    From Marovi AI
    Other languages:
    SummarySource
    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
    PDF Download PDF

    = Outrageously Large Neural Networks:
    The Sparsely-Gated Mixture-of-Experts Layer =

    Noam Shazeer Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Azalia Mirhoseini Contribuyentes principales por igualTrabajo realizado como miembro del programa Google Brain Residency (g.co/brainresidency) Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Krzysztof Maziarz Universidad Jaguelónica, Cracovia, krzysztof.maziarz@student.uj.edu.pl Andy Davis Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Quoc Le Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Geoffrey Hinton Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com Jeff Dean Google Brain, {noam,azalia,andydavis,qvl,geoffhinton,jeff}@google.com

    Resumen

    La capacidad de una red neuronal para absorber información está limitada por su número de parámetros. La computación condicional, donde partes de la red se activan según el ejemplo, ha sido propuesta en teoría como una forma de aumentar dramáticamente la capacidad del modelo sin un aumento proporcional en cómputo. En la práctica, sin embargo, existen desafíos algorítmicos y de rendimiento significativos. En este trabajo, abordamos estos desafíos y finalmente realizamos la promesa de la computación condicional, logrando mejoras superiores a 1000x en capacidad del modelo con solo pérdidas menores en eficiencia computacional sobre clústeres modernos de GPU. Introducimos una capa Sparsely-Gated Mixture-of-Experts (MoE), compuesta por hasta miles de subredes feed-forward. Una red de gating entrenable determina una combinación dispersa de estos expertos para usar en cada ejemplo. Aplicamos la MoE a las tareas de modelado de lenguaje y traducción automática, donde la capacidad del modelo es crítica para absorber las vastas cantidades de conocimiento disponibles en los corpus de entrenamiento. Presentamos arquitecturas de modelo en las que se aplica una MoE con hasta 137 mil millones de parámetros de manera convolucional entre capas LSTM apiladas. En benchmarks grandes de modelado de lenguaje y traducción automática, estos modelos logran resultados significativamente mejores que el estado del arte con menor costo computacional.

    1 Introducción y trabajo relacionado

    1.1 Computación condicional

    Aprovechar la escala tanto en datos de entrenamiento como en tamaño del modelo ha sido central para el éxito del aprendizaje profundo. Cuando los conjuntos de datos son lo suficientemente grandes, aumentar la capacidad (número de parámetros) de las redes neuronales puede dar una precisión predictiva mucho mejor. Esto se ha mostrado en dominios como texto (Sutskever et al., 2014; Bahdanau et al., 2014; Jozefowicz et al., 2016; Wu et al., 2016), imágenes (Krizhevsky et al., 2012; Le et al., 2012) y audio (Hinton et al., 2012; Amodei et al., 2015). Para los modelos típicos de aprendizaje profundo, donde todo el modelo se activa para cada ejemplo, esto lleva a una explosión aproximadamente cuadrática en los costos de entrenamiento, conforme aumentan tanto el tamaño del modelo como el número de ejemplos de entrenamiento. Lamentablemente, los avances en potencia de cómputo y computación distribuida no logran satisfacer tal demanda.

    Diversas formas de computación condicional han sido propuestas como una forma de aumentar la capacidad del modelo sin un aumento proporcional en costos computacionales (Davis & Arel, 2013; Bengio et al., 2013; Eigen et al., 2013; Ludovic Denoyer, 2014; Cho & Bengio, 2014; Bengio et al., 2015; Almahairi et al., 2015). En estos esquemas, grandes partes de una red están activas o inactivas por ejemplo. Las decisiones de gating pueden ser binarias o dispersas y continuas, estocásticas o deterministas. Se proponen diversas formas de aprendizaje por refuerzo y back-propagation para entrenar las decisiones de gating.

    Aunque estas ideas son prometedoras en teoría, ningún trabajo hasta la fecha ha demostrado mejoras masivas en capacidad del modelo, tiempo de entrenamiento o calidad del modelo. Atribuimos esto a una combinación de los siguientes desafíos:

    • Los dispositivos de cómputo modernos, especialmente las GPUs, son mucho más rápidos en aritmética que en branching. La mayoría de los trabajos anteriores reconocen esto y proponen activar/desactivar grandes bloques de la red con cada decisión de gating.

    • Los batch sizes grandes son críticos para el rendimiento, ya que amortizan los costos de transferencias y actualizaciones de parámetros. La computación condicional reduce los batch sizes para los bloques activos condicionalmente de la red.

    • El ancho de banda de la red puede ser un cuello de botella. Un clúster de GPUs puede tener una potencia computacional miles de veces mayor que el ancho de banda agregado entre dispositivos. Para ser computacionalmente eficiente, la demanda relativa de cómputo frente a red de un algoritmo debe exceder esta razón. Las capas de embedding, que pueden verse como una forma de computación condicional, sufren precisamente este problema. Dado que los embeddings generalmente necesitan enviarse por la red, el número de interacciones (ejemplo, parámetro) está limitado por el ancho de banda de red en lugar de por la capacidad computacional.

    • Dependiendo del esquema, pueden ser necesarios términos de pérdida para alcanzar el nivel deseado de sparsity por bloque y/o por ejemplo. Bengio et al. (2015) usan tres de tales términos. Estos asuntos pueden afectar tanto la calidad del modelo como el balanceo de carga.

    • La capacidad del modelo es más crítica para conjuntos de datos muy grandes. La literatura existente sobre computación condicional trata con conjuntos de datos de reconocimiento de imágenes relativamente pequeños, de hasta 600.000 imágenes. Es difícil imaginar que las etiquetas de estas imágenes proporcionen una señal suficiente para entrenar adecuadamente un modelo con millones, mucho menos miles de millones de parámetros.

    En este trabajo, por primera vez abordamos todos los desafíos anteriores y finalmente realizamos la promesa de la computación condicional. Obtenemos mejoras superiores a 1000x en capacidad del modelo con solo pérdidas menores en eficiencia computacional, y avanzamos significativamente los resultados de estado del arte en conjuntos de datos públicos de modelado de lenguaje y traducción.

    1.2 Nuestro enfoque: la capa Sparsely-Gated Mixture-of-Experts

    Nuestro enfoque para la computación condicional es introducir un nuevo tipo de componente de red neuronal de propósito general: una capa Sparsely-Gated Mixture-of-Experts (MoE). La MoE consta de un número de expertos, cada uno una red neuronal feed-forward simple, y una red de gating entrenable que selecciona una combinación dispersa de los expertos para procesar cada entrada (ver Figura 1). Todas las partes de la red se entrenan conjuntamente por back-propagation.

    Refer to caption

    Aunque la técnica introducida es genérica, en este artículo nos centramos en tareas de modelado de lenguaje y traducción automática, que son conocidas por beneficiarse de modelos muy grandes. En particular, aplicamos una MoE convolucionalmente entre capas LSTM apiladas (Hochreiter & Schmidhuber, 1997), como en la Figura 1. La MoE se invoca una vez para cada posición en el texto, seleccionando una combinación potencialmente diferente de expertos en cada posición. Los distintos expertos tienden a especializarse mucho según sintaxis y semántica (ver Apéndice E Tabla 9). Tanto en benchmarks de modelado de lenguaje como de traducción automática, mejoramos los mejores resultados publicados con una fracción del costo computacional.

    1.3 Trabajo relacionado sobre Mixtures of Experts

    Desde su introducción hace más de dos décadas (Jacobs et al., 1991; Jordan & Jacobs, 1994), el enfoque mixture-of-experts ha sido objeto de mucha investigación. Se han propuesto distintos tipos de arquitecturas de expertos, como SVMs (Collobert et al., 2002), procesos gaussianos (Tresp, 2001; Theis & Bethge, 2015; Deisenroth & Ng, 2015), procesos de Dirichlet (Shahbaba & Neal, 2009) y redes profundas. Otros trabajos se han centrado en distintas configuraciones de expertos, como una estructura jerárquica (Yao et al., 2009), números infinitos de expertos (Rasmussen & Ghahramani, 2002) y la adición secuencial de expertos (Aljundi et al., 2016). Garmash & Monz (2016) sugieren un modelo de ensamble en el formato de mixture of experts para traducción automática. La red de gating se entrena sobre un modelo NMT de ensamble preentrenado.

    Los trabajos anteriores se refieren a mixtures of experts de nivel superior. La mixture of experts es el modelo entero. Eigen et al. (2013) introducen la idea de usar múltiples MoEs con sus propias redes de gating como partes de un modelo profundo. Es intuitivo que este último enfoque es más potente, ya que problemas complejos pueden contener muchos sub-problemas, cada uno requiriendo expertos distintos. También aluden en su conclusión al potencial de introducir sparsity, convirtiendo a las MoEs en un vehículo para computación condicional.

    Nuestro trabajo se basa en este uso de las MoEs como componente de red neuronal de propósito general. Mientras que Eigen et al. (2013) usa dos MoEs apiladas que permiten dos conjuntos de decisiones de gating, nuestra aplicación convolucional de la MoE permite decisiones de gating distintas en cada posición del texto. También realizamos sparse gating y demostramos su uso como una forma práctica de aumentar masivamente la capacidad del modelo.

    2 Estructura de la capa Mixture-of-Experts

    La capa Mixture-of-Experts (MoE) consta de un conjunto de $ {\textstyle n} $ "expert networks" $ {\textstyle E_{1},\cdots,E_{n}} $ y una "gating network" $ {\textstyle G} $ cuya salida es un vector disperso de dimensión $ {\textstyle n} $. La Figura 1 muestra una visión general del módulo MoE. Los expertos son ellos mismos redes neuronales, cada uno con sus propios parámetros. Aunque en principio solo requerimos que los expertos acepten entradas del mismo tamaño y produzcan salidas del mismo tamaño, en nuestras investigaciones iniciales en este artículo nos restringimos al caso en que los modelos son redes feed-forward con arquitecturas idénticas, pero con parámetros separados.

    Denotemos por $ {\textstyle G\hspace{0pt}{(x)}} $ y $ {\textstyle E_{i}\hspace{0pt}{(x)}} $ la salida de la red de gating y la salida de la $ {\textstyle i} $-ésima red experta para una entrada dada $ {\textstyle x} $. La salida $ {\textstyle y} $ del módulo MoE se puede escribir como sigue:

    $ {\displaystyle y = {\sum\limits_{i = 1}^{n}{G\hspace{0pt}{(x)}_{i}\hspace{0pt}E_{i}\hspace{0pt}{(x)}}}} $ (1)

    Ahorramos cómputo basándonos en la sparsity de la salida de $ {\textstyle G\hspace{0pt}{(x)}} $. Donde $ {\textstyle {G\hspace{0pt}{(x)}_{i}} = 0} $, no necesitamos calcular $ {\textstyle E_{i}\hspace{0pt}{(x)}} $. En nuestros experimentos, tenemos hasta miles de expertos, pero solo necesitamos evaluar unos pocos para cada ejemplo. Si el número de expertos es muy grande, podemos reducir el factor de branching usando una MoE jerárquica de dos niveles. En una MoE jerárquica, una red de gating primaria elige una combinación dispersa ponderada de "expertos", cada uno de los cuales es a su vez una mixture-of-experts secundaria con su propia red de gating. A continuación nos centramos en MoEs ordinarias. Damos más detalles sobre MoEs jerárquicas en el Apéndice B.

    Nuestra implementación está relacionada con otros modelos de computación condicional. Una MoE cuyos expertos son matrices de pesos simples es similar a la matriz de pesos parametrizada propuesta en (Cho & Bengio, 2014). Una MoE cuyos expertos tienen una capa oculta es similar al block-wise dropout descrito en (Bengio et al., 2015), donde la capa con dropout queda intercalada entre capas completamente activadas.

    2.1 Red de gating

    Softmax Gating:

    Una elección simple de función de gating no dispersa (Jordan & Jacobs, 1994) es multiplicar la entrada por una matriz de pesos entrenable $ {\textstyle W_{g}} $ y luego aplicar la función $ {\textstyle S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x} $.

    $ {\displaystyle {G_{\sigma}\hspace{0pt}{(x)}} = {S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({x \cdot W_{g}})}}} $ (2)

    Noisy Top-K Gating:

    Añadimos dos componentes a la red de gating Softmax: sparsity y ruido. Antes de aplicar la función softmax, añadimos ruido gaussiano ajustable, luego conservamos solo los k valores más grandes, fijando el resto a $ {\textstyle - \infty} $ (lo que hace que los valores de gate correspondientes sean iguales a $ {\textstyle 0} $). La sparsity sirve para ahorrar cómputo, como se describió arriba. Aunque esta forma de sparsity crea algunas discontinuidades teóricamente preocupantes en la salida de la función de gating, no hemos observado todavía que esto sea un problema en la práctica. El término de ruido ayuda al balanceo de carga, como se discutirá en el Apéndice A. La cantidad de ruido por componente se controla mediante una segunda matriz de pesos entrenable $ {\textstyle W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}} $.

    $ {\displaystyle {G\hspace{0pt}{(x)}} = {S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({K\hspace{0pt}e\hspace{0pt}e\hspace{0pt}p\hspace{0pt}T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{({H\hspace{0pt}{(x)}},k)}})}}} $ (3)
    $ {\displaystyle {H\hspace{0pt}{(x)}_{i}} = {{({x \cdot W_{g}})}_{i} + {{{S\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}d\hspace{0pt}a\hspace{0pt}r\hspace{0pt}d\hspace{0pt}N\hspace{0pt}o\hspace{0pt}r\hspace{0pt}m\hspace{0pt}a\hspace{0pt}l\hspace{0pt}{()}} \cdot S}\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}p\hspace{0pt}l\hspace{0pt}u\hspace{0pt}s\hspace{0pt}{({({x \cdot W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}})}_{i})}}}} $ (4)
    $ {\displaystyle {K\hspace{0pt}e\hspace{0pt}e\hspace{0pt}p\hspace{0pt}T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{(v,k)}_{i}} = \begin{cases} v_{i} & {\text{if~}v_{i}\text{~is in the top~}k\text{~elements of~}v\text{.}} \\ {- \infty} & \text{otherwise.} \end{cases}} $ (5)

    Entrenamiento de la red de gating

    Entrenamos la red de gating mediante back-propagation simple, junto con el resto del modelo. Si elegimos $ {\textstyle k > 1} $, los valores de gate para los top k expertos tienen derivadas no nulas con respecto a los pesos de la red de gating. Este tipo de comportamiento ocasionalmente sensible se describe en (Bengio et al., 2013) con respecto a noisy rectifiers. Los gradientes también se propagan hacia atrás a través de la red de gating hacia sus entradas. Nuestro método difiere aquí de (Bengio et al., 2015) que usan gates booleanos y un enfoque de tipo REINFORCE para entrenar la red de gating.

    3 Abordando los desafíos de rendimiento

    3.1 El problema del shrinking batch

    En CPUs y GPUs modernas, los batch sizes grandes son necesarios para la eficiencia computacional, a fin de amortizar la sobrecarga de cargas y actualizaciones de parámetros. Si la red de gating elige $ {\textstyle k} $ de $ {\textstyle n} $ expertos para cada ejemplo, entonces para un batch de $ {\textstyle b} $ ejemplos, cada experto recibe un batch mucho más pequeño de aproximadamente $ {\textstyle \frac{k\hspace{0pt}b}{n} \ll b} $ ejemplos. Esto hace que una implementación ingenua de MoE se vuelva muy ineficiente a medida que aumenta el número de expertos. La solución a este problema del shrinking batch es hacer el batch original tan grande como sea posible. Sin embargo, el tamaño del batch tiende a estar limitado por la memoria necesaria para almacenar las activaciones entre los pases hacia adelante y hacia atrás. Proponemos las siguientes técnicas para aumentar el batch size:

    Mezcla de paralelismo de datos y de modelo:

    En un entorno de entrenamiento distribuido convencional, múltiples copias del modelo en diferentes dispositivos procesan asincrónicamente batches distintos de datos, y los parámetros se sincronizan a través de un conjunto de parameter servers. En nuestra técnica, estos diferentes batches corren sincrónicamente para que puedan combinarse para la capa MoE. Distribuimos las capas estándar del modelo y la red de gating según esquemas convencionales de paralelismo de datos, pero mantenemos solo una copia compartida de cada experto. Cada experto en la capa MoE recibe un batch combinado consistente en los ejemplos relevantes de todos los batches de entrada data-parallel. El mismo conjunto de dispositivos funciona como réplicas data-parallel (para las capas estándar y las redes de gating) y como shards model-parallel (cada uno alojando un subconjunto de los expertos). Si el modelo se distribuye sobre $ {\textstyle d} $ dispositivos, y cada dispositivo procesa un batch de tamaño $ {\textstyle b} $, cada experto recibe un batch de aproximadamente $ {\textstyle \frac{k\hspace{0pt}b\hspace{0pt}d}{n}} $ ejemplos. Así, conseguimos un factor de mejora de $ {\textstyle d} $ en el batch size por experto.

    En el caso de una MoE jerárquica (Sección B), la red de gating primaria emplea paralelismo de datos, y las MoEs secundarias emplean paralelismo de modelo. Cada MoE secundaria reside en un solo dispositivo.

    Esta técnica nos permite aumentar el número de expertos (y por lo tanto el número de parámetros) incrementando proporcionalmente el número de dispositivos en el clúster de entrenamiento. El batch size total aumenta, manteniendo constante el batch size por experto. Los requisitos de memoria y ancho de banda por dispositivo también permanecen constantes, al igual que los tiempos de paso, así como la cantidad de tiempo necesaria para procesar un número de ejemplos de entrenamiento igual al número de parámetros del modelo. Es nuestro objetivo entrenar un modelo de un billón de parámetros sobre un corpus de un billón de palabras. No hemos escalado nuestros sistemas tan lejos en el momento de escribir este artículo, pero debería ser posible añadiendo más hardware.

    Aprovechando la convolucionalidad:

    En nuestros modelos de lenguaje, aplicamos la misma MoE a cada paso temporal de la capa anterior. Si esperamos a que la capa anterior termine, podemos aplicar la MoE a todos los pasos temporales juntos como un gran batch. Hacerlo aumenta el tamaño del batch de entrada a la capa MoE en un factor del número de pasos temporales desplegados.

    Aumentando el batch size para una MoE recurrente:

    Sospechamos que modelos aún más potentes pueden involucrar la aplicación de una MoE de manera recurrente. Por ejemplo, las matrices de pesos de un LSTM u otra RNN podrían reemplazarse por una MoE. Lamentablemente, tales modelos rompen el truco convolucional del párrafo anterior, ya que la entrada a la MoE en un paso temporal depende de la salida de la MoE en el paso temporal anterior. Gruslys et al. (2016) describen una técnica para reducir drásticamente el número de activaciones almacenadas en una RNN desplegada, a costa de recomputar las activaciones hacia adelante. Esto permitiría un gran aumento en el batch size.

    3.2 Ancho de banda de red

    Otra preocupación importante de rendimiento en la computación distribuida es el ancho de banda de red. Dado que los expertos son estacionarios (ver arriba) y el número de parámetros de gating es pequeño, la mayor parte de la comunicación implica enviar las entradas y salidas de los expertos a través de la red. Para mantener la eficiencia computacional, la razón entre el cómputo de un experto y el tamaño de su entrada y salida debe exceder la razón entre la capacidad computacional y la capacidad de red del dispositivo de cómputo. Para GPUs, esto puede ser de miles a uno. En nuestros experimentos, usamos expertos con una capa oculta que contiene miles de unidades activadas con RELU. Dado que las matrices de pesos en el experto tienen tamaños $ {\textstyle i\hspace{0pt}n\hspace{0pt}p\hspace{0pt}u\hspace{0pt}t} $_$ {\textstyle {{s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} \times h}\hspace{0pt}i\hspace{0pt}d\hspace{0pt}d\hspace{0pt}e\hspace{0pt}n} $_$ {\textstyle s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} $ y $ {\textstyle h\hspace{0pt}i\hspace{0pt}d\hspace{0pt}d\hspace{0pt}e\hspace{0pt}n} $_$ {\textstyle {{s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} \times o}\hspace{0pt}u\hspace{0pt}t\hspace{0pt}p\hspace{0pt}u\hspace{0pt}t} $_$ {\textstyle s\hspace{0pt}i\hspace{0pt}z\hspace{0pt}e} $, la razón entre cómputo y entrada/salida es igual al tamaño de la capa oculta. Convenientemente, podemos aumentar la eficiencia computacional simplemente usando una capa oculta más grande, o más capas ocultas.

    4 Balanceo de la utilización de expertos

    Hemos observado que la red de gating tiende a converger a un estado en el que siempre produce pesos grandes para los mismos pocos expertos. Este desbalance se autorrefuerza, ya que los expertos favorecidos se entrenan más rápidamente y por lo tanto son seleccionados aún más por la red de gating. Eigen et al. (2013) describen el mismo fenómeno y usan una restricción dura al inicio del entrenamiento para evitar este mínimo local. Bengio et al. (2015) incluyen una restricción suave sobre el promedio batchwise de cada gate.111Bengio et al. (2015) también incluyen dos pérdidas adicionales. Una controla la sparsity por ejemplo, que no necesitamos ya que la impone el valor fijo de $ {\textstyle k} $. Una tercera pérdida fomenta la diversidad de los valores del gate. En nuestros experimentos, encontramos que los valores del gate naturalmente se diversifican conforme los expertos se especializan (en un ciclo virtuoso), y no necesitamos imponer diversidad de valores del gate.

    Tomamos un enfoque de restricción suave. Definimos la importancia de un experto relativa a un batch de ejemplos de entrenamiento como la suma batchwise de los valores de gate para ese experto. Definimos una pérdida adicional $ {\textstyle L_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $, que se añade a la función de pérdida global del modelo. Esta pérdida es igual al cuadrado del coeficiente de variación del conjunto de valores de importancia, multiplicado por un factor de escala ajustado a mano $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $. Esta pérdida adicional anima a que todos los expertos tengan igual importancia.

    $ {\displaystyle {I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e\hspace{0pt}{(X)}} = {\sum\limits_{x \in X}{G\hspace{0pt}{(x)}}}} $ (6)
    $ {\displaystyle {L_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}\hspace{0pt}{(X)}} = {{w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} \cdot C}\hspace{0pt}V\hspace{0pt}{({I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e\hspace{0pt}{(X)}})}^{2}}} $ (7)

    Aunque esta función de pérdida puede asegurar igual importancia, los expertos aún pueden recibir números muy diferentes de ejemplos. Por ejemplo, un experto puede recibir pocos ejemplos con pesos grandes, y otro puede recibir muchos ejemplos con pesos pequeños. Esto puede causar problemas de memoria y rendimiento en hardware distribuido. Para resolver este problema, introducimos una segunda función de pérdida, $ {\textstyle L_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $, que asegura cargas balanceadas. El Apéndice A contiene la definición de esta función, junto con resultados experimentales.

    5 Experimentos

    5.1 1 Billion Word Language Modeling Benchmark

    Conjunto de datos:

    Este conjunto de datos, introducido por (Chelba et al., 2013) consiste en oraciones únicas barajadas de artículos de noticias, totalizando aproximadamente 829 millones de palabras, con un vocabulario de 793.471 palabras.

    Estado del arte previo:

    Los mejores resultados publicados anteriormente (Jozefowicz et al., 2016) usan modelos consistentes en una o más capas Long Short-Term Memory (LSTM) apiladas (Hochreiter & Schmidhuber, 1997; Gers et al., 2000). El número de parámetros en las capas LSTM de estos modelos varía de 2 millones a 151 millones. La calidad aumenta enormemente con el conteo de parámetros, al igual que los costos computacionales. Los resultados para estos modelos forman la línea superior de la Figura 2-derecha.

    Modelos MoE:

    Nuestros modelos consisten en dos capas LSTM apiladas con una capa MoE entre ellas (ver Figura 1). Variamos los tamaños de las capas y el número de expertos. Para detalles completos sobre la arquitectura del modelo, el régimen de entrenamiento, líneas base adicionales y resultados, ver Apéndice C.

    Cómputo bajo, capacidad variable:

    Para investigar los efectos de añadir capacidad, entrenamos una serie de modelos MoE todos con costos computacionales aproximadamente iguales: alrededor de 8 millones de multiply-and-adds por ejemplo de entrenamiento por timestep en el pase hacia adelante, excluyendo la capa softmax. Llamamos a esta métrica (ops/timestep). Entrenamos modelos con MoEs planas que contenían 4, 32 y 256 expertos, y modelos con MoEs jerárquicas que contenían 256, 1024 y 4096 expertos. Cada experto tenía aproximadamente 1 millón de parámetros. Para todas las capas MoE, había 4 expertos activos por entrada.

    Los resultados de estos modelos se muestran en la Figura 2-izquierda. El modelo con 4 expertos siempre activos rindió (sin sorpresa) de manera similar a los modelos base con cómputo equivalente, mientras que el más grande de los modelos (4096 expertos) logró una impresionante perplejidad 24% menor en el conjunto de prueba.

    Refer to caption Refer to caption

    Test Test #Parameters ops/timestep Training TFLOPS
    Perplexity Perplexity excluding embedding Time /GPU
    10 epochs 100 epochs and softmax layers 10 epochs
    Best Published Results 34.7 30.6 151 million 151 million 59 hours, 32 k40s 1.09
    Low-Budget MoE Model 34.1 4303 million 8.9 million 15 hours, 16 k40s 0.74
    Medium-Budget MoE Model 31.3 4313 million 33.8 million 17 hours, 32 k40s 1.22
    High-Budget MoE Model 28.0 4371 million 142.7 million 47 hours, 32 k40s 1.56

    Cómputo variable, alta capacidad:

    Además del modelo más grande de la sección anterior, entrenamos dos modelos MoE más con capacidad similarmente alta (4 mil millones de parámetros), pero con presupuestos de cómputo más altos. Estos modelos tenían LSTMs más grandes, y menos pero más grandes expertos. Los detalles pueden encontrarse en el Apéndice C.2. Los resultados de estos tres modelos forman la línea inferior de la Figura 2-derecha. La Tabla 1 compara los resultados de estos modelos con el mejor resultado publicado previamente en este conjunto de datos. Incluso el más rápido de estos modelos supera el mejor resultado publicado (controlando por el número de épocas de entrenamiento), a pesar de requerir solo el 6% del cómputo.

    Eficiencia computacional:

    Entrenamos nuestros modelos usando TensorFlow (Abadi et al., 2016) en clústeres con 16-32 GPUs Tesla K40. Para cada uno de nuestros modelos, determinamos la eficiencia computacional en TFLOPS/GPU dividiendo el número de operaciones de coma flotante necesarias para procesar un batch de entrenamiento por el tiempo de paso observado y por el número de GPUs en el clúster. Los conteos de operaciones usados aquí son más altos que los que reportamos en nuestras cifras de ops/timestep, ya que incluimos el pase hacia atrás, incluimos el entrenamiento basado en importance sampling de la capa softmax, y contamos un multiply-and-add como dos operaciones separadas. Para todos nuestros modelos MoE, las operaciones de coma flotante involucradas en los expertos representan entre el 37% y el 46% del total.

    Para nuestros modelos base sin MoE, la eficiencia computacional observada osciló entre 1.07-1.29 TFLOPS/GPU. Para nuestros modelos MoE de bajo cómputo, la eficiencia de cómputo osciló entre 0.74-0.90 TFLOPS/GPU, excepto para el modelo de 4 expertos que no aprovechó plenamente el paralelismo disponible. Nuestro modelo MoE de mayor cómputo fue más eficiente con 1.56 TFLOPS/GPU, probablemente debido a las matrices más grandes. Estas cifras representan una fracción significativa del máximo teórico de 4.29 TFLOPS/GPU declarado por NVIDIA. Los resultados detallados están en el Apéndice C, Tabla 7.

    5.2 100 Billion Word Google News Corpus

    Refer to caption

    En el corpus de mil millones de palabras, añadir capacidad adicional parece producir rendimientos decrecientes a medida que el número de parámetros en la capa MoE supera los mil millones, como puede verse en la Figura 2-izquierda. Hipotetizamos que para un conjunto de entrenamiento más grande, capacidades aún mayores producirían mejoras significativas en calidad.

    Construimos un conjunto de entrenamiento similar consistente en oraciones únicas barajadas del corpus interno de noticias de Google, totalizando aproximadamente 100 mil millones de palabras. De manera similar a la sección anterior, probamos una serie de modelos con costos computacionales similares de aproximadamente 8 millones de ops/timestep. Además de un modelo LSTM base, entrenamos modelos aumentados con capas MoE que contenían 32, 256, 1024, 4096, 16384, 65536 y 131072 expertos. Esto corresponde a hasta 137 mil millones de parámetros en la capa MoE. Los detalles sobre arquitectura, entrenamiento y resultados se dan en el Apéndice D.

    Resultados:

    La Figura 3 muestra la perplejidad de prueba en función de la capacidad después de entrenar sobre 10 mil millones de palabras (línea superior) y 100 mil millones de palabras (línea inferior). Al entrenar sobre los 100 mil millones de palabras completos, la perplejidad de prueba mejora significativamente hasta 65536 expertos (68 mil millones de parámetros), cayendo un 39% por debajo de la línea base con cómputo equivalente, pero se degrada con 131072 expertos, posiblemente como resultado de demasiada sparsity. La brecha cada vez mayor entre las dos líneas demuestra (sin sorpresa) que una mayor capacidad del modelo ayuda más en conjuntos de entrenamiento más grandes.

    Incluso con 65536 expertos (sparsity de capa del 99.994%), la eficiencia computacional para el modelo se mantiene en un respetable 0.72 TFLOPS/GPU.

    5.3 Traducción automática (par único de lenguas)

    Arquitectura del modelo:

    Nuestro modelo fue una versión modificada del modelo GNMT descrito en (Wu et al., 2016). Para reducir cómputo, disminuimos el número de capas LSTM en el encoder y decoder de 9 y 8 a 3 y 2 respectivamente. Insertamos capas MoE tanto en el encoder (entre las capas 2 y 3) como en el decoder (entre las capas 1 y 2). Cada capa MoE contenía hasta 2048 expertos, cada uno con aproximadamente dos millones de parámetros, añadiendo un total de aproximadamente 8 mil millones de parámetros a los modelos. Más detalles sobre arquitectura del modelo, procedimiento de prueba y resultados pueden encontrarse en el Apéndice E.

    Conjuntos de datos:

    Hicimos benchmark de nuestro método en los corpora WMT'14 En$ {\textstyle \rightarrow} $Fr y En$ {\textstyle \rightarrow} $De, cuyos conjuntos de entrenamiento tienen 36M y 5M de pares de oraciones respectivamente. Los protocolos experimentales también fueron similares a los de (Wu et al., 2016): newstest2014 se usó como conjunto de prueba para comparar contra trabajos anteriores (Luong et al., 2015a; Zhou et al., 2016; Wu et al., 2016), mientras que la combinación de newstest2012 y newstest2013 se usó como conjunto de desarrollo. También probamos el mismo modelo en los datos de inglés a francés de producción de Google.

    Model Test Test ops/timenstep Total Training
    Perplexity BLEU #Parameters Time
    MoE with 2048 Experts 2.69 40.35 85M 8.7B 3 days/64 k40s
    MoE with 2048 Experts (longer training) 2.63 40.56 85M 8.7B 6 days/64 k40s
    GNMT (Wu et al., 2016) 2.79 39.22 214M 278M 6 days/96 k80s
    GNMT+RL (Wu et al., 2016) 2.96 39.92 214M 278M 6 days/96 k80s
    PBMT (Durrani et al., 2014) 37.0
    LSTM (6-layer) (Luong et al., 2015b) 31.5
    LSTM (6-layer+PosUnk) (Luong et al., 2015b) 33.1
    DeepAtt (Zhou et al., 2016) 37.7
    DeepAtt+PosUnk (Zhou et al., 2016) 39.2
    Model Test Test ops/timestep Total Training
    Perplexity BLEU #Parameters Time
    MoE with 2048 Experts 4.64 26.03 85M 8.7B 1 day/64 k40s
    GNMT (Wu et al., 2016) 5.25 24.91 214M 278M 1 day/96 k80s
    GNMT +RL (Wu et al., 2016) 8.08 24.66 214M 278M 1 day/96 k80s
    PBMT (Durrani et al., 2014) 20.7
    DeepAtt (Zhou et al., 2016) 20.6
    Model Eval Eval Test Test ops/timestep Total Training
    Perplexity BLEU Perplexity BLEU #Parameters Time
    MoE with 2048 Experts 2.60 37.27 2.69 36.57 85M 8.7B 1 day/64 k40s
    GNMT (Wu et al., 2016) 2.78 35.80 2.87 35.56 214M 278M 6 days/96 k80s

    Resultados:

    Las Tablas 2, 3 y 4 muestran los resultados de nuestros modelos más grandes, comparados con resultados publicados. Nuestro enfoque logró puntuaciones BLEU de 40.56 y 26.03 en los benchmarks WMT'14 En$ {\textstyle \rightarrow} $Fr y En$ {\textstyle \rightarrow} $De. Como nuestros modelos no usaron refinamiento por RL, estos resultados constituyen ganancias significativas de 1.34 y 1.12 puntos BLEU sobre las líneas base sólidas de (Wu et al., 2016). Las puntuaciones de perplejidad también son mejores.222Las perplejidades reportadas son relativas a la tokenización usada tanto por nuestros modelos como por GNMT. En el conjunto de datos de producción de Google, nuestro modelo logró una puntuación BLEU de prueba 1.01 más alta incluso después de entrenar solo durante una sexta parte del tiempo.

    5.4 Traducción automática multilingüe

    Conjunto de datos:

    (Johnson et al., 2016) entrenan un único modelo GNMT (Wu et al., 2016) sobre un conjunto de datos combinado muy grande de doce pares de lenguas. Los resultados son algo peores que los de 12 modelos GNMT de un único par entrenados por separado. Esto no es sorprendente, dado que los doce modelos tienen 12 veces la capacidad y doce veces el entrenamiento agregado del único modelo. Repetimos este experimento con un único modelo aumentado con MoE. Ver Apéndice E para detalles sobre la arquitectura del modelo. Entrenamos nuestro modelo sobre el mismo conjunto de datos que (Johnson et al., 2016) y procesamos el mismo número de ejemplos de entrenamiento (alrededor de 3 mil millones de pares de oraciones). Nuestro tiempo de entrenamiento fue más corto debido al menor presupuesto computacional de nuestro modelo.

    Resultados:

    Los resultados para los modelos GNMT de un único par, el modelo GNMT multilingüe y el modelo MoE multilingüe se dan en la Tabla 5. El modelo MoE logra una perplejidad un 19% menor en el conjunto de desarrollo que el modelo GNMT multilingüe. En cuanto a puntuación BLEU, el modelo MoE supera significativamente al modelo GNMT multilingüe en 11 de los 12 pares de lenguas (hasta por 5.84 puntos), e incluso supera a los modelos GNMT monolingües en 8 de 12 pares de lenguas. El pobre rendimiento en inglés $ {\textstyle \rightarrow} $ coreano parece ser resultado de un sobreajuste severo, ya que para los pares de lenguas más raros un pequeño número de ejemplos reales fueron sobremuestreados en gran medida en el corpus de entrenamiento.

    GNMT-Mono GNMT-Multi MoE-Multi MoE-Multi vs.
    GNMT-Multi
    Parameters 278M / model 278M 8.7B
    ops/timestep 212M 212M 102M
    training time, hardware various 21 days, 96 k20s 12 days, 64 k40s
    Perplexity (dev) 4.14 3.35 -19%
    French $ {\textstyle \rightarrow} $ English Test BLEU 36.47 34.40 37.46 +3.06
    German $ {\textstyle \rightarrow} $ English Test BLEU 31.77 31.17 34.80 +3.63
    Japanese $ {\textstyle \rightarrow} $ English Test BLEU 23.41 21.62 25.91 +4.29
    Korean $ {\textstyle \rightarrow} $ English Test BLEU 25.42 22.87 28.71 +5.84
    Portuguese $ {\textstyle \rightarrow} $ English Test BLEU 44.40 42.53 46.13 +3.60
    Spanish $ {\textstyle \rightarrow} $ English Test BLEU 38.00 36.04 39.39 +3.35
    English $ {\textstyle \rightarrow} $ French Test BLEU 35.37 34.00 36.59 +2.59
    English $ {\textstyle \rightarrow} $ German Test BLEU 26.43 23.15 24.53 +1.38
    English $ {\textstyle \rightarrow} $ Japanese Test BLEU 23.66 21.10 22.78 +1.68
    English $ {\textstyle \rightarrow} $ Korean Test BLEU 19.75 18.41 16.62 -1.79
    English $ {\textstyle \rightarrow} $ Portuguese Test BLEU 38.40 37.35 37.90 +0.55
    English $ {\textstyle \rightarrow} $ Spanish Test BLEU 34.50 34.25 36.21 +1.96

    6 Conclusión

    Este trabajo es el primero en demostrar grandes ganancias de la computación condicional en redes profundas. Identificamos cuidadosamente las consideraciones de diseño y los desafíos del cómputo condicional y los abordamos con una combinación de soluciones algorítmicas y de ingeniería. Aunque nos centramos en texto, la computación condicional puede ayudar también en otros dominios, siempre que se disponga de conjuntos de entrenamiento suficientemente grandes. Esperamos ver muchas implementaciones y aplicaciones novedosas de la computación condicional en los años venideros.

    Agradecimientos

    Nos gustaría agradecer a todos los miembros de los equipos de Google Brain y Google Translate que nos ayudaron con este proyecto, en particular a Zhifeng Chen, Yonghui Wu y Melvin Johnson. Gracias también a nuestros revisores anónimos de ICLR por las útiles sugerencias para mejorar este artículo.

    Referencias

    • Abadi et al. (2016) Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Gregory S. Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian J. Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Józefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mané, Rajat Monga, Sherry Moore, Derek Gordon Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul A. Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda B. Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: Large-scale machine learning on heterogeneous distributed systems. CoRR, abs/1603.04467, 2016. URL http://arxiv.org/abs/1603.04467.
    • Aljundi et al. (2016) Rahaf Aljundi, Punarjay Chakravarty, and Tinne Tuytelaars. Expert gate: Lifelong learning with a network of experts. CoRR, abs/1611.06194, 2016. URL http://arxiv.org/abs/1611.06194.
    • Almahairi et al. (2015) A. Almahairi, N. Ballas, T. Cooijmans, Y. Zheng, H. Larochelle, and A. Courville. Dynamic Capacity Networks. ArXiv e-prints, November 2015.
    • Amodei et al. (2015) Dario Amodei, Rishita Anubhai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Jingdong Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Erich Elsen, Jesse Engel, Linxi Fan, Christopher Fougner, Tony Han, Awni Y. Hannun, Billy Jun, Patrick LeGresley, Libby Lin, Sharan Narang, Andrew Y. Ng, Sherjil Ozair, Ryan Prenger, Jonathan Raiman, Sanjeev Satheesh, David Seetapun, Shubho Sengupta, Yi Wang, Zhiqian Wang, Chong Wang, Bo Xiao, Dani Yogatama, Jun Zhan, and Zhenyao Zhu. Deep speech 2: End-to-end speech recognition in english and mandarin. arXiv preprint arXiv:1512.02595, 2015.
    • Bahdanau et al. (2014) Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473, 2014.
    • Bengio et al. (2015) Emmanuel Bengio, Pierre-Luc Bacon, Joelle Pineau, and Doina Precup. Conditional computation in neural networks for faster models. arXiv preprint arXiv:1511.06297, 2015.
    • Bengio et al. (2013) Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.
    • Chelba et al. (2013) Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, Thorsten Brants, Phillipp Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005, 2013.
    • Cho & Bengio (2014) K. Cho and Y. Bengio. Exponentially Increasing the Capacity-to-Computation Ratio for Conditional Computation in Deep Learning. ArXiv e-prints, June 2014.
    • Collobert et al. (2002) Ronan Collobert, Samy Bengio, and Yoshua Bengio. A parallel mixture of SVMs for very large scale problems. Neural Computing, 2002.
    • Davis & Arel (2013) Andrew Davis and Itamar Arel. Low-rank approximations for conditional feedforward computation in deep neural networks. arXiv preprint arXiv:1312.4461, 2013.
    • Deisenroth & Ng (2015) Marc Peter Deisenroth and Jun Wei Ng. Distributed Gaussian processes. In ICML, 2015.
    • Duchi et al. (2010) John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and stochastic optimization, 2010.
    • Durrani et al. (2014) Nadir Durrani, Barry Haddow, Philipp Koehn, and Kenneth Heafield. Edinburgh’s phrase-based machine translation systems for wmt-14. In Proceedings of the Ninth Workshop on Statistical Machine Translation, 2014.
    • Eigen et al. (2013) David Eigen, Marc’Aurelio Ranzato, and Ilya Sutskever. Learning factored representations in a deep mixture of experts. arXiv preprint arXiv:1312.4314, 2013.
    • Garmash & Monz (2016) Ekaterina Garmash and Christof Monz. Ensemble learning for multi-source neural machine translation. In staff.science.uva.nl/c.monz, 2016.
    • Gers et al. (2000) Felix A. Gers, Jürgen A. Schmidhuber, and Fred A. Cummins. Learning to forget: Continual prediction with lstm. Neural Computation, 2000.
    • Gruslys et al. (2016) Audrunas Gruslys, Rémi Munos, Ivo Danihelka, Marc Lanctot, and Alex Graves. Memory-efficient backpropagation through time. CoRR, abs/1606.03401, 2016. URL http://arxiv.org/abs/1606.03401.
    • He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. IEEE Conference on Computer Vision and Pattern Recognition, 2015.
    • Hinton et al. (2012) Geoffrey Hinton, Li Deng, Dong Yu, George E. Dahl, Abdel-rahman Mohamed, Navdeep Jaitly, Andrew Senior, Vincent Vanhoucke, Patrick Nguyen, Tara N. Sainath, et al. Deep neural networks for acoustic modeling in speech recognition: The shared views of four research groups. IEEE Signal Processing Magazine, 2012.
    • Hochreiter & Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 1997.
    • Ioffe & Szegedy (2015) Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167, 2015.
    • Jacobs et al. (1991) Robert A. Jacobs, Michael I. Jordan, Steven J. Nowlan, and Geoffrey E. Hinton. Adaptive mixtures of local experts. Neural Computing, 1991.
    • Johnson et al. (2016) Melvin Johnson, Mike Schuster, Quoc V. Le, Maxim Krikun, Yonghui Wu, Zhifeng Chen, Nikhil Thorat, Fernanda B. Viégas, Martin Wattenberg, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google’s multilingual neural machine translation system: Enabling zero-shot translation. CoRR, abs/1611.04558, 2016. URL http://arxiv.org/abs/1611.04558.
    • Jordan & Jacobs (1994) Michael I. Jordan and Robert A. Jacobs. Hierarchical mixtures of experts and the EM algorithm. Neural Computing, 1994.
    • Jozefowicz et al. (2016) Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv:1602.02410, 2016.
    • Kingma & Ba (2015) Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
    • Kneser & Ney (1995) Reinhard Kneser and Hermann. Ney. Improved backingoff for m-gram language modeling., 1995.
    • Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E. Hinton. Imagenet classification with deep convolutional neural networks. In NIPS, 2012.
    • Le et al. (2012) Quoc V. Le, Marc’Aurelio Ranzato, Rajat Monga, Matthieu Devin, Kai Chen, Greg S. Corrado, Jeffrey Dean, and Andrew Y. Ng. Building high-level features using large scale unsupervised learning. In ICML, 2012.
    • Ludovic Denoyer (2014) Patrick Gallinari Ludovic Denoyer. Deep sequential neural network. arXiv preprint arXiv:1410.0510, 2014.
    • Luong et al. (2015a) Minh-Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. EMNLP, 2015a.
    • Luong et al. (2015b) Minh-Thang Luong, Ilya Sutskever, Quoc V. Le, Oriol Vinyals, and Wojciech Zaremba. Addressing the rare word problem in neural machine translation. ACL, 2015b.
    • Rasmussen & Ghahramani (2002) Carl Edward Rasmussen and Zoubin Ghahramani. Infinite mixtures of Gaussian process experts. NIPS, 2002.
    • Sak et al. (2014) Hasim Sak, Andrew W Senior, and Françoise Beaufays. Long short-term memory recurrent neural network architectures for large scale acoustic modeling. In INTERSPEECH, pp.  338–342, 2014.
    • Schuster & Nakajima (2012) Mike Schuster and Kaisuke Nakajima. Japanese and Korean voice search. ICASSP, 2012.
    • Shahbaba & Neal (2009) Babak Shahbaba and Radford Neal. Nonlinear models using dirichlet process mixtures. JMLR, 2009.
    • Sutskever et al. (2014) Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. Sequence to sequence learning with neural networks. In NIPS, 2014.
    • Theis & Bethge (2015) Lucas Theis and Matthias Bethge. Generative image modeling using spatial LSTMs. In NIPS, 2015.
    • Tresp (2001) Volker Tresp. Mixtures of Gaussian Processes. In NIPS, 2001.
    • Wu et al. (2016) Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V. Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, Jeff Klingner, Apurva Shah, Melvin Johnson, Xiaobing Liu, Łukasz Kaiser, Stephan Gouws, Yoshikiyo Kato, Taku Kudo, Hideto Kazawa, Keith Stevens, George Kurian, Nishant Patil, Wei Wang, Cliff Young, Jason Smith, Jason Riesa, Alex Rudnick, Oriol Vinyals, Greg Corrado, Macduff Hughes, and Jeffrey Dean. Google’s neural machine translation system: Bridging the gap between human and machine translation. arXiv preprint arXiv:1609.08144, 2016.
    • Yao et al. (2009) Bangpeng Yao, Dirk Walther, Diane Beck, and Li Fei-fei. Hierarchical mixture of classification experts uncovers interactions between brain regions. In NIPS. 2009.
    • Zaremba et al. (2014) Wojciech Zaremba, Ilya Sutskever, and Oriol Vinyals. Recurrent neural network regularization. arXiv preprint arXiv:1409.2329, 2014.
    • Zhou et al. (2016) Jie Zhou, Ying Cao, Xuguang Wang, Peng Li, and Wei Xu. Deep recurrent models with fast-forward connections for neural machine translation. arXiv preprint arXiv:1606.04199, 2016.

    Apéndices

    A Pérdida de balanceo de carga

    Como se discutió en la sección 4, a efectos de balanceo de carga, queremos definir una función de pérdida adicional que anime a los expertos a recibir aproximadamente el mismo número de ejemplos de entrenamiento. Lamentablemente, el número de ejemplos recibidos por un experto es una cantidad discreta, por lo que no puede usarse en back-propagation. En su lugar, definimos un estimador suave $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}} $ del número de ejemplos asignados a cada experto para un batch $ {\textstyle X} $ de entradas. La suavidad nos permite back-propagar gradientes a través del estimador. Este es el propósito del término de ruido en la función de gating. Definimos $ {\textstyle P\hspace{0pt}{(x,i)}} $ como la probabilidad de que $ {\textstyle G\hspace{0pt}{(x)}_{i}} $ sea distinto de cero, dada una nueva elección aleatoria de ruido en el elemento $ {\textstyle i} $, pero manteniendo las elecciones ya muestreadas de ruido en los otros elementos. Para calcular $ {\textstyle P\hspace{0pt}{(x,i)}} $, observamos que $ {\textstyle G\hspace{0pt}{(x)}_{i}} $ es distinto de cero si y solo si $ {\textstyle H\hspace{0pt}{(x)}_{i}} $ es mayor que el $ {\textstyle k^{t\hspace{0pt}h}} $-mayor elemento de $ {\textstyle H\hspace{0pt}{(x)}} $ excluyéndose a sí mismo. La probabilidad resulta ser:

    $ {\textstyle P{(x,i)} = Pr\left( {(x \cdot W_{g})}_{i} + StandardNormal{()} \cdot Softplus{({(x \cdot W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e})}_{i})} \right.} $ (8)
    $ {\textstyle \left. > kth\_ excluding{(H{(x)},k,i)} \right)} $

    Donde $ {\textstyle k\hspace{0pt}t\hspace{0pt}h\hspace{0pt}\_\hspace{0pt}e\hspace{0pt}x\hspace{0pt}c\hspace{0pt}l\hspace{0pt}u\hspace{0pt}d\hspace{0pt}i\hspace{0pt}n\hspace{0pt}g\hspace{0pt}{(v,k,i)}} $ denota el k-ésimo componente más alto de $ {\textstyle v} $, excluyendo el componente $ {\textstyle i} $. Simplificando, obtenemos:

    $ {\displaystyle {P\hspace{0pt}{(x,i)}} = {\Phi\hspace{0pt}\left( \frac{{({x \cdot W_{g}})}_{i} - {k\hspace{0pt}t\hspace{0pt}h\hspace{0pt}\_\hspace{0pt}e\hspace{0pt}x\hspace{0pt}c\hspace{0pt}l\hspace{0pt}u\hspace{0pt}d\hspace{0pt}i\hspace{0pt}n\hspace{0pt}g\hspace{0pt}{({H\hspace{0pt}{(x)}},k,i)}}}{S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}p\hspace{0pt}l\hspace{0pt}u\hspace{0pt}s\hspace{0pt}{({({x \cdot W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}})}_{i})}} \right)}} $ (9)

    Donde $ {\textstyle \Phi} $ es la CDF de la distribución normal estándar.

    $ {\displaystyle {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}_{i}} = {\sum\limits_{x \in X}{P\hspace{0pt}{(x,i)}}}} $ (10)

    Ahora podemos definir la load loss como el cuadrado del coeficiente de variación del vector de carga, multiplicado por un factor de escala ajustado a mano $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $.

    $ {\displaystyle {L_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}\hspace{0pt}{(X)}} = {{w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} \cdot C}\hspace{0pt}V\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}^{2}}} $ (11)

    Desbalance inicial de carga:

    Para evitar errores de out-of-memory, necesitamos inicializar la red en un estado de carga de expertos aproximadamente igual (ya que las restricciones suaves necesitan algo de tiempo para funcionar). Para lograr esto, inicializamos las matrices $ {\textstyle W_{g}} $ y $ {\textstyle W_{n\hspace{0pt}o\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}} $ a todo ceros, lo que produce ninguna señal y algo de ruido.

    Experimentos:

    Entrenamos un conjunto de modelos con arquitectura idéntica (el modelo MoE-256 descrito en el Apéndice C), usando distintos valores de $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $ y $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $. Entrenamos cada modelo durante 10 épocas, y luego medimos la perplejidad en el conjunto de prueba. También medimos los coeficientes de variación en $ {\textstyle I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} $ y $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} $, así como la razón entre la carga del experto más sobrecargado y la carga promedio. Este último valor es significativo para fines de balanceo de carga en hardware distribuido. Todas estas métricas se promediaron sobre varios batches de entrenamiento.

    $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e}} $ $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $ Test Perplexity $ {\textstyle C\hspace{0pt}V\hspace{0pt}{({I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e\hspace{0pt}{(X)}})}} $ $ {\textstyle C\hspace{0pt}V\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}} $ $ {\textstyle \frac{m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}}{m\hspace{0pt}e\hspace{0pt}a\hspace{0pt}n\hspace{0pt}{({L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d\hspace{0pt}{(X)}})}}} $
    0.0 0.0 39.8 3.04 3.01 17.80
    0.2 0.0 35.6 0.06 0.17 1.47
    0.0 0.2 35.7 0.22 0.04 1.15
    0.1 0.1 35.6 0.06 0.05 1.14
    0.01 0.01 35.7 0.48 0.11 1.37
    1.0 1.0 35.7 0.03 0.02 1.07

    Resultados:

    Los resultados se reportan en la Tabla 6. Todas las combinaciones que contenían al menos una de las dos pérdidas llevaron a una calidad de modelo muy similar, mientras que no tener ninguna pérdida fue mucho peor. Los modelos con valores más altos de $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d}} $ tuvieron cargas más bajas en el experto más sobrecargado.

    B Hierachical Mixture of Experts

    Si el número de expertos es muy grande, podemos reducir el factor de branching usando una MoE jerárquica de dos niveles. En una MoE jerárquica, una red de gating primaria elige una combinación dispersa ponderada de "expertos", cada uno de los cuales es a su vez una mixture-of-experts secundaria con su propia red de gating.333 No hemos encontrado la necesidad de jerarquías más profundas. Si la MoE jerárquica consta de $ {\textstyle a} $ grupos de $ {\textstyle b} $ expertos cada uno, denotamos la red de gating primaria por $ {\textstyle G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}} $, las redes de gating secundarias por $ {\textstyle (G_{1},G_{2}..G_{a})} $ y las redes expertas por $ {\textstyle (E_{0,0},E_{0,1}..E_{a,b})} $. La salida de la MoE viene dada por:

    $ {\displaystyle y_{H} = {\sum\limits_{i = 1}^{a}{\sum\limits_{j = 1}^{b}{{{{{G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(x)}_{i}} \cdot G_{i}}\hspace{0pt}{(x)}_{j}} \cdot E_{i,j}}\hspace{0pt}{(x)}}}}} $ (12)

    Nuestras métricas de utilización de expertos cambian a las siguientes:

    $ {\displaystyle {I\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e_{H}\hspace{0pt}{(X)}_{i,j}} = {\sum\limits_{x \in X}{{{G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(x)}_{i}} \cdot G_{i}}\hspace{0pt}{(x)}_{j}}}} $ (13)
    $ {\displaystyle {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{H}\hspace{0pt}{(X)}_{i,j}} = \frac{{{L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(X)}_{i}} \cdot L}\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{i}\hspace{0pt}{(X^{(i)})}_{j}}{|X^{(i)}|}} $ (14)

    $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}} $ y $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{i}} $ denotan las funciones $ {\textstyle L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} $ para la red de gating primaria y la $ {\textstyle i^{t\hspace{0pt}h}} $ red de gating secundaria respectivamente. $ {\textstyle X^{(i)}} $ denota el subconjunto de $ {\textstyle X} $ para el cual $ {\textstyle {G_{p\hspace{0pt}r\hspace{0pt}i\hspace{0pt}m\hspace{0pt}a\hspace{0pt}r\hspace{0pt}y}\hspace{0pt}{(x)}_{i}} > 0} $.

    Parecería más simple dejar $ {\textstyle {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{H}\hspace{0pt}{(X)}_{i,j}} = {L\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d_{i}\hspace{0pt}{(X_{i})}_{j}}} $, pero esto no tendría gradiente con respecto a la red de gating primaria, así que usamos la formulación anterior.

    C 1 Billion Word Language Modeling Benchmark - Detalles experimentales

    C.1 Modelos de 8 millones de operaciones por timestep

    Arquitectura del modelo:

    Nuestro modelo consta de cinco capas: una capa de word embedding, una capa recurrente Long Short-Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997; Gers et al., 2000), una capa MoE, una segunda capa LSTM y una capa softmax. La dimensionalidad de la capa de embedding, el número de unidades en cada capa LSTM, y la dimensionalidad de entrada y salida de la capa MoE son todas iguales a 512. Para cada capa distinta de la softmax, aplicamos dropout (Zaremba et al., 2014) a la salida de la capa, descartando cada activación con probabilidad $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $, y de lo contrario dividiéndola por $ {\textstyle ({1 - {D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b}})} $. Después del dropout, la salida de la capa anterior se añade a la salida de la capa. Esta conexión residual fomenta el flujo de gradientes (He et al., 2015).

    Arquitectura de la capa MoE:

    Cada experto en la capa MoE es una red feed forward con una capa oculta activada con ReLU de tamaño 1024 y una capa de salida de tamaño 512. Por lo tanto, cada experto contiene $ {\textstyle {{\lbrack{512 \ast 1024}\rbrack} + {\lbrack{1024 \ast 512}\rbrack}} = {1\hspace{0pt}M}} $ parámetros. La salida de la capa MoE pasa por una función sigmoide antes del dropout. Variamos el número de expertos entre modelos, usando capas MoE ordinarias con 4, 32 y 256 expertos y capas MoE jerárquicas con 256, 1024 y 4096 expertos. Llamamos a los modelos resultantes MoE-4, MoE-32, MoE-256, MoE-256-h, MoE-1024-h y MoE-4096-h. Para las capas MoE jerárquicas, el factor de branching del primer nivel fue 16, correspondiente al número de GPUs en nuestro clúster. Usamos Noisy-Top-K Gating (ver Sección 2.1) con $ {\textstyle k = 4} $ para las capas MoE ordinarias y $ {\textstyle k = 2} $ en cada nivel de las capas MoE jerárquicas. Por lo tanto, cada ejemplo es procesado por exactamente 4 expertos para un total de 4M ops/timestep. Las dos capas LSTM contribuyen 2M ops/timestep cada una para el total deseado de 8M.

    Líneas base con cómputo equivalente:

    El modelo MoE-4 no emplea sparsity, ya que los 4 expertos siempre se usan. Además, entrenamos cuatro modelos base más con cómputo equivalente sin sparsity:

    • MoE-1-Wide: La capa MoE consiste en un único "experto" que contiene una capa oculta activada con ReLU de tamaño 4096.

    • MoE-1-Deep: La capa MoE consiste en un único "experto" que contiene cuatro capas ocultas activadas con ReLU, cada una con tamaño $ {\textstyle 1024} $.

    • 4xLSTM-512: Reemplazamos la capa MoE con dos capas LSTM adicionales de 512 unidades.

    • LSTM-2048-512: El modelo contiene una capa LSTM de 2048 unidades (y ninguna MoE). La salida del LSTM se proyecta a 512 dimensiones (Sak et al., 2014). El siguiente timestep del LSTM recibe la salida proyectada. Esto es idéntico a uno de los modelos publicados en (Jozefowicz et al., 2016). Lo volvimos a ejecutar para tener en cuenta diferencias en el régimen de entrenamiento, y obtuvimos resultados muy similares a los publicados.

    Entrenamiento:

    Los modelos se entrenaron en un clúster de 16 GPUs K40 usando el método sincrónico descrito en la Sección 3. Cada batch consistía en un conjunto de oraciones que totalizaban aproximadamente 300.000 palabras. En aras del tiempo, limitamos el entrenamiento a 10 épocas (27.000 pasos). El entrenamiento tomó 12-16 horas para todos los modelos, excepto para MoE-4, que tomó 18 horas (ya que todo el cómputo de los expertos se realizaba en solo 4 de 16 GPUs). Usamos el optimizador Adam (Kingma & Ba, 2015). La tasa de aprendizaje base se incrementó linealmente durante los primeros 1000 pasos de entrenamiento y se redujo después de manera proporcional al inverso de la raíz cuadrada del número de paso. La capa de salida Softmax se entrenó eficientemente usando importance sampling de manera similar a los modelos en (Jozefowicz et al., 2016). Para cada modelo, realizamos una búsqueda de hiperparámetros para encontrar la mejor probabilidad de dropout, en incrementos de 0.1.

    Para asegurar una utilización balanceada de expertos, fijamos $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} = 0.1} $ y $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} = 0.1} $, como se describe en la Sección 4 y el Apéndice A.

    Resultados:

    Evaluamos nuestro modelo usando perplejidad sobre el conjunto holdout, usado por (Chelba et al., 2013; Jozefowicz et al., 2016). Seguimos el procedimiento estándar y sumamos sobre todas las palabras incluyendo el símbolo de fin de oración. Los resultados se reportan en la Tabla 7. Para cada modelo, reportamos la perplejidad de prueba, el presupuesto computacional, los conteos de parámetros, el valor de $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $ y la eficiencia computacional.

    Model Test Test ops/timestep #Params excluding Total $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p} $- TFLOPS
    Perplexity Perplexity (millions) embed. & softmax #Params $ {\textstyle P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $ per GPU
    10 epochs (final) (millions) (billions) (observed)
    Kneser-Ney 5-gram* 67.6 0.00001 1.8
    LSTM-512-512* 54.1 2.4 2.4 0.8 0.1
    LSTM-1024-512* 48.2 4.7 4.7 0.8 0.1
    LSTM-2048-512* 45.0 43.7 9.4 9.4 0.8 0.1 0.61
    LSTM-2048-512 44.7 9.4 9.4 0.8 0.1 1.21
    4xLSTM-512 46.0 8.4 8.4 0.8 0.1 1.07
    MoE-1-Wide 46.1 8.4 8.4 0.8 0.1 1.29
    MoE-1-Deep 45.7 8.4 8.4 0.8 0.1 1.29
    MoE-4 45.0 8.4 8.4 0.8 0.1 0.52
    MoE-32 39.7 8.4 37.8 0.9 0.1 0.87
    MoE-256 35.7 8.6 272.9 1.1 0.1 0.81
    MoE-256-h 36.0 8.4 272.9 1.1 0.1 0.89
    MoE-1024-h 34.6 8.5 1079.0 1.9 0.2 0.90
    MoE-4096-h 34.1 8.9 4303.4 5.1 0.2 0.74
    2xLSTM-8192-1024* 34.7 30.6 151.0 151.0 1.8 0.25 1.09
    MoE-34M 31.3 33.8 4313.9 6.0 0.3 1.22
    MoE-143M 28.0 142.7 4371.1 6.0 0.4 1.56

    C.2 Modelos más costosos

    Ejecutamos dos modelos adicionales (MoE-34M y MoE-143M) para investigar los efectos de añadir más cómputo en presencia de una capa MoE grande. Estos modelos tienen presupuestos de cómputo de 34M y 143M ops/timestep. De manera similar a los modelos anteriores, estos modelos usan una capa MoE entre dos capas LSTM. La dimensionalidad de la capa de embedding y la dimensionalidad de entrada y salida de la capa MoE se fijan en 1024 en lugar de 512. Para MoE-34M, las capas LSTM tienen 1024 unidades. Para MoE-143M, las capas LSTM tienen 4096 unidades y una proyección de salida de tamaño 1024 (Sak et al., 2014). MoE-34M usa una capa MoE jerárquica con 1024 expertos, cada uno con una capa oculta de tamaño 2048. MoE-143M usa una capa MoE jerárquica con 256 expertos, cada uno con una capa oculta de tamaño 8192. Ambos modelos tienen 4B parámetros en las capas MoE. Buscamos el mejor $ {\textstyle D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} $ para cada modelo y entrenamos cada modelo durante 10 épocas.

    Los dos modelos lograron una perplejidad de prueba de $ {\textstyle 31.3} $ y $ {\textstyle 28.0} $ respectivamente, mostrando que incluso en presencia de una MoE grande, más cómputo sigue siendo útil. Los resultados se reportan al final de la Tabla 7. El más grande de los dos modelos tiene un presupuesto computacional similar al del mejor modelo publicado en la literatura, y los tiempos de entrenamiento son similares. Comparando después de 10 épocas, nuestro modelo tiene una perplejidad de prueba menor en un $ {\textstyle 18\%} $.

    D 100 Billion Word Google News Corpus - Detalles experimentales

    Arquitectura del modelo:

    Los modelos son similares en estructura a los modelos de 8 millones de operaciones por timestep descritos en la sección anterior. Variamos el número de expertos entre modelos, usando una capa MoE ordinaria con 32 expertos y capas MoE jerárquicas con 256, 1024, 4096, 16384, 65536 y 131072 expertos. Para las capas MoE jerárquicas, los factores de branching del primer nivel son 32, 32, 64, 128, 256 y 256, respectivamente.

    Entrenamiento:

    Los modelos se entrenan en un clúster de 32 GPUs Tesla K40, excepto los dos últimos modelos, que se entrenan en clústeres de 64 y 128 GPUs para tener suficiente memoria para todos los parámetros. Para todos los modelos, los batch sizes de entrenamiento son aproximadamente 2.5 millones de palabras. Los modelos se entrenan una sola pasada sobre aproximadamente 100 mil millones de palabras.

    Implementamos varias optimizaciones de memoria para encajar hasta 1 mil millones de parámetros por GPU. Primero, no almacenamos las activaciones de las capas ocultas de los expertos, sino que las recalculamos en el pase hacia atrás. Segundo, modificamos el optimizador en los parámetros de los expertos para requerir menos almacenamiento auxiliar:

    El optimizador Adam (Kingma & Ba, 2015) mantiene estimaciones de primer y segundo momento de los gradientes por parámetro. Esto triplica la memoria requerida. Para evitar mantener un estimador de primer momento, fijamos $ {\textstyle \beta_{1} = 0} $. Para reducir el tamaño del estimador de segundo momento, lo reemplazamos con una aproximación factorizada. Para una matriz de parámetros, en lugar de mantener una matriz completa de estimadores de segundo momento, mantenemos vectores de promedios por filas y por columnas de esa matriz. En cada paso, la matriz de estimadores se toma como el producto exterior de esos dos vectores dividido por la media de cualquiera de ellos. Esta técnica podría aplicarse de manera similar a Adagrad (Duchi et al., 2010).

    Model Test Test ops/timestep #Params excluding Total TFLOPS
    Perplexity Perplexity (millions) embed. & softmax #Params per GPU
    .1 epochs 1 epoch (millions) (billions) (observed)
    Kneser-Ney 5-gram 67.1 45.3 0.00001 76.0
    4xLSTM-512 54.5 47.0 8.4 8.4 0.1 1.23
    MoE-32 48.5 40.4 8.4 37.8 0.1 0.83
    MoE-256-h 42.8 35.3 8.4 272.9 0.4 1.11
    MoE-1024-h 40.3 32.7 8.5 1079.0 1.2 1.14
    MoE-4096-h 38.9 30.9 8.6 4303.4 4.4 1.07
    MoE-16384-h 38.2 29.7 8.8 17201.0 17.3 0.96
    MoE-65536-h 38.2 28.9 9.2 68791.0 68.9 0.72
    MoE-131072-h 39.8 29.2 9.7 137577.6 137.7 0.30

    Resultados:

    Evaluamos nuestro modelo usando perplejidad sobre un conjunto holdout. Los resultados se reportan en la Tabla 8. La perplejidad después de 100 mil millones de palabras de entrenamiento es 39% menor para el modelo MoE de 68 mil millones de parámetros que para el modelo base. Es notable que la eficiencia computacional medida del modelo más grande (0.30 TFLOPS/GPU) es muy baja en comparación con los otros modelos. Esto probablemente sea resultado del hecho de que, a efectos de comparación con los otros modelos, no aumentamos el batch size de entrenamiento proporcionalmente al número de GPUs. Para comparación, incluimos resultados para un modelo base con cómputo equivalente consistente en 4 LSTMs, y para un modelo de 5-gramas no podado con suavizado Kneser-Ney (Kneser & Ney, 1995).444Aunque el tamaño original del corpus era de 130 mil millones de palabras, los modelos neuronales se entrenaron por un máximo de 100 mil millones de palabras. Los modelos Kneser-Ney 5-gram reportados se entrenaron sobre 13 mil millones y 130 mil millones de palabras respectivamente, dándoles una ligera ventaja sobre los otros resultados reportados.

    E Traducción automática - Detalles experimentales

    Arquitectura del modelo para los MoE de un único par de lenguas:

    Nuestro modelo es una versión modificada del modelo GNMT descrito en (Wu et al., 2016). Para reducir cómputo, disminuimos el número de capas LSTM en el encoder y decoder de 9 y 8 a 3 y 2 respectivamente. Insertamos capas MoE tanto en el encoder (entre las capas 2 y 3) como en el decoder (entre las capas 1 y 2). Usamos un mecanismo de atención entre el encoder y el decoder, con el primer LSTM del decoder recibiendo salida de y proporcionando entrada para la atención555Por razones de rendimiento, usamos una función de atención ligeramente diferente de la descrita en (Wu et al., 2016) - ver Apéndice G. Todas las capas de nuestro modelo tienen dimensionalidad de entrada y salida de 512. Nuestras capas LSTM tienen 2048 unidades ocultas, con una proyección de salida de 512 dimensiones. Añadimos conexiones residuales alrededor de todas las capas LSTM y MoE para fomentar el flujo de gradientes (He et al., 2015). De manera similar a GNMT, para tratar eficazmente con palabras raras, usamos sub-word units (también conocidas como "wordpieces") (Schuster & Nakajima, 2012) para entradas y salidas en nuestro sistema.

    Usamos un vocabulario compartido de fuente y objetivo de 32K wordpieces. También usamos la misma técnica de beam search propuesta en (Wu et al., 2016).

    Entrenamos modelos con distintos números de expertos en las capas MoE. Además de un modelo base sin capas MoE, entrenamos modelos con capas MoE planas que contenían 32 expertos, y modelos con capas MoE jerárquicas que contenían 512 y 2048 expertos. Las capas MoE planas usan $ {\textstyle k = 4} $ y los modelos MoE jerárquicos usan $ {\textstyle k = 2} $ en cada nivel de la red de gating. Por lo tanto, cada entrada es procesada por exactamente 4 expertos en cada capa MoE. Cada experto en la capa MoE es una red feed forward con una capa oculta de tamaño 2048 y activación ReLU. Por lo tanto, cada experto contiene $ {\textstyle {{\lbrack{512 \ast 2048}\rbrack} + {\lbrack{2048 \ast 512}\rbrack}} = {2\hspace{0pt}M}} $ parámetros. La salida de la capa MoE pasa por una función sigmoide. Usamos la función de gating estrictamente balanceada descrita en el Apéndice F.

    Arquitectura del modelo para el MoE multilingüe:

    Usamos la misma arquitectura de modelo que para los modelos de un solo par de lenguas, con las siguientes excepciones: Usamos noisy-top-k gating como se describe en la Sección 2.1, no el esquema del Apéndice F. Las capas MoE en el encoder y decoder son MoEs no jerárquicas con $ {\textstyle n = 512} $ expertos y $ {\textstyle k = 2} $. Cada experto tiene una capa oculta más grande de tamaño $ {\textstyle 8192} $. Esto duplica la cantidad de cómputo en las capas MoE, elevando el presupuesto computacional del modelo entero de 85M a 102M ops/timestep.

    Entrenamiento:

    Entrenamos nuestras redes usando el optimizador Adam (Kingma & Ba, 2015). La tasa de aprendizaje base se incrementó linealmente durante los primeros 2000 pasos de entrenamiento, se mantuvo constante durante 8000 pasos adicionales y luego se redujo de manera proporcional al inverso de la raíz cuadrada del número de paso. Para los modelos de un solo par de lenguas, de manera similar a (Wu et al., 2016), aplicamos dropout (Zaremba et al., 2014) a la salida de todas las capas de embedding, LSTM y MoE, usando $ {\textstyle {D\hspace{0pt}r\hspace{0pt}o\hspace{0pt}p\hspace{0pt}P\hspace{0pt}r\hspace{0pt}o\hspace{0pt}b} = 0.4} $. El entrenamiento se hizo sincrónicamente en un clúster de hasta 64 GPUs como se describe en la sección 3. Cada batch de entrenamiento consistía en un conjunto de pares de oraciones que contenían aproximadamente 16000 palabras por GPU.

    Para asegurar una utilización balanceada de expertos, fijamos $ {\textstyle w_{i\hspace{0pt}m\hspace{0pt}p\hspace{0pt}o\hspace{0pt}r\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}c\hspace{0pt}e} = 0.01} $ y $ {\textstyle w_{l\hspace{0pt}o\hspace{0pt}a\hspace{0pt}d} = 0.01} $, como se describe en la Sección 4 y el Apéndice A.

    Métricas:

    Evaluamos nuestros modelos usando la perplejidad y la métrica estándar de puntuación BLEU. Reportamos puntuaciones BLEU tokenizadas calculadas por el script multi-bleu.pl, descargado de la implementación pública de Moses (en Github), que también se usó en (Luong et al., 2015a).

    Resultados:

    Las Tablas 2, 3 y 4 en la Sección 5.3 muestran comparaciones de nuestros resultados con otros métodos publicados. La Figura 4 muestra la perplejidad de prueba en función del número de palabras en (de los datos de entrenamiento) las oraciones fuente procesadas para modelos con distintos números de expertos. Como puede verse en la Figura, conforme aumentábamos el número de expertos hacia 2048, la perplejidad de prueba de nuestro modelo seguía mejorando.

    Refer to caption Refer to caption

    Encontramos que los expertos efectivamente se vuelven altamente especializados por sintaxis y/o semántica, como puede verse en la Tabla 9. Por ejemplo, un experto se usa cuando el artículo indefinido "a" introduce el objeto directo en una frase verbal que indica importancia o liderazgo.

    Expert 381 Expert 752 Expert 2004
    … with researchers , … … plays a core … … with rapidly growing …
    … to innovation . … plays a critical … … under static conditions …
    … tics researchers . … provides a legislative … … to swift ly …
    … the generation of … … play a leading … … to dras tically …
    … technology innovations is … … assume a leadership … … the rapid and …
    … technological innovations , … … plays a central … … the fast est …
    … support innovation throughout … … taken a leading … … the Quick Method …
    … role innovation will … … established a reconciliation … … rec urrent ) …
    … research scienti st … … played a vital … … provides quick access …
    … promoting innovation where … … have a central … … of volatile organic …

    F Gating estrictamente balanceado

    Debido a algunas peculiaridades en nuestra infraestructura que desde entonces han sido corregidas, en el momento en que ejecutamos algunos de los experimentos de traducción automática, nuestros modelos corrían más rápido si cada experto recibía exactamente el mismo batch size. Para acomodar esto, usamos una función de gating diferente que describimos a continuación.

    Recordemos que definimos la función de gating softmax como:

    $ {\displaystyle {G_{\sigma}\hspace{0pt}{(x)}} = {S\hspace{0pt}o\hspace{0pt}f\hspace{0pt}t\hspace{0pt}m\hspace{0pt}a\hspace{0pt}x\hspace{0pt}{({x \cdot W_{g}})}}} $ (15)

    Sparse Gating (formulación alternativa):

    Para obtener un vector de gating disperso, multiplicamos $ {\textstyle G_{\sigma}\hspace{0pt}{(x)}} $ componente a componente con una máscara dispersa $ {\textstyle M\hspace{0pt}{({G_{\sigma}\hspace{0pt}{(x)}})}} $ y normalizamos la salida. La máscara misma es función de $ {\textstyle G_{\sigma}\hspace{0pt}{(x)}} $ y especifica qué expertos se asignan a cada ejemplo de entrada:

    $ {\displaystyle {G\hspace{0pt}{(x)}_{i}} = \frac{G_{\sigma}\hspace{0pt}{(x)}_{i}\hspace{0pt}M\hspace{0pt}{({G_{\sigma}\hspace{0pt}{(x)}})}_{i}}{\sum_{j = 1}^{n}{G_{\sigma}\hspace{0pt}{(x)}_{j}\hspace{0pt}M\hspace{0pt}{({G_{\sigma}\hspace{0pt}{(x)}})}_{j}}}} $ (16)

    Top-K Mask:

    Para implementar top-k gating en esta formulación, dejaríamos $ {\textstyle {M\hspace{0pt}{(v)}} = {T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{(v,k)}}} $, donde:

    $ {\displaystyle {T\hspace{0pt}o\hspace{0pt}p\hspace{0pt}K\hspace{0pt}{(v,k)}_{i}} = \begin{cases} 1 & {\text{if~}v_{i}\text{~is in the top~}k\text{~elements of~}v\text{.}} \\ 0 & \text{otherwise.} \end{cases}} $ (17)

    Batchwise Mask:

    Para forzar a que cada experto reciba exactamente el mismo número de ejemplos, introducimos una función de máscara alternativa, $ {\textstyle M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,m)}} $, que opera sobre batches de vectores de entrada. En lugar de mantener los top $ {\textstyle k} $ valores por ejemplo, mantenemos los top $ {\textstyle m} $ valores por experto a lo largo del batch de entrenamiento, donde $ {\textstyle m = \frac{k\hspace{0pt}{|X|}}{n}} $, de modo que cada ejemplo se envía a un promedio de $ {\textstyle k} $ expertos.

    $ {\displaystyle {M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,m)}_{j,i}} = \begin{cases} 1 & {\text{if~}X_{j,i}\text{~is in the top~}m\text{~values for to expert~}i} \\ 0 & \text{otherwise} \end{cases}} $ (18)

    Como sugieren nuestros experimentos y como también observan (Ioffe & Szegedy, 2015), usar una función batchwise durante el entrenamiento (como $ {\textstyle M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}} $) requiere modificaciones en la inferencia cuando puede que no tengamos un batch grande de ejemplos. Nuestra solución a esto es entrenar un vector $ {\textstyle T} $ de valores umbral por experto para aproximar los efectos de la máscara batchwise. Usamos la siguiente máscara en tiempo de inferencia:

    $ {\displaystyle {M_{t\hspace{0pt}h\hspace{0pt}r\hspace{0pt}e\hspace{0pt}s\hspace{0pt}h\hspace{0pt}o\hspace{0pt}l\hspace{0pt}d}\hspace{0pt}{(x,T)}_{i}} = \begin{cases} 1 & {\text{if~}{x_{i} > T_{i}}} \\ 0 & \text{otherwise} \end{cases}} $ (19)

    Para aprender los valores umbral, aplicamos una pérdida adicional en tiempo de entrenamiento que se minimiza cuando la máscara batchwise y la máscara umbral son idénticas.

    $ {\displaystyle {L_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,T,m)}} = {\sum\limits_{j = 1}^{|X|}{\sum\limits_{i = 1}^{n}{{({{M_{t\hspace{0pt}h\hspace{0pt}r\hspace{0pt}e\hspace{0pt}s\hspace{0pt}h\hspace{0pt}o\hspace{0pt}l\hspace{0pt}d}\hspace{0pt}{(x,T)}_{i}} - {M_{b\hspace{0pt}a\hspace{0pt}t\hspace{0pt}c\hspace{0pt}h\hspace{0pt}w\hspace{0pt}i\hspace{0pt}s\hspace{0pt}e}\hspace{0pt}{(X,m)}_{j,i}}})}\hspace{0pt}{({X_{j,i} - T_{i}})}}}}} $ (20)

    G Función de atención

    El mecanismo de atención descrito en GNMT (Wu et al., 2016) involucra una "Attention Function" aprendida $ {\textstyle A\hspace{0pt}{(x_{i},y_{j})}} $ que toma un "source vector" $ {\textstyle x_{i}} $ y un "target vector" $ {\textstyle y_{j}} $, y debe calcularse para cada paso temporal de la fuente $ {\textstyle i} $ y cada paso temporal del objetivo $ {\textstyle j} $. En GNMT, la función de atención se implementa como una red neuronal feed forward con una capa oculta de tamaño $ {\textstyle n} $. Se puede expresar como:

    $ {\displaystyle {A_{G\hspace{0pt}N\hspace{0pt}M\hspace{0pt}T}\hspace{0pt}{(x_{i},y_{j})}} = {\sum\limits_{d = 1}^{n}{V_{d}\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}h\hspace{0pt}{({{({x_{i}\hspace{0pt}U})}_{d} + {({y_{j}\hspace{0pt}W})}_{d}})}}}} $ (21)

    Donde $ {\textstyle U} $ y $ {\textstyle W} $ son matrices de pesos entrenables y $ {\textstyle V} $ es un vector de pesos entrenable.

    Por razones de rendimiento, en nuestros modelos, usamos una función de atención ligeramente diferente:

    $ {\displaystyle {A\hspace{0pt}{(x_{i},y_{j})}} = {\sum\limits_{d = 1}^{n}{V_{d}\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}h\hspace{0pt}{({({x_{i}\hspace{0pt}U})}_{d})}\hspace{0pt}t\hspace{0pt}a\hspace{0pt}n\hspace{0pt}h\hspace{0pt}{({({y_{j}\hspace{0pt}W})}_{d})}}}} $ (22)

    Con nuestra función de atención, podemos calcular simultáneamente la función de atención sobre múltiples pasos temporales de la fuente y múltiples pasos temporales del objetivo usando multiplicaciones matriciales optimizadas. Encontramos poca diferencia en calidad entre las dos funciones.