PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Attention with Performers

Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, Adrian Weller

arXiv:2009.14794v4cs.LGcs.CLstat.ML

TL;DR

Performers address the quadratic cost and limited guarantees of efficient attention by approximating full-rank softmax attention with linear complexity and no sparsity or low-rankness priors. They use FAVOR+, combining positive and orthogonal random features to obtain efficient, theoretically controlled kernel approximations. Across pixel, text, and protein tasks, Performers achieve competitive results with examined sparse and dense efficient-attention methods.

  • Problem

    Regular attention is costly for long sequences, while many efficient alternatives do not approximate it and impose structural constraints with limited representation guarantees.

  • Method

    Performers use FAVOR+, positive orthogonal random features that approximate softmax and other kernelizable attention mechanisms while enabling linear-space and linear-time computation.

  • Results

    Performers achieve competitive results with examined efficient sparse and dense attention methods across pixel prediction, text modeling, and protein sequence modeling.

  • Takeaways & Limitations

    Performers provide a linear architecture compatible with regular Transformers and theoretical guarantees including unbiased or nearly unbiased attention estimation, uniform convergence, and lower variance.

  • Takeaways & Limitations

    Small attention-approximation errors can propagate across Transformer layers, so very tight approximations may sometimes be required.

Abstract

from arXiv · show

We introduce Performers, Transformer architectures which can estimate regular (softmax) full-rank-attention Transformers with provable accuracy, but using only linear (as opposed to quadratic) space and time complexity, without relying on any priors such as sparsity or low-rankness. To approximate softmax attention-kernels, Performers use a novel Fast Attention Via positive Orthogonal Random features approach (FAVOR+), which may be of independent interest for scalable kernel methods. FAVOR+ can be also used to efficiently model kernelizable attention mechanisms beyond softmax. This representational power is crucial to accurately compare softmax with other kernels for the first time on large-scale tasks, beyond the reach of regular Transformers, and investigate optimal attention-kernels. Performers are linear architectures fully compatible with regular Transformers and with strong theoretical guarantees: unbiased or nearly-unbiased estimation of the attention matrix, uniform convergence and low estimation variance. We tested Performers on a rich set of tasks stretching from pixel-prediction through text models to protein sequence modeling. We demonstrate competitive results with other examined efficient sparse and dense attention methods, showcasing effectiveness of the novel attention-learning paradigm leveraged by Performers.

1 INTRODUCTION AND RELATED WORK

Performers address the quadratic cost and limited guarantees of efficient Transformer alternatives by approximating regular full-rank softmax attention with linear complexity and no sparsity or low-rankness priors. The architecture extends this approach beyond softmax kernels and performs competitively across diverse tasks.

  • Motivation: Regular Transformers model complex sequence dependencies but require quadratic time and space in sequence length, limiting long-sequence use.This cost can make regular attention prohibitively expensive even for moderate sequence lengths.
  • Motivation: Many efficient alternatives simplify attention or impose structural constraints rather than approximating regular attention, with limited rigorous guarantees about representation power.Examples include local attention, sparsity, pooling, clustering, reversible layers, and shared attention weights.
  • Performers: Performers estimate regular full-rank softmax attention with provable accuracy in linear space and time without assuming sparsity or low-rankness.The approach is designed to remain practical while preserving the representational target of regular attention.
  • Guarantees: Performers provide unbiased or nearly unbiased attention estimates, uniform convergence, lower variance, and compatibility with regular Transformers through small amounts of fine-tuning.These properties supply theoretical guarantees while retaining practical interoperability.
  • FAVOR+: FAVOR+ supports efficient modeling of kernelizable attention mechanisms beyond softmax, enabling comparisons among attention kernels on large-scale tasks.The same mechanism is also presented as a scalable replacement for regular attention beyond Transformers.
  • Evaluation: Across pixel prediction, text modeling, and protein sequence modeling, Performers achieve competitive results against examined efficient sparse and dense attention methods.The evaluation spans tasks with substantially different input modalities and modeling demands.

2 FAVOR+ MECHANISM & POSITIVE ORTHOGONAL RANDOM FEATURES

