AlphaFold2源码解析(7)–模型之Evoformer

这篇文章我们主要药讲解AlphaFold2的Evoformer的代码细节。

Evoformer Stack

该网络有一个双塔结构，在MSA堆栈中具有轴向的自我注意；在Pair堆栈中具有三角形的乘法更新和三角形的自我注意；以及外积平均和注意偏置，以允许堆栈之间的通信。网络的主干由N_block=48个Evoformer块组成。每个块都有一个MSA表示{m_si}和一个Pair表示{z_ij}作为其输入和输出，并通过几层处理它们。每一层的输出都通过一个剩余连接添加到当前的表示中。一些层的输出在被添加之前要经过Dropout。最后的Evoformer块提供了一个高度处理的MSA表示{m_si}和一个Pair表示{z_ij}，其中包含了结构模块\和辅助网络head所需的信息。预测模块也在使用 "单一 "序列表示{s_i}，s_i∈R_cs，cs= 384，i∈{1 . . N_res}。这个单一表示是由MSA表示的第一行的线性投影得出的。

MSA row-wise gated self-attention with pair bias

MSA表征是用连续的门控的行与列的自我注意块来处理的。行向量为残基Pair建立注意力权重，并将来自残基Pair表征的信息整合为一个额外的偏置项。这允许从Pair堆栈到MSA堆栈的信息提取，以鼓励它们之间的一致性。

class MSARowAttentionWithPairBias(hk.Module):bias = (1e9 * (msa_mask - 1.))[:, None, None, :] # (N_seq, 1, 1, N_res)assert len(bias.shape) == 4msa_act = hk.LayerNorm(axis=[-1], create_scale=True, create_offset=True, name='query_norm')(msa_act) # (5120, 84, 64)pair_act = hk.LayerNorm(axis=[-1],create_scale=True,create_offset=True,name='feat_2d_norm')(pair_act) # (84, 84, 128)init_factor = 1. / jnp.sqrt(int(pair_act.shape[-1])) # 初始化因子weights = hk.get_parameter('feat_2d_weights',shape=(pair_act.shape[-1], c.num_head), # 注意力头的个数init=hk.initializers.RandomNormal(stddev=init_factor))nonbatched_bias = jnp.einsum('qkc,ch->hqk', pair_act, weights)  # (84, 84, 128) (128, 8) -> (8, 84, 84) LinearNoBias(LayerNorm(zij))attn_mod = Attention(c, self.global_config, msa_act.shape[-1])msa_act = mapping.inference_subbatch(attn_mod,self.global_config.subbatch_size,batched_args=[msa_act, msa_act, bias],nonbatched_args=[nonbatched_bias],low_memory=not is_training)class Attention(hk.Module):key_dim = self.config.get('key_dim', int(q_data.shape[-1])) # key_dim = 64value_dim = self.config.get('value_dim', int(m_data.shape[-1])) # value_dimnum_head = self.config.num_head # 8assert key_dim % num_head == 0assert value_dim % num_head == 0key_dim = key_dim // num_headvalue_dim = value_dim // num_headq_weights = hk.get_parameter('query_w', shape=(q_data.shape[-1], num_head, key_dim),init=glorot_uniform()) # (64, 8, 8)k_weights = hk.get_parameter('key_w', shape=(m_data.shape[-1], num_head, key_dim),init=glorot_uniform()) # (64, 8, 8)v_weights = hk.get_parameter('value_w', shape=(m_data.shape[-1], num_head, value_dim),init=glorot_uniform()) # (64, 8, 8)q = jnp.einsum('bqa,ahc->bqhc', q_data, q_weights) * key_dim**(-0.5) # (4, 84, 64)  (64, 8, 8) -> (4, 84, 8, 8)k = jnp.einsum('bka,ahc->bkhc', m_data, k_weights) # (4, 84, 8, 8)v = jnp.einsum('bka,ahc->bkhc', m_data, v_weights) # (4, 84, 8, 8)logits = jnp.einsum('bqhc,bkhc->bhqk', q, k) + bias # (4, 84, 8, 8) bqhc (4, 84, 8, 8) bkhc-> (4, 8, 84, 84)bhqkif nonbatched_bias is not None:logits += jnp.expand_dims(nonbatched_bias, axis=0)weights = jax.nn.softmax(logits)weighted_avg = jnp.einsum('bhqk,bkhc->bqhc', weights, v) # (4, 8, 84, 84) (4, 84, 8, 8) -> (4, 84, 8, 8)if self.global_config.zero_init:init = hk.initializers.Constant(0.0)else:init = glorot_uniform()if self.config.gating: ## 添加了门机制 gating_weights = hk.get_parameter('gating_w',shape=(q_data.shape[-1], num_head, value_dim),init=hk.initializers.Constant(0.0))gating_bias = hk.get_parameter('gating_b',shape=(num_head, value_dim),init=hk.initializers.Constant(1.0))gate_values = jnp.einsum('bqc, chv->bqhv', q_data,gating_weights) + gating_bias # (4, 84, 8, 8)gate_values = jax.nn.sigmoid(gate_values) # : g^h_si= sigmoid(Linear(msi))weighted_avg *= gate_values # (4, 84, 8, 8)o_weights = hk.get_parameter('output_w', shape=(num_head, value_dim, self.output_dim),init=init)o_bias = hk.get_parameter('output_b', shape=(self.output_dim,),init=hk.initializers.Constant(0.0)) # (8, 8, 64)output = jnp.einsum('bqhc,hco->bqo', weighted_avg, o_weights) + o_bias  # (4, 84, 64)return output

