2021SC@SDUSC基于人工智能的多肽药物分析问题（十二）

基于人工智能的多肽药物分析问题（十二）

2021SC@SDUSC

1. 前言

上周预测模型已经分析到了IterBlock_w_Str这个层，并且已经详细解释了该层中的 str2str模型，以下图片是局部代码总览，本周，我们将继续分析余下部分内容

2. 代码分析

2.1 Str2MSA模型

class Str2MSA(nn.Module):
    def __init__(self, d_msa=64, d_state=32, inner_dim=32, r_ff=4,
                 distbin=[8.0, 12.0, 16.0, 20.0], p_drop=0.1):
        super(Str2MSA, self).__init__()
        self.distbin = distbin
        n_att_head = len(distbin)

        self.norm_state = LayerNorm(d_state)
        self.norm1 = LayerNorm(d_msa)
        self.attn = MaskedDirectMultiheadAttention(d_state, d_msa, n_att_head, d_k=inner_dim, dropout=p_drop) 
        self.dropout1 = nn.Dropout(p_drop,inplace=True)

        self.norm2 = LayerNorm(d_msa)
        self.ff = FeedForwardLayer(d_msa, d_msa*r_ff, p_drop=p_drop)
        self.dropout2 = nn.Dropout(p_drop,inplace=True)
    def forward(self, msa, xyz, state):
        dist = torch.cdist(xyz[:,:,1], xyz[:,:,1]) # (B, L, L)

        mask_s = list()
        for distbin in self.distbin:
            mask_s.append(1.0 - torch.sigmoid(dist-distbin))
        mask_s = torch.stack(mask_s, dim=1) # (B, h, L, L)
        
        state = self.norm_state(state)
        msa2 = self.norm1(msa)
        msa2 = self.attn(state, state, msa2, mask_s)
        msa = msa + self.dropout1(msa2)

        msa2 = self.norm2(msa)
        msa2 = self.ff(msa2)
        msa = msa + self.dropout2(msa2)

该模型中，定义了一些比较常规的神经网络层，我们可以着重看一下该神经网络各个层的顺序

输入的参数经过的第一层首先是LayerNorm这个自定义的模型，该模型对输入数据做了一些处理，使用到了torch.nn.Parameter的方法，将不属于经典深度学习网络模型的参数，变成了模型中可迭代训练的参数通常。因为一般来说我们的参数都是一些常见的结构（卷积、全连接等）里面的计算参数。而当我们的网络有一些其他的设计时，会需要一些额外的参数同样很着整个网络的训练进行学习更新，最后得到最优的值，经典的例子有注意力机制中的权重参数和positional embedding等。在该模型的后面就用到了注意力机制模型的参数迭代学习。然后再经过一个多注意力机制模型，然后是第二个LayerNorm模型，将参数转化为可迭代参数，然后经过dropout层防止过拟合现象。

MaskedDirectMultiheadAttention：

上面的模型有一个未提到过的自定义模型，现在对他进行一个详细的介绍：

class MaskedDirectMultiheadAttention(nn.Module):
    def __init__(self, d_in, d_out, heads, d_k=32, dropout=0.1):
        super(MaskedDirectMultiheadAttention, self).__init__()
        self.heads = heads
        self.scaling = 1/math.sqrt(d_k)
        
        self.to_query = nn.Linear(d_in, heads*d_k)
        self.to_key   = nn.Linear(d_in, heads*d_k)
        self.to_value = nn.Linear(d_out, d_out)
        self.to_out   = nn.Linear(d_out, d_out)
        self.dropout = nn.Dropout(dropout, inplace=True)

    def forward(self, query, key, value, mask):
        batch, N, L = value.shape[:3] 

        q = self.to_query(query).view(batch, L, self.heads, -1).permute(0,2,1,3) 
        k = self.to_key(key).view(batch, L, self.heads, -1).permute(0,2,1,3) 
        v = self.to_value(value).view(batch, N, L, self.heads, -1).permute(0,3,1,2,4) 
        
        q = q*self.scaling
        attention = torch.matmul(q, k.transpose(-2, -1)) 
        attention = attention.masked_fill(mask < 0.5, torch.finfo(q.dtype).min)
        attention = F.softmax(attention, dim=-1) 
        attention = self.dropout(attention) 
        

        out = torch.einsum('bhij,bhnjk->bhnik', attention, v) 
        out = out.permute(0,2,3,1,4).contiguous().view(batch, N, L, -1)
        
        out = self.to_out(out)
        return out

