Brainformers: Intercambiando Simplicidad por Eficiencia

Brainformers: Trading Simplicity for Efficiency

May 29, 2023

Autores: Yanqi Zhou, Nan Du, Yanping Huang, Daiyi Peng, Chang Lan, Da Huang, Siamak Shakeri, David So, Andrew Dai, Yifeng Lu, Zhifeng Chen, Quoc Le, Claire Cui, James Laundon, Jeff Dean

cs.AI

Resumen

Los Transformers son fundamentales para los recientes avances en procesamiento de lenguaje natural y visión por computadora. Los Transformers tienen una estructura principalmente uniforme, donde las capas alternan entre redes feed-forward y auto-atención para construir una red profunda. En este trabajo, investigamos esta elección de diseño y descubrimos que bloques más complejos, con diferentes permutaciones de primitivas de capas, pueden ser más eficientes. Basándonos en esta idea, desarrollamos un bloque complejo, denominado Brainformer, que consta de un conjunto diverso de capas, como capas feed-forward con compuertas dispersas, capas feed-forward densas, capas de atención y diversas formas de normalización de capas y funciones de activación. Brainformer supera consistentemente a los Transformers densos y dispersos de última generación, tanto en calidad como en eficiencia. Un modelo Brainformer con 8 mil millones de parámetros activados por token demuestra una convergencia de entrenamiento 2 veces más rápida y un tiempo por paso 5 veces más rápido en comparación con su contraparte GLaM. En la evaluación de tareas posteriores, Brainformer también muestra un 3% más de puntuación en SuperGLUE con ajuste fino en comparación con GLaM, con un número similar de parámetros activados. Finalmente, Brainformer supera ampliamente a un modelo denso Primer derivado mediante NAS con un cálculo similar por token en evaluaciones de pocos ejemplos.

English

Transformers are central to recent successes in natural language processing and computer vision. Transformers have a mostly uniform backbone where layers alternate between feed-forward and self-attention in order to build a deep network. Here we investigate this design choice and find that more complex blocks that have different permutations of layer primitives can be more efficient. Using this insight, we develop a complex block, named Brainformer, that consists of a diverse sets of layers such as sparsely gated feed-forward layers, dense feed-forward layers, attention layers, and various forms of layer normalization and activation functions. Brainformer consistently outperforms the state-of-the-art dense and sparse Transformers, in terms of both quality and efficiency. A Brainformer model with 8 billion activated parameters per token demonstrates 2x faster training convergence and 5x faster step time compared to its GLaM counterpart. In downstream task evaluation, Brainformer also demonstrates a 3% higher SuperGLUE score with fine-tuning compared to GLaM with a similar number of activated parameters. Finally, Brainformer largely outperforms a Primer dense model derived with NAS with similar computation per token on fewshot evaluations.