生物計算專家超細緻解讀AlphaFold2 論文:模型架構及應用

ScienceAI發表於2021-07-30
本文來源於生物計算領域的學者Mohammed AlQuraishi的部落格。在這篇文章中,Mohammed AlQuraishi主要講了「是什麼使AlphaFold2(AF2)如此令人驚訝」、「AlphaFold2意味著什麼」、「可以用AlphaFold2做些什麼」。
旁白:雖然作者有點囉嗦,但文章邏輯還是很清晰的,預計讀完需要10-15分鐘,請耐心讀下去。
引子
就在一個多星期前(指2021年7 月 15 日),期待已久的AlphaFold2方法的配套論文和相關程式碼終於公佈於世。該成果的釋出,使我不得不面對許多人提出的各種「關於公開披露AlphaFold2的問題」。

相關連結:https://www.nature.com/articles/s41586-021-03819-2

https://github.com/deepmind/alphafold

如今,程式碼已經以各種方式推廣開來。「AlphaFold2被應用於人類蛋白質和其他20個模型生物體」,幾天前(指2021年7月22日),該研究相關成果的配套論文和資料庫被髮布。總之,我對DeepMind公司做的件事表示非常讚賞。

相關連結:https://www.nature.com/articles/s41586-021-03828-1

https://alphafold.ebi.ac.uk/

在過去的幾個月裡,我檢視了相關論文,並有機會仔細考慮AlphaFold2模型架構。(插句題外話:有些人認為DeepMind公開AlphaFold2原始碼是對RoseTTAFold的回應——但實際上公開程式碼是DeepMind的計劃,遠在《RoseTTAFold》被預印之前。)
圖片文章作者Mohammed AlQuraishi 
旁白:請注意!正文開始之前,Mohammed AlQuraishi有話要囑咐。
本篇文章不是AlphaFold2的高層總結——我建議你在閱讀這篇文章之前閱讀論文,包括補充資訊(SI)。在本文中,我將假設你熟悉該模型。我在這裡的重點是結構預測的技術方面,主要著眼於AlphaFold2如何應用於其他更確切的實際問題。
對於「大家可以用AlphaFold2來做些什麼」的問題, DeepMind 給出的明確答案是「……?是的」。自我監督?是的。自蒸餾?是的。新的損失函式?是的。3D細化?是的。精煉後回收?是的。回收後精製?是的。模板?是的。完整的 MSA?是的。捆綁重量?是的。非繫結權重?是的。關注節點?是的。注意邊?是的。注意座標?是的。對於所有問題,答案都是肯定的!
我有些輕率的描述可能會給人一種印象,彷彿AlphaFold2只是好想法的大雜燴——其實事實並非如此。AlphaFold2除了我們可以想到的功能,還具有非常統一整合功能。這就好像團隊中的不同想法,透過一個反覆推敲的過程,最終實現思想的統一輸出(他們的資訊回收方法,也應用於 AlphaFold2 設計過程本身)。這一成果既是技術創新的傑作,又是精美設計的學習機器。
論文中的「一句話」特別能說明問題,這句話中的每個短語基本都與整篇論文相對應。Facebook Research 的一個團隊之前報導過這一句話。

相關連結:https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1.abstract

圖片

「...我們展示了一種聯合嵌入多序列比對(MSA)和成對特徵的新體系結構,一種新的輸出表示和相關損失,能夠實現精確的端到端結構預測;一種新的等變注意體系結構,使用中間損失實現預測的迭代細化,掩蔽MSA損失,與結構聯合訓練;使用自蒸餾對未標記的蛋白質序列進行學習,並自我估計準確度。」

