Brainformers: Trading Simplicity for Efficiency
Authors: name: Yanqi Zhou; name: Nan Du; name: Yanping Huang; name: Daiyi Peng; name: Chang Lan; name: Da Huang; name: Siamak Shakeri; name: David So; name: Andrew Dai; name: Yifeng Lu; name: Zhifeng Chen; name: Quoc Le; name: Claire Cui; name: James Laundon; name: Jeff Dean
Published: 2023-05-29 18:42:01

Abstract: 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.