Brainformers : Échanger la simplicité contre l'efficacité

Brainformers: Trading Simplicity for Efficiency

May 29, 2023

Auteurs: 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

Résumé

Les Transformers sont au cœur des récents succès en traitement du langage naturel et en vision par ordinateur. Les Transformers possèdent une architecture principalement uniforme, où les couches alternent entre des couches feed-forward et des mécanismes d'auto-attention pour construire un réseau profond. Dans cette étude, nous examinons ce choix de conception et constatons que des blocs plus complexes, intégrant différentes permutations de primitives de couches, peuvent être plus efficaces. En nous appuyant sur cette observation, nous développons un bloc complexe, nommé Brainformer, qui combine divers types de couches telles que des couches feed-forward à portes éparses, des couches feed-forward denses, des couches d'attention, ainsi que diverses formes de normalisation de couches et de fonctions d'activation. Brainformer surpasse systématiquement les Transformers denses et éparses de pointe, tant en termes de qualité que d'efficacité. Un modèle Brainformer avec 8 milliards de paramètres activés par token montre une convergence d'entraînement 2 fois plus rapide et un temps par étape 5 fois plus rapide que son homologue GLaM. Dans l'évaluation des tâches en aval, Brainformer obtient également un score SuperGLUE 3 % plus élevé après ajustement fin, comparé à GLaM avec un nombre similaire de paramètres activés. Enfin, Brainformer surpasse largement un modèle dense Primer dérivé par recherche d'architecture neuronale (NAS) avec un calcul similaire par token dans des évaluations en few-shot.

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.