Multi-Head Attention是transformer中的一种机制，虽然文中所示模型是一个自定义的模型，我们仍然可以借鉴transformer中的思想来进行学习分析，对于同一个文本，一个Attention获得一个表示空间，如果多个Attention，则可以获得多个不同的表示空间。基于这种想法，就有了Multi-Head Attention。换句话说，Multi-Head Attention为Attention提供了多个“representation subspaces”。因为在每个Attention中，采用不同的Query / Key / Value权重矩阵，每个矩阵都是随机初始化生成的。然后通过训练，将词嵌入投影到不同的“representation subspaces（表示子空间）”中。

Multi-Head Attention与经典的Attention一样，并不是一个独立的结构，自身无法进行训练。Multi-Head Attention也可以堆叠，形成深度结构。应用场景：可以作为文本分类、文本聚类、关系抽取等模型的特征表示部分。

Multi-Head Attention与Self-Attention的关系是：Multi-Head Attention的Attention可以是Self-Attention,当然也可以是经典的Attention。

而在我们的自定义模型中，我们也可以看到，第16~18行使用全连接层投影得到了相应的Query / Key / Value权重矩阵，之后利用这三个权重矩阵进行相应的数据处理，将attention变量从(B, h, L, L)这样的四元组，依次转化为了(B, h, L1, L2)，最后在第三维加了常量变成五维变量。

2.2 FinalBlock模型

在IterativeFeatureExtractor模型中，最后一层调用的是FinalBlock模型，代码详见下方：

# 
self.final = FinalBlock(n_layer=n_layer, d_msa=d_msa, d_pair=d_pair,
                               n_head_msa=n_head_msa, n_head_pair=n_head_pair, r_ff=r_ff,
                               n_resblock=n_resblock, p_drop=p_drop,
                               performer_L_opts=performer_L_opts, performer_N_opts=performer_N_opts,
                               SE3_param=SE3_param)
class FinalBlock(nn.Module):
    def __init__(self, n_layer=1, d_msa=64, d_pair=128, n_head_msa=4, n_head_pair=8, r_ff=4,
                 n_resblock=1, p_drop=0.1, performer_L_opts=None, performer_N_opts=None,
                 SE3_param={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}):
        super(FinalBlock, self).__init__()
        
        self.msa2msa = MSA2MSA(n_layer=n_layer, n_att_head=n_head_msa, n_feat=d_msa,
                               r_ff=r_ff, p_drop=p_drop,
                               performer_N_opts=performer_N_opts,
                               performer_L_opts=performer_L_opts)
        self.msa2pair = MSA2Pair(n_feat=d_msa, n_feat_out=d_pair, n_feat_proj=32,
                                 n_resblock=n_resblock, p_drop=p_drop, n_att_head=n_head_msa)
        self.pair2pair = Pair2Pair(n_layer=n_layer, n_att_head=n_head_pair,
                                   n_feat=d_pair, r_ff=r_ff, p_drop=p_drop,
                                   performer_L_opts=performer_L_opts)
        self.pair2msa = Pair2MSA(n_layer=n_layer, n_att_head=4, 
                                 n_feat_in=d_pair, n_feat_out=d_msa, r_ff=r_ff, p_drop=p_drop)
        self.str2str = Str2Str(d_msa=d_msa, d_pair=d_pair, SE3_param=SE3_param, p_drop=p_drop)
        self.norm_state = LayerNorm(SE3_param['l0_out_features'])
        self.pred_lddt = nn.Linear(SE3_param['l0_out_features'], 1)
    def forward(self, msa, pair, xyz, seq1hot, idx):

        msa, att = self.msa2msa(msa)

        pair = self.msa2pair(msa, pair, att)

        pair = self.pair2pair(pair)
       
        msa = self.pair2msa(pair, msa)

        xyz, state = self.str2str(msa.float(), pair.float(), xyz.float(), seq1hot, idx, top_k=32)
        
        lddt = self.pred_lddt(self.norm_state(state))
        return msa, pair, xyz, lddt.squeeze(-1)

由于该模型中的各层都已经讲过了，那么我们直接看该模型的调用，从上述代码的27~40行我们可以看到，该模型的调用和之前的模型顺序类似：

该模型接受的输入包括初始的MSA数据，，通过MSA2MSA层处理MSA特征，在经过MSA2Pair层利用处理后的MSA更新残基对数据，之后是一个往返过程，在上一层利用MSA数据更新了残基对数据之后，再分别经过Pair2Pair和Pair2MSA两层最终将MSA的参数进行更新。最终是将msa数据和pair数据均返回。

至此，我们就已经分析完了IterativeFeatureExtractor这个模型，在我们最外层的模型RoseTTAFoldModule中，仅剩下最后一个DistanceNetwork模型，见下方

2.3 DistanceNetwork模型

