利用機械解釋性提取高效演算法

前傳：DNA 語言模型中的生命之樹

去年，Goodfire AI 的研究人員採用了基因組基礎模型 Evo 2，並在其中直觀地發現了「演化樹」。數千個物種之間的系統發育關係被編碼為模型內部激活中的一個彎曲流形（manifold），而沿著該流形的測地線距離追蹤了實際的演化分支長度。數億年前分化的細菌在流形上相距甚遠，而親緣關係近的物種則彼此靠近。

該模型被訓練用於預測下一個 DNA 代幣（token）。沒有人告訴它關於演化的知識，也沒有給它系統發育樹作為訓練信號。但模型為了準確預測 DNA，需要編碼演化關係，因此它在其內部計算中構建了這些關係的結構化幾何表示。這種表示非常出色，以至於你可以用可解釋性工具將其提取出來，並與地面實況（ground truth）進行有意義的比較。

看到這一點後，我決定將同樣的方法應用於另一類生物基礎模型——那些在單細胞數據上訓練的模型。

如果 Evo 2 從原始 DNA 中學習到了生命之樹，那麼 scGPT 對人類細胞的發育學到了什麼？

尋找流形

對於不熟悉生物學方面的讀者：scGPT 是一個在數百萬個單細胞基因表達譜上訓練的 Transformer 模型。你體內的每個細胞都以不同的水平表達數千個基因，而單細胞 RNA 測序實驗可以同時測量潛在數十萬個單個細胞的表達水平。scGPT 以生成方式在這種數據上進行預訓練，學習從上下文預測被遮蓋的基因表達值。

我想回答的問題是：scGPT 是否在其注意力張量（attention tensor）的某處，編碼了某些生物過程的緊湊幾何表示？如果是這樣，我能否在不預先確切知道位置的情況下找到它？

為了系統地解決這個問題，我使用了一個由 AI 執行者-審核者對驅動的兩階段研究循環，並在預註冊的質量門控下運行。第一階段是廣泛的假設搜索：該循環通過改變生物目標（發育排序、調控結構、通訊幾何）、特徵化策略（注意力漂移、原始嵌入、混合算子）和幾何擬合方法（Isomap、測地線 MDS、一種稱為局部歐幾里得變換的技術），探索了候選流形假設的大型組合空間。這一切都應用於完整的 12 層 × 8 頭 scGPT 注意力張量，這意味著要篩選 96 個獨立的注意力單元。

第一階段的結果是一個強健的正向發現：假設 H65。它在特定的注意力頭中識別出一個緊湊的、大約 8 到 10 維的流形，流形上的位置對應於細胞在造血分化過程中的進展程度。幹細胞聚集在一端；終末分化的血細胞類型（T 細胞、B 細胞、單核細胞、巨噬細胞）在另一端沿著不同的分支展開；且分支拓撲結構與已知的發育層級相匹配，具有在嚴格對照下依然成立的顯著分支結構。

然後我轉向第二階段，這是相當手動的調查：方法論閉合測試、置信區間、結構化留出法和外部驗證。我在來自 Tabula Sapiens 的非重疊面板上驗證了該流形，然後通過凍結頭部零樣本遷移（frozen-head zero-shot transfer）到一個完全獨立的多供體免疫面板上進行了確認。你可以在交互式 3D 查看器中親自探索這個流形並比較不同的提取變體。

但提取出的算法真的有效嗎？

我認為在基礎模型中發現具有生物學意義的流形本身就很酷。但我真正關心的問題是：你是否能將這個幾何對象從模型中「取出」，並將其作為一個獨立的方法來完成有用的工作？

為此，我開發了一個三階段提取流水線：

• 我直接導出相關頭部的凍結注意力權重矩陣，無需重新訓練，只是字面意義上讀取學習到的線性算子。
• 我附加了一個輕量級的學習適配器（adaptor），將原始注意力輸出投影到流形的坐標系中。
• 然後我添加了一個特定任務的讀取頭（用於分類或偽時間預測）。

