论文阅读笔记|【Paper Reading】自注意力机制以及在Alphafold2中的应用人工智能|机器学习|深度学习

这期文章我想开一个新的系列，叫读论文Paper Reading，分享初学者的自己对深度学习方面的论文的理解，大家一起学习交流。尤其会关注一些深度学习和医学相结合的领域。

文章目录

Alphafold2做了一件什么事
- Model
从 Self-Attention机制到 Transformer模型
- Attention机制
- Attention机制的计算
- Self-Attention 自注意力机制
- Query-Key-Value (QKV) 模型
- Multi-head self-attention
Self-Attention in Evoformer
- MSA representation信息的编码
- pair representation信息的编码
结语

开年第一篇笔者想讲讲现在非常热门的Alphafold2，作为一个深度学习的模型来精确预测蛋白质分子的结构。论文的Title为“ Highly accurate protein structure prediction with AlphaFold”，由Deepmind团队于2021年7月15日发表在 Nature杂志上。这篇文章的工作被 Nature和 Science同时评选为了2021年科学十大突破之一。
这篇文章涉及到的细节非常多，在文章的补充材料部分提供了近60页的算法描述，涉及到模型的方方面面。毕竟是革命性的工作，且很多算法和团队之前的工作存在前后联系，弄懂所有的算法不是一件容易的事。这次自己就Alphafold2中 Evoformer模块及自注意力机制(Self-attention)的应用写写自己的理解，顺便学习自注意力机制的由来，以及在Transformer及其变种中的应用，一起学习提高。
Alphafold2做了一件什么事简而言之，Alphafold2可根据一串氨基酸的序列，预测出蛋白质的三级结构，并且达到了原子级别的精度。
我们知道蛋白质的基本构成单位是氨基酸，可由一段氨基酸分子的序列表示，借助各种次级键卷曲等形成三级结构（一种3D空间中的结构）。在以往，还没有任何一个方法能根据氨基酸序列预测出原子级精度的蛋白质三级结构，尤其是当没有同源结构可以参考时。而论文表示Alphafold2解决了这个50年来难以解决的问题。
这个模型是一个深度学习的模型，它结合了蛋白质结构的物理性质、生物学知识，利用了多序列比对的方法，将这些信息融入深度学习算法的设计中，大大提高了预测的精度。

论文阅读笔记|【Paper Reading】自注意力机制以及在Alphafold2中的应用

文章图片
我们可以看到这张图显示了模型的精确度达到了1?（10的-10 次方米），远超其他模型。后面的图显示了真实结构和预测结构之间的差异，分别用蓝色和绿色表示，可以看到误差非常小。模型是在PDB蛋白质数据库中进行预测的。
Model

文章图片
上图是模型的整体一览。
我们输入的是一段氨基酸序列（比如我们目标是预测一段人的氨基酸序列），输出的是每个氨基酸在三维空间中的位置。
模型共分为三个部分，第一个部分是特征提取的部分，第二、三部分分别是编码器和解码器。
将提取后的特征以两个矩阵的形式封装并喂到编码器-解码器的神经网络中，预测出氨基酸序列的3D空间构象。
在特征提取的部分，先输入一段人的氨基酸序列，其作为数据第1行。同时，用这个氨基酸序列搜索基因数据库，搜到与该序列相似的片段，列入数据的第2到n行。我们看到这些数据使用了多序列比对 MSA（Multi Sequence Alignment）的方法。这个MSA representation数据就构成了第一个数据矩阵。
其次，我们将每两个氨基酸之间的关系，用表示一个配对关系的矩阵来表示，表中的每一项是一对氨基酸之间的关系。同时，搜索结构数据库，得到一些氨基酸之间已知的位置信息，获得一些模板(Templates)，写入这个矩阵。这个Pair representation这样就构成了第二个矩阵。
第一个矩阵代表了各物种相似氨基酸序列的特征（多序列比对信息），第二个矩阵代表了氨基酸对之间的特征（结构信息）。

文章图片

将这两个矩阵（上图右）传入编码器。
我们来看看这两个矩阵的形状。
第一个代表了各物种相似氨基酸序列的MSA矩阵形状是s×r×c。s行是蛋白质的数目，第一行是人类蛋白序列，后面是数据库中匹配出来的相似序列。r表示序列长度，有r个氨基酸。c用来表示每个氨基酸的特征长度，我们将氨基酸表征为一个向量V = (V1,V2,V3…,Vc)，c表示特征的数量，即向量的长度(维数)。
这里将氨基酸的特征Embedding（嵌入）为一个向量使用了另一个特征提取的算法，在下一小节会对Embedding的概念做些许展开介绍。
第二个氨基酸对之间特征的pair矩阵形状是r×r×c，是一个方阵。r表示这段人蛋白的氨基酸的序列的长度。每个氨基酸对之间的特征信息也是用一个c维向量来表示。

