近年,人工智能的各個領域,包括自然語言處理、計算機視覺、語音處理,借助深度學習的強大威力,都取得了令人嘆為觀止的巨大進步。將深度學習技術應用于傳統的科學領域,如物理、化學、生物、醫學,即所謂的 AI for Science(科學智能),作為一個新的交叉學科,也逐漸興起,孕育著巨大的潛力,受到廣泛的關注。
ByteDance Research 也在進行 AI for Science 的研究,包括機器學習與量子化學、大規模量子化學計算、AI 制藥等領域一些問題的研究,希望跟業界一起推動領域的發展。本文簡要介紹我們這兩年來取得的一些進展。也拋磚引玉,希望與業界進行更多的交流和合作。
在機器學習和量子化學方向,我們提出的 LapNet 算法,比有代表性的 FermiNet 模型訓練速度提高了 10 倍,能計算的化學體系的規模和精度目前是領域最大的。
在大規模量子化學計算方向,我們開發了 Periodic DMET 算法,使用經典和量子混合計算機(實際是在經典計算機上的模擬),用于周期性體系的計算,只用 20 個量子比特就達到了之前方法用近萬個量子比特才能達到的精度。
在 AI 制藥方向,我們開發的 LM-Design 模型,利用大量蛋白質序列數據,以及一定數量的蛋白質結構和序列對應數據,學習從蛋白質結構到序列轉換的模型,達到了目前蛋白質序列設計的最高精度。
機器學習與量子化學
物理學家狄拉克曾說:對大部分物理學和整個化學,進行數學建模所需要的基本定律已完全清楚,困難只在于這些定律的應用,得到的方程一般都太復雜而無法求解。
量子化學是根據量子力學的原理研究化學現象的學科。其重要的問題是用計算的方法求解分子或周期性體系(如固體)的電子薛定諤方程,從而推算出分子或周期性體系的基態能量、電極性等特性。是所謂的從頭計算(ab initio)問題。傳統的方法有密度泛函理論 DFT、耦合簇 CCSD 等。要么計算的精度不夠高,要么能計算的規模不夠大。
近年,用機器學習的方法解決從頭計算問題成為倍受關注的新方向。其基本想法是借助深度學習強大的表示和學習能力,大幅提升從頭計算的精度和規模。其中的一個路徑是 NN-VMC(Neural Network based Variational Monte Carlo) 。用神經網絡近似薛定諤方程的波函數,通過隨機采樣的方式獲得體系中電子在空間中的樣本,這樣可以計算基于薛定諤方程的整個體系的能量。通過最小化能量的上界,優化神經網絡的參數,不斷迭代,最后得到近似最優的神經網絡(波函數),以及體系的近似基態能量(最小能量)。(注:波函數的平方是電子在空間出現的概率密度函數,有了波函數,就可以進行電子在空間中的隨機采樣。)圖 1 顯示 NN-VMC 的基本原理。其核心問題是如何設計神經網絡和學習算法。
圖 1. NN-VMC 方法的基本原理
NN-VMC 中有代表性的方法是 DeepMind 和 ICL 于 2019 年提出的 FermiNet。之后一些研究機構又提出了一些新的方法。ByteDance Research 從 2021 年起,與北京大學合作,進行了一系列相關研究,提出了幾個新的方法。下面對這些方法做一簡單介紹。
NN-VMC+ECP,是我們開發的結合 NN-VMC 和贗勢 ECP(Effective Core Potential)的方法 [1],可以進一步提高計算的效率和體系的規模。計算化學體系的特性時,往往只需要關注原子中外側軌道的電子。將原子中內側軌道的電子的勢能用定量表示,就可以大幅減少所需要的計算量。我們將 ECP 技巧應用于 NN-VMC,得到了這個新方法,取得了很好的效果。
NN-DMC,是我們提出的將神經網絡和擴散蒙特卡洛法 DMC(Diffusoon Monte Carlo)結合的另一個方法 [2]。DMC 與 VMC 不同,不是計算體系基態能量的上界,而是使用虛時演化來計算體系的基態能量。這個方法,相比 FermiNet 等已有方法也能大幅提高計算的精度和規模。
最近開發的 LapNet 也是一種 NN-VMC 方法 [3],特點是在神經網絡學習時使用前向拉普拉斯算子( Forward Laplacian)。基于薛定諤方程計算體系的能量上界的過程中,需要計算哈密頓算子,包括其中的動能部分。之前的方法都是通過計算相關的黑塞矩陣的方式計算動能,其算法復雜度高,成為學習的一個瓶頸。LapNet 在學習的前向傳播中,通過拉普拉斯算子的計算,直接計算動能,以及哈密頓算子,從而省去了黑塞矩陣的計算。這樣可以大幅提高學習的計算效率。相比 FermiNet,LapNet 有平均 10 倍左右的加速。
ECP、DMC 和 Forward Laplace 屬于三種不同的技術改進(簡化勢能計算、優化采樣,提高計算效率),三個技術結合起來原理上可以更大程度上提高計算規模,也是我們正在嘗試的方法。另外,我們還將 NN-VMC 方法應用于固體的薛定諤方程求解 [4],分子體系的力場 [5]、電極化計算 [6] 等問題,證明了 NN-VMC 方法的實用性。
圖 2 顯示目前 NN-VMC 方法中代表性工作的精度和規模,縱軸表示精度,圓的大小表示規模。我們提出的 LapNet 方法能夠以更高的精度計算更大的體系。最大的體系有 116 個電子。
圖 2. NN-VMC 方法的規模和精度
大規模量子化學計算
通過直接求解薛定諤方程計算化學體系特性(比如基態能量)的方法能處理的規模仍然有限。量子嵌入方法(Quantum Embedding Method)被認為是解決這個問題的一條有效路徑。基本想法是通過分而治之和多精度計算實現大規模化。代表性的方法有密度矩陣嵌入理論 DMET(Density Matrix Embedding Theory)。將體系劃分為若干部分(Fragment),對其中的每個 Fragment 及其對應的環境(Bath)進行高精度計算,對其他部分進行低精度計算。而且根據需要對每個 Fragment 進行并行處理。最后再把高精度計算的結果合并起來,不斷迭代逼近原始體系。這樣可以大幅提高可計算的體系的規模。
圖 3 DMET 方法的直觀解釋
圖 3 示意 DMET 方法的過程。首先對原始體系進行劃分,得到一組 Fragment。假設圖中黃色部分是我們關注的 Fragment,例如是兩個原子。圖中藍色的部分包含環境和其他部分。對關注的 Fragment 及其環境用高精度的方法計算,例如 CCSD,對其他部分用低精度的方法計算,例如,Hartree–Fock 法。對所有的 Fragment 做同樣的并行處理。
具體算法如下。首先,通過低精度求解,得到整體(關注的 Fragment、環境和其他部分)的約化密度矩陣,這個低精度解包含參數。其次,對這個矩陣的 Fragment 及其環境進行奇異值分解,構建投影算符 P(該投影算符僅關注 Fragment 及其環境),利用投影算符構建低維度的體系(圖片),并對其進行高精度求解。之后,將所有 Fragment 的計算結果合并,作為體系整體的近似。最后,通過迭代,調節參數,使得原始的低精度解逐漸逼近合并的高精度解(在 L2 norm 的意義下),直到得到最終結果。
我們基于兩種完全不同的計算范式,實現 DMET 及其變種的 SIE,進行大規模量子化學體系的計算。一是使用經典計算機,另一個是使用量子計算機。本文主要介紹后者的相關工作(前者的工作計劃今后有更大進展時介紹),也稱為量子計算化學。我們考慮在量子計算機上的實現,但只在經典計算機上進行模擬,希望為量子化學的發展做出貢獻。無論是哪種計算范式,DMET 方法使得大規模計算成為可能,我們正在嘗試努力實現 DMET,將可計算的體系提升幾個數量級。
物理學家費曼曾說:自然不是古典力學,如果你想模擬自然,你最好用量子力學。開發量子計算技術的驅動力就是用量子層面的計算設備模擬量子現象。換言之,量子化學是量子計算最合適的應用領域之一。
圖 4. Periodic DMET 方法的示意
我們先后開發了兩個量子計算化學方法,結合量子和經典混合計算機和 DMET 的特點,大幅提高了計算體系的精度和規模。基本想法是用量子計算機實現 DMET 高精度計算的部分,用經典計算機實現 DMET 低精度計算的部分。DMET-ESVQE 計算分子體系 [7], Periodic DMET 計算周期性體系 [8]。前者只用 16 個量子比特,就能實現之前的方法用 144 個量子比特的計算。后者只用 20 個量子比特就能實現之前的方法用將近 1 萬個量子比特的計算。
圖 4 顯示在混合計算機上實現的 DMET Periodic 方法。輸入是晶體,輸出是體系的能量。首先對體系進行劃分,然后并行地計算每一個 Fragment。對關注 Fragment 及其環境在量子計算機上用 U-CCSD 求解。對其他部分在經典計算機上用 Hartree–Fock 法求解。
量子蒙特卡羅法,包括 VMC、DMC,是量子化學的最有效的一系列算法 [9]。我們還提出了將量子計算和量子蒙特卡羅法結合的新方法。這個方法可以體現量子計算對量子化學的一些優勢。具體地,量子計算可以部分解決量子蒙特卡洛法中的符號問題。
AI 制藥
使用 AI 技術輔助藥物發現已經成為被業界廣為接受的新范式。近年有大量的研究,也有一些技術應用于實際的場景。我們進行了基于 AI 技術的藥物設計的研究和開發,包括小分子藥物和大分子藥物(抗體藥物)。
小分子藥物設計過程包括蛋白質靶點的發現,小分子藥物候選的生成,候選與靶點(target)的親和性、候選的成藥性、無毒性等的判斷。目前有 AI 技術可以實現這些藥物的開發步驟。我們開發了基于機器學習的小分子藥物候選生成的方法,MARS 是利用打分函數自動生成候選的方法 [10],DESERT 是根據靶點的形狀自動生成候選的方法 [11]。
MARS 從種子分子開始,不斷編輯分子,直到最后得到最優的小分子藥物候選。生成的過程中使用馬爾可夫鏈蒙特卡洛法(MCMC),其平穩分布是由多個打分函數構成的概率分布。打分函數表示小分子藥物候選的親和性、成藥性、無毒性等。建議分布表示基于圖神經網絡(MPNN)進行小分子藥物候選的編輯前后的條件概率分布。圖神經網絡表示小分子化合物的分子式,結點是原子,邊是化學鍵。對小分子的編輯包括增加新的結點和刪除已有的結點。圖神經網絡上可以預測對小分子的可能的編輯操作(增加或刪除),其參數通過學習得到。MARS 只需要有打分函數、小分子藥物的數據庫(分子式)就可以生成全新的和多樣的小分子藥物候選。目前 MARS 已經被用于實際的小分子藥物設計工作中。
DESERT 通過兩個步驟進行小分子藥物候選的生成。Sketching:采樣與蛋白質靶點的口袋(pocket)形狀互補的藥物候選的形狀,Generating:基于藥物候選的形狀自動生成藥物候選的分子式。圖 5 顯示這個過程。
圖 5. DESERT:自動生成小分子藥物候選
候選和靶點結合,其必要條件是兩者的形狀能夠進行很好的對接。在 Sketching 階段,根據蛋白質靶點的形狀,用啟發式的方法產生候選的形狀。在 Generating 階段,用事先學好的形狀到分子的生成模型 Shape2Mol 根據形狀自動生成分子式。可以利用分子庫里的大量的藥物的分子式和形狀,學習這個生成模型。如圖 6 所示,Shape2Mol 中,編碼器將分子 3D 形狀進行編碼,產生中間表示,解碼器根據中間表示生成分子式。3D 形狀使用體素表示,分子式使用符號序列表示。DESERT 是 2022 年小分子藥物候選生成結合性能最好的方法。
圖 6. 形狀到分子的生成模型 Shape2Mol 的示意
最近我們聚焦在大分子藥物設計上,更一般的問題是蛋白質設計。蛋白質設計包括抗體藥物設計,多肽藥物設計等。如果知道了蛋白質序列(氨基酸序列),就可以預測其結構,知道了蛋白質結構也就能預測其功能。這就是著名的 AlphaFold 做的事情。蛋白質設計可以看作是一個反向的過程。一般從功能出發決定對應的蛋白質結構,再根據蛋白質結構決定對應的蛋白質序列。我們開發了從蛋白質結構生成蛋白質序列的模型 LM-Design。
LM-Design 的輸入是蛋白質結構,輸出是對應的蛋白質序列。LM-Design 由結構編碼器和序列解碼器組成。其中,結構編碼器是已訓練好的表示蛋白質結構的圖神經網絡;序列解碼器基于已預訓練好的大規模蛋白質語言模型(Protein Language Model),與 BERT/Transformer Encoder 相似(使用雙向自注意力),在最后一層插入一個結構適配器(Structural Adaptor)。結構適配器的參數是待學習的。圖 7 給出 LM-Design 的模型架構。
圖 7. 蛋白質結構到序列的生成模型 LM-Design 的架構
LM-Design 的學習和預測是掩碼語言建模(Masked Language Modeling),與 BERT 模型的訓練相似,其目標是多次還原被掩碼的序列中的符號(氨基酸)。也就是說,基于已訓練好的蛋白質語言模型中的信息,以及當前的蛋白質結構信息,對蛋白質序列進行多次改寫。LM-Design 基于全局序列信息對其中很少一部分符號(氨基酸)進行改寫,所以對蛋白質遠距離依存關系能夠進行很好的表示和預測。注:蛋白質折疊之后,序列上距離很遠的氨基酸在結構上也可能很近。
現實中有大量經過測序的蛋白質序列數據,但只有少量的蛋白質結構與序列對齊的數據。LM-Design 的一個優勢是可以利用海量的蛋白質序列數據,充分學習和利用蛋白質序列之間的進化中產生的關聯關系,大幅提高從蛋白質結構到序列生成的預測準確率。此外,我們發現增大預訓練蛋白質模型的規模可以進一步提升準確率。如圖 8 所示,LM-Design 是目前效果最好的蛋白質序列生成模型,圖中圓的大小表示模型的參數量。
圖 8. 蛋白質序列生成方法的精度
-
AI
+關注
關注
87文章
31429瀏覽量
269832 -
機器學習
+關注
關注
66文章
8437瀏覽量
132897 -
深度學習
+關注
關注
73文章
5511瀏覽量
121392
原文標題:字節跳動李航:AI for Science的一些探索和進展
文章出處:【微信號:zenRRan,微信公眾號:深度學習自然語言處理】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論