• <xmp id="om0om">
  • <table id="om0om"><noscript id="om0om"></noscript></table>
  • 3 月 19 日下午 2 點,鎖定 NVIDIA AI 網絡中文專場。立即注冊觀看
    數據科學

    NVIDIA cuEquivariance 數學庫助力藥物和材料研發加速

    用于科學的 AI 模型通常經過訓練,可以預測自然界的運作方式,例如預測生物分子的結構或可能成為下一種電池材料的新固態的性質。這些任務需要高精度和準確性。科學領域的 AI 更具挑戰性的原因在于,高度準確和精確的科學數據往往稀缺,這與從多個資源中大量獲取的文本和圖像不同。

    鑒于對解決方案的高需求和有限的資源,研究人員轉而采用創新方法,例如將自然法則嵌入 AI 模型中,提高其準確性以及減少對數據的依賴。

    去年取得成功的一種方法是將科學問題的對稱性嵌入 AI 模型。這些神經網絡架構在等變神經網絡(Equivariant Neural Networks,ENN)下廣受歡迎,使用對稱相關變換下的等變數學概念構建而成。

    簡而言之,ENN 旨在了解問題的基本對稱性。例如,如果 ENN 的輸入進行了旋轉,輸出也將相應旋轉。這意味著,即使以不同的方向呈現,模型也可以識別相同的物體或圖案。

    為了更好地理解這一概念,請考慮如何使用 ENN 來維護 3D 中對稱運算后輸入和輸出之間的關系。例如,如果 ENN 將分子的 3D 模型作為輸入,并將其屬性預測為輸出,那么無需額外的訓練數據或數據增強,就可以預測任何旋轉版本的分子的相同屬性。ENN“理解”分子的旋轉不會改變其基本屬性(圖 1)。

    On the top left corner, there is a sketch of a molecule. An arrow indicates it is an input to the neural network below. Another arrow from the network from a vector below (in the shape of a large arrow) shows that it is the neural network's output. This shows the usual input-network-output pipeline. This same pipeline appears on the right side, this time with the same molecule in another orientation, rotated from the left. It also connects to the network below and its output, which is rotated correspondingly. These two pipelines are connected with the words “Rotated input” at the top and “Rotated output” at the bottom. This shows that with equivariant networks, the rotation or network operation order does not matter; one can either rotate and pass through the network or pass through the network and then rotate; the result is the same as in the bottom right corner of the image.
    圖 1、等變神經網絡經過專門設計,如果輸入數據發生變化以保持整體模式不變 (例如旋轉或翻轉圖像),則輸出會發生可預測且適當的變化,以匹配此轉換。

    在網絡架構中引入這種基本的對稱性,可在更改輸入數據時提高模型的穩健性和數據效率。與將自然法則嵌入神經網絡的其他策略類似,它還提供了一種提高對未見數據的泛化能力的方法。

    所有這些優勢都是有代價的:從理論上講,構建 ENNs 并不簡單,并且生成的網絡在計算上比非等變版本更昂貴。在本文中,我們將介紹新的數學庫 NVIDIA cuEquivariance 如何應對挑戰并加速 AI 科學模型的開發,并提供藥物發現和材料科學應用的示例。

    等變神經網絡的挑戰?

    許多 AI 模型(包括 Tensor Field Networks、LieConv、Cormorant、SE(3)-Transformer、NequIP 以及 DiffDock 和 Equiformer 等其他模型)都使用獨特的方法來確保一致地處理輸入數據中的更改。它們使用對稱組的基本元素,稱為不可約表示(irreps)或這些元素的變體。這些不可約表示在數學上表示為張量,并且以特定的方式組合,通常涉及張量代數(例如張量積),以確保模型的輸出適當反映應用于輸入的任何對稱變換。

    采用使用 irreps 的 ENN 的一個瓶頸是,為給定對稱組構建和處理這些 irrep 對象的理論復雜性。由于缺乏現有基元或可擴展 API,再加上理論復雜性,因此使用 irreps 形式進行 ENN 創新極具挑戰性。在該領域,更容易選擇重復使用現有的實現,即使它們不理想也是如此。

    此外,在使用基于 irreps 的 ENNs 時存在計算復雜性。數學基礎決定了 irreps 的矩陣表示。對于最常用的對稱運算 (例如 3D 旋轉),這些大小對于計算優化來說可能并不常見,例如 5×5 或 7×7 矩陣。這不允許利用現有的優化技術 (例如數學運算中的 tensor cores),將這些對象開箱即用。

    更重要的是,涉及 irreps 的張量積運算遵循源于群論的異常稀疏模式,即使 irreps 本身是密集的。Irreps 具有稱為 Clebsch-Gordan 系數的特殊混合系數,用于確定如何在代數運算中將兩個不同的 irreps 組合為輸出 irrep。

    例如,相乘兩個 irreps 只能得到特定且有限的輸出 irreps 列表,并且此 選擇規則 由群理論決定。事實上,由于選擇規則,不允許進行多種 irreps 組合,這會導致稀疏的 Clebsch-Gordan 系數,其中大多數為零。從計算的角度來看,忽略群理論所決定的稀疏性會導致內存浪費和算法效率低下。

    加速等變神經網絡?

    為應對這些挑戰,NVIDIA 開發了新的 cuEquivariance 數學庫,該庫為等變神經網絡引入了 CUDA 加速的構建塊。cuEquivariance 現已在 GitHub 和 PyPi 上提供公開測試版。

    cuEquivariance Python 前端引入了一個名為分段張量積(Segmented Tensor Product,STP)的統一框架,該框架考慮到前面提到的混合系數的稀疏模式,使用 irrreps 組織代數運算。STP 泛化了等變多線性乘積的計算,使用戶能夠在 irrreps 之間表示各種此類運算。它還允許用戶自由定義不一定等變的運算,這可能有助于研究社區尚未探索的應用。

    cuEquivariance 基于 STP 框架構建,利用專門的 CUDA 內核來加速最常用的 STP 實例。在 ENNs 中,大多數瓶頸運算是多個受內存限制的運算,逐個執行,導致不必要的中間產品加載和存儲。鑒于 irreps 的規模小且數量多,使用不同的內核調用執行每個操作是另一個開銷來源。cuEquivariance 使用內核融合將這些單獨的運算替換為一些特殊用途的 GPU 內核。

    除了內核融合之外,我們還重新組織了特征的內存布局,使內存訪問更好地映射到 NVIDIA GPU 架構的單指令多線程 (SIMT) 范式。這種專門的后端經過優化,可在 NVIDIA GPU 上顯著加快等變神經網絡中的數學運算。

    圖 2 顯示了 cuEquivariance 加速對兩個熱門科學模型 AI 的影響:DiffDock 和 MACE,前者是預測蛋白質 – 配體結合姿態的擴散模型,后者是機器學習的原子間相互作用勢,廣泛用于材料科學和生物學,用于控制分子動力學模擬。

    The figure shows six panels titled “Speed up of NVIDIA cuEquivariance vs Baseline for Different Workloads on H100.” Each panel is dedicated to the comparison of a different model and operation. The top left panel shows cuEquivariance acceleration of 3.1x for DiffDock TP forward operation, and the top right model shows 15.6x for backward of the same. The middle left panel shows the speedup of MACE-OFF Large model symmetric contraction forward operation, with the acceleration of 17x concerning the baseline. The middle right panel shows 9x for the backward of the same. The bottom left model shows a 2.2x speedup for MACE-OFF small model tensor product forward operation. The bottom right shows MACE-OFF medium model tensor product forward operation with a 3.3x speedup.
    圖 2 、使用基準張量乘積實現,將 DiffDock 和 MACE 模型中等變運算的計算性能與特定 GPU 架構的理論最大值 (“光速”) 進行比較。這凸顯了進一步優化的潛力,尤其是通過利用 Clebsch-Gordan 系數中的稀疏性。

    這些等變神經網絡模型具有多個帶 irreps 的張量運算。出于演示目的,為每個模型選擇計算要求最高的運算。對于 DiffDock,這是一個基于 irrep 的張量積運算(TP)。對于 MACE,我們會考慮兩種影響性能的運算:對稱收縮(SC),即 irreps 張量自身的張量收縮;TP,類似于 DiffDock。對于每個運算,都會顯示向前和向后性能。

    圖 3 展示了具有 cuEquivariance 的 MACE-OFF Large 和 MACE-MP Large 模型的端到端性能。最后,圖 4 顯示了在不同 NVIDIA GPU 之間的性能變化。

    The figure shows four panels titled “Training and Inference time for MACE-OFF and MACE-MP Large on A100. The top row shows results for MACE-OFF Large model. The top left shows training time to be 6.1x accelerated for baseline, while inference acceleration is 7.2x. The bottom row shows the results for MACE-MP Large model. The bottom left shows training time to be accelerated by 5.9x, and the bottom right shows inference to be accelerated by 5.9x with respect to baseline.
    圖 3 、MACE-OFF 和 MACE-MP 等效模型的端到端訓練時間,在 FP64 中使用批量大小 32,并實現基準張量積。這表明了通過使用 Clebsch-Gordan 系數中的稀疏性進行進一步優化的可能性。
    Two panels titled “MACE-OFF Large Symmetric Contraction on Different GPUs.” MACE-OFF Large SC forward operation timings on the left panel are shown, where timing reduces from 6 milliseconds on NVIDIA A100 architecture to approximately four milliseconds on H100. In between, the timing gradually reduces for the L40 architecture. The same trend is also visible in the right panel for MACE-OFF Large model SC backward operation: the timings reduce approximately 25 milliseconds for A100 to further down for NVIDIA L40 and then NVIDIA H100 being approximately 20 milliseconds.
    圖 4. MACE-OFF 關閉的大型模型在不同 NVIDIA GPU 架構下的對稱正向和反向收縮運算時間,其中 10K 個輸入批量進行張量乘積運算,并使用 FP32

    結束語?

    cuEquivariance 的開發標志著加速 AI 科學發展的重要一步。通過解決等變神經網絡的理論和計算挑戰,cuEquivariance 使研究人員、科學家和學者能夠為各種科學應用構建更準確、更高效和更通用的模型。正如其成功集成到廣泛使用的模型如 DiffDock 和 MACE 所證明的那樣,cuEquivariance 已準備好在藥物發現、材料科學等領域推動創新并加速發現。

    通過利用對稱性和高效計算的強大功能,cuEquivariance 為 AI 帶來了新的可能性,有助于實現科學突破。將 cuEquivariance 等開源加速計算工具與系統生成的大規模數據集相結合,可以提高 AI 模型的準確性,從而促進研究和企業產品的更廣泛采用和集成。

    開始使用 cuEquivariance

    ?

    0

    標簽

    人人超碰97caoporen国产