Wide Attention Is The Way Forward For Transformers?

TL;DR

This paper demonstrates that wide, shallow Transformer models can outperform deeper ones in NLP tasks, offering benefits in speed, memory, and interpretability, challenging the traditional emphasis on depth.

Contribution

It introduces and systematically evaluates the effectiveness of wide, shallow Transformer architectures as a superior alternative to deep models in NLP.

Findings

Wide models perform 0.3% better than deep models on average across tasks.

Single layer Transformers are 3.1x faster on CPU inference.

Wide models require less memory and are more interpretable.

Abstract

The Transformer is an extremely powerful and prominent deep learning architecture. In this work, we challenge the commonly held belief in deep learning that going deeper is better, and show an alternative design approach that is building wider attention Transformers. We demonstrate that wide single layer Transformer models can compete with or outperform deeper ones in a variety of Natural Language Processing (NLP) tasks when both are trained from scratch. The impact of changing the model aspect ratio on Transformers is then studied systematically. This ratio balances the number of layers and the number of attention heads per layer while keeping the total number of attention heads and all other hyperparameters constant. On average, across 4 NLP tasks and 10 attention types, single layer wide models perform 0.3% better than their deep counterparts. We show an in-depth evaluation and demonstrate how wide models require a far smaller memory footprint and can run faster on commodity hardware, in addition, these wider models are also more interpretable. For example, a single layer Transformer on the IMDb byte level text classification has 3.1x faster inference latency on a CPU than its equally accurate deeper counterpart, and is half the size. We therefore put forward wider and shallower models as a viable and desirable alternative for small models on NLP tasks, and as an important area of research for domains beyond this.