class DistanceNetwork(nn.Module):
    def __init__(self, n_feat, n_block=1, block_type='orig', p_drop=0.0):
        super(DistanceNetwork, self).__init__()
        self.norm = LayerNorm(n_feat)
        self.proj = nn.Linear(n_feat, n_feat)
        self.drop = nn.Dropout(p_drop)
        #
        self.resnet_dist = ResidualNetwork(n_block, n_feat, n_feat, 37, block_type=block_type, p_drop=p_drop)
        self.resnet_omega = ResidualNetwork(n_block, n_feat, n_feat, 37, block_type=block_type, p_drop=p_drop)
        self.resnet_theta = ResidualNetwork(n_block, n_feat, n_feat, 37, block_type=block_type, p_drop=p_drop)
        self.resnet_phi = ResidualNetwork(n_block, n_feat, n_feat, 19, block_type=block_type, p_drop=p_drop)

    def forward(self, x):
        x = self.norm(x)
        x = self.drop(self.proj(x))
        x = x.permute(0,3,1,2).contiguous()
        
        logits_theta = self.resnet_theta(x)
        logits_phi = self.resnet_phi(x)

        x = 0.5 * (x + x.permute(0,1,3,2))
        logits_dist = self.resnet_dist(x)
        logits_omega = self.resnet_omega(x)

        return logits_dist, logits_omega, logits_theta, logits_phi

该模型的主要作用是根据序列对特征预测distance map，使用的是2D ResNet模型。

ResNet的引入：

ResNet是一种残差网络，咱们可以把它理解为一个子网络，这个子网络经过堆叠可以构成一个很深的网络。网络越深，咱们能获取的信息越多，而且特征也越丰富。但是根据实验表明，随着网络的加深，优化效果反而越差，测试数据和训练数据的准确率反而降低了。按道理，给网络叠加更多层，浅层网络的解空间是包含在深层网络的解空间中的，深层网络的解空间至少存在不差于浅层网络的解，因为只需将增加的层变成恒等映射，其他层的权重原封不动copy浅层网络，就可以获得与浅层网络同样的性能。**更好的解明明存在，为什么找不到？找到的反而是更差的解？**显然，这是个优化问题，反映出结构相似的模型，其优化难度是不一样的，且难度的增长并不是线性的，越深的模型越难以优化。**是由于网络的加深会造成梯度爆炸和梯度消失的问题。**目前针对这种现象已经有了解决的方法：对输入数据和中间层的数据进行归一化操作，这种方法可以保证网络在反向传播中采用随机梯度下降（SGD），从而让网络达到收敛。但是，这个方法仅对几十层的网络有用，当网络再往深处走的时候，这种方法就无用武之地了。

为了让更深的网络也能训练出好的效果，何凯明大神提出了一个新的网络结构——ResNet。这个网络结构的想法主要源于VLAD（残差的想法来源）和Highway Network（跳跃连接的想法来源）。

所以，总的来说,ResNet的主要作用是让更深的网络能够避免梯度爆炸和梯度消失。解决的是深度神经网络的“退化”问题。

接着回到我们的模型，调用该模型时接受(B, L, L, C)四维的输入，经过ResNet后返回相应数据。

下面是ResidualNetwork的具体定义代码

class ResidualNetwork(nn.Module):
    def __init__(self, n_block, n_feat_in, n_feat_block, n_feat_out, 
                 dilation=[1,2,4,8], block_type='orig', p_drop=0.15):
        super(ResidualNetwork, self).__init__()


        layer_s = list()
        # project to n_feat_block
        if n_feat_in != n_feat_block:
            layer_s.append(nn.Conv2d(n_feat_in, n_feat_block, 1, bias=False))
            if block_type =='orig': # should acitivate input
                layer_s.append(nn.InstanceNorm2d(n_feat_block, affine=True, eps=1e-6))
                layer_s.append(nn.ELU(inplace=True))

        # add resblocks
        for i_block in range(n_block):
            d = dilation[i_block%len(dilation)]
            if block_type == 'orig':
                res_block = ResBlock2D(n_feat_block, kernel=3, dilation=d, p_drop=p_drop)
            else:
                res_block = ResBlock2D_bottleneck(n_feat_block, kernel=3, dilation=d, p_drop=p_drop)
            layer_s.append(res_block)

        if n_feat_out != n_feat_block:
            # project to n_feat_out
            layer_s.append(nn.Conv2d(n_feat_block, n_feat_out, 1))
        
        self.layer = nn.Sequential(*layer_s)
    
    def forward(self, x):
        output = self.layer(x)
        return output

到此，RoseTTAFoldModule部分结束

3. 总结

本周已经结束了pyrosetta模式的预测分析