文章图片

我们将这两个矩阵都传入编码器，一个名为Evoformer的神经网络架构。
Evoformer是大名鼎鼎的深度学习模型Transformer的变种。在Transformer中（下一小节详述），我们可以在编码器输入一段序列，输出的是另一段序列，在Transformer的神经网络的编码器中，我们计算序列之间的关系，用学习到的参数，来输出相应的序列，这段序列包含了序列每个位置相对于整体的一个期望值信息（即关注度，我们称为自注意力Self-Attention）。
在Evoformer的编码器中，我们传入这两个矩阵，输出的是两个同样形状的矩阵。
这时这两个矩阵经过神经网络的运算，已经带有了自注意力信息，我们将序列的特征信息和配对的信息融入自注意力中。具体在我们这篇文章的例子中，输出的第一个矩阵的信息中，在c这个维度，已经带有了更新过后的氨基酸序列特征；第二个矩阵，在c这个维度，已经带有了更新过后的氨基酸与氨基酸之间的关系的信息。

文章图片
有了这个信息，我们再把这两个矩阵传入右边解码器，即Structure module（这里第一个矩阵只需要传第一行，即人蛋白的氨基酸序列的信息，因为我们的目的就是预测这段氨基酸的3D结构）。解码器会根据更新后的人蛋白的氨基酸序列的特征，以及氨基酸对之间的关系，来预测输出每一个氨基酸在3D空间中的位置，来得到蛋白质的3D结构。
那么现在就会有几个问题：
1.上面这三段话是什么意思？什么是自注意力？
2.这个Evoformer是如何将传入的氨基酸序列信息和氨基酸对的信息建模，编码成我们需要的序列信息和氨基酸对之间的信息来解码3D结构的呢？
3.之前提到的Transformer中计算序列每个元素的自注意力的思想，为什么可以用这个思想来解释更新后的蛋白质特征信息，以及的氨基酸之间的关系信息呢？
我想先从Self-Attention自注意力机制讲起。
从 Self-Attention机制到 Transformer模型这节专门面向小白，讲讲Self-Attention机制与Transformer模型。
在网络上许多对自注意力机制的讲解都是直接讲QKV模型，让人没有一个科学的连贯思维去认识这个机制的来龙去脉，本文会从头讲解自注意力的思路是怎么来的，自注意力机制体现了什么样的思想。
笔者以自然语言处理（NLP）为例讲解，因为Transformer的提出一开始也是广泛地应用在NLP领域。
处理一个序列（Sequence）的数据，最常见的情况就是处理语言文本。
比如我们先将一句话中的词embedding为一个向量，将词汇向量化。这里Embedding可作为一种从高维到低维映射的方法，我们想象词汇是存在于高维空间的东西，具有相似含义的词汇可能位于高维空间的邻近位置。我们抽取它们的特征，映射到一个较低维的向量空间。
但是光有词汇自身的特征还不够，我们还需要建立上下文之间的联系。

文章图片
比如上图，我们输入一句话的文本，下排embedding为词向量后还没有上下文的联系。
我们需要给一个上下文编码器（Contextual Encoder），输出为Contextual embeddings。
有了上下文联系后，就可以进入下游任务，进行情感分类(Sequence-to-one)、翻译(Sequence-to-Sequence, Seq2Seq)等任务。
因此上下文编码器非常重要。
那么，上下文编码器应该如何实现呢？
由上我们可以推出序列常有两个特点：

序列长度不固定，即变长性
序列需要有一个全连接，保证上下文的联系

Attention机制

文章图片
对于以上这两个特点，我们会想到使用CNN（卷积神经网络）和RNN（循环神经网络）来建模，如上图。
但是由于CNN的卷积特性，每次卷积时卷积核的视野是有限的，导致其建模上下文关系时，范围受限；如果堆叠层数，模型又会变得臃肿，训练困难。
而RNN虽然理论上可对长序列进行建模，且对位置信息敏感，但同样会有因时序拖长导致梯度消失，上下文的信息感受范围任然受限。
如果我们使用原始的全连接网络：

文章图片

