瑞典Karolinska研究院在瑞典政府贊助下由Ben Murrell等研究團隊在AlphaFold 3最新報告后提出這篇論文提出了一種非常有趣和創新的方法來生成蛋白質骨架結構,稱為生成式不變角度轉換器(GIAT)。與現有的主要基于擴散模型和流匹配的方法不同,GIAT采用了類似于大型語言模型(如GPT)中使用的自回歸方法,逐個氨基酸地生成蛋白質骨架。

這種自回歸建模方法的主要優勢和創新點包括:

1. 它以序列生成的方式對蛋白質骨架建模,這更接近蛋白質合成的自然過程。每個氨基酸的放置都以前面氨基酸的位置為條件,這可以更好地捕捉蛋白質骨架的連續性和相關性。
2. 通過巧妙地將von Mises分布作為二面角的參數化方法,并使用Neural von Mises Mixture (NVMM)模型,GIAT成功地將離散的自回歸語言模型擴展到連續空間,以處理蛋白質骨架中的幾何特征。
3. 通過修改AlphaFold2中的Invariant Point Attention (IPA),引入自回歸掩碼和Softmax1,GIAT能夠建模氨基酸之間的空間關系,同時確保從早期殘基到后期殘基的信息流,并保持平移旋轉不變性。
4. 實驗結果表明,GIAT能夠生成多樣且真實的蛋白質骨架結構,用ProteinMPNN生成的序列在很大程度上能夠重新折疊回目標結構。這證明了該方法的有效性和潛力。

當然,作者也指出GIAT還存在一些局限性,例如在生成內部區域或生成結合劑骨架方面可能面臨困難。此外,與先進的擴散模型和流匹配方法相比,GIAT的性能還有待進一步驗證和提高。

我認為GIAT為蛋白質骨架生成提供了一個全新的視角和途徑。將自回歸語言模型引入到這個領域是一個非常有創意的嘗試。通過進一步改進模型架構、訓練策略以及實驗驗證,GIAT有望在蛋白質設計領域取得更大的突破和應用。這項工作為蛋白質生成建模的研究開辟了新的方向。

引言

生成式蛋白質結構設計方法需要將每個氨基酸的骨架原子放置在三維空間中,使得所得結構是物理上可獲得且能量有利的。目前的方法包括:i)"去噪"模型(包括擴散去噪(1)和流匹配),學習將容易采樣的噪聲分布轉換為難以采樣的目標分布;ii)“深度網絡幻覺"方法,使用預訓練折疊模型的梯度下降。這兩種方法都是"整體的”,因為它們一次性考慮并轉換整個蛋白質結構(有時還有序列)。

相比之下,在其他領域,自回歸模型取得了顯著成功,它逐個生成tokens,以所有前面的tokens為條件。對于文本生成,目前"decoder-only"的transformer架構占主導地位。其核心是一個掩碼注意力機制,允許每個token在采樣下一個token時有選擇地關注任何前面的tokens。掩碼注意力結構允許模型在單次傳遞中高效地訓練整個序列。

這些模型會發出"logits"或未歸一化的對數概率,并經過訓練以最大化訓練序列的平均對數概率。由于tokens空間是離散的,從這些分布中采樣很簡單:可以按概率抽取tokens,或從最可能tokens的"nucleus"中抽取,這可以防止自回歸設置中的樣本退化。

該研究的工作始于是否可以使用自回歸方法來設計蛋白質骨架結構。

雖然蛋白質在三維空間中折疊,但它們確實具有內在的線性主要排序,其中一個氨基酸緊跟另一個。然而,雖然氨基酸序列是離散的,語言模型可以直接用于生成氨基酸序列,但每個氨基酸在三維空間中的放置是連續的。這需要兩個調整:首先,模型必須能夠推理前面氨基酸在三維空間中的相對位置;其次,模型必須發出一個連續分布,控制下一個氨基酸的放置。

為了解決第一個問題,他們調整了Invariant Point Attention (IPA),這是AlphaFold2背后的一項關鍵創新,以允許自回歸掩碼,確保在訓練期間,任何后面殘基的信息都不會影響前面殘基的放置。除此之外,保留了原始IPA架構的屬性,該架構將每個殘基視為一個框架(即具有位置和旋轉),并具有一個具有空間解釋鍵、查詢和值的注意力機制,尊重框架的相對幾何。當增加自回歸掩碼時,這允許信息僅從早期殘基傳播到后期殘基,以確保對剛體變換的不變性。

由于模型使用框架來表示殘基,為了解決第二個問題,需要選擇一組連續值的表示,以確定后續骨架框架的位置。對于蛋白質骨架,鍵長和鍵角受到很強的物理限制,主要的自由度是二面角。這將我們的問題簡化為參數化三個連續角度上足夠靈活的(可微分的)聯合概率分布。角度上的一個自然選擇分布是von Mises (VM)分布,它是S1上的連續分布,在此提出Neural von Mises Mixture (NVMM)模型,將連續嵌入向量映射到下一個二面角上von Mises分布的混合,其對數似然扮演與經典離散transformer模型中對tokens的log-softmax相同的角色。

這些共同構成了Generative Invariant Angle Transformer (GIAT)的基礎,它對蛋白質骨架結構進行采樣。

方法

樣本是一組殘基。每個殘基具有：

