默克公司發現：換一個"注意力"方式，AI讀懂細胞的能力提升了25%

這項由默克公司（Merck & Co., Inc.）劍橋研究團隊完成的研究，以預印本形式發布于2026年4月29日，論文編號為arXiv:2604.27124，有興趣深入了解的讀者可以通過該編號在arXiv平臺查詢完整論文。

研究的核心問題聽起來非常專業，但背后的邏輯其實很直觀：當一個AI模型在"閱讀"一個細胞的基因信息時，它是怎么決定哪些基因之間的關系更重要的？這個"決定方式"，也就是所謂的"注意力機制"，是整個模型能否準確理解細胞身份的關鍵。研究團隊發現，長期以來被默認使用的那種注意力方式——就像一個人在考場上只能把注意力分給有限幾道題——其實并不適合生物數據。他們改用了另一種方式——允許同時全力關注很多事情——結果發現，不僅模型理解細胞類型的能力提升了25%，訓練速度也加快了近10%，而且還避免了訓練過程中的災難性崩潰。

這項研究的意義不只在于一個技術指標的提升。單細胞RNA測序數據，簡單說就是通過檢測每個細胞里哪些基因在工作來理解細胞的"身份"，是現代生物醫學研究的核心工具之一。基于這類數據訓練的AI模型，可以用來自動識別細胞類型、預測藥物對細胞的影響、研究疾病發展機制，甚至輔助個性化醫療。研究團隊還開源了一個專門為生物數據設計的高效計算內核，讓這種更好的注意力方式真正能在實際研究中跑起來。

一、當AI學會"閱讀"細胞：背景與挑戰

要理解這項研究解決了什么問題，先要知道這類AI模型是怎么工作的。

細胞可以被理解為一個極其復雜的工廠，而基因就是這個工廠里不同車間的開關。一個細胞在某個時刻"表達"了哪些基因、表達量是多少，就像是這個工廠在某個時間節點的運轉狀態快照。科學家通過單細胞RNA測序技術，能夠捕捉到數十萬甚至數百萬個細胞各自的這種快照。

這類AI基礎模型的工作方式，是把一個細胞的基因表達情況當作一段"文字"來讀——每個基因是一個"詞"，整個細胞的基因組合是一個"句子"。模型通過學習海量細胞數據，理解不同基因之間的協同關系，從而學會區分不同類型的細胞。這類模型與ChatGPT這類語言模型在底層結構上高度相似，都依賴一種叫做"自注意力機制"的核心技術。

問題在于，生物數據和語言數據有一個根本性的不同：生物細胞的"句子長度"差異極大。不同細胞表達的基因數量從幾百到一萬七千多個不等，就像有的人說話只說三個字，有的人一口氣說一萬七千字。研究團隊統計了他們所使用的CellxGene數據集（一個包含1.3億多個細胞的龐大數據庫），發現如果把上下文窗口設置為2048個基因，那么有43%的細胞信息會被截斷，就像讀一篇文章只讀前一半就把后面全部扔掉。而如果想要覆蓋96.6%的細胞，窗口至少需要達到8192個基因。

此外，生物序列還有一個特性：每個細胞必須單獨處理，不能像文本那樣把多篇短文拼湊成一篇長文來湊夠固定長度。這就導致批處理時大量的"空白填充"，計算資源被嚴重浪費。

這兩個挑戰——長序列和大量空白填充——正是這項研究著力解決的核心問題。

二、注意力機制的"競爭性"困境：為什么傳統方式在生物數據上表現不佳

現在來理解"注意力機制"到底是怎么回事，以及為什么傳統方式在生物數據上會出問題。

傳統的注意力機制叫做"softmax注意力"。它的工作原理可以用一個課堂場景來理解：假設一個老師在講課，課堂上有100個學生（對應100個基因），老師一次只能把100%的注意力分配給這100個學生。如果老師把60%的注意力給了第一排的同學，那剩下的40%就必須分給其他99個人。這是一個"零和游戲"——關注了這個，就必然減少對那個的關注。數學上說，softmax會把所有的注意力分數歸一化到一個概率分布上，所有分數加起來必須等于1。

這在語言處理中通常沒什么問題，但在基因調控的世界里，情況恰恰相反。一個基因往往同時受到多個轉錄因子（可以理解為基因調控網絡中的"開關管理員"）的協同調控，這些調控關系是并行的、獨立的，而不是互相競爭的。用競爭性的注意力來模擬這種并行協作關系，就像用一把獨木橋來承載一條多車道高速公路的流量，天然不匹配。