FAVOR+ linearizes kernelizable attention by factoring random-feature representations, then improves softmax approximation with positive and orthogonal features. The resulting mechanism reduces complexity while addressing instability and variance problems in low-kernel-value regions.

  • Generalized kernelizable attention: FAVOR+ approximates attention kernels through random feature maps, allowing the attention computation to be reordered into products of feature matrices.The mechanism is described for bidirectional and, via prefix sums, unidirectional attention.
  • Attention preliminaries: The attention formulation covers bidirectional dot-product attention and extends analogously to unidirectional attention used in autoregressive generation.Bidirectional attention explicitly normalizes the kernel-score matrix, while the unidirectional form uses its lower-triangular part.
  • Generalized kernelizable attention: The reordered approximation uses space O(Lr + Ld + rd) and time O(Lrd), replacing regular attention’s O(L2 + Ld) space and O(L2d) time.Here L is sequence length, d the hidden dimension, and r the number of random features.
  • Generalized kernelizable attention: FAVOR+ can model other practical kernels, including Gaussian and other kernelizable attention mechanisms, by varying the feature-map functions and sampling distribution.This makes the framework applicable beyond softmax attention.
  • Positive random features: Positive random features avoid the instability of trigonometric features, whose negative values and high variance near small kernel scores can produce abnormal renormalizers and prevent or degrade training.The issue is especially relevant for low-relevance tokens whose softmax-kernel values are close to zero.
  • Positive random features: In low softmax-kernel-value regions, positive random features can be arbitrarily more accurate than trigonometric random features according to the MSE-ratio utility function.Figure 2 compares the estimators as a function of feature-vector angle and length, including a fixed-length slice and a low-kernel-value example.
  • Orthogonal random features: Orthogonalizing random samples preserves unbiasedness for isotropic distributions while reducing estimator variance and sharpening large-deviation bounds.The orthogonal random-feature mechanism requires m ≤ d and enables accurate approximation with fewer random features.
  • FAVOR+: Combining positive and orthogonal random features improves attention-matrix approximation and yields an accurate mechanism with reduced feature count and linear space-time efficiency.The paper reports that theoretical results align with experiments for the resulting FAVOR+ mechanism.

3 THEORETICAL RESULTS

Theoretical and empirical results support FAVOR+ as an accurate, low-variance approximation to softmax attention, with orthogonal and positive features improving estimation. The resulting Performer retains competitive modeling behavior while reducing computational costs for long sequences.

  • Kernel approximation: Positive random features can approximate the softmax-kernel, while the regularized softmax-kernel remains a usable proxy under stated conditions.The regularized kernel is also described as a universal lower bound for the softmax-kernel.
  • Kernel approximation: Orthogonal features provably reduce positive-feature estimation mean squared error and provide sharper concentration bounds than non-orthogonal features.For the regularized softmax-kernel, the bounds are exponentially small and strictly better for every d > 0.
  • Approximation guarantees: Uniform convergence bounds show that the approximated attention matrix can approach the exact matrix with projection count independent of sequence length L.The stated approximation guarantee is ∥bA − A∥∞ ≤ ϵ with constant probability.
  • Computational costs: The Performer reaches nearly linear time and sub-quadratic memory consumption because it avoids explicitly storing the O(L2) attention matrix.This combination supports larger-batch training and lower wall-clock time per gradient step at large L.
  • Empirical approximation: Experiments validate the PORF design: orthogonal features outperform IID features, and positive features outperform trigonometric sin/cos features in approximation error.The comparison varies the number of random samples m using L = 4096 and d = 16.
  • Applications: On protein sequences, the Performer’s softmax approximation matches exact-softmax Transformer accuracy, while Performer-RELU achieves the highest accuracy in both evaluated cases.The protein benchmark used a 36-layer model trained on TrEMBL sequences.

6 BROADER IMPACT

Performers are presented as potentially useful beyond standard Transformer settings, particularly for longer biological sequences, lower compute use, and broader attention applications.

  • Biology and Medicine: Performers may enable Transformer-based biological sequence analysis on much longer sequences without constraining attention structure.The paper identifies proteome-scale protein-interaction prediction as an initial application.
  • Environment: Performers with FAVOR+ are described as having lower compute costs and substantially lower space complexity than regular Transformers.The paper connects these reductions to lower energy consumption and CO2 emissions.
  • Research on Transformers: The authors position Performers as a step toward efficient Transformer research with strong mathematical foundations and scalable handling of sequences of thousands of tokens or more.They connect this direction to potential advances in protein language modeling and bioinformatics.
  • Backward Compatibility: FAVOR+ is presented as compatible with regular pre-trained Transformers and usable for fast inference without loss of accuracy.The paper describes this backward compatibility as an additional practitioner-oriented feature.
  • Attention Beyond Transformers: FAVOR+ can approximate exact attention outside Transformers, including in hierarchical, graph, image-processing, and reinforcement-learning applications.These applications are presented as potential extensions of the method.

A.2 TRAINING HYPERPARAMETERS

