三個月前,我們開始重振 Papers with Code(請參閱相關公告推文)。其目標是讓開放的 AI 研究更易於存取和理解,以便人們能輕鬆找到與論文相關的產物,探索 AI 各領域的最先進(SOTA)技術,分享有趣的研究並在彼此的成果上繼續發展。換句話說,其目標是推動引領下一個 Transformer 誕生的研究浪潮。
當然,讓 AI 研究易於存取需要一個強大的搜尋引擎,這樣人類和代理程式才能透過網站或 pwc search CLI 指令快速找到相關工作。代理程式可以透過 Skill 使用這些指令。
值得注意的是,搜尋研究與搜尋一般文字有所不同。一個有用的論文搜尋引擎應該能找到確切的標題或 arXiv 識別碼,但它也應該能理解「用於程式碼生成的輕量級語言模型」這樣的查詢,即使這些詞語並未同時出現在論文中。它需要辨識出「原始 BERT 論文」是一個導航請求,容忍不完整的標題或錯字,並且在模型服務冷啟動或暫時不可用時仍能快速回應。
對於 Papers with Code,我們將其建構為一個混合式搜尋系統。這也基於我們在 ML6 的過往經驗,當時我們為客戶開發了基於 RAG 的系統。結果顯示,混合式搜尋通常優於關鍵字和向量搜尋系統,因為它結合了兩者的優點(更多資訊請參閱此部落格)。
關鍵字搜尋能找到精確的提及,而向量搜尋則能找到更模糊、語義相似的術語。請注意,重排序器(也稱為交叉編碼器)可以進一步改善結果,儘管它們也會帶來額外的開銷和延遲。
Papers with Code 依賴 PostgreSQL 資料庫,因此其全文搜尋功能提供了快速的詞彙基準。對於稠密嵌入,我們使用 pgvector 來增加語義召回,並透過倒數排序融合(RRF)演算法將兩者結合。我們使用三個 Hugging Face 服務來處理稠密嵌入:
Hugging Face Jobs 為我們提供了可突發的 GPU 運算,用於嵌入論文語料庫。
Hugging Face Storage Buckets 提供了資料庫、實驗和 Jobs 之間持久性的資料交接。
Hugging Face Inference Endpoints 提供低延遲嵌入,用於即時查詢和增量更新。
目前,該系統維護著來自 arXiv 和 Daily Papers 超過 110,000 篇現有論文的嵌入。這篇文章將解釋其架構、背後的設計決策以及我們將其投入生產時所學到的經驗。
簡而言之,我們刻意將搜尋分為離線語料庫建構和線上搜尋服務:
耗時且注重吞吐量的工作以 Jobs 形式執行。持久性產物儲存在 Bucket 中。只有小型查詢嵌入步驟位於請求路徑上,由受保護的 Inference Endpoint 提供支援,以驅動線上搜尋。如果該端點冷啟動、繁忙或不健康,搜尋會立即回退到全文檢索。這種分離使系統既強大又快速。
嵌入管線經常以不易察覺的方式失敗:模型版本變更、查詢和文件提示詞混淆、向量截斷方式不同,或者更新後的摘要不再與其儲存的向量匹配。
我們透過將嵌入格式視為版本化 API 來避免這種情況。每篇論文都編碼為:標準化標題 + "\n\n" + 標準化摘要。對於每次向量生成,我們記錄:模型儲存庫和確切版本;輸出維度;輸入格式版本;輸入是查詢還是文件;標準化方法;以及原始標題和摘要的內容雜湊。
我們的生產生成使用 Qwen/Qwen3-Embedding-0.6B,釘選到特定版本,並使用 256 維的 L2 標準化向量。我們在 MTEB 排行榜的幫助下選擇了該模型,這是比較嵌入模型的首選基準。請注意,像 Qwen3 這樣的新型嵌入模型允許兩個新功能:
可以指定動態嵌入大小,這允許在品質與速度/儲存成本之間進行權衡。Qwen 模型稱之為「MRL」,即 Matryoshka Representation Learning 的縮寫。您可以在這裡了解所有相關資訊。我們選擇 256 的嵌入大小以加快搜尋速度。
可以提供指令提示詞。Qwen 嵌入模型支援文件提示詞(我們用來嵌入論文)和即時搜尋使用其查詢提示詞(用來嵌入使用者查詢)。
這個合約追蹤嵌入從匯出、透過 GPU 推論、進入 PostgreSQL,最終進入線上檢索的整個過程。
全語料庫嵌入是典型的批次工作負載。它需要 GPU 在相對較短的時間內執行,受益於高吞吐量,並且在執行之間不應消耗資源。Hugging Face Jobs 非常適合這種形式:Job 由指令、硬體規格以及可選的 Docker 映像檔定義,並且可以執行帶有內聯聲明依賴項的 uv 腳本。
我們的語料庫建構始於從可重複讀取的 PostgreSQL 快照匯出每篇論文的最新版本。匯出器串流資料列而不是將目錄載入記憶體,寫入有界限的 JSONL 分片,並創建包含資料列計數和 SHA-256 校驗和的清單。
我們將該不可變的執行目錄同步到私有 Storage Bucket,並將 Bucket 直接(使用 hf-mount)掛載到 l4x1 Job(一個具有 24GB VRAM 的 NVIDIA L4 GPU)中。從工作者的角度來看,它只是一個檔案系統:
hf jobs uv run \
--flavor l4x1 \
--timeout 6h \
--volume hf://buckets/OWNER/pwc-paper-embeddings:/bucket \
embed_papers_job.py \
--input /bucket/runs/RUN_ID/input \
--output /bucket/runs/RUN_ID/output \
--model Qwen/Qwen3-Embedding-0.6B \
--revision MODEL_REVISION \
--dimensions 256 \
--allow-matryoshka
工作者會執行以下操作:驗證輸入清單和每個分片校驗和;載入釘選的模型版本;按長度排序文本以減少填充;批次呼叫 encode_document(如模型卡所述);如果 GPU 記憶體不足,則自動減少批次大小;將 Matryoshka 表示截斷為 256 維並進行標準化;原子性地寫入 float16 Parquet 分片;並記錄吞吐量、套件版本、硬體、峰值 VRAM、資料列計數和輸出校驗和。
每個完成的分片都有自己的標記,因此重新啟動的 Job 可以跳過已驗證的工作。這對於大型語料庫很有用:重試應該只是恢復工作,而不是覆蓋現有的嵌入。在我們的 5,000 篇論文試驗中,Qwen Job 在 L4 GPU 上以 1024 維度每秒編碼約 75 篇論文。相同的過程可以以 512 和 256 維度確定性地實現,因此我們可以在不支付更多推論費用的情況下比較儲存和檢索的權衡。
Storage Buckets 是 Hub 上可變的、類似 S3 的物件儲存,專為 AI 工作負載優化。它們可以透過 hf://buckets/... 路徑存取,並以讀寫方式掛載到 Jobs 中,無需建立單獨的儲存整合。
對我們來說,Bucket 不僅僅是放置向量的地方。它是三個具有不同生命週期的系統之間的邊界:生產資料庫匯出原始記錄;短暫的 Jobs 消耗這些記錄並產生向量;匯入器在觸及搜尋索引之前驗證結果。我們將產物組織在不可變的執行前綴下:
runs/<run-id>/
├,input/
│ ├,manifest.json
│ └,papers-*.jsonl
└,output/
├,manifest.json
├,embeddings-*.parquet
└,embeddings-*.complete.json
Buckets 本身是刻意可變的,因此不可變性是應用程式層級的規則:執行 ID 永遠不會被覆蓋,每個產物都由清單和校驗和覆蓋。這為我們提供了幾個有用的特性:可重現性:我們可以將資料庫生成追溯到確切的語料庫快照、模型版本和一組產物。安全重試:Jobs 可以從同一執行前綴中已完成的分片恢復。
低成本實驗:多個模型或維度可以重複使用一個已驗證的輸入快照。受控發布:匯入生成不會立即啟用它。我們首先驗證覆蓋範圍並建立其索引。簡單回溯:前一個生成及其產物保持可用,直到新的生成被證明穩定。
只有在匯入器重新檢查綱要、校驗和、維度、標準化、唯一論文 ID 和當前內容雜湊後,我們才會將向量載入 PostgreSQL。然後,我們為新的生成建立一個單獨的 HNSW 索引,並且只有當每個符合條件的當前論文都被覆蓋時,才原子性地將其標記為活躍(HNSW 是基於圖的演算法,可實現快速向量搜尋)。
批次嵌入解決了檢索的文件端問題。使用者查詢仍需要在請求時使用相同的模型合約進行嵌入。我們將釘選的模型部署為由 Text Embeddings Inference (TEI) 支援的經過驗證的 Inference Endpoint。該端點接受查詢文本並使用模型的查詢提示詞返回一個標準化的 256 維向量。請注意,這裡也可以利用 vLLM 或 SGLang。
然後 API 對活躍的 pgvector 生成執行餘弦距離搜尋:
SELECT paper_id,
embedding <=> CAST(:query_vector AS halfvec(256)) AS distance
FROM paper_embeddings
WHERE generation_id = :active_generation
ORDER BY embedding <=> CAST(:query_vector AS halfvec(256))
LIMIT 50;
HNSW 索引保持此查詢快速。在我們的 5,000 篇論文試驗中,256 維 Qwen 索引在精確搜尋中實現了 0.9955 的 Recall@20,p50 和 p95 HNSW 查詢延遲分別為 1.31 毫秒和 2.21 毫秒。其表格和索引使用的儲存空間約為 1024 維版本的 27%,同時在該測試中基本保持了相同的近似最近鄰(ANN)召回率。
該 Endpoint 配置為最多一個副本,並在閒置時可以縮小至零。這是一個有用的成本槓桿,因為這表示在沒有使用時無需付費。然而,這也意味著冷啟動必須是應用程式設計的一部分,而不是被視為例外事件,因為端點啟動並處理流量需要一些時間。
因此,我們的查詢客戶端具有刻意嚴格的行為:一秒的生產環境逾時;非阻塞式併發限制;回應維度、有限性和範數驗證;以查詢和嵌入生成為鍵的短期快取;重複失敗後的斷路器;以及日誌中不包含原始查詢文本,只有標準化指紋。如果端點正在擴展、逾時、返回格式錯誤的向量或沒有可用的併發,我們會立即跳過語義分支。
使用者仍會收到詞彙結果,而不是等待不可靠的依賴。Inference Endpoints 運行非常可靠,並包含一個漂亮的儀表板,您可以快速查看關鍵分析數據。
對於每個查詢,詞彙分支使用加權 PostgreSQL 全文搜尋檢索最多 50 個候選。語義分支從 pgvector 檢索最多 50 個候選。我們使用加權倒數排序融合(RRF)組合它們的排名:
score(d)=∑r ∈ {lexical, semantic}wrk+rankr(d)