更麻煩的是，當序列變長、基因數量達到幾千甚至上萬時，softmax注意力會出現一種叫做"注意力熵坍縮"的現象——通俗說就是，模型的注意力會越來越極端地集中在少數幾個基因上，其他基因幾乎被完全忽視。這種極端集中會導致梯度（訓練過程中用來調整模型參數的信號）急劇膨脹，最終引發訓練崩潰。在之前的單細胞基礎模型研究中，這種訓練失敗的情況相當普遍，造成了大量的計算資源浪費。

研究團隊提出的替代方案是"sigmoid注意力"。它的工作原理完全不同：每個基因之間的注意力分數是獨立計算的，不需要與其他基因競爭。還是用課堂比喻來說，這就好比老師可以同時全力關注每一個學生，給每個學生的關注度都可以獨立達到100%，互不干擾。數學上，sigmoid函數把每對基因之間的關聯分數獨立映射到0和1之間，而不進行跨基因的歸一化。

這種獨立性帶來了兩個直接好處：一是能更真實地模擬基因的并行協同調控關系；二是梯度的傳播更加穩定，因為sigmoid函數的導數（可以理解為"信號放大倍數"）永遠不會超過0.25，而softmax的信號放大倍數會隨著注意力分數的增大呈指數級膨脹。研究團隊對此做了嚴格的數學推導，證明sigmoid注意力的雅可比矩陣（一種衡量函數對輸入變化敏感程度的數學工具）是對角結構的，而softmax的雅可比矩陣是密集耦合的——前者相當于每條電路獨立運行，互不干擾；后者相當于所有電路共用一根總線，一處過載全盤崩潰。

三、讓理論落地：專為生物數據設計的高效計算內核

理論上sigmoid注意力更好，但如果實際跑起來慢得像烏龜，那再好的理論也沒用。這就引出了研究的第二個核心貢獻：一個叫做TritonSigmoid的高效GPU計算內核。

要理解為什么需要專門開發這個內核，先要理解現有工具的局限性。目前最流行的高效注意力計算工具叫做FlashAttention，它專門為softmax注意力優化設計，無法直接用于sigmoid注意力。雖然之前有研究者開發了一個叫FlashSigmoid的工具，但它有兩個致命缺陷：第一，不支持序列填充，也就是說同一批次里所有序列必須等長，這在生物數據中幾乎不可能實現；第二，不兼容最新的NVIDIA GPU架構。用普通PyTorch（一個深度學習框架）直接實現sigmoid注意力雖然支持填充，但速度極慢，在H100 GPU上只能跑到41 TFLOPS的前向計算速度，相當于只用到了硬件理論算力的一小部分。

研究團隊用Triton（一種GPU編程語言，可以理解為專門給GPU寫高效程序的工具）從頭設計了TritonSigmoid，核心創新點有幾處。

第一個創新是"稀疏塊計算"：對于完全由填充組成的空白塊，內核會直接跳過，完全不做任何計算。這就好比考試閱卷時，看到整張答題紙是空白的，直接給0分跳過，而不是逐字逐句去檢查有沒有答案。這使得在有25%填充的情況下，計算效率損失僅為9.3%。

第二個創新是"融合運算"：傳統方式需要把注意力矩陣先寫到內存里，再讀出來進行后續計算；TritonSigmoid將整個注意力計算流程融合成一個連續操作，避免了反復讀寫內存的開銷。對于sigmoid函數本身，內核使用了一個基于tanh（雙曲正切）的硬件加速近似公式，利用現代GPU內置的快速tanh運算單元。

第三個創新是"反向傳播分解"：在訓練神經網絡時，除了前向計算（模型做預測），還需要反向傳播（計算如何調整模型參數）。研究團隊將反向傳播拆分成兩個獨立的內核：一個專門計算查詢矩陣（Q）的梯度，另一個專門計算鍵矩陣（K）和值矩陣（V）的梯度。這種分解消除了并行計算中的"原子操作沖突"，讓不同計算單元能更高效地協同工作。此外，反向傳播時需要用到前向傳播中間結果，但這些中間結果不需要保存在內存里——內核會在反向傳播時重新計算一遍，以此換取內存效率。

最終的性能數據非常突出。在NVIDIA H100 GPU上，TritonSigmoid在16384個token長度、128維度的配置下，前向傳播達到了515.6 TFLOPS，后向傳播達到373.5 TFLOPS。相比之下，FlashSigmoid是439.7/341.6 TFLOPS，FlashAttention-2是360.6/312.5 TFLOPS，普通PyTorch實現只有92.8/204.8 TFLOPS。換算成相對速度，TritonSigmoid比FlashAttention-2快43%，比FlashSigmoid快17%，比普通PyTorch快5.6倍。

在有25%填充的真實生物數據場景下，TritonSigmoid的優勢更加明顯，前向傳播相比普通PyTorch提速了14.58倍，相比FlashAttention-2也快了29%。由于TritonSigmoid用Triton實現，它天然支持JIT（即時編譯）到不同GPU架構，未來面對新一代GPU也不需要重新手寫代碼，具有很好的前向兼容性。