在传统的全连接网络学习完成后，参数是固定值（即权重W）。但是固定形状的权重W不能处理变长序列。比如把“我爱我家”输成“啊我爱我家”，句意没有太多的变化，但是由于原句向右移位1格，原本有效的权重会产生错位，形成错误的预测。因此我们需要动态地根据当前输入，生成网络权重，这种根据输入变化的动态权重即为注意力（Attention）的来源。
在注意力机制中，全连接的思想是不变的。
拓展一下，全连接还可以扩展为图（Graph）：

文章图片

每个词可以是一个节点Node，各个节点之间全连接。每个节点的更新依赖于其他所有节点。从而还可以用图神经网络的方法解决问题。
综上所述，我们不再让权重W成为固定可学习的参数，而让它动态生成，我们将这个权重表示为：α；
α的产生源自于对变长序列的连接操作，是一种动态的关注度(即注意力)，任何向量都会对整个句子产生一个注意力的分布。
Attention机制的计算 【论文阅读笔记|【Paper Reading】自注意力机制以及在Alphafold2中的应用】有了上面的理解，我们就可以理解Attention机制的计算过程了：

文章图片
我们假设此时有一个外来的查询向量q（query），需要了解q在整个句子中的关注度，即注意力。
一共分为2步：
第一步，q乘以每个词向量x，即做向量内积，内积体现了两个向量之间的相关程度。
内积得到各自的分数s，相当于q对各个词进行一次打分（score），接着对score进行归一化softmax处理。softmax后，所有打分加和为1，可视作一种α的概率分布。
α n = p ( z = n ∣ X , q ) = softmax ? ( s ( x n , q ) ) = exp ? ( s ( x n , q ) ) ∑ j = 1 N exp ? ( s ( x j , q ) ) \begin{aligned} \alpha_{n} &=p(z=n \mid \boldsymbol{X}, \boldsymbol{q}) \\ &=\operatorname{softmax}\left(s\left(\boldsymbol{x}_{n}, \boldsymbol{q}\right)\right) \\ &=\frac{\exp \left(s\left(\boldsymbol{x}_{n}, \boldsymbol{q}\right)\right)}{\sum_{j=1}^{N} \exp \left(s\left(\boldsymbol{x}_{j}, \boldsymbol{q}\right)\right)} \end{aligned} αn??=p(z=n∣X,q)=softmax(s(xn?,q))=∑j=1N?exp(s(xj?,q))exp(s(xn?,q))??
第二步，有了α的分布，就可以对所有的输入x做一个期望：
即各自对应的α乘以x，加和，就得到了最终查询词q之于整句话X的注意力Attention。
att ? ( X , q ) = ∑ n = 1 N α n x n = E z ～ p ( z ∣ X , q ) [ x z ] \begin{aligned} \operatorname{att}(\boldsymbol{X}, \boldsymbol{q}) &=\sum_{n=1}^{N} \alpha_{n} \boldsymbol{x}_{n} \\ &=\mathbb{E}_{z \sim p(z \mid \boldsymbol{X}, \boldsymbol{q})}\left[\boldsymbol{x}_{z}\right] \end{aligned} att(X,q)?=n=1∑N?αn?xn?=Ez～p(z∣X,q)?[xz?]?
注意这个加和会随着query向量q变化，因此查询的东西一变，这个期望就可能变化。
Self-Attention 自注意力机制我们由此看到Attention机制解决了两个问题：
1.解决了变长的问题。之前的全连接方式，如果输入向量是变长会出问题；
2.解决了一句话内两个词之间的依赖关系，这个依赖关系更多的是依靠内容，而不是依靠前后的位置关系。因为最后的注意力值是一个权重加和，是一个标量。
那么如果查询向量query(q)是来自自身输入的向量的情况是怎么样的呢：

文章图片

在上图的例子中，我们用第一个词“The”来查询整句话进行打分评估。查询向量query(q)变成了来自自身输入的向量，即自己注意到自己，这就是 Self-Attention。只要有预先学习好的权重W，就可以进行自注意力的计算了。
因此，这个Self-Attention可以作为神经网络模型的单独一层，具有汇聚所有节点的信息的特性，并算得每个节点在所有信息中的综合期望（关注度），即可以理解为每个节点在整体中的自注意力是多少。
Query-Key-Value (QKV) 模型自注意力机制的应用，最著名的就是Transformer模型了，在Transformer模型中，自注意力机制特化成了Q K V模型：