MSA column-wise gated self-attention

逐列关注让属于同一目标残基的元素交换信息。在这两个关注块中，头的数量N_heads=8，k、q和v的尺寸c=32。

这部分的算法与MSA row-wise gated self-attention with pair bias相似，唯一的区别是输入特征是按照列来取的。
还有需要注意的是 Extract Evoformer Stack 的 column-wise部分使用的是MSAColumn Global Attention,但是Evoformer Stack 的 column-wise部分使用的是MSAColumn Attention

class MSAColumnGlobalAttention(hk.Module):。。。。。。bias = (1e9 * (msa_mask - 1.))[:, None, None, :]assert len(bias.shape) == 4msa_act = hk.LayerNorm(axis=[-1], create_scale=True, create_offset=True, name='query_norm')(msa_act)attn_mod = GlobalAttention(c, self.global_config, msa_act.shape[-1],name='attention')# [N_seq, N_res, 1]msa_mask = jnp.expand_dims(msa_mask, axis=-1)msa_act = mapping.inference_subbatch(attn_mod,self.global_config.subbatch_size,batched_args=[msa_act, msa_act, msa_mask],nonbatched_args=[],low_memory=not is_training)msa_act = jnp.swapaxes(msa_act, -2, -3)return msa_act

MSA transition

在逐行和逐列注意力之后，MSA栈包含一个2层MLP作为过渡层。中间的通道数将原来的通道数扩大了4倍。

class Transition(hk.Module):def __call__(self, act, mask, is_training=True):...........num_intermediate = int(nc * self.config.num_intermediate_factor) # 256mask = jnp.expand_dims(mask, axis=-1) #[N_seq, N_res, 1]act = hk.LayerNorm(axis=[-1],create_scale=True,create_offset=True,name='input_layer_norm')(act)#[N_seq, N_res, 64]transition_module = hk.Sequential([common_modules.Linear(num_intermediate, # 256initializer='relu',name='transition1'), jax.nn.relu,common_modules.Linear(nc, # 64initializer=utils.final_init(self.global_config),name='transition2')])act = mapping.inference_subbatch(transition_module,self.global_config.subbatch_size,batched_args=[act],nonbatched_args=[],low_memory=not is_training)return act #[N_seq, N_res, 64]

Outer product mean

"外积平均 "块将MSA表示转换为Pair表示的更新。所有的MSA条目都通过两个独立的线性变换被线性地投射到一个较小的维度c=32。来自两列i和j的这些向量的外积在序列上被平均化，并被投射到维度c_z，以获得Pair表示中的条目ij的更新。

