TurboPairFormer：具有优化的三角形注意内核的快速稳定的蛋白质折叠模型训练

三角注意力是 AlphaFold3 式生物分子模型中的核心计算，其立方成本为token计数。它的共享对偏差在通常的查询和键减少的基础上增加了注意力切片上的梯度减少。我们检查的开源后端通过重复概率重新计算、浮点原子或满分梯度存储来处理这些减少。另外，根据 BF16 舍入的前向输出计算 softmax 后向校正会丢失数值精度。我们推出了 TurboPairFormer，这是一种针对 NVIDIA Hopper GPU 的三角注意力实现，可以解决这两个问题。我们的键瓦片并行后向算法针对查询、键、值和成对偏差梯度重新计算每个概率瓦片一次，使用有序部分缩减进行确定性累积，无需浮点原子或完整分数梯度存储。输出残差补偿保留输出舍入残差的 BF16 近似值，以便在 FP32 中更准确地计算后向校正，而无需更改 BF16 输出。使用裁剪尺寸为 384、640 和 768 以及头部尺寸为 16 和 32 的 BF16 输入，TurboPairFormer 在本文比较的实现中相对于 FP64 参考实现了最低的平均查询、键和对偏差梯度 RMSE。在受控消融中，残余补偿将这些 RMSE 值降低了 28-47%。在固定执行条件下，所有 600 个输入案例中的 5 次重复调用中，所有四个梯度都是按位相同的。通过我们的三角形乘法内核集成到 OpenFold3 中，TurboPairFormer 在 16 个 H100 GPU 上评估的后端配置中实现了每个优化器步骤的最低 GPU 计算时间，在裁剪尺寸 768 时，$1.73\times$ 比 OpenFold3 的 Triton 后端加速，$1.13\times$ 比 cuEquivariance 加速。