華為MindSpore架構師王紫東：當生物計算領域「遇災」，MindSpore做了何解？

機器之心報道

演講：王紫東

編輯：華衞

疾病治療和新藥研發當中，多尺度建模發揮着重要作用。而當傳統科學計算方式遇到生物醫藥這類沒有方程式可尋的問題時，亟待 AI 技術的加持。

9 月 2 日，在 2022WAIC 上海生物計算論壇上，華為中央軟件院昇思MindSpore 開源項目架構師王紫東發表主題演講《MindSpore AI + 科學計算實踐》。演講中，他主要介紹了傳統生物計算領域中的發展瓶頸以及 MindSpore 在這方面所做的工作。

以下為王紫東在 2022WAIC 上海生物計算論壇上的演講內容，機器之心進行了不改變原意的編輯、整理：

大家下午好！很高興有機會和大家一起分享 MindSpore 在 AI + 生物計算方面做的實踐。MindSpore 之前是傳統的 AI 框架，目前正向 AI + 科學計算融合的通用計算架構演進。

總體上講，我們的實踐是在高毅勤老師課題組指導下，共同合作完成的。研究中，我們主要從基於力場的分子模擬入手。

分子模擬的應用與挑戰

藉助分子模擬技術，我們可以模擬化學反應、蛋白質摺疊以及同一物質在不同狀態下的相變等；在行業之中，分子模擬也是比較重要的技術，如材料設計、藥物設計、化工設計等場景都離不開分子模擬。

同時，基於力場的分子模擬也是微觀世界探索中的關鍵問題。針對這樣的問題，研究者做了很多的探索，目標是在更大的空間尺度和時間尺度上得到精確的分子模擬，以揭示微觀規律。

針對微觀世界，最有效的武器即為量子力學。保羅狄拉克曾説，「在有了量子力學之後，大部分物理化學現象的基本定律是完全已知的，困難在於實際應用中方程過於複雜無法求解。」

量子力學本質上是求解薛定諤方程，但薛定諤方程沒辦法處理較複雜的體系，只要粒子數大於 2 就沒有辦法求解。所以只能進行近似求解。學界為了針對更大體系進行分子模擬，發展出了很多方法，如分子動力學、粗粒化動力學和連續介質力學。

雖然能模擬的分子體系越來越大，但是精度在逐漸下降。以分子動力學為例，它利用第一性原理數據或者實驗數據進行參數擬合，得到體系勢能函數，也就是我們常説的分子力場。這種方法算得快，但精度很受限制，沒有辦法模擬精確的化學反應。

事實上，分子模擬一直處於模擬精度和效率無法兼得的狀態，我把它歸結為維度災難。在這樣的瓶頸下，現在已經有一些方法在解決這個問題，最主要的就是 AI。尤其是機器學習在解維度災難問題上很有經驗，也產生了一定結果，所以傳統科學計算領域以後可能會成為 AI 的主戰場。

「AI + 分子模擬」的成就

現在，AI 和分子模擬已經開始進行融合，並且產生了有突破性的工作。這裏我們看兩個案例。

第一個案例 DeePMD，它主要是用 AI 神經網絡去擬合分子力場，訓練數據來自於第一性原理計算數據。DeePMD 兼顧精度和性能，並且獲得了 2020 年的戈登貝爾獎。

案例二是 AlphaFold2，這是非常有突破性的成果，它更偏向 AI 思路，可以從海量數據中直接學習序列至結構的高維映射，效率比較高。令人驚喜的是，它的精度可以與實驗精度相媲美，並被評為十大科技進步之首。

新一代分子模擬軟件

AI 方法已經改變了傳統計算的範式，成為一種新的可能性，軟件上也要有相應的支持。我們來看一下現有軟件能不能較好地支持 AI 新的範式，其實答案可能並不是很樂觀。

主流分子模擬軟件有這麼幾個特點：第一是開發時間長，用户依賴性強；第二是幾乎都由西方發達國家開發，由中國人開發的模擬軟件極少且用户有限；第三是模擬框架老舊，靈活性差，如需添加新算法，往往需要對程序代碼進行大幅改動；第四是程序多用 C/C++，甚至 Fortran 語言編寫，難以兼容目前主流以 Python 為前端語言的 AI 框架。

要想支持這種新範式，可能還是有很多路要走，現在做得並不是特別好。

在這樣的背景下，MindSpore 和高毅勤老師課題組一起開發 MindSPONGE，我們把它稱之為新一代分子模擬軟件。

體系架構如何？

這是我們已知的第一個根植於 AI 計算框架中的分子軟件，打破了傳統 AI 和傳統分子模擬計算的界限，並對它們的架構做了統一化的建模，取得了較好的效果。

我們發現，分子模擬和深度學習的訓練雖然看着好像沒什麼關係，但實際上有相似的計算邏輯。比如，深度學習訓練是要優化一個損失函數，分子模擬可能跟它是差不多的，只不過固定的是參數，參數指的是分子力場，再去優化以空間座標為主的樣本。