class OuterProductMean(hk.Module):"""Computes mean outer product.
计算平均外积。
"""def __call__(self, act, mask, is_training=True):mask = mask[..., None]act = hk.LayerNorm([-1], True, True, name='layer_norm_input')(act) #(N_seq, N_res, hz)left_act = mask * common_modules.Linear(c.num_outer_channel, #32initializer='linear',name='left_projection')(act)right_act = mask * common_modules.Linear(c.num_outer_channel, #32initializer='linear',name='right_projection')(act)if gc.zero_init:init_w = hk.initializers.Constant(0.0)else:init_w = hk.initializers.VarianceScaling(scale=2., mode='fan_in')output_w = hk.get_parameter('output_w',shape=(c.num_outer_channel, c.num_outer_channel,self.num_output_channel),init=init_w)output_b = hk.get_parameter('output_b', shape=(self.num_output_channel,),init=hk.initializers.Constant(0.0))def compute_chunk(left_act):# This is equivalent to## act = jnp.einsum('abc,ade->dceb', left_act, right_act)# act = jnp.einsum('dceb,cef->bdf', act, output_w) + output_b## but faster.left_act = jnp.transpose(left_act, [0, 2, 1]) # (N_seq, N_res, hz)-> (N_seq, hz, N_res) act = jnp.einsum('acb,ade->dceb', left_act, right_act) ##(N_seq, hz, N_res) acd  [N_seq, N_res, hz]ade -> [N_res, hz, hz, N_res]act = jnp.einsum('dceb,cef->dbf', act, output_w) + output_b # [N_res, hz, hz, N_res] dceb [hz, hz, ohz] ->[N_res,N_res, ohz]return jnp.transpose(act, [1, 0, 2])act = mapping.inference_subbatch(compute_chunk,c.chunk_size,batched_args=[left_act],nonbatched_args=[],low_memory=True,input_subbatch_dim=1,output_subbatch_dim=0)epsilon = 1e-3norm = jnp.einsum('abc,adc->bdc', mask, mask)act /= epsilon + normreturn act # (N_res, N_res, 128)

Triangular multiplicative update

三角形乘法更新通过结合每个三角形图边ij、ik和jk内的信息来更新Evoformer块中的Pair表示。每条边ij从所有三角形的其他两条边接收更新，其中涉及到它。有两个对称的版本，一个用于 "传出 "的边，一个用于 "传入 "的边。两者的区别以黄色标出。

这里的Al 11算法与AL 12不同之处是一个按行，一个按列

class TriangleMultiplication(hk.Module):def __call__(self, act, mask, is_training=True):"""Builds TriangleMultiplication module."""........	mask = mask[..., None]act = hk.LayerNorm(axis=[-1], create_scale=True, create_offset=True, name='layer_norm_input')(act)input_act = act # (N_res, N_res, 128)left_projection = common_modules.Linear(c.num_intermediate_channel, # 128name='left_projection')left_proj_act = mask * left_projection(act) # (N_res, N_res, 128)right_projection = common_modules.Linear(c.num_intermediate_channel, name='right_projection')right_proj_act = mask * right_projection(act) # (N_res, N_res, 128)left_gate_values = jax.nn.sigmoid(common_modules.Linear( c.num_intermediate_channel, bias_init=1.,initializer=utils.final_init(gc), name='left_gate')(act)) # (N_res, N_res, 128)right_gate_values = jax.nn.sigmoid(common_modules.Linear(c.num_intermediate_channel, bias_init=1.,initializer=utils.final_init(gc), name='right_gate')(act)) # (N_res, N_res, 128)left_proj_act *= left_gate_valuesright_proj_act *= right_gate_valuesact = jnp.einsum(c.equation, left_proj_act, right_proj_act) # (N_res, N_res, 128)act = hk.LayerNorm(axis=[-1], create_scale=True, create_offset=True, name='center_layer_norm')(act)# (N_res, N_res, 128)output_channel = int(input_act.shape[-1])act = common_modules.Linear(output_channel, initializer=utils.final_init(gc), name='output_projection')(act) # (N_res, N_res, 128)gate_values = jax.nn.sigmoid(common_modules.Linear(output_channel, bias_init=1.,initializer=utils.final_init(gc), name='gating_linear')(input_act))act *= gate_valuesreturn act # (N_res, N_res, 128)

Triangular self-attention

三角形自注意力更新Evoformer块中的Pair表示。起始节点 "版本用来自共享同一起始节点i的所有边的值来更新边ij。决定边ij是否会收到来自边ik的更新不仅由它们的查询键相似度决定，而且还由从这个三角形的第三条边jk衍生的信息bjk调制。此外，我们还用一个从边ij衍生出来的额外的门控gij来扩展更新。这个模块的对称Pair在结束节点周围的边上操作。差异以黄色显示。

