1. 從“黑盒”到“白盒”為什么我們需要新一代生物分子力場如果你在生物物理、計算化學(xué)或者藥物設(shè)計領(lǐng)域工作過哪怕只是淺嘗輒止也一定被“力場”這個詞折磨過。傳統(tǒng)上我們依賴經(jīng)驗力場如AMBER、CHARMM來模擬蛋白質(zhì)、核酸等生物大分子的動力學(xué)行為。這些力場本質(zhì)上是參數(shù)化的物理方程比如用彈簧模型描述化學(xué)鍵用庫侖定律描述靜電相互作用。它們在過去幾十年里功勛卓著但瓶頸也日益明顯參數(shù)化過程極度依賴專家經(jīng)驗和有限的實驗數(shù)據(jù)對復(fù)雜的非共價相互作用如π-π堆積、鹵鍵和反應(yīng)過程描述不準(zhǔn)更別提那些奇奇怪怪的化學(xué)修飾或金屬離子了。這就像一個經(jīng)驗豐富的老師傅靠著一本祖?zhèn)鞯摹憹M了經(jīng)驗公式的手冊在干活遇到手冊外的新零件就只能靠猜。于是基于機器學(xué)習(xí)的力場ML-FF應(yīng)運而生試圖用數(shù)據(jù)驅(qū)動的方式解決這個問題。早期的ML-FF比如ANI、SchNet等確實在精度上實現(xiàn)了飛躍能夠逼近量子化學(xué)計算的水平。但它們大多有個通病像個計算能力超強但“不講道理”的黑盒。你輸入原子坐標(biāo)它輸出能量和力中間發(fā)生了什么不知道。更重要的是這些模型往往缺乏物理歸納偏置比如等變性——一個分子在空間旋轉(zhuǎn)平移后其能量應(yīng)該不變作用力應(yīng)該協(xié)同變換。沒有這個約束模型可能學(xué)到一些“捷徑”在訓(xùn)練集上表現(xiàn)完美但泛化到新體系時一塌糊涂。這就是“LiTEN-FF”這個工作出現(xiàn)的背景。它不僅僅是一個“更好的”力場而是試圖構(gòu)建一個可擴展、高精度且物理可解釋的“基礎(chǔ)模型”?；A(chǔ)模型這個詞最近在AI界很火指的是在廣泛數(shù)據(jù)上預(yù)訓(xùn)練、能通過微調(diào)適配多種下游任務(wù)的模型。把這種思想用到力場上野心不小我們能否訓(xùn)練一個通用的“分子物理理解模型”讓它既能精確計算小分子的能量又能推演蛋白質(zhì)折疊的路徑甚至預(yù)測藥物與靶點的結(jié)合強度LiTEN-FF給出的技術(shù)答案是線性張量四邊形注意力。這個名字聽起來很唬人但拆解開來它核心解決的就是如何高效、等變地建模原子之間那些復(fù)雜、多體、方向性的相互作用。2. 拆解核心引擎線性張量四邊形注意力如何工作要理解LiTEN-FF必須搞懂它的核心創(chuàng)新點。我們一步步來。2.1 傳統(tǒng)注意力機制在分子建模中的困境Transformer架構(gòu)中的注意力機制通過計算所有節(jié)點原子對之間的關(guān)聯(lián)權(quán)重擅長捕捉長程依賴。但在分子體系中直接套用會遇到幾個死結(jié)標(biāo)量問題標(biāo)準(zhǔn)注意力輸出的是標(biāo)量權(quán)重但分子相互作用中的力和扭矩是向量甚至張量。用標(biāo)量去加權(quán)向量無法保證結(jié)果的等變性。計算復(fù)雜度原子數(shù)為N時標(biāo)準(zhǔn)注意力計算所有原子對復(fù)雜度是O(N2)。對于一個中等大小的蛋白質(zhì)幾千個原子這幾乎是不可承受的。方向性信息缺失化學(xué)鍵有方向氫鍵有方向π軌道也有方向。標(biāo)準(zhǔn)注意力對原子位置的編碼如正弦位置編碼難以有效捕獲這種各向異性的幾何信息。2.2 從“標(biāo)量對”到“張量對”四邊形注意力的直覺LiTEN-FF的“四邊形注意力”靈感很可能來源于一個經(jīng)典的化學(xué)概念二面角。二面角由四個原子連續(xù)定義i-j-k-l它描述了分子骨架的扭轉(zhuǎn)狀態(tài)是決定蛋白質(zhì)三維結(jié)構(gòu)的關(guān)鍵。四邊形注意力將這種思想泛化不再局限于連續(xù)的化學(xué)鍵而是考慮所有可能的原子四元組quadruplet。為什么是四元組因為三點確定一個平面四點才能定義一個具有手性的三維體積和方向。通過考慮四元組模型能自然捕獲到原子環(huán)境的體積、形狀和手性等精細(xì)幾何特征這是三元組三角形無法提供的。例如一個手性碳中心的周圍環(huán)境必須通過四元組才能完整描述其立體構(gòu)型。2.3 “線性”與“張量”實現(xiàn)可擴展與等變的關(guān)鍵“四邊形”解決了“看什么”的問題但直接計算所有四元組復(fù)雜度是恐怖的O(N?)?！熬€性”二字就是來解決這個 scalability可擴展性難題的。這里的“線性”并非指模型是線性的而是指其計算復(fù)雜度通過巧妙的分解最終可以做到與原子數(shù)N成線性關(guān)系O(N)。這是如何實現(xiàn)的其核心是一種因式分解的思想。我們可以想象不直接計算所有四元組 (i, j, k, l) 的相互作用而是將其分解為兩兩原子對相互作用的某種組合。具體來說模型可能通過以下步驟近似先計算所有原子對 (i, j) 的“二階”特征比如距離、方向向量。將這些成對特征進(jìn)行組合來“模擬”或“投票”出四元組的整體特征。這個過程可以通過高效的線性代數(shù)操作如外積、張量收縮來實現(xiàn)避免顯式枚舉四元組?！皬埩俊眲t解決了“輸出什么”的問題。在整個計算流水線中特征不再是標(biāo)量而是幾何張量標(biāo)量、向量、二階張量等。網(wǎng)絡(luò)中的每一層操作都設(shè)計為等變操作例如球諧函數(shù)卷積、張量場網(wǎng)絡(luò)中的操作等。這意味著當(dāng)你旋轉(zhuǎn)或平移整個分子系統(tǒng)時網(wǎng)絡(luò)中間層的特征張量會按照其對應(yīng)的幾何類型如向量會跟著一起旋轉(zhuǎn)同步變換最終保證輸出的能量標(biāo)量不變力向量協(xié)同旋轉(zhuǎn)。這就把物理定律直接“烘焙”進(jìn)了模型架構(gòu)極大地提升了泛化能力和物理合理性。簡單類比傳統(tǒng)的ML-FF像在用一堆像素點原子畫圖而LiTEN-FF則像在用樂高積木原子對、原子四元組搭建筑。樂高積木本身帶有形狀和連接方向張量特征并且有一套標(biāo)準(zhǔn)的拼接規(guī)則等變操作這樣搭出來的建筑預(yù)測的分子構(gòu)象不僅結(jié)構(gòu)穩(wěn)固而且無論從哪個角度看旋轉(zhuǎn)各個部件的關(guān)系都是正確的。3. LiTEN-FF模型架構(gòu)與訓(xùn)練實戰(zhàn)解析光有核心思想不夠我們來看看LiTEN-FF具體是怎么搭建和訓(xùn)練的。這部分我會結(jié)合常見的等變網(wǎng)絡(luò)實踐來還原其可能的架構(gòu)細(xì)節(jié)。3.1 輸入編碼從原子坐標(biāo)到幾何張量模型的輸入是原子的類型如C N O和三維坐標(biāo)。第一步是將其轉(zhuǎn)化為豐富的、等變的初始特征。原子嵌入每個原子類型通過一個查找表映射為一個高維標(biāo)量特征向量。相對幾何編碼對于每一對原子 (i, j)計算標(biāo)量部分原子間距離 r_ij。通常會通過一組高斯徑向基函數(shù)將其展開得到一個向量用于編碼距離的“軟”信息。向量部分歸一化的相對位置向量 r_ij / ||r_ij||。這是最基礎(chǔ)的向量特征。高階張量部分可選為了更精細(xì)地描述方向可以將相對位置向量投影到球諧函數(shù)基上生成更高階的張量特征如二階球諧函數(shù)對應(yīng)著“啞鈴形”的方向偏好。構(gòu)建初始相互作用對將原子i的標(biāo)量嵌入、原子j的標(biāo)量嵌入以及它們之間的幾何編碼標(biāo)量向量...結(jié)合起來形成原子對 (i, j) 的初始張量特征。這個特征同時包含了化學(xué)身份信息和空間信息。3.2 線性張量四邊形注意力層信息傳遞的核心這是模型的堆疊層。每一層接收上一層的原子/原子對特征并輸出更新后的特征。消息計算基于四邊形對于中心原子i我們需要聚合來自其他原子的信息。傳統(tǒng)圖神經(jīng)網(wǎng)絡(luò)只考慮鄰居j。在這里L(fēng)iTEN-FF考慮的是以i和另一個原子j為邊的“四邊形背景”。對于每一對 (i, j)模型會查找一個“上下文”原子k可能通過最近鄰選擇然后考慮四元組 (i, j, k, ...) 的特征。計算消息時不僅使用i和j的特征還融入了k的特征從而感知局部三維幾何環(huán)境。實現(xiàn)上這可以通過將 (i, k) 和 (j, k) 的對特征進(jìn)行張量乘積如外積或注意力加權(quán)來實現(xiàn)其輸出是一個新的張量消息 m_ij。由于k的選擇是局部的且計算經(jīng)過線性化分解整個過程的理論復(fù)雜度可以控制在線性級別。消息聚合與更新對于原子i將所有來自j的消息 m_ij 進(jìn)行聚合例如求和。然后將這個聚合后的消息與原子i自身的特征進(jìn)行融合通過可學(xué)習(xí)的非線性函數(shù)如門控機制更新原子i的特征。等變性保持上述所有涉及向量/張量特征的操作如張量乘積、加權(quán)求和都必須使用等變線性層如球諧卷積核、Clebsch-Gordan系數(shù)約束的線性層來處理確保輸入輸出張量的變換規(guī)則一致。3.3 輸出頭與損失函數(shù)經(jīng)過多層這樣的注意力網(wǎng)絡(luò)后每個原子都獲得了包含豐富長程和幾何上下文的特征。能量預(yù)測將每個原子最終的標(biāo)量特征求和再通過一個簡單的多層感知機MLP得到整個分子的總勢能。求和操作本身是等變的對平移旋轉(zhuǎn)不變。力預(yù)測力是能量對原子坐標(biāo)的負(fù)梯度。得益于等變架構(gòu)我們可以利用自動微分直接對預(yù)測的能量求坐標(biāo)的梯度來得到力。這是最直接且物理自洽的方式。另一種方式是直接用原子特征回歸出力向量。訓(xùn)練損失損失函數(shù)通常是能量均方誤差和力均方誤差的加權(quán)和。Loss λ_E * MSE(E_pred, E_QM) λ_F * MSE(F_pred, F_QM)其中QM代表高精度的量子化學(xué)計算參考值如DFT。力的權(quán)重λ_F通常設(shè)得更大因為力的準(zhǔn)確對分子動力學(xué)模擬的穩(wěn)定性至關(guān)重要。注意數(shù)據(jù)來源的挑戰(zhàn)。訓(xùn)練這樣一個基礎(chǔ)模型需要海量、多樣的量子化學(xué)數(shù)據(jù)。這通常來自公共數(shù)據(jù)集如QM9 ANI-1x SPICE和自建計算。數(shù)據(jù)質(zhì)量理論方法、基組和覆蓋的化學(xué)空間元素、成鍵類型、構(gòu)象直接決定了模型的上限。4. 超越基準(zhǔn)測試LiTEN-FF在實際場景中的潛力與挑戰(zhàn)在論文里看它在標(biāo)準(zhǔn)測試集上“吊打”前輩當(dāng)然爽但我們更關(guān)心這東西到底能用來干嘛用起來又會遇到什么坑4.1 潛在應(yīng)用場景拆解高精度分子動力學(xué)模擬這是最直接的應(yīng)用。傳統(tǒng)力場搞不定的體系比如涉及電荷轉(zhuǎn)移、激發(fā)態(tài)、化學(xué)反應(yīng)路徑的可以用LiTEN-FF進(jìn)行納秒甚至微秒級的模擬獲得接近量子化學(xué)精度的軌跡用于研究酶催化機制、藥物結(jié)合動力學(xué)等。蛋白質(zhì)結(jié)構(gòu)預(yù)測與設(shè)計補充雖然AlphaFold2已經(jīng)解決了單鏈蛋白結(jié)構(gòu)預(yù)測的主要矛盾但在預(yù)測蛋白質(zhì)-蛋白質(zhì)復(fù)合物、蛋白-小分子結(jié)合構(gòu)象、以及涉及構(gòu)象變化的別構(gòu)調(diào)節(jié)時物理力場仍然不可或缺。LiTEN-FF可以作為這些“對接”或“柔性對接”問題的精修工具。材料發(fā)現(xiàn)不止生物分子有機半導(dǎo)體、電解質(zhì)、金屬有機框架等材料的性質(zhì)也高度依賴于分子間相互作用。一個通用的力場基礎(chǔ)模型可以快速篩選候選材料預(yù)測其堆積模式、載流子遷移率等。作為下游任務(wù)的預(yù)訓(xùn)練模型你可以把訓(xùn)練好的LiTEN-FF模型“凍結(jié)”將其原子特征提取器作為一個強大的3D分子表示生成器用于下游的分子性質(zhì)預(yù)測溶解度、毒性、活性、反應(yīng)預(yù)測等任務(wù)可能只需要少量數(shù)據(jù)微調(diào)就能取得好效果。4.2 實操中的挑戰(zhàn)與應(yīng)對思路然而把論文模型搬到自己的研究項目中絕不會一帆風(fēng)順。挑戰(zhàn)一計算成本依然高昂。盡管是線性復(fù)雜度但涉及張量操作和四元組信息其單步計算開銷仍遠(yuǎn)大于傳統(tǒng)力場。一次能量/力評估可能需要毫秒到秒級對于需要數(shù)億步的長時間模擬這依然是天文數(shù)字。應(yīng)對策略混合方案采用QM/ML/MM分層策略。用LiTEN-FF處理核心感興趣的區(qū)域如酶的活性口袋用傳統(tǒng)力場處理外圍的溶劑和蛋白骨架。主動學(xué)習(xí)與蒸餾用LiTEN-FF生成高質(zhì)量數(shù)據(jù)去訓(xùn)練一個更輕量級的、專門針對目標(biāo)體系的“學(xué)生”模型如消息傳遞神經(jīng)網(wǎng)絡(luò)用于生產(chǎn)模擬。硬件依賴這類模型極度依賴GPU加速。確保你的代碼充分利用CUDA和大型顯存。挑戰(zhàn)二化學(xué)空間的泛化能力邊界。模型在訓(xùn)練數(shù)據(jù)覆蓋的化學(xué)空間內(nèi)表現(xiàn)良好但對于全新的元素如過渡金屬、極端成鍵高張力環(huán)、自由基、或罕見的化學(xué)環(huán)境預(yù)測可能失效。應(yīng)對策略系統(tǒng)性驗證在新體系上使用前務(wù)必在少量代表性構(gòu)象上做量子化學(xué)計算與模型預(yù)測進(jìn)行對比驗證。不確定性量化先進(jìn)的ML-FF會附帶不確定性估計。關(guān)注模型預(yù)測方差大的區(qū)域這些地方可能就是其認(rèn)知的“盲區(qū)”需要特別小心或補充數(shù)據(jù)。持續(xù)微調(diào)如果研究聚焦于某一類特定分子比如某一類金屬酶收集該體系的高質(zhì)量量子化學(xué)數(shù)據(jù)對預(yù)訓(xùn)練的LiTEN-FF進(jìn)行微調(diào)是提升其在該領(lǐng)域表現(xiàn)的最有效方法。挑戰(zhàn)三軟件集成與工作流搭建。LiTEN-FF是一個研究型模型其代碼可能基于PyTorch或JAX。而主流的分子動力學(xué)軟件如GROMACS AMBER LAMMPS有自己的一套插件接口。應(yīng)對策略尋找官方/社區(qū)接口關(guān)注作者是否提供了例如torchscript或ONNX格式的導(dǎo)出工具以及主流MD軟件對應(yīng)的插件如DeePMD-kit之于LAMMPSTorchANI之于AMBER。自定義集成如果沒有現(xiàn)成方案你可能需要編寫一個“客戶端”程序。該程序從MD軟件讀取坐標(biāo)調(diào)用LiTEN-FF模型計算能量和力再寫回MD軟件。這需要熟悉MD軟件的能量/力計算接口如GROMACS的pull code或自定義勢能函數(shù)接口。這個過程調(diào)試起來很繁瑣要做好心理準(zhǔn)備。5. 從理論到實踐一個簡化的概念驗證工作流假設(shè)我們現(xiàn)在想用LiTEN-FF或類似架構(gòu)的模型來研究一個小分子抑制劑與靶點蛋白的結(jié)合模式。以下是一個高度簡化的概念性工作流展示了可能涉及的步驟和決策點。體系準(zhǔn)備獲取蛋白-抑制劑復(fù)合物的初始晶體結(jié)構(gòu)來自PDB。用分子建模軟件如MDAnalysis,OpenBabel處理結(jié)構(gòu)加氫、優(yōu)化質(zhì)子化狀態(tài)、補充缺失殘基。將體系放入一個水盒子中并添加生理離子濃度。此時你得到了一個包含數(shù)萬原子的初始體系。模型準(zhǔn)備與驗證選擇模型評估LiTEN-FF官方提供的預(yù)訓(xùn)練模型看其訓(xùn)練數(shù)據(jù)是否覆蓋了你的體系中的所有元素C H O N S 可能還有金屬離子。局部驗證從復(fù)合物中提取出抑制劑分子以及結(jié)合口袋內(nèi)的關(guān)鍵氨基酸殘基約50-100個原子生成幾個不同的構(gòu)象如輕微扭轉(zhuǎn)二面角。對這些小體系進(jìn)行高精度DFT計算作為基準(zhǔn)然后用LiTEN-FF計算能量和力對比誤差。確保在核心區(qū)域誤差在可接受范圍如能量誤差1 kcal/mol/atom力誤差幾kcal/mol/?。模擬設(shè)置混合方案示例分區(qū)將體系分為三層QM/ML區(qū)抑制劑分子 結(jié)合口袋內(nèi)直接相互作用的殘基側(cè)鏈約100-200原子。此區(qū)用LiTEN-FF。MM區(qū)蛋白質(zhì)骨架、其余殘基、水分子和離子。此區(qū)用傳統(tǒng)力場如AMBER ff19SB。邊界處理在QM/ML區(qū)和MM區(qū)邊界需要小心處理。通常會將邊界處的化學(xué)鍵“切斷”并用鏈接原子如氫原子來飽和價鍵或者使用更復(fù)雜的機械/靜電嵌入方案。運行與采樣由于LiTEN-FF計算慢長時間平衡模擬不現(xiàn)實?？梢圆捎迷鰪姴蓸臃椒ㄈ缭獎恿W(xué)。定義一兩個關(guān)鍵的反應(yīng)坐標(biāo)如抑制劑與蛋白關(guān)鍵殘基的距離、某個二面角。在反應(yīng)坐標(biāo)上施加偏置勢驅(qū)動體系在結(jié)合態(tài)與非結(jié)合態(tài)之間快速轉(zhuǎn)換從而在較短的計算時間內(nèi)幾十到幾百納秒等效時間采集到足夠的構(gòu)象樣本用于計算結(jié)合自由能。結(jié)果分析與迭代分析模擬軌跡觀察穩(wěn)定的結(jié)合模式、關(guān)鍵相互作用氫鍵、疏水接觸。計算結(jié)合自由能如通過MM/PBSA或TI方法。如果結(jié)果與實驗不符需要排查。是初始結(jié)構(gòu)問題質(zhì)子化狀態(tài)錯了還是力場在某個特定相互作用如鹵鍵上描述不準(zhǔn)可能需要回到第2步針對有疑問的局部相互作用補充量子化學(xué)數(shù)據(jù)對模型進(jìn)行微調(diào)。核心心得從“能用”到“用好”。使用這類前沿模型最大的陷阱是“黑箱信任”。絕不能因為它在基準(zhǔn)測試中表現(xiàn)好就認(rèn)為它在你的特定問題上一定可靠。必須建立一套驗證-質(zhì)疑-再驗證的流程。先從最小的、可驗證的子系統(tǒng)開始確保模型在你的化學(xué)空間里是“清醒”的然后再逐步應(yīng)用到更大的復(fù)雜體系中。計算成本高就更要追求每一次模擬的設(shè)計質(zhì)量用增強采樣等方法“把錢花在刀刃上”。