文章图片
在上图的QKV模型中，X矩阵代表了原数据，列是序列长度，行是各个词向量。
查询向量变成了一个矩阵Q（Query），待打分数据变成了矩阵K（Key），QK相乘归一化得到权值分布矩阵，再乘以一个待计算权值矩阵V（Value），则得到了一个相同形状的自注意力矩阵H。
矩阵Q、K、V都源自于原数据X，只是做了一个线性变换，因此，线性变换的参数Wq、Wk、Wv（参数矩阵）是可学习参数，把原始数据X线性映射到一个不同的空间中。
Multi-head self-attention 在Transformer中，我们还涉及到多头自注意力机制。
即同时并行计算多个QKV模型，每个模型的参数(Wq、Wk、Wv)都不共享，单独学习，最后放到一起，用一个权值矩阵W还原成最初的形状。这样分别计算每个自注意力模型的H，最后拼接起来并以系数W加权求和。
设想自注意力机制存在一个缺点，当模型在对当前位置的信息进行编码时，可能会过度的将注意力集中于自身的位置，因此Transformer作者提出了通过Multi-head的机制来解决这一问题。同时，Multi-head还能够给予注意力层的输出包含有不同子空间中的编码表示信息，从而增强模型的表达能力。

文章图片

在Transformer中，我们还涉及到编码位置信息算法（位置向量+词的embedding向量）、Layer-Normalization、残差连接（网络深层化）、逐行的FFN（用全连接层实现，类似窗口为1的卷积）、Encoder、Decoder等算法，这里就不详述了。
我们可以看到Transformer模型的可解释性相对来说是比较强的，它使用了能让人们感同身受的注意力的概念，可以对序列内部的相关性做多方面的解释。
我们注意到，通过自注意力机制，我们可以发现在序列中最相关的那个节点，并且可以了解注意力的强弱。在自然语言处理中，我们将上下文的语义相关性赋予了注意力，帮助我们解决语言信息的分类、翻译等问题。那么同为处理序列信息，如果我们在这个注意力中赋予了氨基酸序列的各种信息比如：氨基酸的物理性质、生物学性质、多序列比对信息、氨基酸对之间的互相作用信息、氨基酸序列结构模板信息等，如果能对这些信息进行编码建模，融入神经网络中，如果它对氨基酸序列三维空间结构有决定性的作用，那就能拿来预测3D结构了。这里也需要一些对生物信息学、生物化学、物理学等等的一个先验的知识进行假设了。
Self-Attention in Evoformer 回到Alphafold2的整体模型，之前提到了将第一个矩阵MSA（多序列比对信息）和第二个矩阵pair（氨基酸对的结构信息）传入Evoformer编码器，对各类氨基酸的信息进行建模。
整个Evoformer编码器如图：

文章图片

MSA representation信息的编码观察整个Evoformer编码器中，有许多红色的Self-Attention模块，只是冠以了“Row-wise”、“Column-wise”、“gated”、“with pair-bias”等词缀。我们来看看这些Self-Attention模块是如何对氨基酸各类信息进行建模的。
首先看左上角第1个模块“Row-wise gated self-attention with pair-bias”（按行的、带门控的、带对偏移的自注意力层）：

文章图片
按行（Row-wise）表示我们先对每一个物种的氨基酸序列特点做面向自注意力机制的信息编码。
左上角的第2、3、4个线性变换，就是之前讲到的QKV模型的三个变换矩阵Wq、Wk、Wv，得到query、key、value，并且用到了多头自注意力机制，对h层进行并行运算。
由于是按行，每一行我们都单独算注意力的α分布，其等于这一行的query与key的点乘加上pair bias。为什么加上来自第二个配对信息矩阵pair的数据作为pair bias（对信息偏倚）呢？因为我们通过向量的点乘算得了两个氨基酸之间的相关性，但是我们假设氨基酸之间的相关性应该有一部分信息来自于氨基酸对信息的矩阵中，所以作者加上了这个偏倚项，这样算得的自注意力会带上氨基酸对的结构信息。
乘以value之后，再乘以一个门控gate，也就是第1个线性变换做sigmoid，加上门控可以控制我们的自注意力哪些部分可以进行输出，这个门控也可能对氨基酸之间的结构的决定能力产生影响。
我们可以看到，这个小模块中所有的算法设计都蕴含了可能影响基于氨基酸结构因素的先验假设，虽然看似非常的绕，但结果非常的出彩。
接着左上角第2个模块，“column-wise gated self-attention”（按列的、带门控的自注意力层）：

