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

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

圖 1 用于分子性質預測的典型 SE （ 3 ） – transformer 的結構。

在/PyTorch/DrugDiscovery/SE3Transformer存儲庫中， NVIDIA 提供了在QM9 數(shù)據(jù)集上為分子性質預測任務訓練優(yōu)化模型的方法。 QM9 數(shù)據(jù)集包含超過 10 萬個有機小分子和相關的量子化學性質。

訓練吞吐量提高 21 倍

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

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

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

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

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

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

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

內存占用率降低 9 倍

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

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

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

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

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

對于處理以氨基酸殘基為節(jié)點的蛋白質的研究人員來說，這意味著你可以輸入更長的序列并增加每個殘基的感受野。

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

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

融合鍵與值計算

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

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

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

TFN 合并

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

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

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

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

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

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

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

這兩種情況都是正確的

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

圖 6 每個輸出度的部分熔融 TFN 。

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

圖 7 每個輸入度的部分熔融 TFN 。

在最后一種情況下，當兩個條件都滿足時，使用完全融合的卷積。這些卷積作為輸入融合特征，輸出融合特征。這意味著每個 TFN 層只需要一個子層。內部 TFN 層使用此融合級別。

圖 8 全熔合 TFN

基預計算

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

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

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

關于作者

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

審核編輯：郭婷