网站制作主要公司,哪个软件购物便宜又好,推广平台哪个好,站长工具怎么关闭AlphaFold3 的 AtomAttentionEncoder 类中,init_pair_repr 方法方法负责为原子之间的关系计算成对表示(pair representation),这是原子转变器(atom transformer)模型的关键组成部分,直接影响对蛋白质/分子相互作用的建模。
init_pair_repr源代码: def init_pair_repr(…AlphaFold3 的AtomAttentionEncoder类中,init_pair_repr方法方法负责为原子之间的关系计算成对表示(pair representation),这是原子转变器(atom transformer)模型的关键组成部分,直接影响对蛋白质/分子相互作用的建模。
init_pair_repr源代码: def init_pair_repr(self,features: Dict[str, Tensor],atom_cond: Tensor,z_trunk: Optional[Tensor],) - Tensor:"""Compute the pair representation for the atom transformer.This is done in a separate function for checkpointing. The intermediate activations due to theatom pair representations are large and can be checkpointed to reduce memory usage.Args:features:Dictionary of input features.atom_cond:[bs, n_atoms, c_atom] The single atom conditioning from init_single_reprz_trunk:[bs, n_tokens, n_tokens, c_trunk] the pair representation from the trunkReturns:[bs, n_atoms // n_queries, n_queries, n_keys, c_atompair] The pair representation"""# Compute offsets between atom reference positionsa = partition_tensor(features['ref_pos'], self.n_queries, self.n_queries) # (bs, n_atoms // 32, 32, 3)b = partition_tensor(features['ref_pos'], self.n_queries, self.n_keys) # (bs, n_atoms // 32, 128, 3)offsets = a[:, :, :, None, :] - b[:, :, None, :, :] # (bs, n_atoms // 32, 32, 128, 3)# Compute the valid maskref_space_uid = features['ref_space_uid'].unsqueeze(-1) # (bs, n_atoms, 1)a = partition_tensor(ref_space_uid, self.n_queries, self.n_queries) # (bs, n_atoms // 32, 32)b = partition_tensor(ref_space_uid, self.n_queries, self.n_keys) # (bs, n_atoms // 32, 128)valid_mask = a[:, :, :, None] == b[:, :, None, :] # (bs, n_atoms // 32, 32, 128, 1)valid_mask = valid_mask.to(offsets.dtype) # convert boolean to binary# Embed the atom offsets and the valid masklocal_atom_pair = self.linear_atom_offsets(offsets) * valid_mask# Embed pairwise inverse squared distances, and the valid masksquared_distances = offsets.pow(2).sum(dim=-1, keepdim=True) # (bs, n_atoms // 32, 32, 128, 1)inverse_dists = torch.reciprocal(torch.add(squared_distances, 1))local_atom_pair = local_atom_pair + self.linear_atom_distances(inverse_dists) * valid_masklocal_atom_pair = local_atom_pair + self.linear_mask(valid_mask) * valid_mask# If provided, add trunk embeddingsif self.trunk_conditioning:local_atom_pair = local_atom_pair + map_token_pairs_to_local_atom_pairs(self.proj_trunk_pair(z_trunk),features['atom_to_token'])# Add the combined single conditioning to the pair representationa = partition_tensor(self.linear_single_to_pair_row(F.relu(atom_cond)), self.n_queries, self.n_queries)b = partition_tensor(self.linear_single_to_pair_col(F.relu(atom_cond)), self.n_queries, self.n_keys)local_atom_pair = local_atom_pair + (a[:, :, :, None, :] + b[:, :, None, :, :])# Run a small MLP on the pair activationslocal_atom_pair = self.pair_mlp(local_atom_pair)return local_atom_pair
init_pair_repr代码解读:
1. 函数定义与注释
def init_pair_repr(self,features: Dict[str, Tensor],atom_cond: Tensor,z_trunk: Optional[Tensor],
) - Tensor:"""Compute the pair representation for the atom transformer.Args:features: Dictionary of input features.atom_cond: [bs, n_atoms, c_atom] The single atom conditioning from init_single_reprz_trunk: [bs, n_tokens, n_tokens, c_trunk] the pair representation from the trunkReturns:[bs, n_atoms // n_queries, n_queries, n_keys, c_atompair] The pair representation"""功能描述: 方法用于计算原子之间的成对表示(pair representation),描述原子对之间的相互关系。通过输入特征和条件化单原子表示(atom_cond)生成成对表示。如果有 trunk 模块输出(z_trunk),进一步将其纳入建模。 输入参数: features: 包含输入原子特征的字典,例如参考位置、掩码等。atom_cond: 由init_single_repr生成的单原子条件表示,提供单原子特征。z_trunk: 可选的 trunk 模块输出,用于加入全局上下文信息。 输出: 返回形状为[bs, n_atoms // n_queries, n_queries, n_keys, c_atompair]的成对表示张量。2. 计算原子间的位移偏移量
a = partition_tensor(features['ref_pos'], self.n_queries, self.n_queries) # (bs, n_atoms // 32, 32, 3)
b = partition_tensor(features['ref_pos'], self.n_queries, self.n_keys) # (bs, n_atoms // 32, 128, 3)
offsets = a[:, :, :, None, :] - b[:, :, None, :, :] # (bs, n_atoms // 32, 32, 128, 3)功能: 通过分块操作,将原子的三维参考位置(ref_pos)分为 query 和 key 的两个集合,计算原子对的位移向量offsets。理论基础: 原子间的位移向量是物理意义上的距离关系的基础,直接影响距离计算和相互作用建模。细节: partition_tensor将输入张量按块划分,便于后续处理。offsets形状为[bs, n_atoms // n_queries, n_queries, n_keys, 3]。原理解读:
什么是features['ref_pos']?
features['ref_pos']是原子在 3D 空间中的参考坐标,形状为(bs, n_atoms, 3)。 bs是批量大小(batch size)。n_atoms是蛋白质中的原子数量。每个原子的坐标由 3 个值(x, y, z)表示。为什么使用partition_tensor?
partition_tensor将输入张量按滑动窗口分区,使得可以对局部子集进行高效计算。作用:通过滑动窗口对原子的参考坐标进行局部划分: 第一次划分a:窗口大小为n_queries,滑动步长为n_queries,即每次取 32 个原子的局部坐标。第二次划分b:窗口大小为n_keys,滑动步长为n_queries,即每次取 128 个原子的局部坐标。分区后的结果: a:形状为(bs, n_atoms // 32, 32, 3),表示每个滑动窗口内的原子局部坐标(32 个)。b:形状为(bs, n_atoms // 32, 128, 3),表示每个滑动窗口内的原子扩展区域(128 个)。为什么计算offset