四、實驗驗證：sigmoid注意力訓練出的模型真的更好嗎

計算效率解決了，接下來的核心問題是：用sigmoid注意力訓練出來的模型，在理解生物細胞方面真的更好嗎？

研究團隊訓練了四個1.6億參數規模的模型，分別是：2K上下文窗口+softmax注意力、2K上下文窗口+sigmoid注意力、4K上下文窗口+softmax注意力、4K上下文窗口+sigmoid注意力。所有模型都在相同的CellxGene數據集上訓練，使用完全相同的超參數（除了注意力機制本身），都訓練到完全收斂，訓練過程中都使用了梯度裁剪（一種防止訓練不穩定的標準技術）。

評估則在六個完全沒有參與訓練的獨立數據集上進行，覆蓋了大腦、血液、結腸、肺和心臟等不同組織，涵蓋胚胎期到老年期的不同發育階段，包含健康、發育和疾病等不同生物學背景。這種多樣化的評估設計是為了考察模型的泛化能力，而不僅僅是在訓練數據上表現好。

評估指標包含四個維度。第一是驗證損失，用來衡量模型預測被遮蓋基因的準確程度，類似于"填空題得分"，分數越低越好。第二是scIB生物學保守性指標，這是單細胞生物學領域的標準評估框架，包含細胞類型輪廓系數（同類細胞在表征空間中是否聚集在一起）、Leiden聚類的NMI和ARI（無監督聚類結果與真實細胞類型標簽的吻合程度），所有這些指標越高越好。第三是UMAP可視化，一種把高維空間投影到二維平面的技術，用于直觀觀察不同細胞類型的分布格局。第四是最大均值差異（MMD），用來量化不同細胞類型在表征空間中的分離程度，可以理解為"不同細胞類型之間的距離"，距離越大越好。

結果顯示了兩條一致的規律。規律一：在所有六個數據集、兩種上下文長度下，sigmoid注意力的驗證損失均低于softmax注意力。規律二：4K上下文窗口的表現系統性地優于2K上下文窗口，兩種注意力機制均如此，這驗證了更長的上下文窗口對捕捉基因關系的重要性。

在生物學保守性指標上，sigmoid在全部六個數據集上都取得了更好的細胞類型凝聚度（輪廓系數），在六個數據集中的四個上取得了更高的綜合生物學保守性得分。在心臟流出道（Heart OFT）數據集上進行的MMD分析最為亮眼：研究團隊計算了8種細胞類型之間所有28對兩兩比較的MMD值，sigmoid注意力在全部28對比較中均取得了更高的MMD，平均提升幅度達到25%。這意味著，sigmoid模型學到的細胞表征空間中，不同細胞類型之間的"距離"更大、更容易區分，這對下游的細胞類型分類任務非常有利。

為什么sigmoid能在預測損失相近的情況下學到更好的表征？研究團隊的解釋是：softmax的競爭性歸一化迫使模型在關注一個基因時減弱對其他基因的關注，這種約束可能使模型傾向于關注那些最有預測性的少數基因，而忽略了定義細胞身份所需的復雜基因共表達模式。sigmoid的獨立注意力機制允許模型同時、充分地關注多個相關基因，從而更好地捕捉細胞類型的多維度特征，即使這些特征對預測被遮蓋基因本身的貢獻相對有限。

五、極端壓力測試：訓練崩潰時，sigmoid能否力挽狂瀾

除了正常訓練條件下的性能比較，研究團隊還設計了一個"極限壓力測試"，專門用來暴露softmax注意力的穩定性缺陷。

測試條件故意設置得非常苛刻：8192個token的長上下文窗口，完全去掉梯度裁剪保護，其他所有條件與正常訓練完全相同，唯一的變量是注意力機制。這就好比把兩輛汽車放上賽道，但把安全帶和防滾架都拆掉，看誰能跑完全程。

softmax模型的表現是這樣的：前40000步訓練一切正常，損失從約10穩步降到約3，看起來學習進展順利。但從第40000步開始，情況急轉直下。到第55600步，訓練徹底崩潰：損失從3急劇飆升回10以上，全局梯度范數從約100膨脹到1.6×10?，足足增長了四個數量級（也就是膨脹了10000倍）。與此同時，第0層注意力的最大分數從約20暴增到2.3億。一旦崩潰，訓練就再也無法恢復，陷入永久性的發散。

sigmoid模型則從頭到尾平穩運行完全部80000步。損失單調下降，從約10穩步降到約3，沒有任何異常波動。在第55600步——也就是softmax模型崩潰的那個時間點——sigmoid模型完全沒有任何異常，繼續穩定訓練。全程梯度范數保持在10到100之間，還呈現下降趨勢。注意力分數也始終保持在1到5的范圍內，與softmax的2.3億相比，穩定程度天壤之別。

