An Attention Free Transformer

Shuangfei Zhai; Walter Talbott; Nitish Srivastava; Chen Huang; Hanlin Goh; Ruixiang Zhang; Josh Susskind

Source-linked AI summary

An Attention Free Transformer

Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, Josh Susskind

arXiv:2105.14103v2May 28, 2021cs.LGcs.CLcs.CV

TL;DR

Transformers face quadratic attention costs as context size grows. The paper introduces AFT, which combines keys and values using learned position biases before element-wise query interaction, and reports competitive performance across benchmarks with excellent efficiency.

Problem
Quadratic time and space complexity in context size makes Transformers difficult to scale to large contexts.
Method
AFT replaces standard dot-product attention by combining keys and values with learned position biases and then multiplying the reduced context element-wise with the query.
Results
AFT provides competitive performance across image autoregressive modeling, character-level language modeling, and image classification while providing excellent efficiency.
Takeaways & Limitations
AFT maintains global connectivity while offering memory complexity linear in both input and model sizes, and its locality-based variants improve efficiency and performance.

Abstract

from arXiv · show

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.

1 Introduction

Transformers capture long-term dependencies through direct pairwise interactions, but dot-product attention scales quadratically with context size. AFT avoids standard dot-product attention while preserving global interaction, and its local variants improve efficiency and reported performance.

Transformers enable direct interaction between every pair of sequence elements, supporting the capture of long-term dependencies.
Quadratic time and space complexity in context size makes standard Transformers difficult to scale to large contexts.
AFT combines keys and values with learned position biases, then combines the reduced context with queries through element-wise multiplication.
AFT preserves direct interaction between context points while achieving memory complexity linear in both input and model sizes.
AFT-local constrains position biases locally and AFT-conv adds spatial weight sharing while maintaining global connectivity.
Across image autoregressive modeling, character-level language modeling, and image classification, AFT reports competitive performance with excellent efficiency.

2 Multi-Head Attention

Multi-Head Attention applies scaled dot-product attention across multiple heads and concatenates their outputs along the channel dimension. The formulation uses learned linear transformations for queries, keys, and values, with softmax as the default nonlinearity.

Multi-Head Attention performs scaled dot-product attention separately for each head and concatenates the resulting outputs.
Each head uses linear transformations for query, key, and value representations, with softmax as the default row-wise nonlinearity.
Under the stated default, query and key dimensions match within each head, and the concatenated output has feature dimension hd_v.

3 Methodology

AFT replaces dot-product self-attention with learned position-biased key–value aggregation followed by element-wise query gating. Its variants impose locality or spatial weight sharing while preserving global connectivity, and factorized position biases reduce complexity and parameters.

Attention Free Transformer: AFT linearly transforms input X into query, key, and value representations before applying its attention-free operation.It is designed as a plugin replacement for multi-head attention without changing other Transformer components.
Attention Free Transformer: For each target position, AFT weights values using keys and learned pair-wise position biases, then combines the reduced context with the query element-wise.The query nonlinearity is defaulted to sigmoid, and ⊙ denotes element-wise multiplication.
Attention Free Transformer: AFT implicitly performs attention with one attention vector per feature dimension, while avoiding explicit attention-map computation and retaining global query–value interactions.The factorized attention matrices are interpreted as having as many heads as feature dimensions.
AFT-local: AFT-local restricts learned relative position biases to a window of size s while maintaining global connectivity and reducing parameter, time, and space costs.Unlike local Transformers, its connectivity remains global regardless of s.
AFT-conv: AFT-conv adds spatial weight sharing, yielding a CNN-like variant with a global receptive field and specialized convolutional properties.The design combines global connectivity, non-negative convolutional weights, and multiplicative gating.
Parameterization: Factorized position biases reduce parameter counts from T^2 to 2Td′ and empirically improve training and testing performance.The factorization uses a small embedding dimension d′, such as 128.

4 Related Work

Efficient-Transformer research addresses the quadratic cost of standard attention through approximation, sparsity, locality, context compression, or alternative attention operations. AFT differs by eliminating dot-product attention while combining global connectivity, gating, and CNN-like efficiency properties.

Motivation: Efficient Transformer methods target the quadratic time and space cost of attention to support larger contexts and more efficient implementations.The surveyed approaches include approximation, sparsity, locality, hashing, low-rank methods, and kernel approximation.
Approximating the dot product: Linearized attention approximates the exponential kernel but has complexity O(Td^2), making scaling with model dimension difficult compared with AFT.Reformers instead approximate dot-product attention using locality-sensitive hashing, whereas AFT removes the dot product.
Sparse, local attention: Sparse and local attention impose fixed context patterns, while AFT-local uses locality as a bias and retains access to the full context.AFT-local and AFT-conv therefore differ from methods relying only on a selected subset of elements.
Context compression and alternative operations: Other approaches compress context or alter attention computation using adaptive spans, routing, key–value compression, predicted weights, or dynamic convolutions.These methods reduce comparisons, shorten context representations, or replace dot-product interactions with different operations.
MLPs for vision: Vision MLP approaches replace attention, whereas AFT adds key- and position-bias weighting with normalized non-negative values and supports plug-in Transformer replacement.AFT-conv additionally inherits CNN properties including parameter efficiency, performance, and variable-sized input handling.

5 Experiments

The experiments evaluate AFT across image autoregressive modeling, character-level language modeling, and ImageNet-1K classification, emphasizing competitive performance alongside efficiency. Results support locality, factorization, global connectivity, compatibility with pretrained Transformers, and variable-size inputs as useful design properties.