The appendix specifies shared training settings and outlines FAVOR+ as a positive-random-feature approach that supports bidirectional and unidirectional attention without explicitly constructing the full attention matrix.

  • A.2 TRAINING HYPERPARAMETERS: Performer and Transformer runs generally share gradient clipping, weight decay, dropout, learning rate, Adam settings, and memory-maximized batch sizes.The defaults are 0.5 gradient clipping, 0.1 weight decay, 0.1 dropout, learning rate 10^-3, β1 = 0.9, β2 = 0.98, and ϵ = 10^-9.
  • A.2 TRAINING HYPERPARAMETERS: Protein experiments use 16x16 TPU-v2 hardware with 8GB per chip, while concatenated Transformer comparisons use 8x8 TPU-v2 hardware.Using 16x16 instead of 8x8 did not change accuracy for the smaller Transformer models.
  • Unidirectional Attention: The unidirectional algorithm computes masked attention through prefix sums rather than explicitly storing the lower-triangular L × L attention matrix.Prefix sums can be computed in O(L) total steps and O(log L) parallel time.
  • FAVOR+: Orthogonal random features preserve sample marginals while enforcing orthogonality, reducing estimator variance and improving approximation and downstream results.Regular ORFs provide unbiased estimation, whereas Hadamard/Givens ORFs provide small bias tending to zero as d increases.

B.3 TIME AND SPACE COMPLEXITY - DETAILED ANALYSIS

FAVOR+ reduces attention resource requirements by avoiding explicit storage of the L × L attention matrix, yielding linear-time computation and lower space complexity for long sequences.

  • Space Complexity: O(md + Ld + mL) space is achievable for bidirectional and aggregated unidirectional FAVOR+, instead of Θ(L^2 + Ld) baseline space.The unidirectional prefix-sum tensor can otherwise require O(mLd) space before aggregation.
  • Space Complexity: O(m + Ld + mL) = O(Ld + mL) space is achievable with Hadamard ORFs, while Givens ORFs require O(m log(d) + Ld + mL).These variants avoid explicitly storing the attention matrix when m and d are much smaller than L.
  • Time Complexity: O(Lmd) time replaces the baseline O(L^2d) time when sequence length L is much larger than the number of random features m.The algorithm constructs Q′ and K′ within the same O(Lmd) time bound.
  • Implementation Trade-offs: H/G-ORFs reduce feature-construction time to O(L log(m)d), improving the leading constant without changing asymptotic time complexity.R-ORFs instead incur a one-time O(md^2) preprocessing cost, negligible relative to O(Lmd) when L ≫ d.
  • Approximation Trade-offs: Increasing the number of random features m raises computation costs while lowering the variance of the attention estimate.The theoretical analysis indicates that M = Θ(d log(d)) can be used in practice.

C.1 TREMBL DATASET

The TrEMBL evaluation uses large protein-sequence data with separate single-sequence and long-sequence tasks, including held-out protein families and sequence-length processing choices.

  • Dataset: The TrEMBL dataset contains 139,394,261 sequences, including 106,030,080 unique sequences.The dataset consists of computationally analyzed but not manually curated UniProtKB sequences.
  • Dataset Splits: The evaluation uses both an OOD-Test split with held-out Pfam families and an IID split with random train, validation, and test partitions.The OOD-Test split is designed around families excluded from evaluation training data.
  • Sequence Tasks: Standard sequence modeling clips individual sequences to maximum length L = 1024, leaving some sequences significantly truncated.The clipping choice follows the reported sequence-length statistics.
  • Sequence Tasks: The long-sequence task concatenates sequences with end-of-sequence tokens and groups them into non-overlapping sequences of length L = 8192.This creates the long-context training and validation examples.
  • Baselines: Random amino-acid baselines achieve 5% accuracy for the 20 standard amino acids and 4% when five anomalous amino acids are included.An empirical-frequency baseline is also considered because amino-acid frequencies are not uniform.

C.3 TABULAR RESULTS

Table 2 reports accuracy and perplexity for single protein sequence modeling at sequence length L = 1024.

  • Table 2 evaluates single protein sequence modeling with sequences of length L = 1024.

C.4 ATTENTION MATRIX ILLUSTRATION

The Performer’s attention matrices exhibit both local and global structure across protein sequences, including diagonal, block, and vertical patterns. Aggregated attention also recovers recognizable amino-acid similarities.

  • The analysis extracts Performer attention matrices by applying the attention mechanism to original and modified Q, K, V inputs.
  • Many heads show diagonal attention to neighboring tokens, while others attend vertically to fixed positions.
  • Aggregating attention across 800 sequences produces an amino-acid similarity matrix that recognizes pairs such as (D, E) and (F, Y).
  • Detailed visualizations identify block attention in Head 1-2 and vertical attention to tokens ‘M’ and ‘C’ in Head 4-1.
  • Attention patterns include both local and global structure when visualized over the first 25 tokens without renormalization.

D.1 BACKWARDS COMPATIBILITY - ERROR PROPAGATION