進行模塊化拆分後，我們把 MindSPONGE 分為三大模塊：系統建立、力場配置和迭代更新。我們有這麼幾個優勢：1. 易用，無需推導原子受力的解析表達式；2. 高效，天然支持高通量模擬；3. 模塊化，兼容 AI 勢能函數。

有哪些案例實踐？

這裏舉幾個分子動力學的案例。分子動力學主要是知道分子下一步怎麼動，這裏面的核心就是分子力場。不管是構建分子力場，還是使用分子力場，MindSPONGE 都提供了非常方便的使用接口。

比如，克萊森重排有七元環與三元環轉換的動力學過程，用我們自己訓練的 AI 力場，以第一性原理精度去做動力學的迭代，精度及效率會優於傳統方法。

另外，我們對新冠病毒 Delta 變異毒株也進行了模擬。模擬結果發現，Delta 變異時，其氫鍵數目和接觸面積相對野生型毒株有一定增加，這可能就是 Delta 毒株傳播力強的原因。

我們整體依託昇思MindSpore，使能 MindSpore 核心技術圖算融合。MindSPONGE 的性能相較於其它框架有一定的優勢。

DeepMind 開源了 AlphaFold2，這對於業界在該領域發展具有巨大的推進作用，但是其並沒有開源自己的訓練代碼，也沒有提供對應的數據集。同時，對於開源出來的網絡權重，其有着非商用的 License。對於科研而言沒有風險，但對於商業應用的話，存在一定的侵權的問題。

算法層面主要有兩個問題，問題一：AlphaFold2 依賴於共進化信息，當共進化信息缺失嚴重時，蛋白結構預測精度會大幅下降。問題二：當存在多個模型推理出來的蛋白結構時，沒有統一的算法評估這些結構的質量，這會導致無法選擇最優的結構。

軟件層面的問題主要為算法的內存消耗，算法佔用的顯存與序列長度是立方關係，如果基礎軟件內存優化做得不好，很容易導致顯存爆炸。同時，推理時性能主要消耗在共進化信息檢索。

針對上述不同層面的問題，我們協同高毅勤老師課題組一起進行了一些探索。

首先，我們復現了 AlphaFold2 的推理和訓練，還構建了數據集，並把它開源出來，大家可以基於此數據集訓練自己的模型。對於處理性能，如最耗時的前端搜索 MSA，我們都進行了優化，這項工作我們是在 4 月份完成的，當時去參加了 CAMEO 比賽，取得了較好的結果。

第二步，針對算法上依賴共進化信息的問題，我們利用了人工智能領域生成模型的思路，同時借鑑了 diffusion model 的模型構建思路，針對一些孤兒序列，構建了對應的 MSA 信息，取得了較好的效果。

最後，我們在蛋白質結構如何評估、如何挑出最好的蛋白質方面做了一些工作，表現也是良好的，比賽成績優異。這項工作會在近期進行開源，如果大家感興趣可以關注一下。

蛋白質結構預測和分子動力學這兩個案例有意思的地方在於，可以聯動。很多時候，結構並不符合物理的規律，可能很多蛋白質是可以繼續優化的，很多位置是有衝突不合理的，這時需要分子動力學模擬來做進一步的動作，也就是蛋白質結構鬆弛。

當前 AlphaFold2, 在網絡推理使用 JAX，分子動力學過程則使用 OpenMM 進行。由於 MindSPONGE 有統一建模的體系和工具，能夠把蛋白質結構預測做一體化的推理，會給出更合理的結構。

還能提供什麼能力？

此外，模型的內存和性能還有很多值得優化的地方，MindSPONGE 也提供了相應的解決方案，我們稱之為 SOMAS 技術。該技術主要是並行拓撲內存約束優化器，可以對所有內存進行統一統籌管理，效果還是比較明顯的。

性能方面，MindSPONGE 具有圖算融合的特性，可以把小的細碎算子融成大算子，模型的整體性能會因此有提升。

傳統優化改進到元優化

這裏，對 MindSPONGE 的未來進行一些展望。

基於分子力場的模擬，分子力場肯定是最重要的。傳統優化方式主要還是靠人工經驗，比如先有力場，就可以做環境的模擬，模擬出的分佈很可能不符合實驗結果，這樣就會產生誤差；傳統方式是通過人進行分析，是不是有些參數可以改一改，再去嘗試一下。

相對來講，這樣比較盲目。MindSPONGE 採用的則是元優化方式。

現在，我們還有很多東西沒有做完，後續希望把流程都打通，也希望和大家一起合作構建。這裏感謝高毅勤老師課題組，也感謝所有社區貢獻的小夥伴們！

最後， MindSpore 不只針對生物計算方面有相應的實踐，我們在汽車、能源、氣象、航空航天、EDA、材料、金融等方面都有規劃。希望通過支撐更多行業應用，幫助合作伙伴商業落地，完善我們國產基礎軟硬件的能力。目前已經有的，主要是電磁仿真、流體仿真、分子模擬。

我們並不懂行業的技術，如果沒有行業專家的參與，我們根本不知道做的這些東西對不對，所以希望和大家多多合作。如果大家感興趣，可以隨時聯繫我，謝謝大家！

© THE END 

轉載請聯繫本公眾號獲得授權

投稿或尋求報道：[email protected]