AlphaFold 2 来了：结构预测奇迹的背后是什么

经过八个月的漫长等待，Nature 现在发布了 AlphaFold 2 论文。正文或多或少地报告了我们近一年来所知道的内容，并添加了一些花絮，尽管在补充信息中伴随着对架构的细致描述。也许更重要的是，作者已经在 Github 上发布了完整的代码，包括运行管道的所有细节。这次没有小字：您可以对任何蛋白质进行推理（我已经检查过！）。你没有听到这个消息吗？让我刷新你的记忆。 2020 年 11 月，来自谷歌 DeepMind 的人工智能科学家团队无可争议地赢得了第 14 届结构预测关键评估竞赛，这是一项两年一度的盲测，计算生物学家试图预测几种蛋白质的结构，这些蛋白质的结构已通过实验确定但未公开发布。他们的结果如此惊人，而且这个问题对生物学来说如此重要，以至于整个世界都大吃一惊，整个学科——计算生物学——都想知道刚刚发生了什么。现在文章已经上线，兴奋之情溢于言表。我们有 70 多页期待已久的答案，以及数千行代码，毫无疑问，它们将成为计算生物学的基本组成部分。然而，与此同时，我们也有许多新问题。新闻爆出前的秘诀是什么，为何如此有效？它是一段普通用户可以实际运行的代码吗？ AlphaFold 2 的缺点是什么？而且，最重要的是，这对计算生物学意味着什么？对于我们所有人？在这篇评论中，我的目标是作为我八个月前博客文章的延续，我试图解决其中的一些问题。首先，我提供了 AlphaFold 2 架构的鸟瞰图。这不是一个技术说明（SI 尽可能详细，甚至代码也引用了其中的不同部分），而是侧重于架构背后的直觉。我希望这能覆盖那些没有深度学习和生物信息学背景但想知道发生了什么的人；以及那些可能具有正确背景但希望在深入研究之前对全文进行概述的人。根据冰冷、坚硬的事实，我对建筑背后的想法进行了完全个人的评估。也就是说，我会解释我认为哪些想法是 AlphaFold 2 成功的关键，并推测哪些因素使这个团队取得了成功，而其他许多人都没有做到这一点。我是一个有强烈意见的人，但仍然很高兴地宣布我的想法可能完全朝着错误的方向发展。尽管如此，我认为 AlphaFold 2 的故事提出了很多我们作为社区尚未解决的问题，值得在某处适当考虑。最后，我重新审视了我八个月前提出的一些问题。论文或代码已经回答了其中一些问题（例如，运行代码的限制是什么）。还有一些没有明确解决，但我有机会更深入地思考它们，我认为我有一些新颖的见解。还有一些是由新信息引起的问题，我认为我们必须一起回答。我向自己保证，这次我会更简洁——毕竟几个月后我应该写我的博士论文，我真的没有太多时间可以腾出时间。让我们看看我是否管理。

直到周四早上，我们得到的最佳答案是一张图片，早在 11 月，DeepMind 的新闻稿中就已发布。这种模式在当时在互联网上流传，并且从那时起就出现在众多会议和讨论组中。但是，遗憾的是，它缺乏细节，即使是知识最渊博的深度学习专家也只能做出有根据的猜测。 Nature 文章提供了一个非常相似但稍微详细的图表，概述了架构的不同部分。总体思路非常简单，所以我将尝试用几行来勾勒它。如果您不熟悉深度学习，以下内容可能听起来有点抽象，但完全没问题。稍后我将带您了解详细信息。不过，现在让我们尝试获取网络的示意图。为清楚起见，我将图像分成三份，代表 AlphaFold 2 系统的三个主要部分。首先，AlphaFold 2系统利用输入的氨基酸序列查询多个蛋白质序列数据库，构建多序列比对（MSA）。简而言之，MSA 可识别在生物体中已识别的相似但不相同的序列。这可以确定序列中更可能发生变异的部分，并允许我们检测它们之间的相关性。 AlphaFold 2 还尝试识别可能与输入（“模板”）具有相似结构的蛋白质，并构建该结构的初始表示，称为“对表示”。从本质上讲，这是一种氨基酸可能相互接触的模型。在图表的第二部分中，AlphaFold 2 将多序列对齐和模板，并将它们传递给一个转换器。稍后我们将讨论 Transformer 需要什么，但现在您可以将其理解为“oracle”，可以快速识别哪些信息的信息量更大。这部分的目标是改进 MSA 和配对交互的表示，同时也在它们之间迭代地交换信息。更好的 MSA 模型将改进网络对几何的表征，同时有助于改进 MSA 模型。这个过程被组织成块，迭代地重复直到指定数量的循环（已发布模型中的 48 个块）。该信息被带到图表的最后一部分：结构模块。这条复杂的管道采用精炼的“MSA 表示”和“对表示”，并利用它们来构建结构的三维模型。与之前最先进的模型不同，该网络不使用任何优化算法：它只需一步即可生成静态的最终结构。最终结果是一长串笛卡尔坐标，代表蛋白质每个原子的位置，包括侧链。因此，回顾一下：AlphaFold 2 找到与输入相似的序列，使用特殊的神经网络架构提取信息，然后将该信息传递给另一个生成结构的神经网络。