只想說,看完這篇論文後,我對AlphaFold2的印象比以前更深刻。這是一件藝術作品,其中有許多概念創新,這些概念與「計算能力」或「工程」無關;這是一幅比最初細節釋出時我所預期的更為複雜的畫卷。在CASP14之後不久,我與一位朋友進行過一次討論,討論的主題是「學術界在多大程度上落後於DeepMind」。當時我說可能需要十年的學術研究才能完成DeepMind的目標,而朋友認為這大概落後兩年。在他這麼說之後,我回過頭來想,也許是我太苛刻了。
現在我已經讀了這篇論文,我認為學術界可能至少需要努力5-6年才能達到預期的水平(在這一點上,要測量的正確增量不是 RoseTTAFold 與 trRosetta,而是 CASP14 的 trRosetta 與 CASP13 的 Rosetta;即便如此,差距在很大程度上是由於第一個 AlphaFold太過優秀。)
在這裡,我要做的是透過一系列的觀察,然後從個人角度總結AlphaFold2最有趣和令人驚訝的方面。它們(在某種程度上)是按照我所認為的重要性順序列出的。
作者的免責宣告:這是一個複雜的系統,有許多運動部件,而我的評論依據僅僅是《自然》論文中所提供的內容。因此,下面所說的一切都基於我對 AlphaFold2 及其工作原理的理解,這可能並不完全正確。
正文開始了……
目錄
  • 條條道路通 {s i}

  • 資訊流是關鍵

  • 作物是你需要的一切

  • 永遠要精煉

  • SE(3)-等變推理的原因

  • SE(3)-等變推理的方法

  • MSA Transformer

  • DeepMind 的魔力不在於蠻力

