AF3 shaped_categorical函数解读

AlphaFold3  data_transforms 模块的 shaped_categorical 函数的作用是根据给定的概率分布 probs 生成随机样本。这个函数主要用于在掩码语言模型（Masked Language Model，MLM）中对 MSA 数据进行随机替换。

源代码：

def shaped_categorical(probs, epsilon=1e-10):
    ds = probs.shape
    num_classes = ds[-1]
    distribution = torch.distributions.categorical.Categorical(
        torch.reshape(probs + epsilon, [-1, num_classes])
    )
    counts = distribution.sample()
    return torch.reshape(counts, ds[:-1])

源码解读：

函数签名

def shaped_categorical(probs, epsilon=1e-10):

• probs：输入的概率分布张量，形状为 [batch_size, sequence_length, num_classes]。

• epsilon：一个小的常数，用于避免概率为零的情况，确保数值稳定性。默认值为 1e-10。

1. 获取概率分布的形状和类别数

ds = probs.shape
num_classes = ds[-1]

• 获取输入概率分布 probs 的形状 ds。

• 获取最后一维的大小 num_classes，即类别数。

2. 创建分类分布对象

distribution = torch.distributions.categorical.Categorical(
    torch.reshape(probs + epsilon, [-1, num_classes])
)

• 将概率分布 probs 展平为二维张量 [batch_size * sequence_length, num_classes]。

• 在概率分布中加入一个小的常数 epsilon，以避免概率为零的情况。

• 使用 torch.distributions.categorical.Categorical 创建一个分类分布对象 distribution。

3. 采样并恢复原始形状

counts = distribution.sample()
return torch.reshape(counts, ds[:-1])

• 使用 distribution.sample() 从分类分布中生成随机样本。生成的样本形状为 [batch_size * sequence_length]。

• 将生成的样本恢复为原始形状 [batch_size, sequence_length]。

• ds[:-1] 表示取 ds 中从第一个元素到最后一个元素之前的切片，即去掉最后一个元素。

函数输出

• 函数返回一个形状为 [batch_size, sequence_length] 的张量，其中每个位置的值是从对应的概率分布中采样得到的类别索引。

---

总结

• shaped_categorical 函数的作用：根据给定的概率分布生成随机样本，并确保样本的形状与输入概率分布的形状一致。

• 核心逻辑：

  1. 获取概率分布的形状和类别数。

  2. 将概率分布展平为二维张量，并加入一个小的常数以避免概率为零。

  3. 使用 torch.distributions.categorical.Categorical 创建分类分布对象。

  4. 从分类分布中生成随机样本，并将样本恢复为原始形状。

• 输出：一个形状为 [batch_size, sequence_length] 的张量，表示从概率分布中采样得到的类别索引。

这个函数在处理掩码语言模型（MLM）任务时非常关键，它用于生成掩码位置的随机替换值。