五億分子、兩萬餘個蛋白口袋,61 秒全部檢索完畢,一項新的研究把藥物虛擬篩選的時間門檻從 24 小時直接降到了一分鐘。
尋找作用於疾病相關蛋白的小分子化合物,是藥物發現的早期環節。蛋白質上供小分子結合的局部區域稱為“蛋白口袋”。面對成億規模的化合物庫,研究者需先挑出可能與蛋白口袋結合的候選小分子,再開展後續計算與實驗。
傳統的虛擬篩選需要逐一模擬化合物與口袋的相互作用,耗時巨大。2023 年,清華大學團隊提出 DrugCLIP,通過對比學習,分別將蛋白質口袋和小分子對應到同一向量空間,再依據向量相似度檢索候選分子。分子向量可以預先計算,供不同口袋重複查詢,從而減少逐一進行分子對接的計算負擔。
2026 年發表於 Science 的後續研究,將這一方法推進到全基因組規模篩選。團隊用八張 A100 GPU,針對約一萬個人類蛋白的兩萬多個口袋,檢索五億多個類藥分子,在約 24 小時內完成超過十兆組口袋與分子的匹配計算。但對需要反覆調整靶點、比較口袋的研究者,一天的檢索等待仍然不短。
近日,亞利桑那州立大學曹志超團隊提出的 SweepANN,進一步加速了這一過程中的大規模向量檢索。在論文的主實驗中,約五億個分子的向量被預先存入四張 GPU。針對兩萬六千余個蛋白質口袋,系統成功為每一個口袋提取出得分排名前 50 萬的候選分子。整個檢索耗時僅約 61 秒,平均召回率高達 99.63%。
曹志超本科畢業於清華大學,現任亞利桑那州立大學計算與增強智能學院(SCAI)助理教授。他曾任 Meta RocksDB 團隊研究科學家,獲美國國家科學基金會 CAREER 獎,以及多個頂級會議 best paper 等。他主持的智能資料基礎設施實驗室(ASU-IDI)聚焦 AI 基礎設施、記憶體、和資料基礎設施的設計與最佳化。SweepANN 成果已在 arXiv 以預印本形式發佈。
SweepANN 的工作建立在 DrugCLIP 之上。DrugCLIP 通過雙塔分別將蛋白質口袋與小分子對應到同一高維向量空間,按匹配分數排序。SweepANN 沿用公開模型生成的向量和打分方式,重點改進向量庫在 GPU 中的記憶體、批次檢索計算,以及大規模候選名單的維護。
通過系統層面的改進,SweepANN 讓虛擬篩選從一次性大規模計算轉變為研究者可頻繁互動的探索工具。曹志超認為,一條科學路徑先回答“能不能”,系統研究隨後回答“能做多快、成本能降到什麼程度”。
24 小時篩選,還不夠快
過去,研究者常用分子對接軟體,模擬化合物進入蛋白質口袋時的結合構象及相互作用,再依據計算結果挑選候選。每個分子都要單獨進行分子對接(docking)計算。但分子庫擴大到數億規模後,逐一處理便會耗費大量時間。
DrugCLIP 改變了這項初篩的計算方式。團隊利用已知的口袋與分子關係訓練兩個編碼器,分別處理蛋白質口袋和小分子,將兩者轉成同一高維空間中的向量。可能匹配的對象由此獲得較高分數,研究者再按分數篩出一批候選分子。
在 2026 年發表於 Science 的研究中,DrugCLIP 團隊用八張 A100 GPU,在約 24 小時內完成了超過五億個分子與兩萬多個蛋白質口袋的向量檢索。團隊為每個口袋保留約 10 萬個分子,再進行聚類和分子對接,使原本規模極大的初篩任務可以在一台計算節點上運行。
出於系統研究者特有的“算帳”直覺,曹志超團隊第一次看到這組資料,便按口袋數量粗略估算了檢索速度。基於向量搜尋研究的經驗,他們發現平均每個口袋仍需等待數秒,認為 24 小時“還是不夠快”,篩選速度仍有提升空間。
分子庫的規模也在增長。目前,面向藥物虛擬篩選的公開化合物資料庫 ZINC-22 已收錄超過 370 億個可購買化合物,其中超過 45 億個有可直接對接的三維結構。為了維持初篩的覆蓋範圍,則需要按固定比例保留候選。這意味著,分子庫越大,每個口袋要留下的前 K 個分子也越多。
一次 24 小時的檢索看似可以接受,但靶點、口袋和研究重點可能改變。另外,每提出一批新的口袋,或分子庫加入新化合物,都會帶來更多的工作量。
面對不斷擴大的分子庫和候選名單,團隊首先考察了向量搜尋常用的加速辦法。這類方法通常借助索引減少需要讀取的資料,其中圖索引預先記錄向量之間較近的關係,倒排索引則先將向量分組,查詢時優先查看更可能相關的部分。
曹志超將這種思路比作“尋找鄰居”:如果只找最近的十個人,就沒有必要逐個比較所有人。若查詢只需返回前十或前一百個結果,這條路徑通常有效。
但 SweepANN 面對的名單長得多。在主實驗中,系統需要為每個口袋從五億個分子中取出排名前 50 萬的候選,並儘量保留精確計算所選出的分子。
團隊測試發現,要把如此長的名單儘量找全,索引需要查看更多分組或節點,導致其優勢縮小。同時,建立和保存索引、維護候選名單本身都要耗費大量算力資源。
面對這一困境,團隊也比較了 CPU 與 GPU 路線,發現 CPU 的平行計算能力相對有限,處理如此規模的口袋與分子匹配需要更多時間。最終,團隊決定用 GPU 掃描分子庫,為每個口袋計算與所有分子的匹配分數。
SweepANN 怎樣讓全量掃描更高效
DrugCLIP 為 SweepANN 提供了前置模型,將蛋白質口袋和小分子轉成向量,並按兩者的匹配程度打分。基於此,團隊將工作集中在分子向量的存放、GPU 的批次打分和候選名單的維護上。
整個檢索過程可以分為以下環節:分子庫要先放進 GPU,隨後在查詢蛋白質口袋時讀取分子向量、計算匹配分數;得到分數後,系統還要從海量結果中選出並記錄每個口袋排名前 50 萬的分子。
首先需要應對的是記憶體壓力。SweepANN 使用的每個分子向量有 768 個維度,如果每個數都按 32 位浮點格式(FP32)保存,五億條向量將佔用約 1.54 TB,超過實驗所用四張 GPU 的視訊記憶體總量,對記憶體提出了較高的要求。
因此,團隊採用量化方法,每 256 個維度共用一個縮放係數,向量中的數值改用佔 1 字節的 INT8 保存。五億條向量由此縮至約 387 GB,分片存入四張 GPU。計算時,系統按縮放係數轉換當前讀入的資料並直接打分,無須事先恢復整個分子庫。
分子庫留在視訊記憶體後,更換口袋便可直接再次查詢。曹志超解釋,量化雖會損失少量數值精度,但分數變化主要影響排名靠近第 50 萬名邊界的分子,使它們進出候選名單,靠前的結果則較為穩定。與未壓縮向量算出的精確名單相比,最終名單質量基本不受影響。
解決了記憶體問題,下一步便是打分。若一次只處理一個口袋,GPU 每換一個口袋,就得再讀一遍分子資料。為減少這種重複讀取,SweepANN 改為批次處理口袋,讓一次讀取服務更多查詢。
團隊據此將 64 個口袋組成一批。四張 GPU 各存放約 1.25 億個分子,接收同一批口袋,並將各自負責的分子資料分塊讀取。每讀入一塊,就通過矩陣乘法同時為整批口袋計算匹配分數。壓縮資料在打分過程中按需還原,無須另存一份解壓後的分子庫。
所有口袋與分子的組合仍會被計算,但一次讀取的資料服務了更多口袋,計算單元等待資料的時間也隨之減少。單卡測試中,批次從 1 個口袋增至 64 個,吞吐量由每秒 18.2 個查詢升至 453.3 個。
打分完成後,新的壓力落在候選名單上。四張 GPU 各自保存一部分分子,而每個口袋最終都需要全庫得分最高的 50 萬個結果。如果保存一批口袋對所有分子的分數,再統一排序,分數矩陣會佔用大量視訊記憶體;如果每處理一小塊分子就篩選一次,又要反覆整理和合併長名單。
團隊因此設計了全分數、逐塊、分組和窗口四種維護名單的方案,依據剩餘視訊記憶體與任務規模,選擇能夠運行且實測速度更快的方案。
主實驗最終選擇了“窗口計畫”,先設定門檻,再篩選需要保留的候選分子。具體而言,系統先處理少量選出的“種子”資料區塊,從中為每個口袋選出暫時排在前 50 萬名的分子,並以名單中的最低分作為篩選門檻。
接下來掃描其他資料區塊時,每個分子仍會得到分數,但低於門檻的分子無須進入候選緩衝區,高於門檻的結果則暫時留下。
其中,種子資料區塊的選擇將決定篩選分數線的高低,但不會影響按壓縮後分數得到的最終名單。種子塊已選出 50 萬個分子,得分低於名單末位的分子不會進入全庫前 50 萬名。因此,種子選得不理想,只會增加後續整理候選的時間。
另外,為減少反覆發起計算的開銷,團隊還採用 CUDA Graphs 回放,將一批查詢中反覆執行的 GPU 操作預先記錄下來,供後續批次重複執行。
經過這些環節的配合,分子庫載入視訊記憶體後,SweepANN 用四張 RTX PRO 6000 Blackwell GPU,在 61.04 秒內完成兩萬六千余個口袋的檢索。同樣硬體和資料下,團隊早期的逐塊方案耗時 360.56 秒。
從 24 小時到 1 分鐘,能改變什麼
相較於 DrugCLIP 報告的約 24 小時,SweepANN 將同量級的大規模向量檢索推進到分鐘級。同樣面對約五億個分子和兩萬多個蛋白質口袋,前者使用八張 A100 GPU,後者使用四張 RTX PRO 6000 Blackwell GPU,進一步壓縮了(關於卡的使用的評價)。
論文中的約 61 秒指 GPU 駐留檢索時間,從壓縮後的分子庫已經載入視訊記憶體開始計時,到全部口袋的查詢結果生成並留在視訊記憶體中結束。此前對分子庫進行量化並將其載入視訊記憶體,另耗時約 502 秒。
在實際研發中,只要分子庫無需重新編碼,已載入視訊記憶體的向量就能供多批口袋反覆查詢,量化和載入所花的時間將由後續的多次高頻檢索分攤,曹志超補充。
在這項任務中,SweepANN 的平均召回率達到 99.63%。這意味著,與用未壓縮向量逐個口袋掃描全庫後選出的精確前 50 萬名相比,SweepANN 的候選名單平均保留了其中 99.63% 的分子,即排名靠前的高分化合物幾乎不會出現遺漏。
從虛擬篩選到實驗驗證,候選分子還要經過更細緻的計算,以及細胞、動物實驗。物理世界的實驗周期難以跨越,一次失敗往往意味著回到前期步驟重新篩選,因此前期判斷的精準性與迭代效率顯得極為重要。
曹志超認為,速度的提升也有望改變研究者與 AI 協作提出、檢驗假設的節奏。研究者可以針對一種疾病嘗試不同靶點或口袋,迅速取得候選,再交給分子對接等後續計算。如果結果不理想,就及時調整方向並重新檢索,讓一次大型篩選逐漸變成可以反覆進行的探索。
SweepANN 瞄準的是整條虛擬篩選流水線中的中間一環。團隊認為,上游的向量編碼和下游的分子對接也有提速空間,幾個環節更快銜接,才能讓初篩提速更充分地作用於整個探索過程。
此外,作為面向大結果集的向量檢索方案,這套方法的適用範圍並不侷限於藥物篩選。論文除了在藥物相關模型上驗證,還測試了公開的圖像與文字向量資料。曹志超判斷,若一項任務同樣需要從龐巨量資料集中一次找出幾萬乃至幾十萬個相似對象,其處理大結果集的思路便有借鑑價值。
一條科學路徑先回答“能不能做”,系統研究隨後回答“能做多快、成本能降到什麼程度”。未來,團隊還將在更大規模的分子庫及上下游計算環節中進一步追問,如何用更少資源、更短時間完成濕實驗前的探索。
亞利桑那州立大學電腦系博士生張振宇是本論文第一作者,合作者包括該系博士生林奇,以及藥物研發專家張楠。(DeepTech深科技)