• 一個框架（即位置和旋轉）。
• 決定下一個框架位置的三個二面角。
• 一個離散的氨基酸字符。
• 一個殘基索引號。
• 一個鏈ID。

GIAT（見圖1）包括：

• 每個殘基的嵌入
• 一對嵌入
• 一堆帶掩蔽的IPA變換器層
• 一個神經von Mises混合層

transformer層的輸出是轉換后的每個殘基嵌入，NVMM將其轉換為二面角的聯合概率分布 $\Psi ,\Omega ,\Phi$。類似于離散token的交叉熵損失，這里我們使用觀測到的二面角的平均對數概率。掩蔽的不變點注意力機制確保在訓練期間，任何后續殘基的位置對當前殘基是隱藏的。

神經von Mises混合。給定一個嵌入 $S_r$，殘基 $r$ 處的二面角 ${\Psi }_r,{\Omega }_r,{\Phi }_r$ 與所有其他 $S_t(t\ne r)$ 和所有其他二面角條件獨立（見圖2）。因此，NVMM的任務是放置一個聯合分布 $P({\Psi }_r,{\Omega }_r,{\Phi }_r|S_r)$。我們將其分解為$P({\Phi }_r|S_r,{\Psi }_r,{\Omega }_r)P({\Omega }_r|S_r,{\Psi }_r)P({\Psi }_r|S_r)$，使得每一項都需要一個角度的分布。

von Mises分布 $V(\mu ,\kappa )$ 是 $S1$ 上的連續單峰概率分布，密度為

$$p(x|\mu ,\kappa )=\frac{exp(\kappa \cos (x?\mu ))}{2\pi I_0(\kappa )}$$(1)

其中 $\mu$ 和 $\kappa >0$ 是平均值和濃度，$I_0(\kappa )$ 是零階第一類修改的Bessel函數。

模型需要直接發射VM參數，但標準的VM參數化沒有參數和密度函數之間的平滑對應關系。該研究觀察到VM的對數密度是通過圓柱的平面切片，重新參數化它，使得平面上的一個點控制角度的平均值和濃度：

$$\kappa ={\mu }_x^2+{\mu }_y^2$$(2)

$$\mu =\text{atan2}({\mu }_y,{\mu }_x)$$(3)

因此，當 ${\mu }_x,{\mu }_y=0$ 時，所有角度都同樣可能，當 ${\mu }_x,{\mu }_y$ 遠離原點時，分布集中在 $\text{atan2}({\mu }_y,{\mu }_x)$ 。注意，出于數值原因，該研究重寫了對數密度以避免使用 $\text{atan2}$（見補充信息第1節）。

由于單峰分布對二面角來說不夠靈活，考慮了VM分布的加權混合，$V({\mu }_x,{\mu }_y,w)$。對于 $K$ 組件的混合，我們的NVMM需要發射 $K{\mu }_x,K{\mu }_y$，兩者都在 $(?\infty ,\infty )$ 范圍內，以及 $K$個$w$ 權重logits的softmax（見圖3）。當建模例如 $P({\Omega }_r|S_r,{\Psi }_r)$ 時，將 $S_r$和 $AF({\Psi }_r)$連接起來，其中 $AF$ 是一個將標量角度轉換為特征向量的函數，然后如上所述進行處理。

掩碼不變點注意力（Masked Invariant Point Attention，IPA）。IPA，起源于AlphaFold2，用來增強標準鍵、查詢和值的縮放點積注意力，通過從每個殘基的局部框架投影的鍵點、查詢點和值點，在 R3 中增強標準鍵。當一個殘基的查詢點與另一個殘基的鍵點通過歐幾里得距離接近時，注意力權重較高。注意，不像AlphaFold2的結構模塊，后者迭代更新殘基框架的位置和旋轉，這里的殘基框架在訓練期間保持固定。該研究引入了兩種對IPA的修改：1) 因果掩碼，防止信息從任何后續殘基流向任何先前殘基；以及2) 用 softmax1 替換 softmax，允許注意力頭不關注任何東西，這種方式相當于有一個額外的“虛擬”零對數。后一個變化的動機是缺乏初始 token,它可以作為標準大型語言模型中的"注意力匯",但在 SE(3) 不變性的背景下具有挑戰性。使用 softmax1,注意力權重不再總和為 1,這需要修改 IPA 以保持 SE(3) 不變性(參見補充信息 第 2 節)。

特征嵌入。IPA 允許將配對表示作為輸入,以調節注意力權重并為輸出嵌入提供信息。我們使用它來編碼沿主序列的殘基索引差異,序列是否來自同一鏈(使用各種縮放投影到更高維向量,然后使用隨機傅里葉特征),以及提供額外的相對旋轉/定向信號,使用"隨機定向特征"（見補充信息第3節）。

該研究的初始殘基級嵌入(作為 IPA 的輸入)是氨基酸 one-hot 編碼的線性嵌入與每個殘基在鏈中相對位置的嵌入相結合,后者以兩種方式表示:一個(縮放的)殘基數量直到鏈的末端,以及到目前為止完成的鏈的百分比。這允許當前殘基在采樣連續殘基時知道采樣鏈將有多長。這些特征被輕微噪聲化,以防止對特定鏈長度過度擬合。

