Brainformers: Trocar Simplicidade por Eficiência

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

Resumo

Os Transformers são centrais para os recentes sucessos em processamento de linguagem natural e visão computacional. Os Transformers possuem uma estrutura principal majoritariamente uniforme, onde as camadas alternam entre feed-forward e self-attention para construir uma rede profunda. Aqui, investigamos essa escolha de design e descobrimos que blocos mais complexos, com diferentes permutações de primitivas de camadas, podem ser mais eficientes. Com base nessa percepção, desenvolvemos um bloco complexo, chamado Brainformer, que consiste em um conjunto diversificado de camadas, como camadas feed-forward com portas esparsas, camadas feed-forward densas, camadas de atenção e várias formas de normalização de camadas e funções de ativação. O Brainformer supera consistentemente os Transformers densos e esparsos de última geração, tanto em qualidade quanto em eficiência. Um modelo Brainformer com 8 bilhões de parâmetros ativados por token demonstra convergência de treinamento 2x mais rápida e tempo de passo 5x mais rápido em comparação com sua contraparte GLaM. Na avaliação de tarefas subsequentes, o Brainformer também demonstra um aumento de 3% na pontuação SuperGLUE com ajuste fino em comparação com o GLaM, com um número semelhante de parâmetros ativados. Por fim, o Brainformer supera amplamente um modelo denso Primer derivado com NAS (Neural Architecture Search) com computação semelhante por token em avaliações 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.