AF3 AtomAttentionEncoder类的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