Experimental setup: AFT experiments replace baseline Transformer attention modules across image modeling, language modeling, and image classification tasks.The first two tasks use causal AFT models, while image classification uses the encoding model.
Image autoregressive modeling: AFT-local outperforms all Transformer baselines on CIFAR10 while running faster than standard and Image Transformers and using half the memory.AFT-simple also remains competitive, outperforming Image Transformer with strong speed and memory efficiency.
Ablations and properties: Factorized position biases reduce parameters and improve training and testing performance, while learned position biases can produce sparse local patterns without removing global connectivity.AFT-conv’s locality-based variants preserve global interaction regardless of local kernel size.
Character-level language modeling: AFT achieves the lowest training bpc on Enwik8, slightly trails the basic Transformer in testing, and outperforms the other Transformer variants.The deeper, narrower AFT architecture provides the best balance across parameter count, speed, memory, and performance.
Character-level language modeling: AFT-local reaches its best Enwik8 performance at local window size 32, while larger sequence sizes consistently reduce training and testing loss.The window-size results form a U-shape, and increasing T to 2048 and 4096 improves both losses.
Image classification: AFT-conv improves ImageNet-1K top-1 accuracy over AFT-full and maintains global connectivity, with the default 80.8 versus 79.9 for a local-only baseline at kernel size 7.AFT-conv also supports pretrained DeiT initialization and variable-size inputs, reaching 81.6 accuracy on 384 crops versus 81.0 on 224 crops.

6 Conclusions

The paper concludes that AFT replaces dot product attention with an efficient operation and achieves strong benchmark results. It presents AFT as a basis for further Transformer-like model designs.

AFT replaces dot product attention with a new efficient operation and achieves strong results on standard benchmarks.
The authors position AFT as opening a new design space for Transformer-like models.

7 Additional Ablations

Additional ablations examine parameterization, reparameterization, kernel size, query contributions, key behavior, and learned position-bias patterns. These studies connect architectural choices to accuracy, sparsity, and convolutional behavior.

Factorization: Factorized position-bias parameterization improves AFT-full training and test performance while reducing parameter counts.
Reparameterization: AFT-conv reparameterization improves model performance for a kernel size of 7 × 7.
Kernel size: AFT-conv remains comparable to the DeiT reference even with a 3 × 3 kernel.
Query contribution: Removing the query term from AFT-conv causes significant performance drops.
Key visualization: AFT-conv keys gradually evolve into object detectors as network depth increases.
Position-bias variants: The exponentiated position-bias visualizations compare standard, sparsity-regularized, and Gumbel-softmax AFT-conv variants, whose reported top-1 accuracies are 80.8%, 80.9%, and 79.9%.

8 Sparsity

AFT-conv position biases exhibit sparsity, motivating entropy regularization and an extreme one-position-per-head variant. The latter preserves strong ImageNet performance while enabling simpler context reduction.

Sparsity: AFT-conv position biases show sparsity patterns that motivate quantization and pruning.The patterns are observed with an 11 × 11 kernel.
Sparsity: 80.9 vs 80.8 top 1 accuracy follows entropy regularization of AFT-conv with 384 heads and an 11 × 11 kernel.The regularizer minimizes per-head entropy and produces visibly sparser position biases.
Sparsity: The extreme variant assigns each head a learned relative position bias for a single position using Gumbel softmax during training.During inference, Gumbel softmax is replaced with hard max, returning a one-hot vector.
Sparsity: 79.9 top 1 accuracy results from the extreme variant, with less than 1 point drop versus the unregularized model.Its K, V context reduction can use global average pooling and indexing, matching AFT-simple complexity while retaining strong performance comparable to the standard Transformer.

11 Additional Ablations

Additional ablations examine parameterization, reparameterization, kernel size, query contribution, and learned key behavior in AFT variants. They report benefits from factorization and reparameterization, competitive small-kernel performance, and evolving key visualizations.

Factorization of w: The non-factorized AFT-full parameterization performs worse on training and test performance than the factorized version.The ablation is reported in Table 1.
Reparameterization of w: Reparameterization effectively improves AFT-conv performance.The comparison uses a 7 × 7 kernel.
Kernel size: AFT-conv achieves comparable performance to the Deit reference with a 3 × 3 kernel.This result is reported for the 384-head AFT-conv small setting.
Contribution of the query: Removing the query term from AFT-conv produces significant performance drops.The ablation is reported for models with 384 heads and 11 × 11 or 15 × 15 kernels.
Visualizing the key: AFT-conv keys gradually evolve into “object detectors” as the layer level increases.The keys are visualized on randomly sampled ImageNet-1K validation images across layers and heads.

212 Sparsity

The sparsity experiments progressively constrain AFT-conv position biases, from entropy regularization to one-hot selection. These constraints yield sparser patterns and retain strong ImageNet accuracy while simplifying context reduction.

Sparsity: Entropy regularization minimizes each head’s softmax entropy over position-bias logits.The regularizer is combined with cross-entropy loss using weighting 0.001.
Extreme sparsity: The extreme variant assigns one learned relative position bias per head by multiplying biases with Gumbel softmax samples during training.The Gumbel softmax temperature is set to 0.5, and hard max is used during inference.
Extreme sparsity: 79.9 top 1 accuracy is achieved by the extreme variant, with less than 1 point drop compared with the unregularized model.Its position biases are visualized in Figure 8.
Extreme sparsity: Global average pooling and indexing implement K, V context reduction for the extreme variant with the same complexity as AFT-simple.The variant maintains strong performance comparable to the standard Transformer.