華為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]
- 從AI為寵物投保這個「小事」講起,保險變革開始了
- 華為MindSpore架構師王紫東:當生物計算領域「遇災」,MindSpore做了何解?
- 從產線到星空,這些AI工程師帶著演算法上天入地
- 數實融合驅動數字未來 影譜科技攜沉浸式數字技術亮相2022服貿會
- 驅動虛實融合,影譜科技在WAIC推出多款智慧視覺及元宇宙新品
- 英偉達和Ampere共同推出雲遊戲Arm伺服器:每臺最多支援160個使用者
- 擴散模型背後數學太難了,啃不動?谷歌用統一視角講明白了
- 悟空搜尋“降世”,能否大鬧搜尋引擎“天宮”?
- 蘑菇車聯與北京通州達成合作,打造車路雲一體化智慧交通系統
- LeCun沒看好的AGI,傳奇程式設計師John Carmack:已創業,融資兩千萬美元
- 蘑菇車聯與無錫梁溪區簽約,將打造數字孿生城市交通底座
- 這場比賽,讓上百個AI智慧體「捲起來了」
- 大模型時代,那些一起訓練AI模型的企業是怎麼應對資料顧慮的?
- 數字人賽道熱度正高,AI技術如何創造不一樣的“靈魂”?
- Unity否決175億美元併購提議:準備自己收購廣告公司
- 消費級GPU可用,文字轉影象開源新模型Stable Diffusion生成宇宙變遷大片
- 全球遊戲引擎商業化進入新階段,國內頭號玩家Cocos“三分天下”
- 墨芯獲評大灣區高成長性企業TOP50
- 圖神經網路發Nature子刊,卻被爆比普通演算法慢104倍,質疑者:灌水新高度?
- 亞馬遜欲17億美元收購iRobot,掃地機器人“含金量”幾何?