隨著 AlphaFold2 等模型的出現，對加速推理和訓練分子 AI 模型的需求激增。對速度的需求帶來了獨特的計算挑戰，包括算法復雜性、內存效率和嚴格的準確性要求。為解決這一問題，NVIDIA 與合作伙伴合作提供加速解決方案，例如更快的 equivariant 操作和更快的 MSA 生成。
今天，我們發布了 cuEquivariance 中的新內核和 NVIDIA NIM 微服務，以加速分子 AI 模型的訓練和推理，例如由 MIT 和 Recursion 開發的開源基礎模型 Boltz-2。這些加速有助于開發更復雜的分子 AI 系統，并更快地大規模了解分子結構。
NVIDIA cuEquivariance 擴展以加速新一代蛋白質結構模型
NVIDIA cuEquivariance 是一個 CUDA-X 庫，旨在加速 MACE、Allegro、NequIP 和 DiffDock 等幾何感知神經網絡的苛刻計算。它提供高度優化的 CUDA 內核和全面的 API，可顯著加速核心等方差運算，例如涉及 Segmented Tensor Products 的運算。
從 cuEquivariance v0.5 開始，該庫現在包括加速的 Triangle Attention 和 Triangle Multiplication 核函數，這對于獲得諾貝爾獎的蛋白質結構預測模型 (如 AlphaFold2) 來說至關重要且獨一無二。隨著加速三角形運算的加入，cuEquivariance 的影響擴展到蛋白質折疊、RNA/DNA 結合、盲對接、蛋白質復合體預測和親和力評分等應用。
了解蛋白質的 3D 結構至關重要，因為它揭示了蛋白質的工作原理。然而，細胞和生命本身的真正復雜性源于生物復合體內的動態相互作用。這些復合體不僅僅由蛋白質組成，它們是蛋白質、核酸 (如 DNA 和 RNA) 、脂類、碳水化合物和各種小分子的復雜組合，它們協同工作。
預測這些單個分子和配合物的結構和動態行為是分子 AI 的下一個前沿領域。下一個科學突破可以揭示細胞通路、確定疾病機制，并設計能夠精確向特定分子相互作用的藥物。
蛋白質、RNA 和 DNA 都是由重復單元構建的長分子：用于 DNA 和 RNA 的核酸，以及用于蛋白質的氨基酸。當這些序列在細胞內生成時，其特定的構建塊排列會導致它們折疊成復雜的三維結構。這些 3D 形狀至關重要，因為它們決定了分子的功能及其與其他細胞成分的相互作用。
在最先進的幾何感知神經網絡 (如 AlphaFold3、Proteina、Chai-1、Neo-1 和 Boltz-2) 中，三角形乘法和三角形注意力是兩個基本的計算密集型運算。在此類模型中，這些組件通常在最耗時的組件中排名靠前。
配對注意力機制通過引入 Transformers 而流行起來，其工作原理是計算 token (Transformers 詞匯表中的構建塊) 與所有其他 token 的相關性，例如，允許模型理解一個單詞在句子中所有其他單詞的上下文中的相關性。
由于分子 AI 模型的任務是根據 2D 表征預測 3D 結構，因此配對關系無法提供所有上下文信息。在這種情況下，“Triangular Relationships”可以充當捕捉空間關系的強大代理。例如，如果構建塊 i 接近 k，而 k 接近 j，則 i 和 j 可能在空間上相關，即使沒有強直接成對信號也是如此。
對于具有 N 個構建塊的分子，這些運算簡單地顯示出 O (N3) 復雜性。這種計算強度給大分子和復雜的多分子組件帶來了重大挑戰，導致巨大的計算成本和對 AI 模型擴展程度的硬性限制。
cuEquivariance 加速 Triangle Operations
在這里，我們將討論 cuEquivariance forward Triangle Attention 模塊與原版 PyTorch 實現相比的性能。這僅衡量模塊運行時，而非完整的端到端推理或訓練加速。我們稍后將在博文中討論端到端性能基準測試。
?
同樣，具有 BF16 精度的 cuEquivariance 三角形乘法內核可在沒有任何精度回歸的情況下提供高達 5 倍的模塊級加速。
麻省理工學院研究員 Gabriele Corso 表示：“這些 kernel 備受期待，并將成為 Boltz 系列模型不可或缺的一部分，幫助解決速度和內存消耗方面的瓶頸問題。
在下一代版本的Boltz-1x 上，我們比較了基于 PyTorch、Trifast 和 cuEquivariance 實現的各種精度 (TF32、FP32 和 BF16)，并比較了它們的端到端推理運行時間。這些運行使用由Boltz-1x 作者發布的默認測試數據集。在保持精度不變的情況下，從 PyTorch BF16 到 cuEquivariance BF16，我們可將性能提升高達 1.75 倍。如果從 PyTorch FP32 升級到 cuEquivariance BF16，則在使用Boltz-1x 時，性能可提升高達 2.5 倍。
從 PyTorch FP32 到使用 Boltz-1x 的 cuEquivariance BF16 配置，端到端訓練速度最高可提升 1.35 倍。端到端加速可能因模型架構而異。
VantAI 首席技術官 Luca Naef 表示：“這種對 cuEquivariance 的擴展非常有價值，我們已經看到訓練速度和推理速度分別提高了 2 倍和 3 倍以上，大大縮短了模型迭代周期，并對更大的分子實現了一個數量級的推理。
cuEquivariance 提供的加速得到了MIT、VantAI、Molecular Glue Labs (MGL) 、Dyno、Peptone、Genesis 和 Xaira 等多家合作伙伴的好評，他們能夠測試早期版本并提供反饋。我們很高興社區中的其他人能夠從這些加速中受益，并提供關鍵反饋，幫助我們改進工作，從而推動科學創新的發展。
借助適用于數字生物學的 Boltz-2 NIM 實現企業級協同折疊
麻省理工學院的Boltz團隊與 Recursion 合作開發的新一代Boltz-2 模型以Boltz-1 等模型的成功為基礎，代表著向前邁出的重要一步。“Boltz-2”旨在打造一個更大、功能更強大的模型，將推理時間優化從Boltz-1x中合并，并融合了獨特、先進的親和力預測功能。NVIDIA 將 Boltz-2 封裝為 NVIDIA NIM，以提供這種先進模型的易用版本。
NIM 是易于使用的預構建容器，可為先進的 AI 模型提供經過優化的生產就緒型推理。Boltz-2 NIM 將為研究人員和開發者提供對其強大功能的簡化訪問，為要求嚴苛的藥物研發工作流程實現實時預測和高效的測試時擴展。這種方法普及了先進的分子 AI，允許更廣泛的用戶利用 Boltz-2 的預測能力。
面向下一個分子 AI 前沿的加速計算
由 cuEquivariance 驅動的更高計算效率至關重要。對于訓練，這些加速的內核使研究人員能夠構建更大的 基礎模型，這些模型可以進一步利用預訓練 scaling laws，其中計算吞吐量的增加通常與模型性能的提高相關。此外，由此帶來的計算時間和成本效率為更多的模型開發周期釋放了資源，進一步突破了新一代功能的極限。在測試時，加速有助于在 silico 實驗中實現更廣泛的應用，使 virtual screening 活動能夠擴展到數十萬甚至數百萬個 inferences。
Recursion 首席技術官 Ben Mabey 表示：“NVIDIA 的 cuEquivariance 庫提供了顯著的加速，這對于像 Boltz-2 這樣的結構感知生物分子模型至關重要。通過解決關鍵的計算瓶頸，這將加快制藥行業部署這些強大的模型用于藥物發現的研發周期。”
除了這些庫級加速之外，NVIDIA 還通過 NVIDIA NIM 微服務等產品增強了對先進模型的訪問。例如，將 Boltz-2 等高級模型封裝為 NIM，可為研究人員和開發者提供精簡的生產就緒型解決方案，以部署這些強大的功能，高效擴展要求嚴苛的藥物研發工作流程。
NVIDIA 與更廣泛的科學社區合作，開發和完善 cuEquivariance 庫等基礎軟件，并通過 NIM 提供優化的模型訪問。這些產品由 NVIDIA 計算平臺提供支持，使科學界能夠突破研究界限，加速從計算洞察到現實世界對更廣泛的藥物研發和生物學的影響的過程。
立即試用 cuEquivariance
目前，這些加速可通過 Apache 2.0 許可下的 PyTorch API 前端提供。單擊此處詳細了解加速、支持的精度和示例。
?
