DARTFormer: Finding The Best Type Of Attention

Jason Ross Brown, Yiren Zhao, Ilia Shumailov, Robert D. Mullins

Keywords: transformer, attention, nas, nueral, architecture, search, homogeneous, heterogeneous, accuracy, xformer

TL;DR: NAS can be used to identify the best efficient Transformer attention mechanism for a given task. This leads to the idea of using mixed attention in a Transformer, however, this does not perform as well as one might expect.

Abstract: Given the wide and ever growing range of different efficient Transformer attention mechanisms, it is important to identify which attention is most effective when given a task. In this work, we are also interested in combining different attention types to build heterogeneous Transformers. Focussing on NLP-based tasks, we first propose a DARTS-like Neural Architecture Search (NAS) method to find the best attention for a given task. In this setup, all heads use the same attention (homogeneous models). Our results suggest that NAS is highly effective at this, and it identifies the best attention mechanisms for IMDb byte level text classification and Listops. We then extend our framework to search for and build Transformers with multiple different attention types, and call them heterogeneous Transformers. We show that whilst these heterogeneous Transformers are better than the average homogeneous models, they cannot outperform the best. We explore the reasons why heterogeneous attention makes sense, and why it ultimately fails.