SE(3)-Transformers 是在NeurIPS 2020上推出的多功能图形神经网络。 NVIDIA 刚刚发布了一款开源优化实现,它使用的内存比基线正式实施少9倍,速度比基线正式实施快21倍。
SE(3)-Transformer 在处理几何对称性问题时非常有用,如小分子处理、蛋白质精制或点云应用。它们可以是更大的药物发现模型的一部分,如RoseTTAFold和此 AlphaFold2 的复制。它们也可以用作点云分类和分子性质预测的独立网络(图 1 )。
在/PyTorch/DrugDiscovery/SE3Transformer存储库中, NVIDIA 提供了在QM9 数据集上为分子性质预测任务训练优化模型的方法。 QM9 数据集包含超过 10 万个有机小分子和相关的量子化学性质。
训练吞吐量提高 21 倍
与基线实施相比, NVIDIA 实现提供了更快的训练和推理。该实现对 SE(3)-Transformers 的核心组件,即张量场网络( TFN )以及图形中的自我注意机制进行了优化。
考虑到注意力层超参数的某些条件得到满足,这些优化大多采取操作融合的形式。
由于这些,与基线实施相比,训练吞吐量增加了 21 倍,利用了最近 GPU NVIDIA 上的张量核。
此外, NVIDIA 实现允许使用多个 GPU 以数据并行方式训练模型,充分利用 DGX A100 ( 8x A100 80GB )的计算能力。
把所有东西放在一起,在 NVIDIA DGX A100 上, SE(3)-Transformer现在可以在 QM9 数据集上在 27 分钟内进行训练。作为比较,原始论文的作者指出,培训在硬件上花费了 2 . 5 天( NVIDIA GeForce GTX 1080 Ti )。
更快的培训使您能够在搜索最佳体系结构的过程中快速迭代。随着内存使用率的降低,您现在可以训练具有更多注意层或隐藏通道的更大模型,并向模型提供更大的输入。
内存占用率降低 9 倍
SE(3)-Transformer 是已知的记忆重模型,这意味着喂养大输入,如大蛋白质或许多分批小分子是一项挑战。对于 GPU 内存有限的用户来说,这是一个瓶颈。
这一点在DeepLearningExamples上的 NVIDIA 实现中已经改变。图 3 显示,由于 NVIDIA 优化和对混合精度的支持,与基线实现相比,训练内存使用减少了 9 倍。
除了对单精度和混合精度进行改进外,还提供了低内存模式。启用此标志后,模型在 TF32 ( NVIDIA 安培体系结构)或 FP16 ( NVIDIA 安培体系结构、 NVIDIA 图灵体系结构和 NVIDIA 伏特体系结构)精度上运行,模型将切换到以吞吐量换取额外内存节省的模式。
实际上,在具有 V100 32-GB GPU 的 QM9 数据集上,基线实现可以在内存耗尽之前扩展到 100 的批大小。 NVIDIA 实现每批最多可容纳 1000 个分子(混合精度,低内存模式)。
对于处理以氨基酸残基为节点的蛋白质的研究人员来说,这意味着你可以输入更长的序列并增加每个残基的感受野。
SE(3)-Transformers 优化
与基线相比, NVIDIA 实现提供了一些优化。有关更多信息,请参阅/PyTorch/DrugDiscovery/SE3Transformer存储库中的源代码和文档。
融合键与值计算
在“自我注意”层中,将计算关键帧、查询和值张量。查询是图形节点特征,是输入特征的线性投影。另一方面,键和值是图形边缘特征。它们是使用 TFN 层计算的。这是 SE(3)-Transformer 中大多数计算发生的地方,也是大多数参数存在的地方。
基线实现使用两个独立的 TFN 层来计算键和值。在 NVIDIA 实现中,这些被融合在一个 TFN 中,通道数量增加了一倍。这将启动的小型 CUDA 内核数量减少一半,并更好地利用 GPU 并行性。径向轮廓是 TFN 内部完全连接的网络,也与此优化融合。概览如图 4 所示。
TFN 合并
SE(3)-Transformer 内部的功能除了其通道数量外,还有一个degree,它是一个正整数。程度特征有维度. TFN 接受不同程度的特征,使用张量积组合它们,并输出不同程度的特征。
对于输入为 4 度、输出为 4 度的图层,将考虑所有度的组合:理论上,必须计算 4 × 4 = 16 个子图层。
这些子层称为成对 TFN 卷积。图 5 显示了所涉及的子层的概述,以及每个子层的输入和输出维度。对给定输出度(列)的贡献相加,以获得最终特征。
NVIDIA 在满足 TFN 层上的某些条件时,提供多级融合以加速这些卷积。通过创建尺寸为 16 倍的形状,熔合层可以更有效地使用张量核。以下是应用熔合卷积的三种情况:
- 输出功能具有相同数量的通道
- 输入功能具有相同数量的通道
- 这两种情况都是正确的
第一种情况是,所有输出特征具有相同数量的通道,并且输出度数的范围从 0 到最大度数。在这种情况下,使用输出融合特征的融合卷积。该融合层用于 SE(3)-Transformers 的第一个 TFN 层。
第二种情况是,所有输入特征具有相同数量的通道,并且输入度数的范围从 0 到最大度数。在这种情况下,使用对融合输入特征进行操作的融合卷积。该融合层用于 SE(3)-Transformers 的最后一层 TFN 。
在最后一种情况下,当两个条件都满足时,使用完全融合的卷积。这些卷积作为输入融合特征,输出融合特征。这意味着每个 TFN 层只需要一个子层。内部 TFN 层使用此融合级别。
基预计算
除了输入节点特性外, TFN 还需要基矩阵作为输入。每个图边都有一组矩阵,这些矩阵取决于目标节点和源节点之间的相对位置。
在基线实现中,这些矩阵在前向传递开始时计算,并在所有 TFN 层中共享。它们依赖于球形 h ARM ,计算起来可能很昂贵。由于输入图不会随着 QM9 数据集而改变(没有数据扩充,没有迭代位置细化),这就引入了跨时代的冗余计算。
NVIDIA 实现提供了在培训开始时预计算这些基础的选项。整个数据集迭代一次,基缓存在 RAM 中。前向传递开始时的计算基数过程被更快的 CPU 到 GPU 内存拷贝所取代。
结论
我鼓励您在/PyTorch/DrugDiscovery/SE3Transformer NVIDIA GitHub 存储库中检查 SE ( 3 ) – transformer 模型的实现。在评论中,分享您计划如何采用和扩展此项目。