OmniNet: Omnidirectional Representations from Transformers

TL;DR

OmniNet introduces omnidirectional representations by allowing tokens to attend to all tokens across all layers, guided by an efficient meta-learner attention mechanism. It combines this with a standard Transformer and uses IndexSort plus a pooling step to produce final representations, with partitioning to bound compute. Across autoregressive language modeling, machine translation, Long Range Arena, and Vision Transformer tasks, OmniNet achieves state-of-the-art results on LM1B and WMT'14 En-De/En-Fr and improves vision tasks in few-shot and fine-tuning scenarios, while maintaining favorable compute-performance trade-offs. The work demonstrates the value of cross-layer information flow and provides several efficient attention variants (kernel-based, low-rank, block-based) to enable scalable omnidirectional attention.

Abstract

This paper proposes Omnidirectional Representations from Transformers (OmniNet). In OmniNet, instead of maintaining a strictly horizontal receptive field, each token is allowed to attend to all tokens in the entire network. This process can also be interpreted as a form of extreme or intensive attention mechanism that has the receptive field of the entire width and depth of the network. To this end, the omnidirectional attention is learned via a meta-learner, which is essentially another self-attention based model. In order to mitigate the computationally expensive costs of full receptive field attention, we leverage efficient self-attention models such as kernel-based (Choromanski et al.), low-rank attention (Wang et al.) and/or Big Bird (Zaheer et al.) as the meta-learner. Extensive experiments are conducted on autoregressive language modeling (LM1B, C4), Machine Translation, Long Range Arena (LRA), and Image Recognition. The experiments show that OmniNet achieves considerable improvements across these tasks, including achieving state-of-the-art performance on LM1B, WMT'14 En-De/En-Fr, and Long Range Arena. Moreover, using omnidirectional representation in Vision Transformers leads to significant improvements on image recognition tasks on both few-shot learning and fine-tuning setups.

Figures (9)

• Figure 1: Overview of OmniNet. In the diagram, the omnidirectional module, when partition size is $P=L$, combines the information across all positions ($1:N$), across all layers ($1:L-1$), and for each position selects the best of all layers via a pooling operation to generate the final representations.
• Figure 2: Performance of pre-trained OmniNet and equivalent ViT models in few-shot learning setup on downstream tasks, when transferred using only few images (1, 5, 10, and 25) per class.
• Figure 3: Performance of ViT and OmniNet (with different partition sizes) in terms of top-1 accuracy on ImageNet 5-shot linear, versus their computational costs in terms of number of FLOPs.
• Figure 4: Contribution of different layers in Omnidirectional representations for a given set of examples. On top, we plot the omnidirectional attention maps (using OmniNet$_{B/16}$-P$12$ ) of one of the heads, over all layers, when CLS token in the last layer is used as query. On the bottom, we show the contribution of each layer to the pooling operation of the Omnidirectional module.
• Figure 5: Contributions of different layers in omnidirectional representations for Example #1.