The appendix examines error propagation, feature redrawing, attention kernels, and comparisons with Linear Transformer and other scalable methods. It reports both approximation and training-stability considerations across these experiments.

  • D.1 BACKWARDS COMPATIBILITY - ERROR PROPAGATION: Approximation errors can propagate through multiple Transformer layers, so tight attention approximations may be necessary.The stated sources of propagation include MLPs and multiple heads.
  • D.2 APPROXIMATE SOFTMAX - EXTENDED PROPERTIES: Feature redrawing helps correct training degradation caused by an unlucky random-feature set during large-scale multilayer training.
  • D.2 APPROXIMATE SOFTMAX - EXTENDED PROPERTIES: After 100K steps on ImageNet64 (U), Performer-ReLU, Performer-Softmax, and Performer-Softmax (SMREG) achieve 3.67, 3.69, and 3.67 BPD, respectively.
  • D.3 GENERALIZED ATTENTION: Kernel comparisons on TrEMBL vary with renormalization and effective batch size, which slightly changes kernel rankings.
  • D.4 COMPARISON WITH LINEAR TRANSFORMER: The Linear Transformer exhibited exploding gradients and NaN errors in both unidirectional and bidirectional ProGen settings.
  • D.5 COMPARISON WITH OTHER METHODS: Performers obtained the largest Long Range Arena score among the tested scalable Transformer methods.

E COMPUTATION COSTS - EXTENDED RESULTS

Extended results describe Performer scaling across layers and compare its computational costs with standard attention and Transformer models. The theoretical appendix further reports concentration and mean-squared-error advantages for orthogonal estimators.

  • E COMPUTATION COSTS - EXTENDED RESULTS: The Performer setup scales to, but is not necessarily limited to, 20 layers.
  • E COMPUTATION COSTS - EXTENDED RESULTS: FAVOR is compared with standard attention in attention-time complexity, while model-level timing compares Transformer and Performer costs.
  • E COMPUTATION COSTS - EXTENDED RESULTS: The Performer is reported as nearly optimal relative to the maximum possible speedup where attention simply returns the V-vector.
  • E COMPUTATION COSTS - EXTENDED RESULTS: The extended benchmarks include a Small configuration in which attention is the dominant source of computation and memory.The Regular configuration uses (nheads, nlayers, dff, d) = (8, 6, 2048, 512).
  • F.4.1 ORTHOGONALITY UNIVERSALLY IMPROVES CONCENTRATION: The appendix defines broader “beautiful functions” whose concentration results apply to softmax and Gaussian kernels through the proposed estimators.
  • F.4.1 ORTHOGONALITY UNIVERSALLY IMPROVES CONCENTRATION: Orthogonal random features provide exponentially small tail-probability bounds that are strictly better than those from unstructured features for any dimensionality d.
  • F.4.1 ORTHOGONALITY UNIVERSALLY IMPROVES CONCENTRATION: Orthogonal estimators also improve mean-squared-error concentration for every d > 0, rather than only asymptotically.

F.5 PROOF OF THEOREM 4

The theorem establishes approximation guarantees for attention matrices using orthogonal random features under bounded-query and key assumptions. The required projection count depends on approximation and dimensional parameters rather than sequence length.

  • Approximation guarantee: The optimal projection count for the trigonometric mechanism is independent of sequence length and scales as Θ(d log(d)).With this choice, the attention matrix can be approximated to arbitrary precision in O(Ld^2 log(d)) time.
  • Theorem assumptions: Under bounded query and key norms, the theorem considers RBF kernels with spectral distributions such as the Gaussian kernel.Queries and keys are assumed to lie within a radius-R ball, while the attention matrix is formed from functions g and h applied to query-key products.
  • Approximation guarantee: The uniform-convergence result makes the required number of projections a function of dimensionality d, error tolerance ϵ, and radius R.The theorem expresses this dependence as m = Ψ(ϵ, d, R).
  • Approximation mechanism: The approximation uses trigonometric orthogonal random features to estimate a generalized attention matrix with concentration guarantees.For RBF kernels, the relevant bounded functions are sine and cosine features, enabling the cited orthogonal-feature result.

F.6 DISCUSSION OF THEOREM 4

The projection budget is controlled by dimensionality, target error, and query-key radius, but fixed projections cannot preserve approximation quality as query and key norms grow indefinitely.

  • Parameter dependence: Larger dimensionality requires more random projections, on the order of d log(d), to achieve a given approximation error.The required count also depends on the target error ϵ.
  • Practical limitation: A fixed projection count cannot approximate hard attention on sequences of unlimited length because query and key lengths must grow for concentrated softmax selection.The authors report that this limitation does not appear at the sequence lengths tested experimentally.
Loading 2009.14794v4…