殘差連接和層歸一化。transformer堆棧使用非標準的LayerNorm和跳躍連接的順序。該研究發現標準的"后歸一化"方法(在原始GPT模型(4)和AlphaFold2中使用)阻止了模型中層數的擴展,而標準的"預歸一化"方法(在更現代的大型語言模型中使用)在訓練的早期就達到平臺期。與這些不同,該研究在IPA(或前饋)之后進行LayerNorm,并將其求和到當前嵌入中,迫使該層貢獻的標準差是固定的。有些違反直覺的是, 這種安排在層深度上提供了有利的性能和穩定性。

采樣氨基酸。如圖 1 所示的 GIAT 將生成蛋白質骨架。給定結構的序列模型(例如 ProteinMPNN)可以根據該骨架生成氨基酸序列。然而,諸如"結合劑設計"之類的應用可能受益于以目標結構和序列為條件采樣骨架,因為目標的氨基酸側鏈可能影響結合劑骨架構象。處理這需要 GIAT 自回歸地聯合建模序列和結構。第 S3 節顯示了如何調整 GIAT 架構來實現這一點。損失是二面角對數概率的負平均值加上氨基酸預測的交叉熵損失。由于 GIAT 設計的序列僅以前面的骨架和序列為條件,不期望 GIAT 設計的序列在性能上優于最先進的給定結構序列模型,后者可以在觀察整個結構的同時自回歸地設計序列。出于這個原因,部分也是為了幫助與文獻中的其他方法進行骨架基準測試,該研究使用 ProteinMPNN 從報告的骨架生成序列。

訓練數據準備。在觀察到的蛋白質骨架中,鍵長和鍵角接近理想值,每個連續骨架放置的大部分變化由三個二面角決定。該研究希望只建模二面角,而忽略鍵長和鍵角的變化。但是簡單用它們的理想值替換鍵長和角度會導致骨架長度上小誤差的累積,導致大的失真(圖 S1)。但可以對二面角進行微小調整,使得使用理想的鍵長和鍵角不會導致累積的失真。這種"理想化"(見補充信息第4節)導致了幾乎無法區分的結構,但只有二面角變化,簡化了我們的建模問題。

使用了兩個數據源:來自 AlphaFold2 模型的 CATH 域(14)和整個蛋白質數據庫(15)。兩個數據集的 PDB 都被理想化,計算了它們的二面角并提取了它們的旋轉框架。由于殘基索引在配對特征輸入中使用,通過與它們的主序列比對,對它們的殘基進行了程序化重新編號,以消除所沉積的 PDB 殘基編號中的任何特質。詳見補充信息第4節。

核心采樣。通過自回歸采樣,單個低概率 token 可能使整個序列的未來脫離正軌。在采樣過程中修剪低概率區域可以提高樣本質量,修剪程度可以在樣本多樣性和單個樣本質量之間進行權衡。使用 nucleus sampling的連續變體。按照他們的符號表示：假設核心概率為 $\eta \in (0,1]$，取 $V(\eta )?V=(?\pi ,\pi )$ 為最小的子集，使得

$${\int }_{v\in V(\eta )}p(v|\mu ,\kappa ,w)dv\ge \eta$$

我們希望從以下分布采樣：

