[论文解读] 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 Bioinformatics被引用 40

一句话总结

引入 Performer，一种 Transformer 变体，使用 Fast Attention Via Orthogonal Random features (FAVOR) 来实现线性尺度的注意力，从而在具备理论保证和向后兼容性的前提下实现长上下文蛋白质序列建模。

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.

研究动机与目标

• 在不依赖稀疏先验的前提下，推动对长生物序列（如蛋白质）的 Transformer 模型进行扩展。
• 开发具有理论保证且与预训练 Transformer 兼容的线性时间注意力机制。
• 在蛋白质序列建模和 ImageNet64 上展示有效性，同时分析复杂度和收敛性。

提出的方法

• 给出 Generalized Attention (GA)，其中注意力可以基于核函数。
• 用 Fast Attention Via Orthogonal Random features (FAVOR) 取代常规注意力，以获得注意力矩阵的无偏、低秩近似。
• 使用随机特征映射来估计查询与键之间的核相似度，从而得到计算形式 Q'K'ᵀ，而无需形成完整的 A 矩阵。
• 结合正交随机特征 (ORFs) 以降低方差并提高近似质量。
• 提供时间/空间复杂度分析，显示在空间上为 O(Ld log d)，在时间上为 O(Ld² log d)，相对于常规注意力的 O(L²d)。
• 通过仅替换注意力机制而保持其他 Transformer 组件不变，展示向后兼容性。

实验结果

研究问题

• RQ1 FAVOR 能否在广义核基注意力下，对标准注意力矩阵给出无偏估计的近似？
• RQ2 基于 FAVOR 的注意力是否能随着序列长度 L 线性扩展，同时在长蛋白质序列上保持准确性？
• RQ3 在长格式蛋白质序列建模和 ImageNet64 上，Performer 相较于常规 Transformer 和稀疏注意力方法的表现如何？
• RQ4 在 GA 框架中，FAVOR 的理论保证（无偏性和一致收敛性）有哪些？

主要发现

• FAVOR 在 token 数量上实现线性时间，空间复杂度亚平方，取代了二次注意力，成为低秩近似。
• 该方法提供对注意力矩阵的无偏估计以及一致收敛保证。
• 在对蛋白质序列任务进行微调后，使用 FAVOR 训练的 Performer 能恢复或接近 Transformer 的性能，并且在蛋白质任务上可以超出一些稀疏注意力基线，如 Reformer。
• 正交随机特征相比非结构化特征降低近似误差，提升下游性能。
• 该方法与预训练的常规 Transformer API 兼容，可以作为常规注意力的可扩展替代品使用。