SE（3）-Transformers 是在NeurIPS 2020上推出的多功能圖形神經(jīng)網(wǎng)絡(luò)。 NVIDIA 剛剛發(fā)布了一款開(kāi)源優(yōu)化實(shí)現(xiàn)，它使用的內(nèi)存比基線正式實(shí)施少9倍，速度比基線正式實(shí)施快21倍。

SE（3）-Transformer 在處理幾何對(duì)稱(chēng)性問(wèn)題時(shí)非常有用，如小分子處理、蛋白質(zhì)精制或點(diǎn)云應(yīng)用。它們可以是更大的藥物發(fā)現(xiàn)模型的一部分，如RoseTTAFold和此 AlphaFold2 的復(fù)制。它們也可以用作點(diǎn)云分類(lèi)和分子性質(zhì)預(yù)測(cè)的獨(dú)立網(wǎng)絡(luò)（圖 1 ）。

圖 1 用于分子性質(zhì)預(yù)測(cè)的典型 SE （ 3 ） – transformer 的結(jié)構(gòu)。

在/PyTorch/DrugDiscovery/SE3Transformer存儲(chǔ)庫(kù)中， NVIDIA 提供了在QM9 數(shù)據(jù)集上為分子性質(zhì)預(yù)測(cè)任務(wù)訓(xùn)練優(yōu)化模型的方法。 QM9 數(shù)據(jù)集包含超過(guò) 10 萬(wàn)個(gè)有機(jī)小分子和相關(guān)的量子化學(xué)性質(zhì)。

訓(xùn)練吞吐量提高 21 倍

與基線實(shí)施相比， NVIDIA 實(shí)現(xiàn)提供了更快的訓(xùn)練和推理。該實(shí)現(xiàn)對(duì) SE（3）-Transformers 的核心組件，即張量場(chǎng)網(wǎng)絡(luò)（ TFN ）以及圖形中的自我注意機(jī)制進(jìn)行了優(yōu)化。

考慮到注意力層超參數(shù)的某些條件得到滿(mǎn)足，這些優(yōu)化大多采取操作融合的形式。

由于這些，與基線實(shí)施相比，訓(xùn)練吞吐量增加了 21 倍，利用了最近 GPU NVIDIA 上的張量核。

圖 2 A100 GPU 上的訓(xùn)練吞吐量。批次大小為 100 的 QM9 數(shù)據(jù)集。

此外， NVIDIA 實(shí)現(xiàn)允許使用多個(gè) GPU 以數(shù)據(jù)并行方式訓(xùn)練模型，充分利用 DGX A100 （ 8x A100 80GB ）的計(jì)算能力。

把所有東西放在一起，在 NVIDIA DGX A100 上， SE（3）-Transformer現(xiàn)在可以在 QM9 數(shù)據(jù)集上在 27 分鐘內(nèi)進(jìn)行訓(xùn)練。作為比較，原始論文的作者指出，培訓(xùn)在硬件上花費(fèi)了 2 。 5 天（ NVIDIA GeForce GTX 1080 Ti ）。

更快的培訓(xùn)使您能夠在搜索最佳體系結(jié)構(gòu)的過(guò)程中快速迭代。隨著內(nèi)存使用率的降低，您現(xiàn)在可以訓(xùn)練具有更多注意層或隱藏通道的更大模型，并向模型提供更大的輸入。

內(nèi)存占用率降低 9 倍

SE（3）-Transformer 是已知的記憶重模型，這意味著喂養(yǎng)大輸入，如大蛋白質(zhì)或許多分批小分子是一項(xiàng)挑戰(zhàn)。對(duì)于 GPU 內(nèi)存有限的用戶(hù)來(lái)說(shuō)，這是一個(gè)瓶頸。

這一點(diǎn)在DeepLearningExamples上的 NVIDIA 實(shí)現(xiàn)中已經(jīng)改變。圖 3 顯示，由于 NVIDIA 優(yōu)化和對(duì)混合精度的支持，與基線實(shí)現(xiàn)相比，訓(xùn)練內(nèi)存使用減少了 9 倍。

圖 3 SE （ 3 ） – transformer s 的基線實(shí)現(xiàn)和 NVIDIA 實(shí)現(xiàn)之間的訓(xùn)練峰值內(nèi)存消耗比較。在 QM9 數(shù)據(jù)集上每批使用 100 個(gè)分子。 V100 32-GB GPU 。

除了對(duì)單精度和混合精度進(jìn)行改進(jìn)外，還提供了低內(nèi)存模式。啟用此標(biāo)志后，模型在 TF32 （ NVIDIA 安培體系結(jié)構(gòu)）或 FP16 （ NVIDIA 安培體系結(jié)構(gòu)、 NVIDIA 圖靈體系結(jié)構(gòu)和 NVIDIA 伏特體系結(jié)構(gòu)）精度上運(yùn)行，模型將切換到以吞吐量換取額外內(nèi)存節(jié)省的模式。

實(shí)際上，在具有 V100 32-GB GPU 的 QM9 數(shù)據(jù)集上，基線實(shí)現(xiàn)可以在內(nèi)存耗盡之前擴(kuò)展到 100 的批大小。 NVIDIA 實(shí)現(xiàn)每批最多可容納 1000 個(gè)分子（混合精度，低內(nèi)存模式）。

對(duì)于處理以氨基酸殘基為節(jié)點(diǎn)的蛋白質(zhì)的研究人員來(lái)說(shuō)，這意味著你可以輸入更長(zhǎng)的序列并增加每個(gè)殘基的感受野。

SE（3）-Transformers 優(yōu)化

與基線相比， NVIDIA 實(shí)現(xiàn)提供了一些優(yōu)化。

