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

By Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Jared Davis, Tamás Sarlós, David Belanger, Lucy J. Colwell, Adrian Weller. 2020

96 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Materials Chemistry, Molecular Biology, Large Language Models, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Materials Chemistry

In Molecular Biology

In Large Language Models

In Efficiency & Compression

Citations in the corpus

Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth (2021, background): “One such variant relies on low-rank or kernel-based approximations to the full attention matrix (Katharopoulos et al., 2020; Wang et al., 2020; Choromanski et al., 2020), in which case the paths likely converge even faster to rank one due to the imposed low-rankness.”
Deformable DETR: Deformable Transformers for End-to-End Object Detection (2020, background): “Katharopoulos et al. (2020); Choromanski et al. (2020) rewrite the calculation of self-attention through kernelization approximation.”
Long Range Arena: A Benchmark for Efficient Transformers (2020, background): “Aside from the standard vanilla Trans-former (Vaswani et al., 2017) and a simple local attention baseline, we compare Sparse Trans-formers (Child et al., 2019), Longformers (Beltagy et al., 2020), Linformers (Wang et al., 2020), Reformers (Kitaev et al., 2020), Sinkhorn Transformers (Tay et al., 202”
Efficient Transformers: A Survey (2020, methodology): “…et al., 2020) O ( n ( k + m )) FP+M ETC (Ainslie et al., 2020) O ( N 2 g + NN g ) FP+M Synthesizer (Tay et al., 2020a) O ( N 2 ) LR+LP Performer (Choromanski et al., 2020a) O ( N ) KR Funnel Transformer (Dai et al., 2020) O ( N 2 ) FP+DS Linformer (Wang et al., 2020c) O ( N ) LR Linear…”
FNet: Mixing Tokens with Fourier Transforms (2021, methodology): “While Fourier Transforms have previously been used to approximate or speed up computations in Convolutional Neural Networks (El-Bakry and Zhao, 2004; Mathieu et al., 2014; Highlander and Rodriguez, 2015; Pratt et al., 2017; Lin et al., 2018; Chitsaz et al., 2020; Goldberg et al., 2020), Recurrent Ne”
A Survey of Transformers (2021, methodology): “The first version of Performer [18] is inspired from the random Fourier feature map [105] that was originally used to approximate Gaussian kernel.”
MambaOut: Do We Really Need Mamba for Vision?* (2024, unknown): “To address this issue, a variety of token mixers with linear complexity to token length have been introduced [78], such as dynamic convolution [44, 90, 92], Linformer [84], Longformer [7], Big Bird [105], and Per-former [14].”
Self-supervised learning in medicine and healthcare (2022, methodology): “Techniques for natural language processing have been applied directly as a pre-training task for protein data by framing them as supervised tasks for predicting the structure or functions of proteins with partially masked sequences 40,41 .”
Sub-Linear Memory: How to Make Performers SLiM (2020, methodology): “A recent line of work aimed to address this restriction, using either structured sparsity [8], truncated back-propagation [12], clustering [20, 31] or linear attention methods [18, 10, 11, 33, 23].”
Hybrid Random Features (2021, methodology): “The new application of random features for softmax kernel in Transformers proposed in (Choromanski et al., 2020; 2021b) led to fruitful research on the extensions and limitations of these methods.”