7.1 將矩陣化為坐標:剛體群組(Rigid Bodies)的概念#
一個反直覺的難題:為什麼不直接預測座標#
若要向一位工程師描述 AlphaFold 2 的結構模組,最直覺的說法是:「輸入序列與演化資訊,輸出每個原子的三維座標。」這個說法聽起來毫無問題——事實上,AlphaFold 1 正是照著這個直覺設計的。但 DeepMind 在 2018 到 2020 年間放棄了這條路,而放棄的理由不是預測不準,而是預測在數學上從來沒有被設定成一件「對」的題目。
問題出在損失函數。假設我們用均方誤差直接回歸原子的絕對座標:
這裡藏著一個幾何上的殘缺。同一個蛋白質的正確結構,可以被任意剛體運動搬動:整體平移三公釐,或繞任意軸旋轉九十度,摺疊的物理本質完全不變,但上面的損失函數卻會從近乎零跳到一個極大的數值。也就是說, \(L\) 不具平移不變性,也不具旋轉不變性。它對「答案的姿態」與「答案的形狀」一視同仁地索求精確,於是網路被迫花費大量參數去記憶一個本質上任意、且每個蛋白質都不同的整體朝向。
這還只是第一層,更嚴重的是第二層:當損失對某個方向的擾動不敏感時,梯度就會消失。想像模型已經學會了正確的相對構形,卻把整個蛋白質放在離原點很遠的地方,或隨便轉了一個角度。此時真正的形狀誤差已趨近於零,但 \(L\)
物理學家與數學家早就給出了標準答案:不要用絕對座標描述物體,要用量在群作用下保持某種性質的方式來描述。但在深度學習裡,這個答案的落腳點不是「把座標換成不變特徵」那樣簡單——因為網路最終必須交出一組實際的原子座標。於是 AlphaFold 2 選擇了一條更精妙的路:讓網路在相對幾何上思考,但在輸出時重新組裝成絕對座標。
框架表示:把一個殘基視為一個三維座標系#
具體的做法,是為每個殘基 \(i\) 引入一組參數:一個旋轉矩陣 \(R_i \in \mathbb{R}^{3 \times 3}\) 與一個平移向量 \(t_i \in \mathbb{R}^3\) 。這一對參數合稱一個框架,在數學上就是三維歐氏空間中特殊正交群與平移群的直積:
也就是說,每個殘基都配備了一個「自己的座標系」:原點在 \(t_i\) ,三條軸由 \(R_i\) 的三個列給出。用這個框架,一個殘基的骨幹原子(氮原子 \(N\) 、碳原子 \(C\) 、以及 \(\alpha\) 碳 \(C_\alpha\) )就不再需要獨立的座標變數,而可以由一組固定的局部常數 \(a_k\) 決定:
其中 \(k\) 遍歷 \(N\) 、 \(C_\alpha\) 、 \(C\) 三個骨幹原子。 \(a_k\) 是常數,它的數值取自理想化的肽鍵幾何,也就是相鄰殘基之間 \(C_\alpha\) 至 \(C_\alpha\) 距離約 3.8 Å 、相鄰肽鍵平面夾角約 111 度這樣的統計平均值。這個常數化的處理有一個極重要的後果:同一個殘基的骨幹形狀被凍結了,網路不能任意拉伸肽鍵、不能把肽鍵扭成不可能的角度,它能操縱的只有整塊剛體的姿態。這正是物理上正確的約束——肽鍵確實是剛性的。
於是,一個長度為 \(L\) 的蛋白質的完整三維結構,其描述從「 \(3L\) 個原子、每個原子三個座標、共 \(9L\) 個連續自由度」被重新參數化成「 \(L\) 個 \(SE(3)\) 元素、合計 \(6L\) 個自由度」。這不是資訊的損失,而是資訊的重組:所有可行的蛋白質構形依然被完整覆蓋,但覆蓋的方式變得對幾何友善得多。
等變而非不變:為什麼這個重組是關鍵#
這個重組帶來一個可以寫成定理的性質,稱為剛體運動等變性:若把輸入的整個結構用某個 \((G_R, G_t) \in SE(3)\) 變換,那麼結構模組的輸出也會被同一個剛體運動變換:
這跟前面那個失敗的損失函數形成鮮明對比。等變性意味著模型對整體姿態的輸出會跟著姿態一起轉,所以訓練目標裡那些與姿態無關的誤差分量被徹底移除,只剩真正的形狀誤差在驅動梯度。整體旋轉與平移的資料擴增因此變得不再必要——它已經被架構本身免費提供了。
這也解釋了為什麼「讓網路直接回歸絕對座標」與「讓網路操作相對幾何」是兩種本質不同的設計。前者是要求網路內部記住一個任意的參考系;後者是讓網路的每一層運算都在一個有物理意義的局部座標系內進行,再把結果送回全局。前者把幾何學推給資料,後者把幾何學寫進架構。
剛體群組:把幾個殘基綁成一塊#
還有一層精化。真實蛋白質中,相鄰殘基之間的相對構形並非任意——存在螺旋、伸展、彎折、超螺旋等等二級結構型態。若讓網路為每個殘基各自獨立地更新一個 \(SE(3)\) 框架,等於要求它獨自學會「這三個殘基組成一段 \(\alpha\) 螺旋」這種關係,資訊利用率很低,精度也受限。
AlphaFold 2 的做法是引入剛體群組:把鏈上連續的 \(b\) 個殘基綁成一個剛體,這個群組共享同一組 \((R, t) \in SE(3)\) 參數,而群組內每個殘基仍然保有自己的局部框架相對於群組框架的偏移量,這些偏移量在訓練中也可以學習。整體結構由多個剛體群組的組合描述,而每個殘基的側鏈則另外以六個可旋轉鍵描述,這是另一種「把局部幾何凍結」的技巧。
剛體群組帶來三個好處。第一是自由度大幅下降:一個 \(b = 4\) 的群組,其內部約束了 \(3(b-1)\) 個相對自由度,網路不必再學這些平凡的局部幾何。第二是結構的生成變得階層化:網路先決定「這一段要彎成什麼角度」,再決定「這一段相對於全局放在哪裡」,這與蛋白質實際的建構過程(先形成二級結構元素、再折疊排布)在概念上是同構的。第三是精度顯著提升:在 AlphaFold 2 的消融實驗中,把獨立的逐殘基框架換成剛體群組,是結構模組中最有效的單項改進之一。
不過為了精確起見,需要補一句實作層面的說明。在 AlphaFold 2 的官方程式碼中,剛體群組是以偏移量來索引的:骨幹部分最基本的是偏移為 0 的逐殘基群組,每個殘基有自己的骨幹框架;側鏈部分則藉由六個可旋轉鍵的角度,額外引入偏移量為 1 到 4 的衍生群組,例如 \(\chi_1\) 的旋轉會帶動 \(\beta\) 與 \(\gamma_2\) 等後續原子一起轉動。換句話說,AF2 真正的剛體群組體系是「骨幹逐殘基加上側鏈逐扭轉」的混合體,而「把連續 \(b\) 個殘基綁成一個剛體」則是這個思想的一般化形式,在文獻與後續工作中被廣泛沿用。理解這一點很重要:剛體群組不是單純的壓縮技巧,而是把「哪些原子必須一起動」這條物理知識,寫進了參數化的拓撲結構裡。
指數映射:把網路的輸出變成一次形變#
框架表示只是靜態的描述,真正精巧的部分是網路如何更新它。結構模組並不直接回傳一個新的 \(R_i\) 或 \(t_i\) ,而是回傳一個增量:一個旋轉向量 \(u \in \mathbb{R}^3\) 與一個平移向量 \(v \in \mathbb{R}^3\) 。這個旋轉向量被轉換為旋轉矩陣的方式,是李群 \(SO(3)\) 的指數映射:
其中 \(\hat{u}\) 是 \(u\) 的反對稱矩陣, \(\theta = \|u\|\) 是旋轉角。這個形式的妙處在於: \(u\) 的長度直接就是旋轉角度,方向就是轉軸,因此網路輸出的是一個連續的、可微的、無奇異的旋轉參數化——沒有歐拉角在 \(\pm \pi\) 附近的萬向節死鎖,也沒有四元數的雙重覆蓋問題。
更新規則因此簡潔地寫成:
這裡採用的是「全域增量」的慣例: \(u_i\) 與 \(v_i\) 都是在全域座標系中表達的變換量。若採用「局部增量」的慣例,則寫成 \(R_i \leftarrow R_i \operatorname{Rot}(u_i)\) 、 \(t_i \leftarrow t_i + R_i v_i\) ,兩者在數學上等價,只是增量被解讀為在殘基自身座標系中的變化。在實作上,AlphaFold 2 內部以四元數儲存旋轉矩陣以兼顧穩定性,但數學內容完全相同。
這一步是整個設計中最關鍵的一次範式轉移。結構模組的一次迭代,等價於對結構施加一次連續的形變。網路不是在一個迴圈裡算出答案,而是反覆地修正草稿:初始框架是一個退化的構形,經過多層結構模組層的逐步更新,逐漸舒展成摺疊好的三維結構。這正是摺疊漏斗模型在計算上的可微分版本——第 2 章描述的那個「從高熵的展開鏈走向低熵的天然構形」的漏斗,在這裡被實作成一個可以沿梯度下降的優化過程,而網路學會的正是如何執行這個下降過程。
初始化的選擇:從「黑洞」到「展開的鏈」#
框架表示還帶來一個前面沒有交代的問題:既然結構模組是對框架做形變,那麼初始框架是什麼?AlphaFold 2 採用的是黑洞初始化:所有殘基的框架被設在原點附近,旋轉取恆等矩陣,於是初始構形是一團完全退化的、所有原子擠在原點的東西。這個選擇乍看之下極不合理,但有一個非常實際的優點:它不引入任何假的偏置。如果初始化時把序列攤開成一條直鏈,網路就必須額外學會「把直鏈彎起來」,而且會不自覺地沿著直鏈的方向偏置整個預測。
另一種在直覺上更自然、但被實驗證明有害的初始化,是由序列組成的展開鏈,也就是用理想化的鍵長鍵角把序列展成直鏈。這個做法看似提供了一個合理的起點,卻反而有害:AlphaFold 團隊發現它會造成顯著的精度損失。原因很可以理解——展開鏈是一個錯誤的局部極小值,一旦網路從這個起點出發,摺疊的早期動態就被鎖死在攤開的樣子上,而摺疊漏斗的入口本來就允許非常廣泛的構形分布。這個對照有一個一般性的啟示:在深度學習的迭代式生成模型中,初始分佈的選擇往往比任何單一網路元件更決定成敗。
結構模組的資料流#
結構模組在 AlphaFold 2 中由 8 個堆疊的區塊組成,輸入是三樣東西:從 Evoformer 來的序列表示(384 維)、配對表示(128 維),以及一組初始框架。每一區塊內部包含不變點注意力、一個逐點 MLP、以及一個框架更新模組;八個區塊的末端還有一個精細的骨幹更新模組與側鏈扭轉角預測頭。
flowchart TD
MSA[MSA 表示 256 通道] --> EVO[Evoformer 48 區塊]
Z[配對表示 z_ij 128 通道] --> EVO
EVO --> S[序列表示 s_i 384 通道]
EVO --> Z2[更新後配對表示 z_ij]
S --> SM[結構模組 8 區塊]
Z2 --> SM
INIT[初始框架 R_i, t_i] --> SM
SM --> IPA[不變點注意力 IPA<br/>4 個頭 每頭輸入 128 輸出 16]
IPA --> UPD[框架更新<br/>輸出 u_i, v_i]
UPD --> EXP[指數映射<br/>R 新 = Rot(u) R, t 新 = t + v]
EXP --> FR[新框架 R_i, t_i]
FR --> ATOM[骨架座標 x_ik = t_i + R_i a_k]
ATOM --> SM
FR --> TOR[骨架扭轉角 與 側鏈扭轉角]注意圖中的回饋箭頭:更新後的框架會成為下一個結構模組區塊的輸入框架。這意味著結構模組的八個區塊並非八次獨立的預測,而是八次連續的形變,每一層都在前一層的基礎上繼續收斂。
這解決了 AlphaFold 1 的什麼問題#
現在可以回頭精確地回答本章開頭的問題:這套設計究竟修補了 AlphaFold 1 的哪個缺陷?
AlphaFold 1 的流程是:網路輸出一個距離圖譜,把殘基對之間的距離分箱到 37 個區間;然後在推論階段,用一個基於距離約束的優化器從隨機初始化出發,把距離圖譜反演成一組三維座標。這個流程有兩個結構性的問題。
其一是資訊的量化損失。距離圖譜只有 37 個分箱,資訊在網路輸出那一刻就被捨棄了;更關鍵的是,距離本身不含手性資訊——一個鏡像的結構與原結構有完全相同的距離圖譜。這意味著即使網路完美預測了所有距離,你仍然面對一個幾何上二義的逆問題,必須靠外部先驗去挑選正確的那一支。而框架表示一次解決了這兩個問題:因為網路輸出的是相對幾何的參數化增量,距離與手性不再需要被中距離圖譜傳遞,手性被編碼在 \(R_i \in SO(3)\) 之中,而 \(SO(3)\) 本身就不是鏡像對稱的群,其行列式恆為 \(+1\) ,所以鏡像構形在表示層面根本不可達。
其二是推論時的搜尋失敗。梯度下降的座標優化是局部方法,摺疊漏斗裡充滿了局部極小值;對長鏈蛋白質,這個優化經常收斂到錯誤構形。也就是說,AF1 把最難的部分(從二維約束到三維結構)留給了一個不可學習、不可微、充滿遺漏解的推論時優化器。
| 面向 | AlphaFold 1 | AlphaFold 2 的框架表示 |
|---|---|---|
| 中間表示 | 37 箱距離圖譜 | \(SE(3)\) 框架集合 |
| 手性資訊 | 遺失,距離對鏡像不變 | 內建於 \(R_i \in SO(3)\) |
| 幾何對稱性 | 需外部處理,損失函數不具不變性 | 架構層級等變性 |
| 座標產生方式 | 推論時非可微梯度下降 | 網路內部可微的連續形變 |
| 失敗模式 | 局部極小值、幾何二義性 | 形變過大導致的幾何失真 |
| 每個殘基對的輸出 | 一個純量,也就是分箱距離 | 一個 \(SE(3)\) 增量,也就是 \(u\) 與 \(v\) |
值得補充的是,最終輸出並非只有骨幹三原子。AlphaFold 2 的結構模組末端還有一條側鏈分支:由框架推導出各殘基的六個扭轉角,再由扭轉角重建完整的側鏈原子座標,並輸出每個原子的 PDB 格式座標以及一個以每個 \(C_\alpha\) 為代表的位置置信度。這條分支與骨幹框架共享同一套「局部常數加上剛體變換」的表示哲學,只是變換的參數從 \((R, t)\) 換成了扭轉角。
小結:從「猜座標」到「操作幾何」#
本章的核心可以濃縮成一句話:AlphaFold 2 沒有讓網路更會猜絕對座標,而是讓網路操作相對幾何。具體機制是三層:第一層是框架表示,把每個殘基變成一個 \(SE(3)\) 元素,使得骨架原子由 \(x_{ik} = t_i + R_i a_k\) 生成;第二層是剛體群組,把連續殘基綁成共享參數的剛體,以對應蛋白質的二級結構階層;第三層是指數映射更新,使得每一次迭代都是一次可微的連續形變。
這三層共同把摺疊漏斗從一個描述性的比喻,變成一個可計算、可微分、可學習的優化地形。而要讓這個地形上的每一步都走在物理與幾何允許的方向上,注意力機制必須本身懂得三維幾何——那正是下一節的主題:不變點注意力如何在每一個注意力頭上,把剛體運動等變性硬編碼進去。
值得留意的還有一層方法論上的意義。第 7 章之前,深度學習在蛋白質結構上的每一次進展,本質上都是「找到更好的輸入表示」:用 MSA 取代單一序列、用配對表示取代距離圖譜。而從本章開始,進步的來源變成了「把物理約束寫進架構本身」。這個轉向的重要性遠大於任何單一技術改進——它意味著結構預測從此不再是純粹的資料驅動問題,而成為可以讓物理學家在模型裡寫下假設的混合問題。第 13 章談到後 AlphaFold 時代的挑戰時,這一點會再次浮現:模型可以內建物理,卻始終無法為尚未被理解的生物物理現象提供正確的先驗。