條條道路通{Si}
整篇論文中最重要的一行,IMO,也就是SI中「演算法20」的第10行;在下圖中,它與的不變點注意(IPA)的消融結合時,它充當了結構模組的核心元件。如果對我來說有一個賣點,那就是這條線;因為我認為這是 AlphaFold2 的「魔力」癥結所在。我這樣說不是因為這條線本身,而是因為它所暗指的部分。
圖片兩個目標集的消融結果:CASP14 域集(N=87 個蛋白質域)和 PDB 鏈的測試集,模板覆蓋率 ≤ 30%,同一性為 30%(N=2,261 個蛋白質鏈)。(來源:論文)
在 3D 空間中摺疊的「長 1D 聚合物」,是構建用於蛋白質推理的 ML 系統的挑戰之一。這裡的關鍵是「1D」、「3D」和「長」。一方面,鑑於蛋白質的序列性質,使用基於序列的架構對它們進行編碼是很自然的。在我開發 RGN 模型時,使用了當時領先的序列問題架構 (LSTM)。然而,這種結構難以推理長程相互作用,這也是蛋白質中的常見現象。儘管如此,一維表示仍然很方便,因為它們很容易對映到物理物件,允許引數化(例如,內部座標)以直接自洽的方式預測蛋白質結構。
大多數領域,直到第一個 AlphaFold,都採用了一種不同的方法;該方法依賴於 2D 矩陣使用殘基間距離對結構進行編碼。這更好地捕獲了遠端互動(相對於順序模型),但引入了「將 2D 表示對映到基本上是 3D 空間中的 1D 曲線的物件」的尷尬。通常,這種對映是在後處理步驟中使用基於物理的鬆弛或 SGD(第一個 AlphaFold)完成的,但仍有很多需要改進之處。
DeepMind 在 CASP14 上首次展示 AlphaFold2 時,他們描述了一種看似混合的 2D/3D 方法;其中結構最初以 2D 表示進行編碼,然後轉換為 3D 表示。關於這是如何完成的,有很多猜測(稍後會詳細介紹)。但請記住,關於「1D-2D 不匹配」的尷尬,是一個微妙的點。 這裡的問題不是關於如何將 2D 距離矩陣對映到 3D 座標集(存在許多方法);相反,它是關於物件之間的結構不匹配;(在資料/張量型別方面,而不是蛋白質結構方面) 基本上是一維(蛋白質本身)以及其在二維矩陣中的分散式表示。人們可以在平均行和列方面提出各種技巧,以將 2D 表示轉換為 1D 表示,但它們的技巧總是讓我感到震驚。
出乎意料的是,我發現AF2架構非常優雅的是MSA表示的中心作用,特別是該表示的第一行;它最初對原始輸入序列進行編碼,但最終投影到它們所表示的物件上{s i}在「AF2交響樂」中扮演著重要角色。從某種基本意義上說,AF2中蛋白質的中心編碼完全是一維的——正是這個一維物件{s i},最終捕獲了預測三維結構所需的內容。這就是演算法20的第10行。從這個簡陋的{s i}開始,整個結構被投影到3D空間上,沒有任何明確的3D推理。誠然,在完整的AF2模型中,IPA模組確實在3D中執行,並從2D表示的單向(傳入)通訊中獲益;但值得注意的是,當該模組被燒蝕時,AF2的效能與完整版本相比幾乎保持不變(只要以後在這方面啟用更多迴圈)。
這是這篇論文最讓我興奮的方面,也是最“神奇”的部分。他們構建它的方式非常聰明。確實,使用一維表示進行學習是很困難的,尤其是對於遠端互動。所以 AF2 所做的就是使用一個一維物件來表示蛋白質,但是將它耦合到一個二維物件(成對的 {z ij})來克服一維表示的缺點。我懷疑 {z ij} 隱含地充當 {s i} 的記憶體,幫助它儲存可以在訓練和預測期間參考和迭代的資訊(我的意思不是{s i} ,因為這是最後一個投影表示,但我在這裡過載符號 - 我實際上是指 {m 1i},evoformer 中 MSA 表示的第一行。)更廣泛地說,{z ij} 可能透過為模型提供更豐富的表示來促進學習過程。在過程的每一步,{s i} 都保持更新,與 {z ij} 來回通訊,以便 {z ij} 中建立的任何內容都可以訪問 {s i}。因此,{s i} 在所有主要模組中處於最前沿和中心位置。最後,在結構模組中,最終是 {s i},而不是 {z ij} 對結構進行編碼(提取四元數以生成結構)。這避免了必須將 2D 表示投影到 3D 空間的尷尬。
最初閱讀手稿並進入演算法 20 時,我對「最終將如何提取結構」的預期感到非常興奮。當發現真正重要的是 {s i} 而不是 {z ij} 時,我感到一種眩暈感。
在某種程度上,DeepMind團隊所做的工作(如果允許我在這裡進行一點推測的話)是開發一種方法;其中「用於學習和推理的表示」與「用於保持系統狀態和預測結構的表示」分離,從而產生兩個方面的最佳結果。在大多數ML體系結構中,使用一個表示來完成這兩個任務;而在AF2中,它被拆分。當DeepMind第一次宣佈他們的混合配對/MSA表示時,我並沒有真正理解其中的「要點」;但現在我已經詳細瞭解了論文內容,我認為上述是他們的主要動機之一。
旁註:當他們合併結構模板時,他們將「扭轉角」作為行新增到 {s i}(實際上也是 {m 1i}),因此結構資訊在早期就明確嵌入到 {s i} 中。當然,不能保證在推理期間模板可用,但它們在訓練期間偶爾可用;這意味著 evoformer 必須從一開始就學習如何使用 {s i} 中的此類幾何資訊進行推理。
資訊流是關鍵
如果論文有一個統一的主題,一個廣泛的原則,即AlphaFold2的設計是為了最大限度地提高其元件之間的資訊流。此外,它在某種程度上反映了我們對蛋白質的理解。我將用幾個具體的例子來說明這一點。
首先是對錶示 {z ij} 和 MSA 表示 {m si} 之間的通訊。在每次迭代中,兩者相互更新,確保持續同步。然而,更新是不對稱的。MSA→pair 方向比較繁重,使用「OuterProductMean」的每一個元件。這也許再次反映了 {s i} 最終扮演的核心角色。然而,MSA 表示並非不受配對錶示的影響。後者也更新了前者;但透過輕觸,偏向了 MSA 表示的不同列如何相互關注。事實上,它是一種輕觸,而不是我預期的重型交叉注意機制。然而, {z ij} 影響 {m si} 的逐列注意力這一事實具有完美的生物學意義;因為配對錶示應該編碼殘基位置之間的耦合強度。這也避免了我之前提到的「hack」,必須平均行或平均列來提取一維資訊——它永遠不會完成,因為 {z ij} 的各個元素直接偏置 {m si}(演算法 7 的第 3 行和第 5 行)。
另一個例子是對錶示內的通訊。它利用了在殘基對之間執行的兩種形式的新型「三角」注意力——這是一種非常昂貴的計算操作。這裡也透過使用僅偏向注意力的輕觸方法來實現效率;更重要的是透過將注意力限制在具有一個共同殘基的成對上,直覺是三角不等式應該限制這些數量(因為三個餘數形成一個三角形)。這個很有趣,因為它說明了一個重要的原則:「幾何約束不需要從字面上實現」。因為許多人,包括我的團隊,一直試圖這樣做。例如,透過數學強制三角不等式,而是資訊上,即在注意力機制的資訊流模式中。實際上,AF2 所做的是將幾何直覺轉化為允許這些距離對(或者編碼它們的表示)輕鬆交流的架構特徵。對我來說,這是 ML 架構工程的一般原則。它說明了如何將先驗資訊整合到可學習模型中:不是透過硬約束甚至軟約束,而是作為學習過程本身的一個方面。(相對於機率程式設計或貝葉斯建模,我認為深度學習大致如此。知識很少直接注入 DL;而是透過設計注入架構中,使其更易於學習。在這方面從某種意義上說,硬 SE(3) 等方差實際上是一個例外,但稍後會詳細介紹。)
最後一個例子是在 IPA 結構模組中。當啟用它時,可以從演算法 22 的第 7 行中的對錶示中獲得輕觸,其中 {z ij} 偏向於 3D 空間中殘基之間的注意力,以及重 touch 在第 8 行和第 11 行,其中 {z ij}  可直接用於 3D 推理引擎。輕觸再次具有完美意義,因為配對錶示應該編碼哪些殘基在空間上接近,這是空間推理時的寶貴資訊。
「作物」是你需要的一切
AF2(以及它之前的第一個 AF,雖然以一種看似奇怪的方式進行訓練:不是在整個蛋白質上,而是在蛋白質片段上,或者 AF2 團隊所說的「多肽作物」。他們不短, 通常有幾百個殘基。但對於更長的蛋白質,這些作物只捕獲整個序列的一小部分。此外,通常製作兩種不連續的作物(相同的蛋白質),然後縫合在一起。值得注意的是,雖然 AF2 主要針對多達 256 個殘基的多肽進行訓練(後來對 384 個殘基進行了微調),但它可以預測超過 2000 個殘基的蛋白質結構,這是一項驚人的壯舉。不僅因為它絕對是一個非常困難的問題,而且因為 AF2 是針對更短的多肽進行訓練的。
全域性背景對蛋白質很重要。在兩種不同的蛋白質中,相同的氨基酸序列通常具有不同的結構;否則蛋白質結構預測早就解決了!(當然,隨著序列長度的增加,這種情況逐漸變得不那麼真實。)將有數百個殘基分開的無實體「作物」,並要求模型預測它們的相對方向,且它們的分離長度是未知的(如果它們相距超過 32 個殘基),這似乎是一項不可能完成的任務。但這裡至少有兩個改善因素在起作用。首先,AF2 正在使用 MSA/協同進化模式,這些模式對這些資訊進行編碼,而不考慮線性鏈分離。其次,這僅在訓練時完成,而不是在推理時完成。在推理期間,AF2 確實可以訪問整個序列,因此上下文敏感性問題沒有實際意義。在訓練期間,AF2 可能會從一種蛋白質中的一對作物中得到一個訊號;而從不同蛋白質中的一對相似作物中得到一個相互矛盾的訊號,這完全沒問題。這種情況可能會告訴模型,當它以某種未知的間隔看到這兩種作物時,存在固有的不確定性。回想一下,在訓練期間,不需要模型能夠準確預測結構。重要的是透過梯度更新將有用的資訊傳遞給模型。
這是解耦的另一個例項,通常是耦合的(前一個是學習和保持狀態的合適表示)。在大多數ML模型中,訓練和推理任務保持非常相似,其思想是訓練任務與推理任務越相似越好。但這並不是真的必要,因為訓練是為了從資料中獲取有用的訊號,而推理是為了做出準確的預測。這個想法肯定不是AF2獨有的;生成模型在生成和訓練過程中通常涉及完全不同的任務,當然,訓練過程中使用的可微損失函式通常與真實目標函式不同。儘管如此,AF2在有監督的學習環境中展示了這一思想的相當強大的應用;幾乎支援故意解耦任務,然而通常人們傾向於將這種解耦視為建模的失敗。
我不知道上述是否是AF2團隊的意圖,因為使用目前的TPU幾乎不可能在全長蛋白質上訓練像AF2這樣大的東西,所以記憶效率可能在他們的思維中起到很大的作用。儘管如此,從生物物理學的角度來說,最初可能是一種計算技巧,但最終卻成為了一個好主意。
順便說一下,這解決了 CASP14 期間關於 AF2 的一大謎團——推理和訓練時間之間的差距。由於三次縮放,以及推理是對全長蛋白質進行的,而訓練是對作物進行的,因此兩者之間(可能)存在非常大的計算時間差距。
永遠要精煉
AF2 的另一個顯著特點是其始終開啟的最佳化模式。即,它總是能夠在距原生狀態一定距離處獲取初步結構,並將其改進為更接近該原生狀態。在多個模組和多個粒度級別中都是如此,這使得系統在利用不同型別的資料時非常強大和通用
它在結構模組中最為明顯和自然,其中迭代 IPA 過程的權重被繫結,因此重複應用相同的操作。這是有道理的,因為 IPA 的預期功能是改進來自 evoformer 的結構。但是,evoformer 本身也始終處於細化模式。這並沒有在架構本身中明確編碼(evoformer 的 48 層的權重是解開的),但在它的訓練方式中很明顯,鼓勵它以這種方式行事。例如,原始輸入可以包括同源結構的模板,其中一些可能與所尋求的結構相似,從而為 evoformer 的第一層提供一個結構(編碼成對錶示);該結構基本上是完整的,並且不應該被搞砸了。這是關鍵並概括了這一點的難點,因為 AF2 也可能獲得一個沒有結構同源物的序列,從而為 evoformer 的第一層提供幾乎沒有結構資料——在這兩種情況下,evoformer 必須學會正確行為。MSA 的重複子取樣加強了這一點,因為每個樣本提供不同程度的序列覆蓋。
同樣的現象也發生在回收利用上。首先,回收本身是一種細化形式,因為整個網路具有繫結的權重,最多可以再重新應用 3 次。但是回收的行為也使 evoformer 成為一個「精煉者」,因為在稍後的回收迭代中相同的 evoformer 可以呈現出比第一次(預回收)迭代中的 evoformer 更遠的結構。
另一個鼓勵改進的機制是在結構模組和回收中使用中間損失。即,在訓練期間,模型被最佳化以最小化最終預測結構的損失以及預測透過系統的某些部分的中間結構的損失。這鼓勵 AF2 不僅正確預測結構,而且在系統的早期迭代中快速預測。在結構模組中,這是非常明確地完成的;它的損失函式實際上是所有迭代的平均值。在回收中,它有點微妙。僅用一次迭代的損失用於反向傳播,但是因為迭代次數是隨機取樣的,所以效果是一樣的;鼓勵模型在早期的回收迭代中獲得正確的結構。
AF2 在不同階段對不同任務的魯棒性在隨論文提供的動畫中顯而易見。LmrP (T1024) 的影片顯示了在 evoformer 的第一層之後基本完整的結構,而 Orf8 (T1064) 的影片一直持續到最後,幾乎類似於分子動力學摺疊模擬(它不是一個,很明顯。)順便說一句,這些動畫也暗示了 AF2 在不同 MSA 深度序列方面的行為。對於深度 MSA,它的作用可能類似於嚴重依賴協同進化訊號的 pre-AF2 方法,或多或少完全基於該訊號推斷結構,僅使用單個 evoformer 層。對於具有非常淺的 MSA 的序列,它屬於 evoformer 和結構模組的後期階段才能真正摺疊蛋白質,我懷疑該模型正在學習和應用有關蛋白質結構的一般物理知識。下圖的「無 IPA 和無回收」皮膚顯示當回收和 IPA 關閉時,淺層 MSA 的 AF2 效能顯著下降,支援這一假設。此外,儘管 Sergey Ovchinnikov 有一個引人注目的替代理論,但 AF2 在預測來自未配對 MSA 的蛋白質複合物方面的明顯成功可能是由於這種一般的物理學知識。
圖片外部產品的平均值。尺寸:s:序列,r:殘基,c:通道。(來源:論文)
總的來說,我發現不斷改進的想法強大且廣泛有用,尤其是當它可以應用時,而不必在整個模型中回溯。我不確定這是否是新穎的迭代監督學習語境中(它可能在生成計算機視覺方面有相似之處)。它非常類似於 RL,這顯然是 DeepMind 的強項。
SE(3)-等變推理的原因
可能我讀這篇論文最大的驚喜是「消融研究,尤其是 IPA 的消融」,AF2 大肆宣傳的 SE(3) 等變變壓器。雖然我沒有像其他人在釋出前那樣認為這個模組有那麼多的價值,但刪除它似乎做的很少的事實仍然令人震驚。如果貢獻如此之少,所有這些工作和機器的意義何在?
這裡其實有兩個驚喜。首先是沒有 IPA,AF2 只是簡單地吐出 3D 座標,沒有任何明確的 SE(3)-invariant 轉換對錶示中的「距離」到 3D 空間。事實上,正如我之前提到的,AF2 的無 IPA 版本完全依賴於 1D {si} 進行結構生成。這意味著它選擇了一個特定的全域性參考框架,在其中生成結構,我認為對於包括我自己在內的許多人來說,這似乎太粗糙而無法工作。但顯然它做得很好。
第二個驚喜是,3D 中的推理,即在結構的初始版本在全域性參考系中具體化之後的推理,似乎並不是非常重要,除非回收並刪除。這與我們的直覺背道而馳,即某些空間模式,尤其是分佈在蛋白質的多個不連續元素上的空間模式,在 3D 空間中更明顯,因此應該受益於 3D 推理。從實際的角度來看,它似乎也排除了進入等變網路的所有方法論研究,至少就它應用於蛋白質而言(下一節將詳細介紹)。
這當然是一種解釋,但我不認為它是完全準確的。關鍵在於,只要保留回收,去除 IPA 就可以了。當兩者都被消融時,效能顯著下降。此外,如果去除回收但保留 IPA,則 AF2 的效能幾乎不會受到干擾。這對於 IPA 來說是一個相當令人印象深刻的展示,因為回收基本上是 48 層 evoformer 的四倍(由於各種技巧,計算成本方面並非如此),而 IPA 只有 8 個權重繫結層。從這個角度來看,IPA 層比 evoformer 層更有效,至少對後期細化方面而言。隨著現在重點從單鏈預測轉移到複合物和高階組裝,空間推理的重要性只會增加,我預計 IPA 及其未來的衍生物將繼續發揮重要作用。
SE(3)-等變推理的方法
撇開 SE(3)-等方差的效用不談,它在 AF2 中如何執行的問題可能是論文發表之前最受期待的問題。退一步思考這個子領域在過去幾年中的發展方向會有所幫助。最近在等變神經網路(equivariant neural networks)中的一系列活動可以說是從 Thomas 等人的一篇論文開始的,儘管有一些先行工作。該論文依賴於群論機制,採用(卷積)濾波器,將球諧函式用作其基組。該公式在數學上很簡單,並且已經在許多論文中進行了闡述,特別是 Fuchs 等人的 SE(3)-equivariant Transformer,它不僅概括了從卷積到自注意力的方法,而且還共享了一個相同的 DeepMind 團隊在 CASP14 期間使用的名稱來描述他們現在所說的 IPA。這自然會讓人們猜測 AF2 使用了與這種方法非常相似的東西,包括我之前關於 AF2 的帖子。回想起來,這沒什麼好處,特別是因為這些方法大約是在同一時間開發的,所以沒有理由相信它們會相互影響。

相關連結:https://arxiv.org/abs/1802.08219

https://arxiv.org/abs/2006.10503

https://moalquraishi.wordpress.com/2020/12/08/alphafold2-casp14-it-feels-like-ones-child-has-left-home/

在發展群論方法的同時,也出現了一系列圖論方法來解決分子等變推理的問題。基於圖的方法不依賴於球諧函式,而是將分子嵌入為圖,空間資訊編碼在連線節點的邊緣中,這些節點依次對原子進行編碼。這一系列研究已應用於小分子和蛋白質,但可以說,後者最有用。與所有聚合物一樣,蛋白質本質上允許在每個原子上構建明確的參考系,並且這一事實已被基於圖的方法所利用,產生了巨大的影響。Ingraham 等人的 Structured Transformer 是在機器學習和蛋白質的背景下使用這種結構的第一個(也許是第一個)方法,似乎這項工作是 IPA 的靈感來源。對於蛋白質,使用基於圖的結構比基於組的結構有很多優勢,但這是一個較長的討論,不會影響 IPA 的工作原理。可以說 IPA 完全屬於基於圖形的陣營,對於蛋白質而言,IMO 最有意義。

相關連結:https://papers.nips.cc/paper/2019/file/f3a4ff4839c56a5f460c88cce3666a2b-Paper.pdf

IPA 的工作原理非常有趣,因為它可以說是整個 AF2 中最新穎的神經原語,結合了多種空間推理機制,可能會為未來的分子建模提供資訊。我將重點介紹幾個。我要強調的是,我的大部分想法都是高度投機的,尤其是在元推理上。
首先是 IPA 注意力機制,它有點像野獸(演算法 22 中的第 7 行)。一般來說,IPA 做了很多這種「噴射」,其中殘基會生成許多 3D 向量,以控制注意力並相互傳送值。這些向量在每個殘基的區域性參考系中生成,然後轉換為全域性參考系,因此它們都是等變的並且能夠在 3D 空間中關注整個蛋白質。因為幾何匹配項是負數,它會減弱而不是增加註意力。這使得它相當保守——導致大多數殘基對彼此忽略——因為它要求查詢/鍵恰好匹配以避免衰減,這不太可能偶然發生。衰減程度是一個可學習的人均引數 (γh),IPA 有可能學會將其關閉,但幾何匹配永遠不會對注意力產生積極的貢獻。綜上所述,這些方面可能會導致空間區域性性偏差——每個殘基更容易為查詢/鍵噴射通用區域性向量,當轉換為全域性參考系時,將關注附近的物件,而不是遠處的精確定位向量從殘基的中心開始,當轉換到全域性參考框架時,它恰好落在蛋白質的正確部分。
其次是值,特別是 3D 向量形式的幾何值,如何在殘基之間進行傳遞(演算法 22 的第 3 行和第 10 行)。每個殘基再次噴射多個向量,所有向量都轉換為全域性參考系,其中對於任何接收殘基,在將所有其他殘基的向量轉換回接收殘基的區域性參考系之前,所有其他殘基的向量在整個蛋白質的歐幾里德空間中進行平均。我想這可以實現複雜的幾何推理能力,這些能力可以反映蛋白質生物物理學的一般方面,並且與 MSA 中包含的遺傳/共同進化資訊的聯絡較少。
例如,考慮一個「催化三聯體」,其中一個殘基必須檢測其他兩個殘基相對於自身的定位和取向。它以適當的位置和方向噴射查詢向量,並且由於它們是在其本地參考系中指定的,因此它們形成了 IPA 可以反覆學習和應用的通用空間模式。如果查詢向量指向正確的方向,我們可以想象關鍵向量簡單地說明每個殘基的位置和方向,本質上返回 T j。當轉換到全域性參考系時,鍵將匹配查詢,避免衰減。大多數其他注意力對會減弱;過濾噪音。現在相關的殘基相互關聯,每個殘基都返回一個再次編碼 T j 的值。該資訊隨後在第 10 行和下游層進行處理,以推動殘基更好地符合 IPA 對活性催化位點的預期。
以上是否是對 IPA 工作原理的準確描述當然是未知的,除非仔細檢查學習的權重和模型的行為,否則可能無法知道。
第三,也是迄今為止最具推測性的,是 IPA 可能會結合其元件來執行元推理。我的意思是推理不是關於蛋白質結構,而是關於 AF2 本身,即它關於當前推理任務的知識狀態以及它在後續 IPA 迭代中可能如何改進(我指的是在推理過程中實時學習,而不是透過梯度下降。)想象一下,AF2 期望一個蛋白質片段與另一個相互作用,但不確定它的位置。在第一次迭代期間,它噴射多個間隔很寬的查詢向量,分佈在它認為可能包含所尋找片段的蛋白質區域的大部分割槽域。幾乎所有查詢都找不到匹配項,但其中一個,如果 AF2 的假設是正確的,則可能。一旦找到並基於獲得的資訊,AF2 可以在隨後的 IPA 迭代中傳送更多定向查詢,以更好地定位片段的確切位置和方向,然後根據其學習的蛋白質知識改進結構。
這裡的關鍵是使用 IPA 的迭代來控制其推理和發現過程。與傳統取樣方法(隨機考慮不同蛋白質構象)相比,IPA(和 evoformer)可能會積極推理如何提高其對當前試圖摺疊的蛋白質的瞭解。可以肯定的是,不能保證這種元推理正在發生,我懷疑即使 DeepMind 對此進行評估也不是微不足道的。如果發生元推理,它可以解釋 Orf8 和其他具有淺 MSA 的結構的摺疊時間長(在到達天然狀態的意義上)。如果 AF2 用向量噴射空間來尋找和定向蛋白質片段,那將是進行一種搜尋形式,可以進行多次迭代才能得出結論。這將解釋為什麼重複應用相同的計算最終可以摺疊蛋白質,以及在消融研究中 IPA 和回收到具有淺 MSA 的蛋白質的特殊重要性。
最後一點:在演算法 22 的第 10 行,所有這些幾何資訊都被編碼回 {s i} 中,這是一個沒有明確幾何但都是“潛在”的物件,這對我來說很了不起。在它被輸送到 {s i}和所有被遺忘之前,許多等變架構都專注於在整個網路中維護其表示的顯式幾何方面,但在 IPA 中,顯式幾何僅被簡要使用,以執行我上面描述的操作。它展示了在 {s i} 中編碼幾何資訊的能力和依賴隱式表示的優勢,避免了在整個網路中形式上保持等變的複雜性(一般來說,等變層很難很好地處理神經網路非線性,特別是群論公式,如果不是數學上的,那麼至少從最佳化的角度來看。)
MSA Transformer
這在方法論上是細微的,但在效能方面卻是非常重要的。MSA Transformer作為AlphaFold2 的輔助損耗之一,下圖顯示了它對於具有淺 MSA 的序列至關重要,如果沒有它,效能會大大降低。
圖片在最近的 PDB 鏈組上,MSA 深度的不同值相對於基線的消融精度,透過模板覆蓋率在 30% 同一性≤ 30% 時(N = 2261 條蛋白質鏈)進行過濾。(來源:論文)
對我來說,AlphaFold2 令人印象深刻的是它對淺層 MSA 的魯棒性。長期以來,對序列(或本例中的 MSA)進行自我監督學習似乎是解決問題的一種自然方式,它在這裡成功了。
有趣的是,當 MSA 非常淺時(例如,只有一個序列),這種增益就會消失;但這也是有道理的,因為那時 MSA 轉換器是無資訊的。
DeepMind 的魔力不在於蠻力
我將以我從論文中得出的最違反直覺的結論作為結束。進入這個領域,我曾預料到,至少在一些元件中(最顯著的是喚起式注意機制),如果不進行「軸向注意」這樣的簡化,這種機制的規模會非常可怕。DeepMind已經運用了他們巨大的計算資源,以暴力手段解決了「其他群體必須透過智慧才能解決的問題」。
具有諷刺意味的是,事實似乎幾乎完全相反。AlphaFold2 給人留下深刻印象的不是一流的硬體,而是一流的軟體和 ML 工程,這使得蠻力擴充套件變得沒有必要。我認為這部分,對於學術界來說將是最難複製的;因為它不是僅僅增加國家對計算資源的投資就能提升的(這無疑會有所幫助);這將是對內部專業軟體工程能力的一個更高要求。這一點在許多方面都很明顯:在IPA和回收利用方面,謹慎使用梯度停止,節省了大量的計算;在裁剪方面,它從根本上改變了系統的誘導偏差;在經過深思熟慮選擇的初始化方案方面……
當然,DeepMind可以魚和熊掌兼得。他們在CASP14中使用的8個整合就證明了這一點;他們將計算需求增加了一個數量級,雖然這似乎是非常邊際的收益(我懷疑他們是在CASP14之後才發現這一點的),但總的來說,他們的傳統似乎是計算節儉。從最好的意義上說,這是我對AlphaFold2團隊最大的讚美,因為他們在許多方面都取得了卓越的成就。
原文連結:https://moalquraishi.wordpress.com/2021/07/25/the-alphafold2-method-paper-a-fount-of-good-ideas/

相關文章