class TriangleAttention(hk.Module):"""Triangle Attention."""def __call__(self, pair_act, pair_mask, is_training=False):.......bias = (1e9 * (pair_mask - 1.))[:, None, None, :]pair_act = hk.LayerNorm(axis=[-1], create_scale=True, create_offset=True, name='query_norm')(pair_act) # (N_res, N_res, 128)init_factor = 1. / jnp.sqrt(int(pair_act.shape[-1]))weights = hk.get_parameter('feat_2d_weights', shape=(pair_act.shape[-1], c.num_head),init=hk.initializers.RandomNormal(stddev=init_factor))nonbatched_bias = jnp.einsum('qkc,ch->hqk', pair_act, weights)attn_mod = Attention(c, self.global_config, pair_act.shape[-1])pair_act = mapping.inference_subbatch( attn_mod, self.global_config.subbatch_size,batched_args=[pair_act, pair_act, bias], nonbatched_args=[nonbatched_bias], low_memory=not is_training)return pair_act # (N_res, N_res, 128)

Transition in the pair stack

Pair堆栈中的过渡层相当于MSA堆栈中的过渡层：一个2层MLP，中间的通道数将原来的通道数扩大4倍。

    act = hk.LayerNorm(axis=[-1],create_scale=True,create_offset=True,name='input_layer_norm')(act)#[N_seq, N_res, 64]transition_module = hk.Sequential([common_modules.Linear(num_intermediate, # 256initializer='relu',name='transition1'), jax.nn.relu,common_modules.Linear(nc, # 64initializer=utils.final_init(self.global_config),name='transition2')])act = mapping.inference_subbatch(transition_module,self.global_config.subbatch_size,batched_args=[act],nonbatched_args=[],low_memory=not is_training)return act #[N_seq, N_res, 64]

Additional inputs

Template stack

成对的模板特征进行线性投影，形成初始模板表征t_stij，t_stij∈R^ct，ct=64，i，j∈{1 . . N_res}，s_t∈{1 . . N_templ}。每个模板表示都是用模板Pair堆栈独立处理的，所有可训练的参数在模板间共享。
输出的表征与模板点对点注意力汇总，其中Pair表征{z_ij}被用来形成查询，并注意力各个模板。这个模块的输出被添加到Pair表征中。

此外，模板扭转角的特征被嵌入一个小的MLP，并与MSA表征相连接，作为额外的序列行。这些额外的行参与所有的MSA堆栈操作，但不参与掩盖的MSA损失。虽然模板扭转角和MSA特征在概念上是不同的量，但它们被嵌入了不同的权重集，因此，学习过程大概会促使嵌入具有可比性，因为它们是由相同的下游模块以相同的权重处理。

if c.template.enabled: # 是否使用模版template_batch = {k: batch[k] for k in batch if k.startswith('template_')}template_pair_representation = TemplateEmbedding(c.template, gc)(pair_activations,template_batch,mask_2d,is_training=is_training)

Unclustered MSA stack

未聚类的MSA序列特征被线性投影以形成初始表征{e_s_{e}i}，e_s_{e}i∈Rce，c_e=64，s_e∈{1 . . . N_extra_seq}，i∈{1 . . N_res}。这些表示是用包含4个块的Extra MSA堆栈处理的。它们与主要的Evoformer块高度相似，显著的区别是使用了全局的列式自我关注和较小的表示尺寸，以允许处理大量的序列。最终的Pair表征被用作主Evoformer堆栈的输入，而最终的MSA激活是未使用的。


该模块的细节与Evoformer Stack一致，都是使用EvoformerIteration类来做为模型，具体入口如下：

  extra_msa_feat = create_extra_msa_feature(batch) #(5120, N_res, 25)extra_msa_activations = common_modules.Linear(c.extra_msa_channel,name='extra_msa_activations')(extra_msa_feat) # (5120, N_res, 64) # Extra MSA Stack.# Jumper et al. (2021) Suppl. Alg. 18 "ExtraMsaStack"extra_msa_stack_input = {'msa': extra_msa_activations, # (5120, N_res, 64)'pair': pair_activations, # (N_res, N_res, 128)}extra_msa_stack_iteration = EvoformerIteration(c.evoformer, gc, is_extra_msa=True, name='extra_msa_stack')