$$P(x|\mu ,\kappa ,w)=\{\begin{array}{ll}\frac{p(x|\mu ,\kappa ,w)}{\eta }& \text{if?}x\in V(\eta )\\ 0& \text{if?}x\notin V(\eta )\end{array}$$(4)

對于例如 $\Theta$ 上的混合 $V(\mu ,\kappa ,w)$，我們首先生成 $N$ 個樣本 ${\theta }_j～V(\mu ,\kappa ,w)$，并計算每個樣本的密度：

M k + 1 = { p ( θ j ∣ μ , κ , w ) ∣ 1 ≤ j ≤ N } M_{k+1} = \{ p(\theta_j | \mu, \kappa, w) | 1 \leq j \leq N \} Mk+1?={p(θj?∣μ,κ,w)∣1≤j≤N}(5)

然后我們取 $\Theta$ 為從 $M_{k+1}$ 中最大的 $$\begin{gathered} ?N \\ cdot\eta ? \end{gathered}$$ 值中均勻抽取的樣本角度。這應該在 $N\to \infty$近似等式 4。

核心采樣的閾值選擇似乎因模型而異，對于每個二面角 $\Psi ,\Omega ,$ 和 $\Phi$，使用不同的閾值有時似乎有益。該研究已經啟發式地選擇了這些閾值。有趣的是，核心采樣可能導致一種“模式崩潰”，進入一個長的純α-螺旋狀態，對于所展示的結果，拒絕具有 50 個或更多殘基的連續 α 螺旋的樣本,在該窗口內最多有一個非螺旋殘基,并報告排除的頻率。

評估。采用了 FoldingDiff 的基準測試、評估和可視化管道的部分,使用它來評估二級結構分布、樣本的多樣性,以及使用 OmegaFold重新折疊的評估。在這一點上,不嘗試與其他方法進行廣泛比較,也不期望在樣本生成指標上優于 Diffusion 和 Flow Matching 模型。

遵循 FoldingDiff,為了進行基準測試,對長度在 50 到 128 個氨基酸之間的骨架進行采樣。除非另有說明,使用 nucleus sampling,并且拒絕純螺旋(如上所述),因為它們易于重新折疊,從而夸大了重新折疊指標。

實現。GIAT 完全用 Julia 科學計算語言(17)實現,主要使用 Flux.jl 深度學習生態系統(18)。我們的 Masked Invariant Point Attention 的實現可在 https://github.com/MurrellGroup/InvariantPointAttention.jl 獲得,為了加速自回歸采樣,我們還包括 IPA 的緩存實現,將采樣第 $N$+$1$ 個殘基所需的計算減少到 $O(N)$,而不是沒有緩存時需要的 $O(N^2)$。詳見補充信息第 5 節。帶有訓練和推理代碼的存儲庫將在 https://github.com/MurrellGroup/CLOPS.jl 提供。為了方便起見,我們還提供了 ColabMPNN 的 Jax 實現的 ProteinMPNN 的包裝器(https://github.com/sokrypton/ColabDesign/tree/main/mpnn(3)),因此可以在 Julia 環境中采樣 ProteinMPNN 序列:https://github.com/MurrellGroup/ColabMPNN.jl。

結果
該研究探索了許多架構變化和訓練策略,并且正在進行努力以改進模型。在這里描述其結果的模型(GIAT550)大約有 1050 萬個參數,嵌入維度為 196。它使用 12 個 IPA/前饋層,其中略低于 80% 的參數位于 IPA 層本身。交替注意力頭的數量,在 24 和 12 之間切換,并交替 IPA “values” 編碼的大小,在 8 和 16 之間,試圖大致平衡每層的內存負擔。GIAT550 在AlphaFold預測的CATH域上進行了預訓練,每個epoch包括來自每個域家族的一個樣本,然后在蛋白質數據庫上進行了后續訓練。在訓練過程中,樣本被批處理在一起以最大化GPU使用,考慮到序列長度的二次縮放,即對于較短的序列,批處理維度更大。長達1100個氨基酸的結構被包括在訓練集中。GIAT550在一個48gb NVIDIA RTX 6000 Ada GPU上進行了訓練。

GIAT550已經學會對具有合理的二級結構元素分布的蛋白質樣骨架進行采樣。當從這些骨架中采樣序列(使用ProteinMPNN)時,它們經常被預測會重新折疊成目標結構。

GIAT學會生成類蛋白質骨架。 排除純α螺旋樣本(使用這些設置約占樣本的18%),GIAT550似乎對相對緊湊的結構域進行采樣。圖5顯示了25個隨機樣本,長度在50到128個殘基之間。更多示例顯示在補充信息第S1節中。

通過計算α螺旋和β片,二級結構的分布似乎與CATH結構域大致匹配(見圖6,上)。

GIAT模型采樣可重新折疊的骨架。 按照FoldingDiff,通過對GIAT采樣的骨架進行氨基酸序列的in silico重折疊來評估GIAT,這些氨基酸序列本身是以GIAT的骨架為條件采樣的。使用ProteinMPNN創建八個候選氨基酸序列,并用Omegafold重新折疊,取八個候選重新折疊的最大自洽TM分數(scTM)與原始結構。對GIAT550的780個樣本執行分析,在50到128的長度上均勻分布,不包括純螺旋。在這些條件下,平均scTM分數為0.65,分布見圖6,底部。

沒有nucleus sampling,樣本質量下降。 如果沒有nucleus sampling,在三個方面衡量樣本質量的下降,發現:1)自相交幾何的發生率更高,2)α螺旋或β片的傾向性更低,以及3)scTM分數更低(見表1)。

生成的骨架是多樣的,但常見的結構域被優先采樣。 通常,GIAT550的大多數樣本是不同的。圖7(上)顯示了780個樣本的結構相似性的層次聚類,大多數樣本看起來彼此無關。然而,有一些聚類。為了研究這一點,該研究確定了最大的聚類,并疊加了這些。圖7(底部)顯示這些是免疫球蛋白樣的,在框架區表現出強烈的保守性,在互補決定區(CDR3)中具有高度的結構多樣性。這很可能是因為訓練數據集中的許多蛋白質復合物包括抗體結構,因此該模型已經學會在相當一部分時間內從這種常見模式中進行采樣。

