基于线性可扩展长上下文变换器的蛋白质掩码语言建模
Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers
Krzysztof Choromański, Valerii Likhosherstov, David Dohan, Xingyou Song, Gane, Andreea, Tamás Sarlós, Hawkins, Peter, Jared Quincy Davis, David Belanger, Lucy J. Colwell, Adrian Weller
arXiv (Cornell University) · 2020-06-05
摘要
变换器模型在众多领域均取得了顶尖的研究成果。然而，训练注意力机制以学习远距离输入间的复杂依赖关系所产生的成本问题日益受到关注。为此，利用已学习注意力矩阵的结构与稀疏性特征的解决方案层出不穷。不过，涉及长序列的现实应用场景，如生物序列分析，往往难以满足这些方案的假设条件，从而限制了此类模型的应用探索。为应对这一挑战，我们提出了一种基于正交随机特征快速注意力机制（FAVOR）的新型变换器架构——Performer。该架构的计算复杂度随序列中标记数量呈线性增长而非平方级增长，其空间复杂度低于二次方水平，且无需预设任何稀疏模式先验。此外，它还具备坚实的理论保障：能够对注意力矩阵进行无偏估计并实现一致收敛。该架构还能与预训练的常规变换器向后兼容。我们在极具挑战性的蛋白质序列建模任务中验证了它的有效性，并提供了详尽的理论分析。
本摘要由 AI（Hy-MT2-Pro）翻译生成，仅供参考，请以原文为准。
AI 综述
基于全文生成- 主要观点
- 论文针对标准 Transformer 注意力随输入 token 数 L 二次增长、长序列训练成本过高，而稀疏或结构化注意力先验在生物序列分析等真实长序列任务中可能不成立的问题，提出 Performer 架构与 FAVOR 机制。FAVOR 以正交随机特征近似常规注意力矩阵，使注意力计算随 token 数线性缩放，具有亚二次空间复杂度，且不引入稀疏模式先验。作者将常规点积注意力纳入更广的 Generalized Attention（GA），标准 softmax 注意力被表述为其中一个特例，并主张 FAVOR 可处理可由随机特征估计的核驱动注意力。论文进一步主张 FAVOR 对注意力矩阵提供无偏估计与一致收敛保证，并可与预训练常规 Transformer 向后兼容，经少量微调恢复性能。在蛋白质序列建模任务中，Performer 在单向与双向设置下被验证，Performer-RELU 取得最高准确率，softmax 近似达到与精确 softmax Transformer 相同准确率，论文还报告了 ImageNet64 上的效果。
- 研究方法
- 论文先定义标准双向与单向点积注意力，将注意力矩阵 A 分解为 DQ B DK，其中 B 为高斯核矩阵，再通过随机特征映射 φ 对 B 进行低秩无偏近似。为降低蒙特卡洛估计方差，作者引入正交随机特征，包括 R-ORFs 与 H/G-ORFs，并比较其空间、时间与偏置特性。FAVOR 算法不显式构造 L×L 注意力矩阵，双向情形通过 ˆD^{-1}(Q′((K′)⊤V)) 计算，单向情形利用前缀和（prefix-sum）处理 tril 注意力与 C=[V 1_L]。理论分析使用负相关理论证明 FAVOR 的一致收敛，给出 Mopt=Θ(d log d) 不依赖 L，并可在 O(Ld^2 log d) 时间近似 A 至任意精度。实验在 Jax 中实现，仅替换注意力组件，其余 Transformer 组件保持不变，比较 backward pass 速度与内存，默认配置为 (nheads,nlayers,dff,d)=(8,6,2048,512)，并在 V100 GPU 16GB 上测到显存上限。蛋白质实验训练 36 层模型，使用 TrEMBL 2019 年 1 月发布数据，涵盖单向与双向设置，以 Reformer 作为稀疏注意力基线；还进行预训练 BERT 在 LM1B 上向 Performer 的迁移微调，以及 L=4096、d=16 时随 M 变化的近似误差测试，标准差基于 10 个样本。
- 学术贡献
- 论文提出 FAVOR 机制作为常规注意力的可替换模块，给出 O(Ld^2 log d) 时间与 O(Ld log d) 空间量级，避免显式存储 O(L^2) 注意力矩阵。论文提出 Generalized Attention 类，将标准注意力作为特例，并把 FAVOR 扩展到可由随机特征估计的核注意力，拓展了注意力机制的理论表述范围。理论层面的贡献是证明 FAVOR 对注意力矩阵的无偏估计与一致收敛，给出最优随机特征数 Mopt=Θ(d log d) 且不依赖序列长度 L。领域拓展上，论文在蛋白质序列建模中展示长上下文 Transformer 的有效性，并与 Reformer 稀疏注意力对比，同时报告 ImageNet64 结果。工程与兼容性方面，Jax 实现保持与常规 Transformer 的 API 兼容，支持从预训练 BERT 迁移并经微调恢复准确率。
- 存在局限
- 蛋白质实验主要围绕 TrEMBL 序列上的掩码语言建模与准确率评估，材料未明确给出对多蛋白拼接、蛋白质相互作用或四级结构预测等论文动机中任务的直接端到端评估。向后兼容并非零成本：预训练 Transformer 权重迁入 Performer 后初始准确率仅为 0.07，需要微调；近似误差会经 MLP、多头、多层传播，限制直接迁移。理论收敛保证针对 RBF 核、R-ORFs、查询与键行位于球 B(R) 等条件，且为常数概率，材料未明确覆盖所有 Generalized Attention 核与 H/G-ORFs 的完整有限样本保证。实验比较稀疏注意力时主要使用 Reformer 作为基线，材料未明确系统比较局部窗口、LSH、聚类或分桶等其他长序列注意力在蛋白质任务上的表现。随机特征数 M 与重采样策略被设为超参数，材料未明确给出在不同序列长度、不同蛋白家族或分布外数据上的完整精度—成本权衡。
- 进一步研究方向
- 在论文提到的多蛋白拼接与蛋白质相互作用或四级结构预测任务上直接评估 Performer，并比较材料中提及的进化序列统计模型，以检验全局注意力对跨蛋白交互的增益。对 H/G-ORFs 与 R-ORFs 在蛋白质掩码语言建模中开展系统的精度—时间—内存消融，固定 L、d 与层数，扫描 M 和重采样周期，建立近似误差与下游准确率的关系。将 Generalized Attention 中的 arc-cosine 或 angular 等点式非线性高斯核与 Performer-RELU 比较，检验非 softmax 核在长蛋白序列上的表示能力。设计分阶段微调协议，研究从预训练常规 Transformer 向 Performer 迁移时如何减少初始性能损失，并测量不同层数、头数与 dff 下的恢复梯度步数。在 L≫d 条件下，对单向 FAVOR 的前缀和实现与双向实现在蛋白质生成任务中进行并行效率与数值稳定性对比，明确何时并行前缀和聚合优于存储完整 GPS 张量。
本综述由 AI（deepseek-flash）阅读开放获取全文自动生成，仅供参考，不代表原文作者观点，请以原文为准。
参考文献
相关论文
话题：Machine Learning in Bioinformatics