用于科學的 AI 模型通常經過訓練，可以預測自然界的運作方式，例如預測生物分子的結構或可能成為下一種電池材料的新固態的性質。這些任務需要高精度和準確性。科學領域的 AI 更具挑戰性的原因在于，高度準確和精確的科學數據往往稀缺，這與從多個資源中大量獲取的文本和圖像不同。
鑒于對解決方案的高需求和有限的資源，研究人員轉而采用創新方法，例如將自然法則嵌入 AI 模型中，提高其準確性以及減少對數據的依賴。
去年取得成功的一種方法是將科學問題的對稱性嵌入 AI 模型。這些神經網絡架構在等變神經網絡（Equivariant Neural Networks，ENN）下廣受歡迎，使用對稱相關變換下的等變數學概念構建而成。
簡而言之，ENN 旨在了解問題的基本對稱性。例如，如果 ENN 的輸入進行了旋轉，輸出也將相應旋轉。這意味著，即使以不同的方向呈現，模型也可以識別相同的物體或圖案。
為了更好地理解這一概念，請考慮如何使用 ENN 來維護 3D 中對稱運算后輸入和輸出之間的關系。例如，如果 ENN 將分子的 3D 模型作為輸入，并將其屬性預測為輸出，那么無需額外的訓練數據或數據增強，就可以預測任何旋轉版本的分子的相同屬性。ENN“理解”分子的旋轉不會改變其基本屬性（圖 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，前者是預測蛋白質 – 配體結合姿態的擴散模型，后者是機器學習的原子間相互作用勢，廣泛用于材料科學和生物學，用于控制分子動力學模擬。
這些等變神經網絡模型具有多個帶 irreps 的張量運算。出于演示目的，為每個模型選擇計算要求最高的運算。對于 DiffDock，這是一個基于 irrep 的張量積運算（TP）。對于 MACE，我們會考慮兩種影響性能的運算：對稱收縮（SC），即 irreps 張量自身的張量收縮；TP，類似于 DiffDock。對于每個運算，都會顯示向前和向后性能。
圖 3 展示了具有 cuEquivariance 的 MACE-OFF Large 和 MACE-MP Large 模型的端到端性能。最后，圖 4 顯示了在不同 NVIDIA GPU 之間的性能變化。
結束語?
cuEquivariance 的開發標志著加速 AI 科學發展的重要一步。通過解決等變神經網絡的理論和計算挑戰，cuEquivariance 使研究人員、科學家和學者能夠為各種科學應用構建更準確、更高效和更通用的模型。正如其成功集成到廣泛使用的模型如 DiffDock 和 MACE 所證明的那樣，cuEquivariance 已準備好在藥物發現、材料科學等領域推動創新并加速發現。
通過利用對稱性和高效計算的強大功能，cuEquivariance 為 AI 帶來了新的可能性，有助于實現科學突破。將 cuEquivariance 等開源加速計算工具與系統生成的大規模數據集相結合，可以提高 AI 模型的準確性，從而促進研究和企業產品的更廣泛采用和集成。
?