最后一点是该模型以迭代方式工作。生成最终结构后，它将获取所有信息（即 MSA 表示、对表示和预测结构）并将其传递回 Evoformer 块的开头，即我们图表的第二部分。这允许模型改进其预测，并生成一些您可以在文章页面中找到的有趣视频。与大多数生物信息学程序一样，AlphaFold 2 配备了“预处理管道”，这是该学科对“调用其他代码的 Bash 脚本”的术语。管道运行多个程序来查询数据库，并使用输入序列生成多序列比对 (MSA) 和模板列表。每个程序的脚本略有不同，但 AlphaFold 2 与您的花园品种蛋白质结构预测预处理管道并没有太大区别。值得解释多序列​​比对的含义。在 MSA 中，我们打算预测其结构的蛋白质的序列在一个大型数据库中进行比较（通常类似于 UniRef，尽管在后来的几年中，使用源自宏基因组学的序列来丰富这些比对是很常见的）。潜在的想法是，如果两个氨基酸紧密接触，其中一个的突变将紧随另一个的突变，以保持结构。考虑以下示例。假设我们有一种蛋白质，其中带负电荷的氨基酸（例如谷氨酸）靠近带正电荷的氨基酸（例如赖氨酸），尽管它们在氨基酸序列中都相距甚远。这种库仑相互作用稳定了蛋白质的结构。现在想象一下，第一个氨基酸突变成带正电的氨基酸——为了保持这种接触，第二个氨基酸将在进化压力下突变成带负电的氨基酸，否则产生的蛋白质可能无法折叠.当然，真实情况很少像这个例子那样明确，但你明白了。查找模板遵循一个完全不同但又密切相关的原则。模板建模背后的哲学可以用格言“太阳底下没有新鲜事”来编码。蛋白质会发生变异和进化，但尽管发生了变化，但它们的结构往往保持相似。例如，在下图中，我展示了对应于不同生物体的四种不同肌红蛋白的结构。您可以欣赏它们看起来几乎相同，但是如果您查看序列，您会发现巨大的差异。例如，右下角的蛋白质与左上角的蛋白质只有约 25% 的氨基酸相同。然而，在大多数情况下，保护发生在较小的范围内，其中蛋白质片段（例如，酶的活性中心）在其周围环境进化时基本保持不变。大小并不重要：使用正确的方法，可以识别其中一些保守片段并将它们用作构建结构的指南。这是结构预测中的一个重要因素，CASP14 中的目标通常根据可用模板的数量进行分类。这里有什么特别的吗？并非如此：CASP14 的大多数参与者都遵循非常相似的策略。使用相关突变从 MSA 中提取结构信息的想法已有数十年历史，而收集其他蛋白质片段来模拟目标结构的想法甚至可能更早。我想说，到目前为止，没有什么新东西。

这是故事真正开始的地方。巨大的 AlphaFold 2 神经网络的第一部分 Evoformer 的任务是从多序列比对和模板中提取每一盎司的信息。听到从多序列比对中提取信息（“协同进化分析”）多年来一直是结构生物信息学的主要追求，您可能不会感到惊讶。人们在 90 年代开始研究它，尽管成功有限。在上个十年之初，几个小组开始确定一些阻碍先前尝试的偏见，并开发了强大的统计机制来纠正它们。几年来取得了一些持续的进展。然后，在 CASP13 (2018) 中，几个小组证明实际上不需要稳健的统计数据：你只需要训练深度残差神经网络。 AlphaFold 2 的 Evoformer 彻底改造了这个过程，并进一步推进了几个步骤。 Evoformer 背后的中心思想是信息在整个网络中来回流动。在 AlphaFold 2 之前，大多数深度学习模型将采用多序列比对并输出一些关于几何接近度的推断。因此，几何信息是网络的产物。相反，在 Evoformer 中，对表示既是产品又是中间层。在每个循环中，该模型利用当前的结构假设来改进对多序列比对的评估，这反过来又会导致新的结构假设，依此类推。表示，序列和结构，交换信息，直到网络达到可靠的推理。这作为示例更容易理解。假设您查看多序列比对并注意到一对氨基酸之间的相关性。让我们称它们为 A 和 B。您假设 A 和 B 很接近，然后将此假设转化为您的结构模型。随后，您检查所述模型并观察到，由于 A 和 B 很接近，因此 C 和 D 很可能应该接近。这导致了另一个基于结构的假设，可以通过在 MSA 中搜索 C 和 D 之间的相关性来确认。通过多次重复此操作，您可以很好地理解结构。网络中的第一步是为 MSA 和模板定义“嵌入”。请记住，多重序列比对最终是来自有限字母表的符号序列：离散变量的主要例子。另一方面，神经网络本质上是连续的设备，依靠差异化从训练集中学习。 “嵌入”是深度学习魔法书中的一个技巧，它允许将离散变量转换为连续空间（“嵌入空间”），从而可以训练网络。这听起来很复杂，实际上非常简单。您只需要定义一层接收离散输入并输出一些连续向量的神经元。嵌入可能是预训练的，这在自然语言处理 (NLP) 中很常见，但更常见的是，它与我们试图学习的任何目标一起训练。在 AlphaFold 2 中，嵌入是普通的密集神经网络。

