Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers

Krzysztof Choromanski; Valerii Likhosherstov; David Dohan; Xingyou Song; Andreea Gane; Tamas Sarlos; Peter Hawkins; Jared Davis; David Belanger; Lucy Colwell; Adrian Weller

Source-linked AI summary

Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers

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

arXiv:2006.03555v3Jun 5, 2020cs.LGcs.CLstat.ML

TL;DR

Long-sequence applications such as biological sequence analysis are constrained by the quadratic cost of standard attention and by alternatives that impose structural assumptions. The paper introduces Performer with FAVOR, which approximates attention with linear scaling while providing theoretical guarantees and broad attention compatibility. The authors report effectiveness on protein sequence modeling and describe compatibility with regular Transformers after fine-tuning.

Problem
Standard Transformer attention scales quadratically with sequence length, while structural approximations may be unsuitable for long biological sequences and long-range dependencies.
Method
Performer uses Fast Attention Via Orthogonal Random features (FAVOR) to approximate regular and generalized kernel-based attention without sparsity-pattern priors.
Results
Performer provides linear scaling in sequence length, unbiased attention estimation, uniform convergence, and empirical effectiveness on protein sequence modeling.
Takeaways & Limitations
FAVOR offers a scalable attention replacement for long-sequence modeling while retaining compatibility with regular Transformer architectures and extending to other attention settings.
Takeaways & Limitations
Approximation error can propagate through Transformer components, so pretrained regular-Transformer weights cannot be transferred directly without fine-tuning.

Abstract

from arXiv · show

Transformer models have achieved state-of-the-art results across a diverse range of domains. However, concern over the cost of training the attention mechanism to learn complex dependencies between distant inputs continues to grow. In response, solutions that exploit the structure and sparsity of the learned attention matrix have blossomed. However, real-world applications that involve long sequences, such as biological sequence analysis, may fall short of meeting these assumptions, precluding exploration of these models. To address this challenge, we present a new Transformer architecture, Performer, based on Fast Attention Via Orthogonal Random features (FAVOR). Our mechanism scales linearly rather than quadratically in the number of tokens in the sequence, is characterized by sub-quadratic space complexity and does not incorporate any sparsity pattern priors. Furthermore, it provides strong theoretical guarantees: unbiased estimation of the attention matrix and uniform convergence. It is also backwards-compatible with pre-trained regular Transformers. We demonstrate its effectiveness on the challenging task of protein sequence modeling and provide detailed theoretical analysis.

1 Introduction and related work

Transformers model complex dependencies but standard attention becomes prohibitively expensive for long sequences, while many alternatives impose structural assumptions that may limit long-range modeling. Performer addresses this with FAVOR, offering scalable attention, theoretical guarantees, and compatibility with regular Transformer components.

Standard Transformer attention scales quadratically with sequence length, making long biological sequences difficult to process efficiently.Approximations based on localized gradients can also fail to capture long-distance correlations.
Protein sequence models show promise for predicting structure and function, but their applicability beyond single-protein design remains limited.
FAVOR replaces quadratic attention with linear-scaling computation, sub-quadratic space complexity, and no sparsity-pattern prior.It estimates regular attention directly rather than simplifying it through structural priors.
FAVOR provides unbiased attention-matrix estimation and uniform-convergence guarantees, while supporting a general kernel-based attention class that includes regular attention.
The mechanism is empirically evaluated on protein sequence modeling and is API-compatible with regular Transformers, preserving other architectural components.The supplied contribution summary also states that pretrained Transformer components can be retained with FAVOR replacing standard attention.

2 Generalized Attention via FAVOR mechanism

FAVOR replaces explicit quadratic attention with unbiased random-feature approximations that support Generalized Attention while reducing long-sequence computation and memory requirements.

2.1 Preliminaries - standard attention mechanism: Regular dot-product attention forms an L×L matrix and therefore requires O(L2d) time, making end-to-end processing of long sequences impractical.The matrix is computed and stored explicitly in the standard formulation.
2.2 Generalized Attention (GA): Generalized Attention extends regular attention to kernel-based mechanisms parameterized by K, g, and h, with regular attention as a special case.FAVOR applies when the corresponding kernel can be effectively estimated by random features.
2.3 Towards FAVOR: approximating attention with random features (RFs): FAVOR approximates the attention matrix with an unbiased low-rank decomposition, avoiding explicit construction and storage of the full matrix.The approximation uses random feature maps and is leveraged without explicitly computing the approximate matrix.
2.4 Towards FAVOR: refinements via orthogonal random features: Orthogonal random features reduce Monte Carlo estimator variance and yield more accurate approximations with substantially better downstream results.The paper uses orthogonal features while preserving their marginal distributions, including blockwise use when M > d.
2.6 Time and space complexity analysis: FAVOR reduces bidirectional space complexity to O(Md + Ld + ML), or O(Ld + ML) with H-ORFs, instead of Θ(L2 + Ld).These variants avoid storing the attention matrix explicitly when M, d ≪ L.
2.6 Time and space complexity analysis: The FAVOR algorithm runs in O(LMd) time, which is much lower than the O(L2d) baseline when L ≫ M.H-ORFs and G-ORFs can improve the leading constant through O(L log(M)d) computation for projected representations.

