[Paper Review] An Attention Free Transformer

Shuangfei Zhai, Walter Talbott+5 more|arXiv (Cornell University)|May 28, 2021

Advanced Neural Network Applications26 references42 citations

TL;DR

Introduces Attention Free Transformer (AFT), which replaces dot-product self-attention with a non-attention, linear-complexity module, plus variants like AFT-local and AFT-conv that preserve global connectivity with locality or convolutional weight sharing.

ABSTRACT

We introduce Attention Free Transformer (AFT), an efficient variant of Transformers that eliminates the need for dot product self attention. In an AFT layer, the key and value are first combined with a set of learned position biases, the result of which is multiplied with the query in an element-wise fashion. This new operation has a memory complexity linear w.r.t. both the context size and the dimension of features, making it compatible to both large input and model sizes. We also introduce AFT-local and AFT-conv, two model variants that take advantage of the idea of locality and spatial weight sharing while maintaining global connectivity. We conduct extensive experiments on two autoregressive modeling tasks (CIFAR10 and Enwik8) as well as an image recognition task (ImageNet-1K classification). We show that AFT demonstrates competitive performance on all the benchmarks, while providing excellent efficiency at the same time.

Motivation & Objective

• Motivate the need for scalable Transformers due to quadratic attention costs.
• Propose a plug-in attention-free module that maintains global interactions without dot-product attention.
• Introduce AFT-local and AFT-conv to exploit locality and spatial weight sharing while preserving global connectivity.
• Demonstrate competitive performance and efficiency on image autoregressive modeling, language modeling, and image classification.

Proposed method

• Replace standard multi-head attention with an interaction among Q, K, V where K and V are combined with learned position biases and the result is element-wise multiplied by the query.
• Formulate AFT as Y_t = sigma_q(Q_t) ⊙ [ sum_{t'} exp(K_{t'}+w_{t,t'}) ⊙ V_{t'} / sum_{t'} exp(K_{t'}+w_{t,t'}) ], with w being learned pairwise position biases.
• Provide a factorized parameterization w_{t,t'} = u_t^T v_{t'} to reduce parameters.
• Introduce AFT-local (local w with s-bounded window) and AFT-conv (spatial weight sharing, akin to CNNs) variants.
• Offer AFT-simple as a highly efficient extreme variant with no position biases (s=0).
• Discuss parameterizations and reparameterizations (e.g., w-factorization, w normalization) to improve training stability and performance.

Experimental results

Research questions

• RQ1 Can an attention-free mechanism match or exceed Transformers on autoregressive image modeling, language modeling, and image classification?
• RQ2 Do locality and weight-sharing variants (AFT-local, AFT-conv) offer practical gains in efficiency and accuracy while preserving global connectivity?
• RQ3 How does factorizing the position-bias parameters impact model performance and parameter count?
• RQ4 What are the comparative complexities and practical runtimes of AFT variants versus standard and other efficient transformers?
• RQ5 Can AFT variants integrate with existing Transformer architectures as plug-ins?

Key findings

• AFT variants achieve competitive performance across CIFAR10, Enwik8, and ImageNet-1K while offering favorable efficiency.
• AFT-local outperforms several Transformer baselines on CIFAR10 image autoregressive modeling in both speed and memory usage.
• Factorized position-bias parameterization dramatically reduces parameters and improves performance compared to non-factorized w.
• AFT-simple delivers strong results with even greater efficiency, closely matching or surpassing some vision Transformer baselines.
• AFT-conv substantially improves ImageNet Top-1 accuracy over baseline configurations, with global connectivity maintained via convolutional-style weight sharing.
• Across tasks, AFT variants demonstrate linear or near-linear time/space behavior relative to sequence length and feature dimension, in contrast to quadratic costs of standard attention.