這個結果與前面的理論分析完全吻合：sigmoid的導數上界為0.25，無論輸入多大，信號放大倍數永遠不超過這個值；而softmax的局部Lipschitz常數（可以理解為最大信號放大倍數）會隨注意力分數的增大呈指數級增長，當序列很長、注意力分數很大時，這個放大倍數會達到天文數字，最終引發訓練崩潰。

六、訓練速度：sigmoid到底快了多少

除了質量更好、更穩定，sigmoid注意力還能讓訓練更快完成。

研究團隊測量了四種模型規模（1.6億、4億、6億、14億參數）在三種上下文長度（2K、4K、8K）下的端到端訓練速度，使用16塊H100 GPU進行測量，然后將吞吐量數據換算成處理1.316億個細胞（完整訓練數據）所需的GPU小時數。

在4K上下文長度下，sigmoid比softmax快：4億參數模型節省5.1%（1739 vs 1832 GPU小時），6億參數模型節省3.0%（2336 vs 2408 GPU小時），14億參數模型節省4.0%（4180 vs 4349 GPU小時）。對于14億參數模型，速度優勢隨上下文長度遞增：2K時快2.1%，4K時快4.0%，8K時快7.5%，僅在8K上下文這一個配置上就節省了645 GPU小時。

對于1.6億參數模型，研究團隊實際完整跑完了訓練。在2K上下文下，sigmoid用了653 GPU小時，softmax用了717 GPU小時，sigmoid快了9%；在4K上下文下分別是896和934 GPU小時，sigmoid快了4%。

速度優勢隨上下文長度增大而增大是符合預期的：注意力計算的復雜度與序列長度的平方成正比，序列越長注意力在總計算量中占的比例就越大，因此sigmoid注意力在計算上的內在簡潔性（不需要做跨token的歸一化）帶來的收益也就越顯著。端到端的加速幅度小于內核級別的TFLOPS提升，是因為總訓練時間還包括數據加載、非注意力層計算和GPU間通信等其他開銷。

歸根結底，這項研究給出了一個清晰的答案：在單細胞生物學基礎模型的訓練中，把softmax注意力換成sigmoid注意力，不是一個"理論上可能更好"的選項，而是一個在多個維度上經過實驗驗證的務實選擇。表征質量更好了，細胞類型識別能力更強了，訓練更快了，在極端條件下也不會崩潰。研究團隊同時解決了讓這種替換真正可行的工程障礙，開發出了支持生物數據特有的長序列和填充需求的高效內核，并將其開源。

生物信息學和AI醫療這兩個領域正在以極快的速度融合，越來越多的藥物研發、疾病診斷和精準醫療方案將依賴于這類能理解細胞語言的AI模型。如何讓這些模型在有限的計算預算內訓練得更好、更穩定，是一個具有直接現實影響的工程問題。這項研究提供的解答——一個改動看似微小，但影響深遠的注意力機制替換——值得這個領域認真對待。對這項研究感興趣的讀者，可以通過arXiv編號2604.27124查閱完整論文，研究團隊的代碼也已在GitHub公開。

Q&A

Q1：sigmoid注意力和softmax注意力的核心區別是什么？

A：softmax注意力是"競爭式"的，關注一個基因就必須減少對其他基因的關注，所有注意力分數加起來必須等于1。sigmoid注意力是"獨立式"的，每對基因之間的關聯分數獨立計算，互不影響，可以同時對多個基因保持高度關注。這種獨立性更符合基因調控中多個轉錄因子并行調控同一目標基因的生物學現實，也使得訓練過程中的梯度更加穩定。

Q2：TritonSigmoid內核為什么比現有工具快？

A：TritonSigmoid主要通過三個機制實現加速：對完全是空白填充的計算塊直接跳過（稀疏塊計算）、將注意力計算融合成單一操作避免反復讀寫內存（融合運算），以及將反向傳播拆分成兩個獨立內核消除并行沖突。此外，sigmoid函數本身比softmax計算更簡單，不需要跨token的歸一化操作，在硬件層面天然更高效。在H100 GPU上，其峰值前向計算速度達到515 TFLOPS，比FlashAttention-2快43%。

Q3：單細胞基礎模型訓練為什么容易出現梯度爆炸崩潰？

A：單細胞數據的序列往往很長（可達數千甚至上萬個基因），在長序列場景下，softmax注意力的分數會變得非常大，而其局部Lipschitz常數（可理解為信號放大倍數）會隨注意力分數的增大呈指數級增長。當這種放大效應在多層網絡中疊加傳播時，梯度會急劇膨脹，最終超出數值表示范圍，導致訓練不可逆地崩潰。sigmoid注意力的導數永遠不超過0.25，從根源上消除了這種指數級放大的可能性。