AF3 make_msa_mask函数解读

qq_273900232025-03-11 19:01

AlphaFold3 data_transforms 模块的 make_msa_mask 函数主要是 为 MSA (多序列比对) 生成初始化的 mask 掩码特征： msa_mask和 msa_row_mask。用于后续的特征处理，确保模型在处理 MSA 时能够识别哪些数据是有效的，哪些需要被忽略（比如填充的 0）。

源代码：

复制代码

def make_msa_mask(protein):
    """Mask features are all ones, but will later be zero-padded."""
    protein["msa_mask"] = torch.ones(protein["msa"].shape, dtype=torch.float32)
    protein["msa_row_mask"] = torch.ones(
        (protein["msa"].shape[0]), dtype=torch.float32
    )
    return protein

源码解读：

复制代码

protein["msa_mask"] = torch.ones(protein["msa"].shape, dtype=torch.float32)

protein["msa"] 是一个 (num_seq, num_res) 形状的张量：
- num_seq：MSA 的序列数。
- num_res：每条 MSA 序列的氨基酸个数。
torch.ones(protein["msa"].shape, dtype=torch.float32) 生成一个全 1 的掩码：
- 形状：(num_seq, num_res)
- 作用：指示哪些 MSA 位置是有效的。
- 后续可能的变化 ：
  - 在 MSA 被 zero-padding（填充 0） 之后，填充部分的 msa_mask 可能会被置为 0。
  protein["msa_row_mask"] = torch.ones(
  (protein["msa"].shape[0]), dtype=torch.float32
  )
protein["msa_row_mask"] 是一个 (num_seq,) 形状的 1D 张量：
- 作用：指示哪些 MSA 序列是有效的。
- 后续可能的变化 ：
  - 在 MSA 被 zero-padding 或者裁剪 之后，可能会有 0 置入其中。

代码意义

目的

msa_mask ：用于 逐个氨基酸 位置的掩码，确保后续计算时忽略填充区域。
msa_row_mask ：用于 整条 MSA 序列 的掩码，可能用于批处理时过滤无效序列。

作用

确保 MSA 计算时不会受 填充值 (zero-padding) 影响。
在后续的 extra_msa_mask 计算、聚类、特征提取时，利用 msa_mask 来过滤无效数据。
在 DataPipeline 处理过程中，MSA 会被填充到固定长度，msa_mask 可以帮助 标记填充区域。

代码关系

变量	形状	作用
`msa_mask`	`(num_seq, num_res)`	逐个氨基酸位置的掩码，标记 MSA 位置是否有效
`msa_row_mask`	`(num_seq,)`	逐行掩码，标记 MSA 序列是否有效

最终，make_msa_mask 主要是 初始化 MSA 掩码，确保后续处理可以正确忽略填充区域 ，在 AlphaFold3 处理中起到 屏蔽无效数据 的作用。