3 Theoretical convergence analysis

FAVOR provides uniform convergence guarantees for approximating generalized attention, including regular attention, with a feature count independent of sequence length.

With Mopt = Θ(d log(d)), FAVOR approximates the attention matrix to any precision in O(Ld^2 log(d)) time regardless of sequence length L.The optimal random-feature count depends on d, not L.
The convergence theorem bounds the approximation error for generalized attention on compact query-key domains with high probability.The result applies to radial basis function kernels and includes regular attention with Gaussian kernels.
FAVOR’s uniform-convergence guarantee applies to generalized attention, whose regular-attention instance is covered by the same analysis.

4 Experiments

Experiments show that FAVOR reduces computation and memory costs, maintains controllable approximation error, supports pretrained-model adaptation, and enables longer protein-sequence training than the tested Transformer baselines.

4.1 Computation costs: The Performer approaches linear time in sequence length, uses sub-quadratic memory, and achieves nearly optimal speedup for large L in backward-pass comparisons.The comparison used the default architecture (8, 6, 2048, 512), with Transformer memory dominated by the explicit O(L^2) attention matrix.
4.2 Approximation error and compatibility with regular Transformer: Orthogonal random features generally yield lower attention-matrix and output approximation error than unstructured features as M varies.Figure 2 evaluates L = 4096 and d = 16 across 10 samples.
4.2 Approximation error and compatibility with regular Transformer: Transferred pretrained Transformer weights initially retain 0.07 accuracy, and fine-tuning the Performer quickly recovers accuracy in a small fraction of the original gradient steps.Approximation error can propagate through MLPs, heads, and layers, so direct weight transfer is not immediately sufficient.
4.3 Multiple layer training: On TrEMBL, the Reformer significantly drops in accuracy, suggesting sparse attention may be insufficient for protein tasks requiring global interactions.The experiments used a 36-layer model in both unidirectional and bidirectional settings.
4.4 Large length training: For concatenated TrEMBL sequences at L = 8192, the standard Performer trained at batch size 8 per chip and reached ≈24%, while the smaller Transformer plateaued at ≈19%.The baseline Transformer overloaded memory at batch size 1 per chip, requiring a smaller architecture.
4.4 Large length training: On ImageNet64, the Performer matches the Reformer under the reported TPU setting and can be 2x faster in the unidirectional case depending on hardware.

5 Conclusion

The paper concludes that FAVOR makes attention scalable while retaining unbiased estimation, enabling Performer to improve Transformer space and time complexity with respect to sequence length.

FAVOR is presented as an unbiased estimator of the original attention algorithm with linear space and time complexity in sequence length L.The conclusion identifies potential applications beyond Transformers, including hierarchical, graph, image, and reinforcement-learning attention.

6 Broader impact

The authors describe broader implications for long biological sequences, lower-compute Transformer use, mathematically grounded architecture research, pretrained-model reuse, and attention beyond Transformers.

Biology and Medicine: For biology and medicine, Performer may enable longer biological-sequence analysis and protein-interaction prediction without constraints on attention structure.The paper connects this to proteome-scale interaction prediction and the bottleneck of large evolutionary alignments.
Environment: Lower compute and space complexity could translate into reduced CO2 emissions and energy consumption relative to regular Transformers.
Research on Transformers: The method may guide efficient Transformer research toward architectures with stronger mathematical foundations and support sequence lengths of thousands or more.The authors frame this as a potential gateway to applications such as protein language modeling.
Backward Compatibility: Performer is backward-compatible with regular pretrained Transformers and can support fast inference without accuracy loss when up-training is unnecessary.
Attention Beyond Transformers: FAVOR can approximate exact attention outside Transformers, including hierarchical attention networks, graph attention networks [47], image processing [23], and reinforcement learning or robotics [44].

A Theoretical results

The proof establishes FAVOR’s approximation guarantee by applying an existing theorem to the default R-ORF setting and then bounding the resulting attention error. The selected Performer configuration uses generalized attention unless explicitly labeled otherwise.

The default FAVOR proof applies Theorem 3 from [32] because the RBF kernel uses the bounded function f = cos and has Lf = 1.
The proof bounds the attention approximation error by propagating the intermediate error through the query and key normalization factors, then chooses δ = ϵ/(g*h*) to complete the guarantee.
Unless explicitly called Performer-SOFTMAX, the paper’s default Performer uses the optimal generalized-attention setting, which also supports approximate softmax comparisons.