一旦我们的 MSA 和模板处于正确的“嵌入空间”，Evoformer 就该发挥它的魔力了。要了解 Evoformer，您首先需要熟悉迄今为止最热门的深度学习架构：transformer。不乏解释这种架构的材料，坦率地说，很多（如果不是大多数）都会比我的要好。如果您对深入分析感兴趣，我会推荐 Jay Alammar 的 The Illustrated Transformer。如果您只想知道最低限度，请阅读以下内容。而且，如果您已经知道像手掌这样的转换器，请按 Ctrl+F 并在下面的几段中找到“回到 Evoformer”这句话。 Transformer 架构于 2017 年由 Google Brain 的一个团队在一篇题为“Attention is all you need”的论文中介绍。正如您可能从显眼的标题中想象的那样，关键要素是一种称为注意力的新颖机制。注意的目标是确定输入的哪些部分对神经网络的目标更重要。换句话说，要确定它应该注意输入的哪些部分。想象一下，您正在尝试训练一个神经网络来生成图像标题。一种可能的方法是训练网络处理整个图像——比如 512×512 图片中的约 250k 像素。这可能有效，但有一些原因表明它不是最好的主意。也许首先，因为这不是我们人类所做的：当我们看一张图片时，我们并没有看到它“作为一个整体”。相反，我们将其划分为不同的模式：孩子、狗、飞盘。幸运的是，事实证明我们可以训练一个设计巧妙的神经网络层来执行这项任务。而且，根据经验，它似乎大大提高了性能。到目前为止，图像字幕并不是唯一的例子。与原始的“Attention is all you need”论文更相关的是机器翻译的情况，其中注意力揭示了序列的哪些部分对当前翻译部分很重要。变压器普遍使用的原因还有很多。例如，在机器翻译中，它们有助于改善“梯度消失问题”，这是训练过程中的一个常见障碍。在基于序列的模型中，相对于经典的循环神经网络 (RNN) 模型，它们可以显着加快训练速度。最重要的是，变压器在各种任务中凭经验证明了卓越的性能。特别是，过去一年中大部分最引人注目的 AI 成就都落后于他们——例如，GPT-3 中的 GPT 代表“生成式预训练变压器”。在我认为我的解释完整之前，我确实觉得我必须提到最后一个免责声明。变压器没有在许多领域得到广泛应用是有原因的。注意矩阵的构建导致了二次内存成本。即使您拥有配备 80 GB 专用内存的上一代 NVIDIA A100，它也可以很快被填满。不出所料，谷歌著名的张量处理单元 (TPU) 的关键优势之一是每个内核的大量内存。这种情况可能会随着新架构（例如 Performer 或 Perceiver）的引入而迅速改变，它们将这种二次成本降低为伪线性成本。但是，无论如何，我已经超越了自己。回到 Evoformer。 Evoformer 架构使用的不是一个，而是两个变压器（“双塔架构”），两者之间有一个清晰的通信通道。每个头部专门用于它正在查看的特定类型的数据，或者是多序列比对，或者是氨基酸之间成对相互作用的表示。它们还包含连续表示的信息，允许定期交换信息和迭代细化。

让我们首先看看负责 MSA 的负责人，我将其称为“MSA 转换器”，以纪念 Facebook AI Research 2 月份发表的论文，该论文实现了类似的想法。 MSA 转换器在一个非常大的蛋白质符号矩阵上计算注意力。为了减少原本不可能的计算成本，注意力被“分解”为“行”和“列”组件。即，网络首先在水平方向计算注意力，让网络识别哪些氨基酸对更相关；然后在垂直方向上，确定哪些序列的信息量更大。 AlphaFold 2 的 MSA 变换器最重要的特点是行（水平）注意力机制结合了来自“对表示”的信息。在计算注意力时，网络会添加一个偏置项，该偏置项是直接根据当前对表示计算得出的。这个技巧增强了注意力机制并允许它精确定位相互作用的残基对。另一个变压器头，作用于对表示的那个，以类似的方式工作，当然，虽然很多细节不同。该网络的关键特征是注意力以残差三角形的形式排列。这里的直觉是加强三角不等式，这是度量空间的公理之一。这是一个非常聪明的想法，因为基于深度学习的结构预测的经典问题之一是距离分布不能嵌入到三维空间中。似乎这个技巧解决了这个问题，然后还有更多。经过多次迭代（论文中的 48 次），网络建立了蛋白质内部相互作用的模型。现在，是时候建立一个结构了。如果你了解这么多，你现在可能知道 AlphaFold 2 处理序列搜索数据以生成两个“表示”：多序列比对 (MSA) 的表示，它捕获序列变异；以及“残基对”的表示，它捕获了哪些残基可能彼此相互作用。现在的问题是：我们如何从这些中获得结构？这是结构模块的工作结构模块背后的想法在概念上非常简单，但是通过......的许多细节变得混乱。