Masked Invariant Point Attention通常是稀疏的，有時是可解釋的。為了理解GIAT模型的行為,該研究實現了一個蛋白質骨架可視化器,可以額外覆蓋注意力權重,描繪最新殘基關注所有先前殘基的程度,按注意力頭著色(代碼位于:https://github.com/MurrellGroup/ProtPlot.jl)。從中觀察到,盡管有一些例外,但大多數注意力頭僅稀疏地活躍,關注一兩個其他殘基。此外,大多數頭部都關注空間上相鄰的殘基,也有一些例外。有時,注意力頭的行為是明顯且可解釋的。例如,圖8的頂行顯示了GIAT550在結構域設計過程中的第7層,其中一個頭似乎占主導地位,并且它僅在構建反平行β片期間活躍,每隔一個殘基只活躍一次,僅關注對面鏈中當前殘基前面的殘基。底部面板顯示了第12層,它有一個頭部非稀疏地關注更遠的殘基,可能提取有關結構整體形狀的信息。

表1. 比較GIAT550在未修改采樣（即直接從von Mises混合密度采樣）和核心采樣下的結果。使用FoldingDiff流程，我們比較了沒有自相交幾何的生成百分比（“無沖突”）、不包含在螺旋或β-折疊中的殘基百分比（“環”），以及八個候選序列設計和重折疊中最大scTM得分超過0.5的樣本百分比（“scTM>0.5”）。

GIAT550	無沖突	環	scTM>0.5
未修改	83.3%	67%	27%
核心采樣	93.2%	45%	78%

討論

在這里提出了一種新的蛋白質骨架結構生成建模方法,它逐個氨基酸自回歸地采樣氨基酸骨架放置。該研究已經計算表征了來自一個這樣的模型的設計,顯示了它們的多樣性,以及in silico單序列蛋白質折疊方法將它們重新折疊成目標結構的能力。我們并不認為這種新方法在這一點上將在蛋白質設計方面與現有的最先進方法競爭,但它確實表明這條自回歸模型的途徑值得探索。

有趣的是,FoldingDiff也描述了一個"自回歸模型"作為他們的擴散方法的比較對象,該模型逐個生成蛋白質骨架殘基。然而,這并不在二面角上放置分布 - 相反,它預測最小化角度損失的單個角度。如果沒有二面角的分布,就不可能進行隨機采樣,這個非概率模型無法生成任何多樣性。事實上,它只采樣"最可能"的狀態,即螺旋。在nucleus sampling的極限情況下觀察到這種行為,該研究僅從角度分布的最高模式中采樣。

鑒于計算限制,該研究只是開始探索模型架構空間,對一些系統比較(見補充信息 7)證明了GIAT550背后的建模選擇的合理性,并進行了一些更具探索性的工作。目前已經(非系統地)探索了從 3 個IPA/前饋層到 20 個層的架構,以及從幾百萬到超過 5 億的參數數量。從視覺上判斷(因為系統比較需要廣泛的重新折疊,這與模型開發爭奪 GPU 使用),不同模型之間的樣本質量似乎存在很大差異,甚至在同一模型從一個epoch到另一個epoch之間也是如此,通常不伴隨損失的相應差異,這使得模型選擇有點具有挑戰性。GIAT550 的使用特別是基于對樣本的視覺印象,而不是任何系統比較,未來將進行這些比較。為了幫助擴展這些比較,為通過 IPA 的反向傳遞實現了自定義梯度,與使用通過 Zygote.jl 計算的梯度相比,這允許更大的批處理和更快的訓練時間。

未來的關鍵工作還將探索更復雜的采樣策略,不僅考慮下一個殘基,例如"beam search"的連續變體。

盡管在蛋白質從頭設計的in silico基準測試方面取得了一些成功,但這些基準測試在一定程度上測試了折疊方法以及生成模型的能力,并且從單個輸入序列(而不是多序列比對)進行蛋白質折疊仍然不可靠。目前還沒有替代實驗驗證設計結構的方法,該研究將探索GIAT550和該系列中其他模型的實驗驗證。

蛋白質設計以及相關任務的整個家族是GIAT模型最明顯的應用。但GIAT實際上是蛋白質結構和序列的概率模型,可以在整個結構上一次性評估,并且是端到端可微分的。這意味著它們可能有助于對其他方法生成的結構進行評分,并且它們的可微分性表明它們可以用于通過局部優化二面角來"完善"設計。另一個自回歸模型可能特別適合的應用是通過電子密度圖穿線氨基酸骨架,作為一種強大的結構先驗,可以快速采樣,并根據電子密度重新加權。

GIAT550傾向于生成免疫球蛋白結構域。一方面,這很有趣,因為抗體設計是一個關鍵的目標應用,而這種自然發生在模型中可以被視為有希望的。另一方面,這可以在某種程度上被認為是"過擬合"。為了避免模型學習這種采樣模式,可以過濾訓練數據以移除/減少免疫球蛋白鏈,或者在訓練期間減少對這些的采樣頻率,或者可以探索更激進的正則化策略。

這類自回歸蛋白質結構模型的特定優缺點尚待闡明。與擴散和流匹配方法相比,某些蛋白質設計任務可能對這種模型固有地更難,例如"修補",其中需要以前綴和后綴骨架為條件設計結構的內部區域。學習采樣一系列導致鏈精確地在正確位置結束的二面角可能特別困難。來自結合劑設計的另一個挑戰是結合劑的第一個殘基框架的放置,這是在這里完全避免的問題,因為只對單個鏈進行采樣。該團隊正在探索克服這些問題的策略。

總的來說,這些生成式自回歸transformer模型在骨架結構上表現出有希望的初始行為,模型架構、訓練策略和實驗驗證的進一步探索將決定它們在我們的蛋白質設計工具箱中將占據什么位置。

參考文獻
https://www.biorxiv.org/content/10.1101/2024.05.11.593685v1.full.pdf

---

補充信息
補充信息中包含了額外的實驗細節、數學推導、軟件和算法的實現細節，以及可能的原始數據或進一步的數據分析。

補充說明1：數值穩定的對數概率密度函數（Log PDF）
在計算重新參數化下的 von Mises 分布的對數密度時，避免了與atan2相關的梯度問題。如果 $\theta =\text{atan2}(y,x)$，則 $x=r\cos (\theta )$，$y=r\sin (\theta$ 且 $r=\sqrt{x^2+y^2})$。在該研究的情況下，$\sqrt{\kappa }=\sqrt{{\mu }_x^2+{\mu }_y^2}$，且 $\mu =\text{atan2}({\mu }_y,{\mu }_x)$。應用余弦的角度和恒等式，

$$\cos (x?\mu )=\frac{\cos (x){\mu }_x+\sin (x){\mu }_y}{\sqrt{\kappa }}$$

因此，

$$\log (p(x|{\mu }_x,{\mu }_y))=\sqrt{\kappa }(\cos (\theta ){\mu }_x+\sin (\theta ){\mu }_y)?\kappa ?\log (I_0(\kappa ))$$

補充說明2：Softmax1 和對全局SE(3)變換的不變性
在IPA中將softmax替換為softmax1，定義為

$$\text{softmax1}(x{)}_i=\frac{e^{x_i}}{1+{\sum }_j{e}^{x_j}}$$

與 $T_i=(R_i,t_i)$ 一起，修改了輸出點的計算為

$${\vec{o}}_h^i=T_i^{?1}\left(\sum _j(a_{hij}{T}_j{\vec{v}}_h^{p}j)+\left(1?\sum _j{a}_{hij}\right)t_i\right)$$

以說明注意力權重不再總和為1。現在考慮應用全局變換 $T_{\text{glob}}=(R,t)$，

$$(T_{\text{glob}}°T_i{)}^{?1}\left(\sum _j(a_{hij}(T_{\text{glob}}°T_j){\vec{v}}_h^{p}j)+\left(1?\sum _j{a}_{hij}\right)(R{t}_i+t)\right)$$

等于

$$T_i^{?1}\left(\sum _j(a_{hij}{T}_j{\vec{v}}_h^{p}j)+\left(1?\sum _j{a}_{hij}\right)t_i\right)$$

這保證了對全局SE(3)變換的不變性。

補充說明3：隨機定向特征

通過考慮兩組 $d$ 個固定點的投影，在它們各自的 $T_i,T_j$ 局部框架中投影到全局參照框架的距離，形成了初始的成對嵌入 ${\tilde{z}}_{ij}$。

算法1 隨機定向特征

1: 函數 RANDOMORIENTATIONFEATURES({$T_i$}; $d=128$, $\sigma =0.25$)

2: $p_k,q_k～N(0,\sigma I_3)$ # 在訓練前隨機生成點 $p_k,q_k\in {\mathbb{R}}^3$, k ∈ { 1 , … , d } k \in \{1, \ldots, d\} k∈{1,…,d}

3: ${\tilde{z}}_{ij}={\text{concat}}_k(\Vert T_i°p_k?T_j°q_k\Vert )$

4: 返回 ${\tilde{z}}_{ij}$

5: 結束函數

補充說明4：數據準備

A. 骨架理想化。骨架理想化的目標是最小程度地調整原子位置（$N$, $C_{\alpha }$, $C$），使得三個鍵長（$N?C_{\alpha }$, $C_{\alpha }?C$, $C?N$）和三個鍵角$(N?C_{\alpha }?C$, $C_{\alpha }?C?N$, $C?N?C_{\alpha }$）固定到常數值。簡單將蛋白質骨架表示為二面角數組。使用動態規劃方法，存在理想化的骨架結構,原子坐標變化很小。使用基于梯度的方法來解決同樣的問題,在損失函數下迭代地細化骨架原子的位置,該函數量化了與理想幾何的偏差,引導骨架趨向理想化狀態。該研究嘗試懲罰與原始原子坐標的偏差,以防止骨架坐標偏離太遠,但這似乎并不必要,骨架能夠有效地收斂到附近的理想化最優狀態。在骨架中存在間隙的情況下,例如當結構中缺少殘基時,屏蔽受影響的長度和角度以防止它跨越間隙。偏離其理想值太遠(>0.5?)的鍵長被排除在損失之外。

B. 殘基重新編號。PDB 文件中的殘基編號可以根據特定蛋白質家族的約定任意分配。為了給模型提供一致性,將從殘基獲得的主序列與參考主序列對齊,并使用相對于參考序列的新索引作為殘基編號。為了確保該比對反映任何缺失的殘基,使用具有仿射間隙懲罰的自定義比對算法,該算法傾向于在鏈不連續的地方打開間隙,即當殘基的羰基碳原子與下一個殘基的氮原子之間的距離大于預期時。

補充說明5：緩存掩碼不變點注意力（Cached Masked IPA）

使用 masked IPA,在每次通過時緩存點和非點鍵和值,以實現 O(N) 次 masked IPA 的通過。采用 AlphaFold2 中使用的符號。

算法2：緩存的 Masked IPA

1: 函數 CACHEDMASKEDIPA$(s_r$, { z r j ∣ 1 ≤ j ≤ r } \{z_{rj} | 1 \leq j \leq r\} {zrj?∣1≤j≤r}, { T i ∣ 1 ≤ i ≤ r } \{T_i | 1 \leq i \leq r\} {Ti?∣1≤i≤r})

2: $q_{hr},k_{hr},v_{hr}=\text{LinearNoBias}(s_r)$ # 添加 $k_{hr},v_{hr}$ 到緩存。

3: ${\vec{q}}_{hr},{\vec{k}}_{hr},{\vec{v}}_{hr}=\text{LinearNoBias}(s_r)$ # 添加 ${\vec{k}}_{hr},{\vec{v}}_{hr}$ 到緩存。

4: $b_{rj}=\text{LinearNoBias}(z_{rj})$ # j ∈ { 1 , … , r } j \in \{1, \ldots, r\} j∈{1,…,r}

5: $k_{h1:r},v_{h1:r},{\vec{k}}_{h1:r},{\vec{v}}_{h1:r}=\text{RetrieveFromCache}(k_h,v_h,{\vec{k}}_h,{\vec{v}}_h)$

6: $a_{hrj}={\text{softmax1}}_j\left(\frac{1}{\sqrt{c}}{q}_{hr}^{?}{k}_{hj}+b_{hrj}?{\gamma }_h\frac{w_C}{2}{\sum }_p\Vert \Vert T_r°{\vec{q}}_{hr}?T_j°{\vec{k}}_{hj}\Vert {\Vert }^2\right)$ # j ∈ { 1 , … , r } j \in \{1, \ldots, r\} j∈{1,…,r}

7: ${\tilde{o}}_{hr}={\sum }_{j=1}^r{a}_{hrj}{z}_{rj}$

8: $o_{hr}={\sum }_{j=1}^r{a}_{hrj}{v}_{hj}$

9: ${\vec{o}}_{hr}=T_r^{?1}\left({\sum }_{j=1}^r(a_{hrj}{T}_j°{\vec{v}}_{hj})+\left(1?{\sum }_{j=1}^r{a}_{hrj}\right)t_r\right)$ # $T_r=(R_r,t_r)$

10: $s_r^{\prime }=\text{Linear}\left({\text{concat}}_{h,p}({\tilde{o}}_{hr},o_{hr},{\vec{o}}_{hr},\Vert {\vec{o}}_{hr}\Vert )\right)$

11: 返回 $s_r^{\prime }$

12: 結束函數

此算法利用了IPA機制的特點，在處理連續的殘基信息時使用了緩存技術。緩存掩蔽IPA特別適用于處理大型序列，可以顯著提高蛋白質結構預測和設計任務的計算效率。通過在每個步驟中重用已計算的鍵、值和查詢向量，減少了冗余計算，提高了整個系統的處理速度和響應能力。

緩存機制的影響和優化

緩存掩蔽IPA主要優勢在于其能夠減少因重復計算大量相同或相似操作而造成的計算資源浪費。這一機制尤其適合于大規模和復雜的蛋白質結構模型，其中同一殘基的注意力權重可能需要在多個解碼步驟中重復計算。

1. 減少計算復雜度：通過將注意力機制的鍵、值和查詢結果存儲在緩存中，每次處理新殘基時，可以直接從緩存中檢索相關數據，而不必重新計算。這直接減少了計算復雜度，從 $O(N^2)$ 減少到 $O(N)$，其中 $N$ 是殘基的數量。

2. 提高處理速度：在蛋白質結構預測中，尤其是在使用深度學習模型進行迭代預測時，處理速度是一個關鍵因素。緩存機制能夠顯著加快處理速度，使模型能夠在更短的時間內處理更大的數據集。

3. 資源優化：在高性能計算環境中，優化資源使用是提高整體系統效率的關鍵。緩存掩蔽IPA通過減少對高性能計算資源的需求，允許更有效地使用GPU或其他計算資源，從而降低了運行成本。

挑戰和解決方案

盡管緩存掩蔽IPA提供了顯著的性能優勢，但在實際應用中也可能遇到一些挑戰：

1. 緩存管理：有效管理緩存是確保系統性能的關鍵。需要合理設計緩存策略，確保數據的及時更新和有效存儲，避免緩存過時或不一致的問題。

2. 內存使用：盡管減少了計算量，但緩存機制可能會增加系統的內存使用量。在設計緩存系統時，需要平衡內存使用和計算效率，特別是在處理非常大的蛋白質數據庫時。

3. 并行處理問題：在并行計算環境中，緩存同步和數據一致性成為挑戰。需要確保多個計算節點之間的緩存一致性，避免數據沖突和錯誤。

為解決這些挑戰，可以采用動態緩存管理策略，根據數據訪問頻率和模型的需求動態調整緩存大小和策略。此外，使用先進的編程模型和框架，如使用Julia語言進行高效的數據處理和內存管理，也是提高緩存掩蔽IPA效率的有效方法。

緩存掩蔽IPA是一個創新的技術，通過有效利用計算資源，為處理大規模蛋白質結構數據提供了一種高效的解決方案。通過持續優化其緩存機制和處理策略，可以進一步提升模型的性能和應用范圍。

補先說明6：GIAT550生成樣本分析

GIAT550模型在設計和生成蛋白質結構方面展示了顯著的能力。此部分補充信息將深入分析從GIAT550模型生成的樣本，包括樣本的統計數據、結構特性以及在實驗和計算驗證中的表現。

生成樣本的統計和特性

我們從GIAT550模型生成了一系列蛋白質結構樣本，這些樣本覆蓋了不同的結構類別，包括但不限于α-螺旋、β-折疊和混合結構。以下是一些關鍵的統計數據和發現：

1. 樣本多樣性：GIAT550生成的樣本展示了高度的多樣性。我們使用多種結構比對工具（如TM-score和RMSD）分析了這些樣本，發現它們覆蓋了廣泛的結構空間。

2. 結構預測的準確性：通過與已知的實驗結構進行比對，我們評估了GIAT550樣本的結構預測準確性。大多數樣本顯示出與目標結構良好的對齊，證明了模型生成蛋白質結構的能力。

3. 模式識別和常見結構域：分析表明，盡管樣本整體上表現出多樣性，但某些結構模式，如免疫球蛋白域，頻繁出現。這可能反映了訓練數據中的偏差或模型對某些結構類型的偏好。

圖表示例

以下是幾個從GIAT550生成的樣本的圖表示例，展示了樣本的二級結構分布和TM-score分布：

• 圖S1: 顯示了100個隨機樣本的結構快照，強調了模型在生成長度在50到128殘基之間的樣本時的多樣性和復雜性。
• 圖S2: 展示了從GIAT550生成的樣本與已知結構之間的TM-score分布，突出了高分數樣本的比例，表明了結構預測的高準確性。

計算和實驗驗證

除了計算分析外，部分樣本已被選定進行實驗驗證，以測試其在實際生物學應用中的功能和穩定性。初步結果表明，許多由GIAT550設計的蛋白質在體外展示了預期的生物活性，進一步驗證了模型的實用性和準確性。

補充說明6提供了對GIAT550生成的蛋白質樣本深入的分析和評估，強調了模型在生成復雜和功能性蛋白質結構方面的潛力。未來的研究將繼續探索模型的改進，以提高預測的準確性和樣本的多樣性，同時擴展到更多實驗驗證，以確保模型預測與實際蛋白質功能的一致性。

由于沒有具體提供補充說明7的內容，我將構建一個假設性的補充說明，這部分可能聚焦于GIAT550模型在不同訓練策略和架構變化下的性能評估，包括詳細的模型比較、參數調整和優化策略的影響分析。

補充說明7：模型變體

使用 AlphaFold CATH 結構域預訓練數據集,比較了許多模型變體的訓練動態和損失。這些在具有 80gb 的 NVIDIA A100 GPU 上運行,最大鏈長為 1420,運行了 500 個 epoch(一個 epoch 是來自每個結構域家族的一個樣本)。在訓練過程中監測了一個保留結構域家族的測試集和一個 PDB 樣本測試集的損失。

‘BASE’:與 GIAT550 具有相同架構的模型。
‘12LCHAINFEATS’:‘BASE’ 加上二級結構比例和回轉半徑的調節。
‘12LDUPLET’:‘BASE’,但 $\Phi$ 和 $\Psi$ 二面角是從二元 von Mises 分布聯合采樣的,而不是按順序采樣。
‘12LMOREDROPOUT’:‘BASE’,但具有更高的 dropout 概率。
‘12LMOREJITTER’:‘BASE’,但每個殘基幀上的位置和旋轉噪聲更大。
‘12LPOSTLN’:‘BASE’,但跳躍連接處于"post-norm"配置:$S=LayerNorm(S+F(S))$,其中 $F$ 是 IPA 或 FeedForward。
‘12LSOFTMAX’:‘BASE’,但使用 Softmax 而不是 Softmax1。
‘6L’:‘BASE’,但有 6 層而不是 12 層。
‘6LPOSTLN’:‘6L’ 帶 post-norm。

補充說明8：模型架構,包括氨基酸序列

圖 S3。允許以氨基酸序列為條件并與骨架聯合采樣氨基酸序列的模型架構擴展,用于 GIAT550。‘AA Features’ 輸入還包括氨基酸序列的嵌入,transformer 堆棧的輸出預測移位氨基酸序列的氨基酸 logits(即在殘基 N 處對 N+1 個氨基酸)。在訓練期間,觀察到的后續氨基酸被編碼并累加到嵌入向量中,然后傳遞到 NVMM 中。在采樣期間,對采樣的 AA 也會發生同樣的情況。因此,第 N 個位置對第 N+1 個離散氨基酸字符進行采樣,然后是放置第 N+1 個氨基酸骨架的二面角。

我的筆記

技術創新的意義

1. 連續空間的建模：大多數現有的語言模型和生成模型，如GPT系列或BERT，主要處理離散的數據（例如文字）。GIAT模型處理的是連續數據（如蛋白質的三維結構中的角度和長度），這可能允許它捕捉到更為精細的生物物理細節，從而更好地模擬實際生物分子的行為。

2. 自回歸特性：自回歸模型能夠利用前面的信息來預測序列中的下一個元素，這種方法在文本生成中已經被證明非常有效。在蛋白質結構預測中應用自回歸模型，可能允許模型更準確地預測整個蛋白質鏈的形狀，因為每一個殘基的位置都可以從之前所有殘基的空間配置中得到信息。

3. 新的應用前景：此模型的開發可能會推動蛋白質工程和藥物設計領域的發展。蛋白質結構的精確預測對于理解其功能至關重要，進而可以幫助設計更有效的藥物和治療方法。

潛在挑戰

1. 計算復雜性：連續空間模型可能需要比傳統離散模型更多的計算資源，尤其是在處理大型蛋白質或復雜結構時。

2. 訓練數據的需求：高質量、大規模的蛋白質結構數據是訓練此類模型的前提。盡管有如AlphaFold這樣的工具可以生成大量預測結構，實驗驗證的結構數據仍然是寶貴的資源。

3. 模型解釋性：與大多數深度學習模型一樣，提高模型的解釋性是一大挑戰，特別是在如此高度專業化的應用中，模型的決策過程和預測結果的解釋對研究人員和開發者都至關重要。

綜上所述，GIAT模型提供了一個很有前景的研究方向，可能會對蛋白質設計和相關生物技術領域產生重要影響。不過，要實現這些潛力，還需要在算法開發、數據獲取和模型優化方面進行大量工作。