融合鍵與值計(jì)算

在“自我注意”層中，將計(jì)算關(guān)鍵幀、查詢(xún)和值張量。查詢(xún)是圖形節(jié)點(diǎn)特征，是輸入特征的線性投影。另一方面，鍵和值是圖形邊緣特征。它們是使用 TFN 層計(jì)算的。這是 SE（3）-Transformer 中大多數(shù)計(jì)算發(fā)生的地方，也是大多數(shù)參數(shù)存在的地方。

基線實(shí)現(xiàn)使用兩個(gè)獨(dú)立的 TFN 層來(lái)計(jì)算鍵和值。在 NVIDIA 實(shí)現(xiàn)中，這些被融合在一個(gè) TFN 中，通道數(shù)量增加了一倍。這將啟動(dòng)的小型 CUDA 內(nèi)核數(shù)量減少一半，并更好地利用 GPU 并行性。徑向輪廓是 TFN 內(nèi)部完全連接的網(wǎng)絡(luò)，也與此優(yōu)化融合。概覽如圖 4 所示。

圖 4 NVIDIA 實(shí)現(xiàn)中的鍵、查詢(xún)和值計(jì)算。鍵和值一起計(jì)算，然后沿通道維度分塊。

TFN 合并

SE（3）-Transformer 內(nèi)部的功能除了其通道數(shù)量外，還有一個(gè)degreed，它是一個(gè)正整數(shù)。程度特征d有維度2d+1. TFN 接受不同程度的特征，使用張量積組合它們，并輸出不同程度的特征。

對(duì)于輸入為 4 度、輸出為 4 度的圖層，將考慮所有度的組合：理論上，必須計(jì)算 4 × 4 = 16 個(gè)子圖層。

這些子層稱(chēng)為成對(duì) TFN 卷積。圖 5 顯示了所涉及的子層的概述，以及每個(gè)子層的輸入和輸出維度。對(duì)給定輸出度（列）的貢獻(xiàn)相加，以獲得最終特征。

圖 5 TFN 層中涉及的成對(duì)卷積，輸入為 4 度，輸出為 4 度。

NVIDIA 在滿(mǎn)足 TFN 層上的某些條件時(shí)，提供多級(jí)融合以加速這些卷積。通過(guò)創(chuàng)建尺寸為 16 倍的形狀，熔合層可以更有效地使用張量核。以下是應(yīng)用熔合卷積的三種情況：

輸出功能具有相同數(shù)量的通道

輸入功能具有相同數(shù)量的通道

這兩種情況都是正確的

第一種情況是，所有輸出特征具有相同數(shù)量的通道，并且輸出度數(shù)的范圍從 0 到最大度數(shù)。在這種情況下，使用輸出融合特征的融合卷積。該融合層用于 SE（3）-Transformers 的第一個(gè) TFN 層。

圖 6 每個(gè)輸出度的部分熔融 TFN 。

第二種情況是，所有輸入特征具有相同數(shù)量的通道，并且輸入度數(shù)的范圍從 0 到最大度數(shù)。在這種情況下，使用對(duì)融合輸入特征進(jìn)行操作的融合卷積。該融合層用于 SE（3）-Transformers 的最后一層 TFN 。

圖 7 每個(gè)輸入度的部分熔融 TFN 。

在最后一種情況下，當(dāng)兩個(gè)條件都滿(mǎn)足時(shí)，使用完全融合的卷積。這些卷積作為輸入融合特征，輸出融合特征。這意味著每個(gè) TFN 層只需要一個(gè)子層。內(nèi)部 TFN 層使用此融合級(jí)別。

圖 8 全熔合 TFN

基預(yù)計(jì)算

除了輸入節(jié)點(diǎn)特性外， TFN 還需要基矩陣作為輸入。每個(gè)圖邊都有一組矩陣，這些矩陣取決于目標(biāo)節(jié)點(diǎn)和源節(jié)點(diǎn)之間的相對(duì)位置。

在基線實(shí)現(xiàn)中，這些矩陣在前向傳遞開(kāi)始時(shí)計(jì)算，并在所有 TFN 層中共享。它們依賴(lài)于球形 h ARM ，計(jì)算起來(lái)可能很昂貴。由于輸入圖不會(huì)隨著 QM9 數(shù)據(jù)集而改變（沒(méi)有數(shù)據(jù)擴(kuò)充，沒(méi)有迭代位置細(xì)化），這就引入了跨時(shí)代的冗余計(jì)算。

NVIDIA 實(shí)現(xiàn)提供了在培訓(xùn)開(kāi)始時(shí)預(yù)計(jì)算這些基礎(chǔ)的選項(xiàng)。整個(gè)數(shù)據(jù)集迭代一次，基緩存在 RAM 中。前向傳遞開(kāi)始時(shí)的計(jì)算基數(shù)過(guò)程被更快的 CPU 到 GPU 內(nèi)存拷貝所取代。

關(guān)于作者

Alexandre Milesi 是 NVIDIA 的深度學(xué)習(xí)算法工程師。他擁有法國(guó) UTC 的機(jī)器學(xué)習(xí)碩士學(xué)位，以及法國(guó)索邦大學(xué)的機(jī)器人和多智能體系統(tǒng)碩士學(xué)位。在加入 NVIDIA 之前， Alexandre 是伯克利實(shí)驗(yàn)室的附屬研究員，使用深度強(qiáng)化學(xué)習(xí)解決電子 CTR ical 網(wǎng)格問(wèn)題。在 NVIDIA ，他的工作集中于藥物發(fā)現(xiàn)和計(jì)算機(jī)視覺(jué)的 DL 算法，包括等變圖神經(jīng)網(wǎng)絡(luò)。

審核編輯：郭婷