這個流水線的關鍵屬性是，核心工作（即實際的生物學知識）完全來自 scGPT 在預訓練期間學習到的凍結注意力權重。適配器和讀取頭體積小且訓練成本低，而且它們從未接觸過模型預訓練時使用的原始數據集。你最終得到的是一個獨立的算法，可以作為文件發送並獨立於 scGPT 運行。

那麼：它的表現如何？

我將提取出的算法與生物學家在實踐中使用的成熟方法進行了基準測試：scVI（單細胞數據的深度生成模型）、Palantir（基於擴散映射和馬可夫鏈的偽時間方法）、Diffusion Pseudotime（Scanpy 實現版本）、CellTypist（基於邏輯回歸、在大型參考圖譜上訓練的細胞類型分類器）、PCA 以及原始表達基線。這些是單細胞生物信息學工具箱中的標準工具，由領域專家歷經多年開發和完善。

在偽時間深度排序（衡量方法恢復從幹細胞到成熟血細胞真實發育進程的能力）方面，提取出的算法表現最優，在配對分層統計中顯著優於所有測試的替代方案。在分類（區分細胞類型）方面，情況雖不如前者明確但依然強勁：提取出的讀取頭在分支平衡準確率以及關鍵亞型區分任務（如 CD4/CD8 T 細胞分離和單核細胞/巨噬細胞區分）上領先。在某些階段級別和分支級別的 macro-F1 指標上，擴散式基線或原始表達具有優勢，因此這並非全勝，但提取出的算法在各方面都穩居頂尖梯隊，並在最具生物學意義的終點指標上佔據主導地位。

現在，你可能會合理地問：這是否只是擁有更高級探針（probe）的結果？也許任何在 scGPT 嵌入之上擬合的足夠靈活的函數都能做得同樣好，而「流形發現」部分並沒有做出任何實質貢獻。我測試了這一點。一個擁有 175,000 個可訓練參數、在凍結的 scGPT 平均池化嵌入上擬合的 3 層 MLP，在 8 個分類終點中的 6 個上表現顯著遜於提取出的 10 維讀取頭。而且提取出的讀取頭在完成這些任務時，在完整的 12 次分割測試中評估速度快了 34.5 倍，且可訓練參數減少了約 1,000 倍。

讓我重申一遍：機械可解釋性在 scGPT 注意力頭中發現的幾何結構，當被提取並直接使用時，優於在模型嵌入之上疊加 MLP 的標準方法。源自可解釋性的方法同時更準確、更快速且更小巧。

到底能縮減到多小？

一旦你擁有了一個有效的提取算法，自然而然的下一個問題是你實際需要其中的多少部分。出於實際原因，壓縮很有趣，但出於可解釋性原因，它更有趣，因為你在保持性能的同時將算法壓縮得越徹底，就越接近理解它實際在做什麼。

最初提取的算子匯集了 scGPT 的三個注意力頭，大小為 17.5 MB。按現代標準來看並不算大，但也不易於直觀檢查。第一個壓縮步驟是問：我們真的需要所有三個頭嗎？還是單個頭就承載了核心幾何結構？我掃描了 scGPT 張量中的所有 96 個注意力單元，發現單個單元（第 2 層，第 5 頭）本身就承載了實質性的、可遷移的發育幾何結構。由這個單頭構建的緊湊算子大小為 5.9 MB，在基準測試套件上與三頭版本相比幾乎沒有損失。

第二個壓縮步驟更具侵略性：對單頭算子進行截斷奇異值分解（SVD）。這將權重矩陣分解為低秩分量，並丟棄低於選定秩閾值的所有內容。在秩為 64 時，生成的替代模型縮小到 0.73 MB，這已經非常微小，且在所有八個匯總分類終點上仍然擊敗了凍結的 scGPT 平均池化 + MLP 基線。與密集單頭算子相比，它確實在 8 個終點中的 5 個上產生了統計學上的顯著損失，因此這並非無損壓縮。但秩為 64 的版本仍然是比標準探針方法更好的算法，且大小僅為幾分之一 MB。

