|
隨著擴(kuò)散模型殺入蛋白質(zhì)設(shè)計(jì)領(lǐng)域，時(shí)隔幾天，Grigoryan和Ingraham等人以Generate Biomedicines公司的名義發(fā)表了一個(gè)新的擴(kuò)散蛋白生成模型Chroma。它的創(chuàng)新點(diǎn)為通過隨機(jī)圖神經(jīng)網(wǎng)絡(luò)的架構(gòu)進(jìn)行低復(fù)雜度的長距離推理，通過一個(gè)基于凸優(yōu)化的結(jié)構(gòu)模塊從預(yù)測的殘基間幾何信息高效地預(yù)測蛋白質(zhì)的三維結(jié)構(gòu)，通過低溫采樣提升擴(kuò)散效果，以及針對特定分布或性質(zhì)的定向生成。基于以上方法，Chroma可以使蛋白質(zhì)設(shè)計(jì)成為外部約束下的貝葉斯推理過程，在生成大量可設(shè)計(jì)、新穎的蛋白基礎(chǔ)上，可以根據(jù)根據(jù)對稱性、motif、性質(zhì)分類器等等進(jìn)行蛋白工程，對于binder、抗體等設(shè)計(jì)應(yīng)用。除此以外本文還進(jìn)行了一些創(chuàng)新實(shí)驗(yàn)包括根據(jù)形狀、自然語言指導(dǎo)生成，類似AIGC中的模型進(jìn)行基于特定描述的圖像生成一樣，進(jìn)行蛋白生成，這展示了Chroma的應(yīng)用潛力。 2. 模型方法
Chroma的主要架構(gòu)如上圖展示，其中多個(gè)模塊都是在已有模型上的改進(jìn)，比如隨機(jī)圖神經(jīng)網(wǎng)絡(luò)是在Ingraham2019年發(fā)表的GNN架構(gòu)[1]上進(jìn)行的改進(jìn)，解碼器借鑒了Terminator，側(cè)鏈網(wǎng)絡(luò)借鑒了Possu Huang 3DCNN的設(shè)置[2]，坐標(biāo)更新方法也仿照了Alphafold2的設(shè)置，但都進(jìn)行了效率上的改進(jìn)。主要的擴(kuò)散模型在a圖展示，包括主鏈網(wǎng)絡(luò)和設(shè)計(jì)網(wǎng)絡(luò)兩部分。蛋白結(jié)構(gòu)在正向擴(kuò)散過程中在被加以鏈和回旋半徑約束的情況下加噪，并在逆向過程中去噪獲得采樣主鏈，并被最后輸入到一個(gè)基于圖的設(shè)計(jì)網(wǎng)絡(luò)中采樣獲得蛋白序列和側(cè)鏈構(gòu)象。b圖展示了在擴(kuò)散中使用的隨機(jī)圖神經(jīng)網(wǎng)絡(luò)和結(jié)構(gòu)模塊，最終從加噪結(jié)構(gòu)中獲得了去噪結(jié)構(gòu)。c圖展示了與擴(kuò)散模型先驗(yàn)結(jié)合的條件模型，其關(guān)鍵在于將給定約束條件的蛋白質(zhì)條件分布的擴(kuò)散過程加入來自外部分類器的梯度，本文這里的外部分類器包含范圍廣泛，指任何從加噪樣本預(yù)測標(biāo)簽的監(jiān)督模型，具體屬性見下文。2.1 結(jié)構(gòu)約束擴(kuò)散模型現(xiàn)有的一些蛋白擴(kuò)散模型方法會在二面角、接觸圖上等低維度信息上進(jìn)行，本文指出由于高斯噪音方法屬于無關(guān)擴(kuò)散，在三維信息上進(jìn)行高斯噪音去噪需要模型從頭學(xué)習(xí)這些基本的結(jié)構(gòu)關(guān)系，而這往往又是不理想的。噪聲損害了蛋白結(jié)構(gòu)都會天然滿足的簡單的鏈和密度約束，另一方面，高噪聲樣本又會因?yàn)樽陨矸植纪獾男再|(zhì)使得在特定數(shù)據(jù)上訓(xùn)練的針對蛋白體系的神經(jīng)網(wǎng)絡(luò)架構(gòu)性能難以發(fā)揮，這可能使得模型會較難學(xué)到基本約束并生成不符合常識的序列結(jié)構(gòu)。本文則回避了這一問題，通過白化去除數(shù)據(jù)中冗余的相關(guān)性關(guān)系，并在白化坐標(biāo)上進(jìn)行無關(guān)擴(kuò)散，并引入了SO(3)不變性的保留鏈和密度約束的線形復(fù)雜度的多變量高斯分布以實(shí)現(xiàn)這些基本約束，這等同于在原始坐標(biāo)上進(jìn)行相關(guān)擴(kuò)散。模型訓(xùn)練使用了以2.6埃，50% identity為閾值篩選得到的28819個(gè)pdb以及1726個(gè)90% identity過濾的抗體，通過PyRosetta填補(bǔ)缺失，8:2:2根據(jù)相似度劃分?jǐn)?shù)據(jù)集。主鏈模型在8個(gè)GPU訓(xùn)練優(yōu)化ELBO損失函數(shù)，設(shè)計(jì)網(wǎng)絡(luò)在1個(gè)GPU優(yōu)化序列和chi角的損失函數(shù)。本文首先指出了似然模型的通病，即最大似然擬合的低容量模型通常會過度分散。解決的思路一般為以減少樣本多樣性為代價(jià)增加對高似然狀態(tài)的采樣，比如貪婪、波束采樣等方法。低溫采樣也是其中一種方法，在原本采樣分布基礎(chǔ)上添加溫度項(xiàng)以實(shí)現(xiàn)多樣性和似然的平衡。對于擴(kuò)散模型也存在同樣的多樣性和質(zhì)量的平衡，Dhariwal等人[3]表示簡單的調(diào)整，如對分?jǐn)?shù)函數(shù)或噪聲的縮放是有問題的，在DALL-E 2、Imagen等模型中往往使用的基于分類器的引導(dǎo)方法。本文提出進(jìn)行了2個(gè)操作并將兩者融合使用：1. 對反向擴(kuò)散的分?jǐn)?shù)函數(shù)進(jìn)行溫度項(xiàng)進(jìn)行縮放，值得注意的是雖然該方法可以起到平衡作用，但不足以在溫度擾動的分布中進(jìn)行適當(dāng)?shù)闹匦录訖?quán)。2. 引入退火Langevin動力學(xué)的思想從低溫分布中采樣，加入平衡率選項(xiàng)，雖然該方法需要有足夠的平衡時(shí)間。本文也發(fā)現(xiàn)如果沒有低溫采樣的過程，生成結(jié)果中會包含不少沒有二級結(jié)構(gòu)的不合理蛋白，上圖展示了在3個(gè)相同的隨機(jī)種子的案例上調(diào)節(jié)溫度帶來的二級結(jié)構(gòu)上的生成質(zhì)量提升。
2.3 隨機(jī)圖神經(jīng)網(wǎng)絡(luò)
與以往的高復(fù)雜度的網(wǎng)絡(luò)相比，本文復(fù)現(xiàn)了高效N-body模擬方法Barnes-Hut算法，并基于消息傳遞網(wǎng)絡(luò)和隨機(jī)圖構(gòu)建了隨機(jī)圖神經(jīng)網(wǎng)絡(luò)。這么做的前提之一是計(jì)算物理學(xué)的主要經(jīng)驗(yàn)，涉及O(N2)密集相互作用的N-body模擬往往只需要O(NlogN)尺度的計(jì)算就可以有效模擬。另一個(gè)前提是在transformer的衍生應(yīng)用中有不少也探索了在降低計(jì)算復(fù)雜度的前提下加入長距離信息。目前的文獻(xiàn)中，大多數(shù)蛋白質(zhì)結(jié)構(gòu)的圖神經(jīng)網(wǎng)絡(luò)或者以基于局部連接，例如k-NN，或全連接圖的方式處理信息。局部圖對于高度空間定位的任務(wù)來說是很自然的，比如結(jié)構(gòu)條件下的序列設(shè)計(jì)和特定殘基環(huán)境的描述，并且計(jì)算負(fù)擔(dān)較小，但是否能夠有效地推理全局結(jié)構(gòu)依舊存疑。本文的方法對短距離和長距離連接混合的隨機(jī)圖進(jìn)行采樣。具體來說，每個(gè)節(jié)點(diǎn)的邊從潛在的邊集合中不放回地加權(quán)抽樣，權(quán)重來自于空間距離，最終圖使用20近鄰邊以及40隨機(jī)逆立方采樣邊（Barnes-Hut方法）。該方法在于統(tǒng)一采樣、指數(shù)采樣相比，該方法沒有丟失局部信息也沒有引入額外的長度分布，保留了短距離和長距離的平衡，邊際采樣概率隨距離緩慢單調(diào)地衰減。2.4 坐標(biāo)更新在Alphafold的工作中，生成蛋白質(zhì)幾何結(jié)構(gòu)的神經(jīng)網(wǎng)絡(luò)對整個(gè)分子系統(tǒng)進(jìn)行了迭代局部坐標(biāo)更新。本文將該問題類比為機(jī)器學(xué)習(xí)中的 '結(jié)構(gòu)化預(yù)測 '問題學(xué)習(xí)勢能函數(shù)，并利用優(yōu)化和幾何方法進(jìn)行解決。關(guān)于Alphafold可以學(xué)習(xí)到勢能函數(shù)也在Sergey等人[4]的論文中得到證實(shí)。本文提出了一種新的凸優(yōu)化方法實(shí)現(xiàn)殘基間幾何預(yù)測和端到端可微調(diào)優(yōu)化，可以將成對的殘基間幾何信息和置信度信息轉(zhuǎn)換為成對的旋轉(zhuǎn)平移信息，將問題變換為在旋轉(zhuǎn)平移空間內(nèi)通過局部分析與反復(fù)迭代求解全局優(yōu)化結(jié)構(gòu)的一個(gè)凸優(yōu)化問題。在上述過程中，局部優(yōu)化時(shí)鄰居是固定的，目標(biāo)函數(shù)被定義為：
其中w為置信度，t代表平移，O代表旋轉(zhuǎn)，Proj函數(shù)為Kabsch 算法的投影以計(jì)算優(yōu)化RMSD的位置，這個(gè)過程被迭代以求得優(yōu)化后的局部pose，并最終使用Alphafold2中相同的函數(shù)再轉(zhuǎn)換為坐標(biāo)。另外本文也進(jìn)行了一些局部信息增強(qiáng)的操作。2.5 可編程性Chroma的另一個(gè)比較大亮點(diǎn)是它的可編程性，可以針對下游任務(wù)的需求進(jìn)行條件性偏向性采樣，比如在訓(xùn)練時(shí)引入的特定的標(biāo)簽。
上圖展示了本文進(jìn)行的條件生成的類型，主要包括：通過指定某些氨基酸對在空間上接近，可以在包括binder設(shè)計(jì)，環(huán)肽設(shè)計(jì)等場景上應(yīng)用。操作方式為兩種，或者訓(xùn)練一個(gè)分類器判別在去噪之后某些位點(diǎn)間距離是否小于某閾值，或者根據(jù)擴(kuò)散過程中的坐標(biāo)信息直接計(jì)算某些位點(diǎn)間的距離以擬合分布。2. 子結(jié)構(gòu)RMSD匹配（motif）這里的應(yīng)用場景是基于具有潛在功能的特定子結(jié)構(gòu)進(jìn)行生產(chǎn)，這個(gè)特定子結(jié)構(gòu)可以是某結(jié)合位點(diǎn)，可以是結(jié)合骨架、也可以是特定的表位。操作方式同樣分兩種，或者為建立一個(gè)分類器以評價(jià)去噪后的結(jié)構(gòu)與motif之間的RMSA是否小于某閾值，或者同樣根據(jù)擴(kuò)散過程計(jì)算在不同時(shí)間節(jié)點(diǎn)的結(jié)構(gòu)匹配RMSD，擬合RMSD變化與噪聲和motif之間函數(shù)，構(gòu)建一個(gè)RMSD的分類器。這里的應(yīng)用場景為不僅給定了子結(jié)構(gòu)給定，也同樣給定了具體的序列位置，即殘基編號。此時(shí)對給定片段之間的缺失部位的采樣變?yōu)榱穗p側(cè)節(jié)點(diǎn)為缺失部位兩側(cè)片段的布朗橋，對分類器的評價(jià)依舊是RMSD。另外本文也設(shè)置了專門的損失函數(shù)避免采樣過程中的碰撞和不連續(xù)。Chroma同樣可以處理對稱設(shè)計(jì)問題，這對于大型蛋白對稱組裝體很有幫助。同樣本文在對可能破壞對稱性的錯(cuò)誤進(jìn)行限制后，提出了2種操作方式，一種稱為對稱廣播，計(jì)算全局梯度更新并將其和噪音投影到所有對稱位置；另一種稱為對稱平均，平均對稱子單元的更新和噪音，這兩種方式都可用于更新采樣過程并引進(jìn)對稱性。除了上述的很貼近應(yīng)用的一些條件采樣之外，作者還探索了一些獨(dú)特的采樣方式?？紤]到蛋白質(zhì)經(jīng)常通過特定的形狀實(shí)現(xiàn)特定的功能，比如受體、細(xì)胞孔隙等，本文探索了利用形狀約束對蛋白質(zhì)進(jìn)行采樣。約束方法來自最優(yōu)傳輸方法[5] ，這種方法可用于比較殘基之間的對應(yīng)關(guān)系和幾何距離。通過對預(yù)測距離矩陣和理想形狀的距離矩陣計(jì)算Wasserstein距離（點(diǎn)云距離）和Gromov- Wasserstein距離（不同domain之間的未對齊距離），添加形狀的損失函數(shù)用于優(yōu)化采樣。
在以往的方法中對于特定領(lǐng)域的遷移往往使用微調(diào)或者自適應(yīng)遷移等方法。本文則另訓(xùn)練了一個(gè)模型ProClass對于加噪的主鏈坐標(biāo)信息進(jìn)行分類，該模型的架構(gòu)如上所示，訓(xùn)練數(shù)據(jù)為CATH、PFAM、二級結(jié)構(gòu)、界面殘基等等信息，這意味這擴(kuò)散模型也可以根據(jù)該分類器向著特定二級結(jié)構(gòu)，界面殘基，家族進(jìn)行生成。
本文最新穎的案例則為仿照AIGC的方法，從特定的自然語言提示出發(fā)進(jìn)行蛋白生成。本文訓(xùn)練了一個(gè)從文本到蛋白質(zhì)骨架的生成模型ProCap，該模型使用在科學(xué)文本上訓(xùn)練的1.25億參數(shù)模型GPT-Neo，期望它能夠理解和蛋白質(zhì)相關(guān)的語言描述，也使用了ProClass的網(wǎng)絡(luò)以編碼主鏈。模型的輸入包括由ProClass編碼的結(jié)構(gòu)和由GPT-Neo編碼的語言描述和任務(wù)字符，這里的任務(wù)字符用于區(qū)分是單鏈任務(wù)還是復(fù)合物任務(wù)。如上圖架構(gòu)所示，ProCap在語言和結(jié)構(gòu)模型外，通過可學(xué)習(xí)的線性層來連接這兩個(gè)模型，一個(gè)可學(xué)習(xí)的語言模型頭來將原始語言模型輸出轉(zhuǎn)換為概率。ProCap在約45000個(gè)由PDB和Uniprot組成的文本描述-結(jié)構(gòu)的數(shù)據(jù)對上使用單張V100訓(xùn)練，其中的文本描述來自網(wǎng)站上的描述性文本或者注釋。在結(jié)果展示部分，文章依次對擴(kuò)散非條件生成和各種可編程性的定向生成結(jié)果進(jìn)行了基于拆分?jǐn)?shù)據(jù)集的分析。生成過程效率很高，只需要使用一般GPU，對于3000個(gè)氨基酸以上的復(fù)合物生成只需要幾分鐘。
首先在非條件生成上，作者比較了低溫采樣帶來的影響，如果沒有低溫采樣，模型會天然結(jié)構(gòu)分配了很高的似然，同時(shí)生成蛋白大多為不合理的結(jié)構(gòu)，比如有很少的二級結(jié)構(gòu)。而加入了低溫采樣則會緩解該問題，圖左側(cè)展示了大量生成出的結(jié)構(gòu)，圖右側(cè)則對一些單鏈性質(zhì)，比如二級結(jié)構(gòu)分布、contact密度、長距離contact，回轉(zhuǎn)半徑等進(jìn)行了生成蛋白和和天然蛋白的比較，兩者的相似度很高，生成蛋白的螺旋比較多，這一定程度上說明生成結(jié)果的合理性。另一方面本文也對生成結(jié)果進(jìn)行了Omegafold評價(jià)、與天然TERMs的比對分析均顯示了較高的可設(shè)計(jì)性。
1. 幾何約束（a圖）：展示了對稱性采樣，上圖展示了更多的對稱生成樣本，包括各種的循環(huán)對稱、二面體對稱、四面體對稱、八面體對稱和二十面體對稱。這樣的對稱性設(shè)置很嚴(yán)格，作者發(fā)現(xiàn)甚至有一些對稱的樣本中間出現(xiàn)了空腔，不具備良好的接觸界面，這也是后續(xù)可能需要改進(jìn)的問題。受益于本文的高效計(jì)算，這里展示的最大案例為包含96000個(gè)殘基的二十聚體。2. 子結(jié)構(gòu)（b圖）：本文展示了5個(gè)案例，包括對二氫葉酸還原酶一半結(jié)構(gòu)的重構(gòu)，抗體對CDR的重構(gòu)，其他預(yù)先指定的motif（αββ packing motif，編碼糜蛋白酶催化三聯(lián)體活性部位的骨架片段，EF-hand Ca-binding motif）的自由生成，結(jié)果顯示可以生成具備多樣性的合理的結(jié)構(gòu)。3. 形狀（c圖），本文以英文字母和阿拉伯?dāng)?shù)字作為形狀約束進(jìn)行生成，從結(jié)果來看雖然無法評價(jià)合理性，但是在二級結(jié)構(gòu)，側(cè)鏈packing方面都較為正常。
上圖展示了通過ProClass和ProCap進(jìn)行條件生成的結(jié)果，所有豎行之間都共享相同的隨機(jī)種子有對應(yīng)關(guān)系，可以看到在二級結(jié)構(gòu)方面有著肉眼可見的對特定二級結(jié)構(gòu)組合的富集，在折疊類型的生成上對于beta桶和Ig折疊的生成效果較好，但是對于Rossman折疊上則只具備一些特征，出現(xiàn)了和非條件生成一樣的很長的單螺旋。右側(cè)則展示了自然語言條件的生成結(jié)果，可以看到條件生成結(jié)果使得模型的困惑度的到了提升，甚至在一個(gè)案例上超越了PDB結(jié)構(gòu)。另外作者也在案例分析時(shí)指出，語言模型條件并不是經(jīng)常生效的，一個(gè)影響很大的因素是語言模型的梯度和結(jié)構(gòu)擴(kuò)散模型梯度的融合，語言模型梯度權(quán)重過小會導(dǎo)致和無條件生成沒有區(qū)別，大了又會導(dǎo)致生成樣本質(zhì)量下降，這些都是后續(xù)需要討論的問題。
另外本文比較了在生成中添加來自ProCap文本描述是否為生成蛋白帶來采樣上的改善，即TM分?jǐn)?shù)的變化和損失函數(shù)的變化。上圖展示了對應(yīng)上文的2個(gè)自然語言案例上加入語言條件與否對于兩個(gè)指標(biāo)的影響，可以看到兩者之間最初都是比較嘈雜的，但慢慢地兩個(gè)指標(biāo)上條件模型均獲得了穩(wěn)定指標(biāo)提升，這也體現(xiàn)了語言模型條件在優(yōu)化蛋白上的潛力。
最后則是對整個(gè)蛋白空間進(jìn)行可視化，對應(yīng)標(biāo)題照亮蛋白空間。本文與之前蛋白空間可視化的語言模型相比，有著明確的結(jié)構(gòu)信息并可以將這些結(jié)構(gòu)信息與空間位置進(jìn)行嚴(yán)格對應(yīng)，正如左側(cè)展示的一樣，相似的類別二級結(jié)構(gòu)比例，大小，在空間中產(chǎn)生了明顯聚集。這里引入了一個(gè)新穎度評價(jià)方式“CATH-可壓縮性”，即在TM分?jǐn)?shù)高于0.5的情況下，實(shí)現(xiàn)對蛋白質(zhì)中80%的殘基的貪婪覆蓋所需的CATH域的數(shù)量，作者意圖說明生成蛋白在新穎度上有著明顯的優(yōu)勢且具備多樣性，一方面，左下角的子圖中展示了天然pdb和Chroma生成蛋白隨著長度增加生成蛋白的新穎度變化，生成蛋白占據(jù)明顯優(yōu)勢。另一方面，本文將蛋白的拓?fù)涿枋龇M(jìn)行Umap投影，并根據(jù)新穎度進(jìn)行上色，以右側(cè)的12個(gè)為代表的生成蛋白不僅覆蓋廣泛，而且有著極高的新穎度。Chroma作為一個(gè)新擴(kuò)散生成模型，在擴(kuò)散模型、AIGC和Alphafold2正值風(fēng)口的當(dāng)下，與bakerlab的RFdiffusion同期發(fā)布。它吸納了眾多模型的優(yōu)點(diǎn)實(shí)現(xiàn)了高效高質(zhì)量的可編程蛋白設(shè)計(jì)，同時(shí)文章也指出了一些可以調(diào)整的設(shè)置作為開放問題，這會進(jìn)一步促使AI蛋白設(shè)計(jì)領(lǐng)域的發(fā)展。該模型潛在應(yīng)用廣泛，包括邦德設(shè)計(jì)、以抗原做條件的抗體設(shè)計(jì)、酶的活性位點(diǎn)嫁接、材料設(shè)計(jì)等等，這些應(yīng)用都亟待改進(jìn)以及實(shí)驗(yàn)驗(yàn)證。Chroma 旨在轉(zhuǎn)移蛋白設(shè)計(jì)的焦點(diǎn)問題，將主鏈采樣問題轉(zhuǎn)向特定任務(wù)特定功能的采樣——即蛋白質(zhì)打算做什么。通過利用大自然積累的數(shù)以億計(jì)的蛋白質(zhì)進(jìn)化實(shí)驗(yàn)，Chroma 等生成模型有望推動人類健康和生物工程的生物分子多樣性的另一次改變。
參考鏈接： Illuminating protein space with a programmable generative model. John Ingraham, Max Baranov, Zak Costello, Vincent Frappier, Ahmed Ismail, Shan Tie, Wujie Wang, Vincent Xue, Fritz Obermeyer, Andrew Beam, Gevorg Grigoryan. bioRxiv 2022.12.01.518682 John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph- based protein design. Advances in neural information processing systems, 32, 2019. Namrata Anand, Raphael Eguchi, Irimpan I Mathews, Carla P Perez, Alexander Derry, Russ B Altman, and Po-Ssu Huang. Protein sequence design with a learned potential. Nature commu- nications, 13(1):1–11, 2022. Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Ad- vances in Neural Information Processing Systems, 34:8780–8794, 2021. State-of-the-art estimation of protein model accuracy using AlphaFold. James P. Roney, Sergey Ovchinnikov. bioRxiv 2022.03.11.484043 Gabriel Peyre ?, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends? in Machine Learning, 11(5-6):355–607, 2019.
|
