單細胞模型的稀疏自編碼器:解鎖基礎模型內部的生物學知識
我認為生物基礎模型並不像大語言模型,目前的評估方式系統性地低估了它們包含的知識,而我透過應用 AI 安全社群的機械解釋性工具,成功映射出這些模型內部隱藏的生物學邏輯與因果電路。
人們正爭相構建越來越大的單細胞基礎模型(基於 RNA 定序數據訓練),但在我看來,我們甚至還沒從現有模型中提取出其內含知識與能力的冰山一角。
為了說明我的觀點,我將在本文中提出三個論點,並展示背後的實證研究工作。
論點 1:生物基礎模型與大型語言模型(LLM)不同,業界習慣以同樣的方式評估它們,導致我們系統性地低估了它們所包含的內容。 當你與 GPT 互動時,表層輸出(生成的文本)是衡量模型能力的良好指標。你可以閱讀它寫的內容並形成合理的評價。生物基礎模型在這方面則根本不同。像 Geneformer 或 scGPT 這樣的模型接收細胞的基因表達譜,並產生嵌入(embeddings)、被遮蓋基因的預測或細胞類型分類。這些表層輸出僅僅是模型內部運作的極小部分。該模型已在數千萬個細胞上進行了訓練,它為了達成訓練目標而構建的表徵中,包含了壓縮的生物學知識,而這些知識永遠不會直接出現在任何你可以看到的輸出中。僅憑細胞類型標註或擾動預測的基準測試表現來評估這些模型,就像是透過讓人類科學家填寫選擇題考卷來評估其水準一樣。
論點 2:人們不斷將生物基礎模型稱為「虛擬細胞」,但這是一個被暗示而非經過測試或驗證的標籤。 這個詞出現在資助申請書、新聞稿甚至一些論文中,彷彿這些模型已經內化了細胞生物學的運作模擬是一個既定事實。也許它們確實做到了,或者它們只是學會了複雜的統計規律,表面上看起來像生物學,但在仔細檢查下就會瓦解。我的研究顯示,這些模型在有意義的層面上確實是細胞的模型,但這是一個需要實證處理的經驗性主張。
論點 3:正確的工具已經存在,它們源自 AI 安全社群在機械解釋性(mechanistic interpretability)方面的工作。 稀疏自編碼器(SAEs)、因果電路追蹤(causal circuit tracing)、特徵消融(feature ablation)、激活補丁(activation patching):這些方法是為了理解語言模型而開發的,很大程度上是受對齊(alignment)問題的驅動。事實證明,它們非常適合生物基礎模型,原因很簡單:在語言模型中,當你發現一個電路時,你往往缺乏關於該電路在深層意義上是否「正確」的基準真相(ground truth),因為模型內部的計算並不需要對應到任何客觀的外部現實。而在生物基礎模型中,你有數十年的分子生物學研究、精選的通路數據庫、基因組規模的擾動篩選以及特徵明確的調節網絡可以進行驗證。生物學提供了語言模型所缺乏的基準真相。這使得生物基礎模型可以說是目前存在的機械解釋性方法最佳(且真實)的試驗場。
接下來是我最近產出的三篇論文的故事,每一篇都建立在前一篇的基礎上。在這些研究中,我將基於 SAE 的解釋性工具包應用於(不久前的)兩個領先單細胞基礎模型(Geneformer V2-316M 和 scGPT whole-human),並逐步繪製出它們知道什麼、如何計算,以及它們知識的邊界在哪裡。
SAE 圖譜 (The SAE Atlas)
第一個問題非常簡單:這些模型內部有什麼?
神經網絡以「疊加」(superposition)的方式編碼信息。這在語言模型的解釋性文獻中已得到證實,但尚未有人系統性地在生物基礎模型中證明或嘗試解決它。
我針對 Geneformer V2-316M(18 層,d=1152)和 scGPT whole-human(12 層,d=512)每一層的殘差流激活訓練了 TopK 稀疏自編碼器。SAE 將密集、疊加的激活分解為稀疏、可解釋的特徵,每個特徵(理想情況下)對應一個單一的生物學概念。結果產生了一對特徵圖譜:Geneformer 有 82,525 個特徵,scGPT 有 24,527 個特徵,跨越 30 層共計超過 107,000 個特徵。
疊加現象是巨大的。SAE 恢復的特徵中有 99.8% 對於 SVD 等標準線性方法是不可見的,這意味著如果你試圖使用 PCA 或類似方法來理解這些模型,你只能看到 0.2% 的表徵結構。僅憑這一點,就足以讓任何認為自己能根據標準降維方法理解模型行為的人深思。
這些特徵具有豐富的生物學意義。針對五個主要數據庫(Gene Ontology, KEGG, Reactome, STRING, 和 TRRUST)的系統性標註顯示,29% 到 59% 的特徵對應到已知的生物學概念,且在不同層之間呈現有趣的 U 型分佈:早期層標註率高(捕捉基礎通路成員身份),中間層下降(模型似乎在此構建更抽象、更難標註的表徵),後期層再次上升(模型在此重建與輸出相關的生物學類別)。這些特徵還組織成共激活模塊(Geneformer 中有 141 個,scGPT 中有 76 個),展現出因果特異性(當你消融一個特徵時,下游影響集中在特定的輸出基因而非廣泛擴散,中位特異性為 2.36 倍),並形成了連接 63% 到 99.8% 特徵的功能管道跨層信息高速公路。
到目前為止,結果令人振奮。模型顯然已經內化了大量有組織的生物學知識:通路、蛋白質相互作用、功能模塊、層次抽象。這看起來非常接近業界喜歡講述的「虛擬細胞」故事。
繪製佈線圖 (Mapping the Wiring)
SAE 圖譜告訴我們模型內部存在哪些特徵。下一個問題是:它們如何相互作用?計算圖譜是什麼樣的?
我為生物基礎模型引入了因果電路追蹤。該方法的工作原理是在源層消融一個 SAE 特徵(將其在殘差流中的激活設置為零),然後測量所有後續層中每個下游 SAE 特徵的反應。這會給你帶方向、帶正負號的因果邊緣:L 層的特徵 A 因果性地驅動 L+k 層的特徵 B,具有效應大小 d 和方向(興奮性或抑制性)。這不是相關性,不是共激活,也不是互信息,而是一種干預。
應用於四種實驗條件後,結果是一個包含 96,892 條顯著邊緣的因果電路圖,這是通過 80,191 次前向傳播計算得出的。
該圖譜的幾個屬性令人驚訝:
抑制性主導。 65% 到 89% 的因果邊緣是抑制性的:消融源特徵會減少下游特徵的激活。這意味著特徵主要編碼必要的信息。移除一個特徵會導致依賴它的下游特徵失去激活,而不是為其他特徵騰出空間(後者會產生興奮性邊緣)。模型的計算結構是相互依賴的,而非競爭性的。約 20% 的興奮性部分可能反映了去抑制作用:移除某些特徵會釋放其他受壓制的特徵。
生物學一致性。 在源特徵和目標特徵都有生物學標註的邊緣中,53% 共享至少一個本體論術語(ontology term)。模型內部超過一半的計算路徑連接了生物學相關的特徵。特定電路可以直接解釋為已知的生物學級聯反應。例如,在 Geneformer 中,L0 的 DNA 修復特徵因果性地驅動 L1 的 DNA 損傷反應特徵(d = -1.87,113 個共享本體論術語),後者又連接到 L6 的著絲點(Kinetochore)特徵(d = -3.47),重現了 DNA 損傷檢測、修復機制激活與有絲分裂檢查點參與之間公認的聯繫。該模型僅透過基因表達數據的訓練,就發現了一個分子生物學家需要數十年實驗工作才能表徵的電路。
跨模型收斂。 當我比較 Geneformer 和 scGPT(具有不同架構、訓練數據組成和訓練目標的模型)的因果佈線時,我發現它們獨立學習到了驚人相似的內部電路。1,142 個生物領域對在不同架構間是保守的,其富集程度超過隨機機率的 10 倍。更具代表性的是,疾病相關領域在這一共識集中超額代表了 3.59 倍,這意味著對人類健康最重要的生物學,正是兩個模型最可靠收斂的生物學。兩個完全不同的神經網絡在獨立訓練後,內部佈線了相同的生物學,且這種收斂在疾病相關通路上最強。
窮盡式研究與發現生物特徵的暗物質
在第三篇論文中,我不再只追蹤 30 個精選特徵,而是追蹤了 Geneformer 第 5 層中每一個活躍的 SAE 特徵(共 4,065 個),產生了 1,393,850 條顯著的因果邊緣。這比論文 2 中的選擇性採樣擴大了 27 倍。
結果推翻了選擇性分析中的幾個結論。完整的電路圖揭示了一個重尾(heavy-tailed)的樞紐架構,僅 1.8% 的特徵就佔據了極高比例的連接性。但有趣的部分在於:前 20 個樞紐特徵中,有 40% 完全沒有生物學標註。它們不對應 GO、KEGG 或 Reactome 中的任何已知通路。這些是模型在計算中最依賴的特徵,而它們恰恰是我們早期帶有標註偏見的採樣中系統性排除掉的。
這具有嚴重的方論意義!如果你只解釋那些已有生物學標籤的特徵,你就是在路燈下找東西:你會找回已知的生物學並得出模型已學會生物學的結論,而模型實際最依賴的特徵則處於黑暗中,未被研究。這些無標籤的樞紐特徵中,有些可能代表了無法整齊對應到現有通路數據庫的新穎生物程序,另一些可能是模型為了以我們尚未概念化的方式壓縮細胞狀態而發明的計算抽象。無論如何,它們正是最有趣的發現可能隱藏的地方,而任何預先過濾標註的解釋性流程在結構上都無法找到它們!
此外,最初的 SAE 圖譜顯示某些特徵與分化狀態相關:有些特徵在成熟細胞中更活躍,有些在祖細胞中更活躍。但那只是相關性,對於「虛擬細胞」主張而言,關鍵問題在於:放大一個與分化相關的特徵是否真的能將細胞狀態推向成熟?
答案是肯定的。後期層特徵(L17)因果性地將細胞推向成熟,而早期層特徵則將其推離成熟。模型學到了一種依賴於層級的分化梯度,而且我們可以引導它:放大一個後期層的分化特徵,細胞的計算狀態就會向更成熟的表型移動。這是第一個因果證據,證明這些模型編碼了類似功能性發育程序的信息,這也是我們目前最接近驗證「虛擬細胞」隱喻的證據。
這一切意味著什麼?
好消息是,生物基礎模型包含的知識遠比任何人提取出來的都要多。超過 107,000 個可解釋特徵,組織成生物通路,由重現已知分子生物學的因果電路連接,並在獨立架構間收斂。「虛擬細胞」的隱喻並非毫無根據;這些模型內部正在發生真實、有結構、具有生物學意義的計算,我們可以識別、繪製甚至引導它。是的,這些知識中有很大一部分是相關性的,但並非全部。而且我們面臨一個大問題,即至少前一代的模型沒有學會調節網絡。更多內容請參閱此處。
這同時也是一個明確的方法論警告:在計算上最重要的特徵,往往是不具備生物學標籤的那些。未來在該領域的任何工作都需要解決標註偏見問題,否則它將不斷產生證實我們已知知識的結果,卻錯失我們未知的領域。
我越來越確信這裡存在著巨大的機會。為 AI 安全開發的機械解釋性工具,事實證明是從基礎模型中提取生物學知識的強大工具。