AlphaFold2代码阅读（三）

2021SC@SDUSC

---

前言

看了Evoformer的代码，一如既往地晦涩难懂，所以在看代码的同时阅读了论文，果然感觉好了很多。

---

一、class EvoformerIteration

1.代码

class EvoformerIteration(hk.Module):
 
  def __init__(self, config, global_config, is_extra_msa,
               name='evoformer_iteration'):
    super().__init__(name=name)
    self.config = config
    self.global_config = global_config
    self.is_extra_msa = is_extra_msa

  def __call__(self, activations, masks, is_training=True, safe_key=None):
 
    c = self.config
    gc = self.global_config

    msa_act, pair_act = activations['msa'], activations['pair']

    if safe_key is None:
      safe_key = prng.SafeKey(hk.next_rng_key())

    msa_mask, pair_mask = masks['msa'], masks['pair']

    dropout_wrapper_fn = functools.partial(
        dropout_wrapper,
        is_training=is_training,
        global_config=gc)

    safe_key, *sub_keys = safe_key.split(10)
    sub_keys = iter(sub_keys)

    msa_act = dropout_wrapper_fn(
        MSARowAttentionWithPairBias(
            c.msa_row_attention_with_pair_bias, gc,
            name='msa_row_attention_with_pair_bias'),
        msa_act,
        msa_mask,
        safe_key=next(sub_keys),
        pair_act=pair_act)

    if not self.is_extra_msa:
      attn_mod = MSAColumnAttention(
          c.msa_column_attention, gc, name='msa_column_attention')
    else:
      attn_mod = MSAColumnGlobalAttention(
          c.msa_column_attention, gc, name='msa_column_global_attention')
    msa_act = dropout_wrapper_fn(
        attn_mod,
        msa_act,
        msa_mask,
        safe_key=next(sub_keys))

    msa_act = dropout_wrapper_fn(
        Transition(c.msa_transition, gc, name='msa_transition'),
        msa_act,
        msa_mask,
        safe_key=next(sub_keys))

    pair_act = dropout_wrapper_fn(
        OuterProductMean(
            config=c.outer_product_mean,
            global_config=self.global_config,
            num_output_channel=int(pair_act.shape[-1]),
            name='outer_product_mean'),
        msa_act,
        msa_mask,
        safe_key=next(sub_keys),
        output_act=pair_act)

    pair_act = dropout_wrapper_fn(
        TriangleMultiplication(c.triangle_multiplication_outgoing, gc,
                               name='triangle_multiplication_outgoing'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))
    pair_act = dropout_wrapper_fn(
        TriangleMultiplication(c.triangle_multiplication_incoming, gc,
                               name='triangle_multiplication_incoming'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))

    pair_act = dropout_wrapper_fn(
        TriangleAttention(c.triangle_attention_starting_node, gc,
                          name='triangle_attention_starting_node'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))
    pair_act = dropout_wrapper_fn(
        TriangleAttention(c.triangle_attention_ending_node, gc,
                          name='triangle_attention_ending_node'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))

    pair_act = dropout_wrapper_fn(
        Transition(c.pair_transition, gc, name='pair_transition'),
        pair_act,
        pair_mask,
        safe_key=next(sub_keys))

    return {'msa': msa_act, 'pair': pair_act}

  看完这段代码，知道这段代码主要是实现了一个Evoformer Block，这个Evoformer Block应该会迭代很多次。
其中__call__函数的参数如下：
参数:
activations: 包含activations的字典:
msa: MSA 激活, shape是[N_seq, N_res, c_m].
pair: pair 激活, shape是[N_res, N_res, c_z].
masks:masks的字典:
msa: MSA mask, shape 是[N_seq, N_res].
pair: pair mask, shape是 [N_res, N_res].
is_training:模块是否处于培训模式
safe_key: prng.SafeKey封装rng密钥
返回值:
输出（outputs）

为了能更深入的理解这部分代码，所以我要结合论文理解了。

2.理解代码

以下是我阅读了这部分论文后做出的总结。首先要从整体的网络结构说起。

网络结构

  AlphaFold2通过结合基于蛋白质结构的进化、物理和几何约束，开发了新型神经网络架构和训练方法，大大提高了结构预测的准确性，其架构如下图所示：

  该网络包括两个主要阶段，即主干模块和结构模块。主干模块由重复堆叠的Evoformer模块构成，以生成已处理的MSA和残基对（配对的氨基酸）的表示。
  其中MSA是指多重序列比对，表示对一组蛋白质序列的演化关系、同源关系的分析结果，例如氨基酸的突变情况，通常用于揭示蛋白质二级结构与三级结构甚至个别氨基酸的保守性。

Evoformer模块

  Evoformer 模块包含了许多新颖的基于注意力和非基于注意力的组件。Evoformer的关键创新是提出了在 MSA 内交换信息的新机制，并且其配对表示允许直接推理空间和进化关系。主干模块之后是结构模块。该模块的输出是蛋白质的每个残基的旋转和平移信息，从而引入了明确的 3-D 结构。这些表示以所有旋转角度和所有位置坐标进行初始化，并能快速收敛至具有精确原子细节的高度准确的蛋白质结构。

这一模块的关键创新包括：

（1）通过打破链原子结构，允许同时对结构的所有部分进行局部细化。

（2）提出一种新颖的等变Transformer架构，允许网络隐式推理未表示的侧链原子。

（3）提出一个损失项，给残基的方向正确性提供关键权重。

在结构模块和整个网络中，通过重复将最终损失应用于输出，然后将输出递归地馈送到相同模块来强化迭代细化。使用整个网络的迭代细化，显着提高了准确性，而额外的训练时间很少。

Evoformer模块关键原理

  Evoformer模块的关键原理是将蛋白质结构预测视为 3-D 空间中的图推理问题，其中图的边缘由邻近的残基定义，配对表示（以数组表示）的元素则编码有关残基之间关系的信息，如下图所示：
  对于MSA 表示数组，其列编码输入序列的各个残基，而行表示这些残基出现的序列。在这个框架内，定义了许多更新操作，这些更新操作应用于每个块中，其中不同的更新操作被串联应用。与之前的工作不同，这些操作将持续应用于每个模块中，而不是只在网络中应用一次，这使得从不断演化的 MSA 表示和配对表示之间可以进行连续通信。

预测过程解释

  AlphaFold2 是如何预测蛋白质结构的呢？AlphaFold2为网络中的 48 个 Evoformer 模块中各自单独训练了一个结构模块，同时保持主网络的所有参数保持不变。在前几个模块之后产生的轨迹非常平滑，表明 AlphaFold2 对蛋白质结构进行了不断的增量改进，直到它不能再改进。

---

总结

结合论文看代码，果然对EvoformerIteration的理解更加深入了。
随着对AlphaFold2理解的越来越多，我觉得我应该系统的学习一下深度学习。