文章图片
之前的Row-wise我们对每一个物种的氨基酸序列特点信息进行了建模。
按列（column-wise）表示序列中同一个位置的氨基酸，在不同物种的蛋白质中的相关性进行编码。
与Row-wise相比，由于研究的是不同物种的蛋白质中的相关性，这里就不再加上pair bias的偏倚信息了。其余部分都是一样的。
通过先对行建模，再对列建模，最后加上一个全连接层（第3个绿色的模块）进行信息的提炼，输出的矩阵中就更新了之前编码的氨基酸位置信息。
pair representation信息的编码在整体模型中，我们看到MSA和pair之间的信息存在一个交互，pair到MSA我们已经在第1个自注意力模块里看到了。
我们看紫色的“Outer product mean”，它进行了MSA到pair的信息回流：

文章图片
我们看到在MSA在接受pair作为第1层的bias信息加和，经历了三层编码后，存在一个“Outer product mean”的支路，将更新过后的带氨基酸位置信息加回到pair信息中，完成了信息的交互。
具体来说，在MSA中，每一列的信息，也就是一个S×Cm的矩阵代表了这一个氨基酸位置的（所有物种的）信息。不同位置的氨基酸由于更新了位置信息，我们将两个向量做外积，形状变为（s,c,c），通过一系列线性变换，加回到pair信息中，长度为Cz。
在后面的模块中，作者对氨基酸对之间的信息按行和列做了自注意力层的编码，对pair的内容进行了更新。具体如何实现留一个坑，有空会在下次Paper Reading会继续讲解。
我们对MSA和pair进行了位置信息的更新后，就可以进入解码器来构造3D结构了。由于我们只需要预测人的氨基酸序列的3D结构，所以只需要把MSA的第一行内容传入解码器。在解码器中，作者同样使用了自注意力机制，也加上了对偏移(pair bias)进行空间信息编码，并为编码构造3D结构的信息(旋转和位移)，而进行了“魔改”。具体算法，包括训练、模型评估等等其他细节也会有空在下次Paper Reading继续讲解。
整个Evoformer编码器的架构总体上运用了多个自注意力层+全连接层的结构，并对MSA信息和pair信息进行建模编码，将信息融入修改过的自注意力层，得到更新氨基酸空间位置信息的MSA和pair矩阵。
结语 Alphafold2是一个复杂的模型，它通过一串氨基酸的序列，预测出原子精度的蛋白质三级结构。它在编码器中使用了Transformer的变种Evoformer，使用了许多改进的自注意力层，许多算法的改进是结合了基于氨基酸空间结构影响因素的先验假设进行了巧妙的设计。
本文是自己对自注意力机制的应用的理解，整理了自己的思路。从自注意力机制在NLP中的运用，可以看到它解决了两个问题：1.解决变长输入，2.上下文相关。追根溯源，它大概解决了以下问题：通过Attention找到最相关的节点，根据先验知识将建模好的信息编码入自注意力层，更新优化Embedding后的数据。
在Alphafold2中，我们希望序列信息MSA和配对信息pair能为我们建立空间3d结构所用，需要了解序列中每个氨基酸的影响力，这个影响力（即自注意力）通过序列信息和配对信息建模，编码，将序列的Embedding数据进行更新，以达到能以此信息输入解码器构造3D结构的能力。
在建模中涉及到的大量参数，通过PDB蛋白质数据库进行训练学习。
我们可以看到，整个模型的大大小小的各个算法设计都蕴含了可能影响基于氨基酸结构因素的先验假设，虽然看似非常的绕，但结果非常的出彩，也是凝聚了Deepmind团队大量心血，在大量前期工作基础上完成的，包括一年半前同样发表在Nature上的Alphafold1代的算法模型。
笔者是深度学习初学者，又是一名医学生，第一次做深度学习方面的的Paper Reading Report，欢迎指正错误和不足。
参考文献：
1.Jumper, J., Evans, R., Pritzel, A. et al. Highly accurate protein structure prediction with AlphaFold. Nature 596, 583–589 (2021).
2.Qiu, X., Sun, T., Xu, Y. et al. Pre-trained models for natural language processing: A survey. Sci. China Technol. Sci. 63, 1872–1897 (2020).
3.Vaswani, Ashish, et al. “Attention is all you need.” Advances in neural information processing systems. 2017.
4.邱锡鹏. A Tutorial of Transformers. VALSE Tutorial. 2021.06.20.
（本文发在我的公众号“飞舞的矩阵”和CSDN博客）