此外，現在可解釋性的回報到來了。我對秩為 64 的替代模型進行了因子消融審計：系統地一次移除 64 個因子中的一個，測量性能下降程度，並按必要性對其進行排名。結果顯示，在 64 個因子中，僅僅四個因子就佔了總匯總消融影響的 66%。接著，當我檢查這四個因子在生物學上對應什麼時，它們解析成了明確的造血基因程序。

因此，機械可解釋性正成為雙用途技術

讓我們暫時從具體結果中退一步，考慮一下這裡的高層次教訓。

使這一結果有趣的特性，也正是讓我對將同樣技術應用於大型語言模型（LLM）持謹慎態度的原因。因為這個論點是雙向的。如果你能提取出模型用來出色完成某事的算法，你也有可能改進它完成該事的方式：通過識別低效組件、擴展相關電路、以新方式組合提取的子程序、用更清晰的提取版本替換模糊的學習版本並釋放模型學習其他內容的能力等。換句話說，機械可解釋性正在成為一種能力增強工具。這是一個眾所周知的理論擔憂，但現在看來它正成為一個現實問題。

考慮幾個場景。你識別出語言模型中負責多步規劃的電路，將其提取出來，發現它在低秩下運行且具有大量冗餘，並發表論文展示如何壓縮它。現在，任何訓練下一代模型的人都可以更有效地初始化該電路，或為關鍵組件分配更多容量。或者：你發現模型的思維鏈推理依賴於一種特定的注意力模式，該模式以可預測的方式通過中間代幣路由信息，並發表了關於其運作機制的詳細說明。現在，構建推理時間擴展流水線的人可以直接優化該路由，而不是依賴模型從頭開始重新發現它。

這是我特意選擇將可解釋性工作重點放在生物基礎模型上的原因之一。雖然我同意推動生物學也可能伴隨風險，但我相信考慮到目前的 AI 風險格局，我們確實需要盡快推動生物學，而推動生物學正是我努力在做的事。

用於新知識發現的機械可解釋性

從更積極和普遍的角度來看：除了作為審計/監測工具外，機械可解釋性還可以作為一種知識發現工具。考慮到：

• 模型學到了關於造血作用的某些知識，而現有的生物信息學方法尚未完全捕捉到這些知識，至少沒有以同樣緊湊的形式捕捉到。
• 可解釋性流水線發現了一種表示，當將其提取並部署為獨立算法時，在最具生物學意義的基準測試中優於成熟工具。
• 從模型內部提取的知識在操作相關的意義上是全新的：以前沒有人擁有這種特定的算法，而且它比人們一直在使用的工具效果更好。

加入我們

如果你喜歡機械可解釋性，我鼓勵你考慮從 LLM 轉向生物基礎模型。

我在此描述的工作是關於生物基礎模型機械可解釋性更廣泛研究計劃的一部分。此前，我發表了關於 scGPT 和 Geneformer 的基於注意力的可解釋性方法的全方位壓力測試。與此同時，我開發了一個稀疏自編碼器圖譜，涵蓋了 Geneformer 和 scGPT 所有層的 107,000 多個特徵。造血流形論文是最新的一篇。

這裡還有很多工作要做，無論是將這些方法應用於其他生物系統和發育過程，還是開發更好的無監督流形發現技術，以實現超越目前半監督方法所能達到的規模。我認為這是目前研究領域中進行可解釋性工作的最佳切入點之一，既具有方法論上的趣味性，對生物醫學具有實際用途（是的，還有人類智能增強），且相對於令我擔憂的 LLM 可解釋性能力外部性而言是安全的。

你可以在 biodynai.com 找到更多關於研究計劃、正在進行的項目以及參與方式的信息。包含所有補充材料的完整論文已發表在 arXiv 上，交互式 3D 流形查看器位於 biodyn-ai.github.io/hema-manifold。

參與討論