[Paper Review] Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers

Krzysztof Choromański, Valerii Likhosherstov+9 more|arXiv (Cornell University)|Jun 5, 2020

Machine Learning in Bioinformatics40 citations

TL;DR

Introduces Performer, a Transformer variant using Fast Attention Via Orthogonal Random features (FAVOR) to achieve linear-scale attention, enabling long-context protein sequence modeling with theoretical guarantees and backward compatibility.

ABSTRACT

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.

Motivation & Objective

• Motivate scaling Transformer models to long biological sequences (e.g., proteins) without relying on sparsity priors.
• Develop a linear-time attention mechanism with theoretical guarantees and practical compatibility with pretrained Transformers.
• Demonstrate effectiveness on protein sequence modeling and ImageNet64 while analyzing complexity and convergence.

Proposed method

• Present Generalized Attention (GA) where attention can be kernel-based.
• Replace regular attention with Fast Attention Via Orthogonal Random features (FAVOR) to obtain an unbiased, low-rank approximation of the attention matrix.
• Use random feature maps to estimate kernel similarities between queries and keys, yielding a computational form Q'K'ᵀ without forming the full A matrix.
• Incorporate orthogonal random features (ORFs) to reduce variance and improve approximation quality.
• Provide time/space complexity analysis showing O(Ld log d) space and O(Ld² log d) time vs O(L²d) for regular attention.
• Show backward compatibility by swapping only the attention mechanism while keeping other Transformer components intact.

Experimental results

Research questions

• RQ1 Can FAVOR approximate the standard attention matrix with unbiased estimates for generalized kernel-based attention?
• RQ2 Does the FAVOR-based attention scale linearly with sequence length L while maintaining accuracy for long protein sequences?
• RQ3 How does the Performer perform on long-form protein sequence modeling and ImageNet64 compared to regular Transformers and sparse attention methods?
• RQ4 What are the theoretical guarantees (unbiasedness and uniform convergence) for FAVOR in the GA framework?

Key findings

• FAVOR achieves linear time in the number of tokens and sub-quadratic space, replacing the quadratic attention with a low-rank approximation.
• The approach provides unbiased estimation of the attention matrix and uniform convergence guarantees.
• Performers trained with FAVOR recover or closely match Transformer performance on protein sequence tasks after fine-tuning and can outperform some sparse-attention baselines like Reformer on proteins.
• Orthogonal random features reduce approximation error compared to unstructured features, improving downstream performance.
• The method is API-compatible with pretrained regular Transformers and can be used as a scalable replacement for regular attention.