B.1 Training Hyperparameters

The experiments use fixed optimization and regularization settings for Performer and Transformer runs, while allocating comparable compute across the main protein experiments. The implementation also evaluates several default FAVOR configurations.

Performer and Transformer runs use gradient clipping of 0.5, weight decay of 0.1, dropout of 0.1, learning rate 10^-3, Adam, and maximized batch size.Adam uses β1 = 0.9, β2 = 0.98, and ϵ = 10^-9.
All 36-layer protein experiments use the same compute budget of 16x16 TPU-v2 devices, while concatenated Performer experiments also use 16x16 TPU-v2 devices.The shorter Transformer models use 8x8 devices without changing accuracy when scaled to 16x16.
The default FAVOR configuration uses attention renormalization, a numerical stabilizer of 10^-6, 256 features, orthogonal features, and zero orthogonal scaling.
The default generalized-attention configuration uses attention renormalization, 256 features, a ReLU kernel, and kernel epsilon 10^-3.

C.1 TrEMBL Dataset

The dataset is built from the large TrEMBL protein corpus using both IID and held-out-family splits, with separate preprocessing for standard and long-sequence modeling. Empirical amino-acid frequencies are also characterized for baseline construction.

TrEMBL contains 139,394,261 sequences, including 106,030,080 unique sequences, and represents approximately 99.5% of UniProtKB sequences.
The evaluation uses an OOD-Test split holding out selected Pfam families and an IID split that randomly divides the remaining sequences into training, validation, and test sets.
Standard sequence modeling clips individual proteins to maximum length L = 1024, while the long-sequence task concatenates proteins into non-overlapping sequences of length L = 8192.
The empirical baseline assigns amino-acid probabilities according to training-set frequencies, whose estimated distribution is visualized for the 20 standard amino acids by class.The visualization is reported as consistent with TrEMBL statistics.

C.3 Tabular Results

The tabular evaluation covers single-protein modeling with accuracy and perplexity metrics, while additional analyses inspect attention structure and amino-acid similarity. The paper also identifies a limitation in directly transferring pretrained Transformer weights to Performer inference.

Tabular Results: Single-protein results at L = 1024 report accuracy and perplexity, with bidirectional models evaluated on positions selected by 15% masking.Table 2 applies its empirical-baseline results to both unidirectional and bidirectional models.
Attention Matrix Illustration: Attention visualizations show both local diagonal patterns and global vertical patterns across protein-sequence positions, including heads attending to fixed tokens.The analysis uses the first four layers and all eight heads for the BPT1_BOVIN sequence.
Amino Acid Similarity: Aggregating attention across 800 sequences yields an amino-acid similarity matrix that recognizes highly similar pairs such as (D, E) and (F, Y).
Backwards Compatibility: Pretrained Transformer weights cannot be used immediately for Performer inference because error propagation from nonattention components is a primary source of the mismatch.

D.2 Generalized Attention

The study evaluates generalized attention kernels across TPU configurations, highlighting accuracy, numerical stability, and kernel selection for larger batches.

Experimental setup: For TrEMBL at L = 512, the experiments compare sigmoid, exponential, ReLU, absolute, gelu, cosine, tanh, and identity kernels.They also vary whether attention renormalization is enabled.
Kernel selection: The generalized ReLU kernel is used by default because it was empirically optimal for large-batch runs.The authors note that behavior differed between 2x2 and 4x4 TPU configurations.
Experimental setup: The kernel comparisons use log-scaled axes to emphasize high-accuracy runs and expose NaN-related early stopping.Runs were performed on 2x2 TPU-v2 systems with batch size 128 per device.
Experimental setup: A corresponding generalized-attention setup was also evaluated on 4x4 TPU-v2 systems.

E Extended computation costs

The extended benchmarks examine Performer depth, attention complexity, and full-model cost, showing scalability toward longer sequences and deeper models.

Extended computation costs: The Performer can scale to at least 20 layers, while its attention complexity is compared directly with standard Transformer attention.The extended benchmarks cover both attention-only and full-model timing scenarios.
Extended computation costs: The Performer is benchmarked against an optimal V-vector baseline to assess how closely its speed approaches the maximum possible attention speedup.The “X” baseline returns the V-vector without performing attention computation.
Benchmark caveat: The benchmark notes that explicit tensor construction in Jax inflates reported memory relative to an actual model.This caveat applies to the attention-complexity benchmark rather than model memory itself.
Model configurations: The study also evaluates a Small configuration where attention is the dominant computational and memory component.The Small version uses (nheads, nlayers, dff, d) = (1, 6, 64, 64), compared with the Regular version’s (8, 6, 2048, 512).