Hugging Face、Papers with Code の検索機能をエンドポイントやジョブで強化

3 ヶ月前、私たちは Hugging Face によって引き継がれた「Papers with Code」の復活プロジェクトを開始しました（詳細は Reddit の投稿 もしくは 公式ツイート をご覧ください）。このプロジェクトの目的は、オープンな AI 研究へのアクセスを容易にし、誰もが関連するアーティファクトを見つけやすくし、AI の各分野における最先端（SOTA）を検索し、興味深い研究成果を共有し、互いの成果の上に新たな構築を行えるようにすることです。一言で言えば、次世代の Transformer を生み出す研究ブームを支えるプラットフォームを目指すものです。

もちろん、AI 研究を誰もが利用可能にするためには、強力な検索エンジンが不可欠です。これにより、人間もエージェントも、ウェブサイト上または pwc search という CLI コマンド を通じて、関連する研究をすばやく見つけることができます。なお、この CLI は Skill を介してエージェントが利用可能です。

研究論文を検索する行為は、通常のテキスト検索とは性質が異なります。有用な論文検索エンジンには、正確なタイトルや arXiv 識別子を見つける能力だけでなく、「コード生成用の小型言語モデル」のような、文中にそのままの形で出現していないクエリも理解できる必要があります。

また、「オリジナルの BERT 論文」といった表現からナビゲーション意図を認識し、不完全なタイトルや入力ミスにも寛容であるべきです。さらに、モデルサービスが起動待ち（コールド）状態にある場合や一時的に利用できない状況でも、迅速に応答できることが求められます。

Papers with Code では、このシステムをハイブリッド検索として構築しました。これは、私たちが ML6 で得た経験に基づいています。ML6 ではクライアント向けに RAG ベースのシステムを開発していました。その結果、ハイブリッド検索はキーワード検索やベクトルベースの検索よりも一般的に優れたパフォーマンスを発揮することが分かりました。これは両者の利点を組み合わせたものだからです（詳細は こちら をご覧ください）。キーワード検索は正確な一致を検出するのに対し、ベクトル検索はより曖昧で意味的に類似した用語を見つけます。なお、rerankers（クロスエンコーダーとも呼ばれます）を使用すると結果をさらに改善できますが、その分オーバーヘッドとレイテンシが増加することにも注意が必要です。

Papers with Code では、全文検索の基盤として PostgreSQL データベースを採用しています。これにより、高速な語彙ベースの検索が可能になっています。一方、意味的な検索を実現するために pgvector を利用し、埋め込みベクトル（dense embeddings）による検索機能を追加しました。そして、両者の結果を統合する手法として、reciprocal rank fusion (RRF) アルゴリズムを採用しています。

埋め込みベクトルの生成には、Hugging Face の 3 つのサービスを活用しています。

• Hugging Face Jobs は、論文コーパスの埋め込み処理に必要な GPU 計算リソースをバースト可能に提供します。
• Hugging Face Storage Buckets は、データベース、実験データ、Jobs の間での永続的なデータ受け渡しを支えます。
• Hugging Face Inference Endpoints は、リアルタイムのクエリや増分更新に対応する低遅延な埋め込み提供を担います。

現在、このシステムでは arXiv や Daily Papers から収集された 11 万 1000 件以上の最新の論文に関する埋め込みデータを保持しています。この記事では、そのアーキテクチャと設計の背景にある判断、そして本番環境への移行過程で得た教訓について解説します。

TL;DR

検索機能は意図的に、「オフラインでのコーパス構築」と「オンラインでの検索サービス」に分割して設計しました。

オフラインコーパス構築とオンラインハイブリッド検索パイプラインのアーキテクチャ

スループット指向でコストのかかる処理は「Jobs」として実行されます。永続的な成果物は「Bucket」に保存され、クエリ埋め込みという小規模なステップのみが、保護された「Inference Endpoint」の背後に配置されてオンライン検索を駆動します。もしこのエンドポイントが起動待ち（コールド）状態や混雑、不健康な状態であれば、検索は即座に全文検索へフォールバックします。この分離により、システムは強力かつ高速を実現しています。

厳格な埋め込み契約から始める

埋め込みパイプラインは、以下のような微妙な理由で失敗することがあります。モデルの改訂が行われたり、クエリとドキュメントのプロンプトが混同されたり、ベクトルの切り捨て方が異なったり、更新された抄録が保存済みのベクトルと一致しなくなったりするケースです。

これらを回避するため、私たちは埋め込みフォーマットをバージョン管理された API として扱います。すべての論文は以下のようにエンコードされます。

normalized title + "\n\n" + normalized abstract

各ベクトル生成においては、以下の情報を記録します。

• モデルのリポジトリと正確なリビジョン番号
• 出力次元数
• 入力フォーマットのバージョン
• 入力がクエリかドキュメントかの区別
• 正規化手法
• ソースとなるタイトルと抄録のコンテンツハッシュ

私たちの本番環境では、`Qwen/Qwen3-Embedding-0.6B` を使用しており、正確なリビジョンに固定し、256 次元の L2 正規化ベクトルを生成しています。なお、Qwen3 に代表される最新の埋め込みモデルでは、以下の 2 つの新機能が利用可能です。

動的な埋め込みサイズを指定することで、精度と速度・ストレージコストのトレードオフを調整できます。Qwen モデルではこれを「MRL（Matryoshka Representation Learning）」と呼びます。詳細は こちら や ブログ記事 で確認できます。今回は検索を高速化するため、埋め込みサイズを 256 に設定しました。

また、指示付きプロンプトも提供可能です。Qwen の埋め込みモデルは、論文を埋め込む際に使用する document プロンプトと、ユーザーのクエリを埋め込む際に検索システムが使用する query プロンプトをサポートしています。

この仕組みは、埋め込みのエクスポートから GPU での推論、PostgreSQL への格納、そして最終的なオンライン検索に至る一連の流れに沿っています。

Jobs を使ってデータベースのスナップショットをベクトルコーパスに変換

全コーパスの埋め込み生成は、典型的なバッチ処理ワークロードです。比較的短期間 GPU リソースが必要で、高いスループットが求められ、実行間中はリソースを消費しないことが理想です。Hugging Face Jobs はまさにこの要件に合致しています。Job はコマンドと ハードウェアの仕様、そして任意の Docker イメージで定義され、依存関係をインラインで宣言した uv スクリプトを実行できます。

Papers with Code 上の検索機能を支える Hugging Face のインフラは、まず PostgreSQL スナップショットからすべての論文の最新バージョンをエクスポートすることから始まります。このスナップショットは再現可能な読み取り保証を提供しており、エクスポート処理ではカタログ全体をメモリにロードするのではなく、行単位でストリーミングして有限サイズの JSONL シャードへ書き出します。同時に、行数と SHA-256 チェックサムを含むマニフェストファイルも生成されます。

作成された不変のランディレクトリは、プライベートな Storage Bucket に同期され、hf-mount を用いて l4x1 Job（NVIDIA L4 GPU 搭載、VRAM 24GB）に直接マウントされます。ワーカー側から見れば、これは単なるファイルシステムとして扱われます。

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

このワーカーは以下の処理を実行します。

• 入力マニフェストと各シャードのチェックサムを検証し、
• ピンされたモデルリビジョンを読み込み、
• テキストを長さに基づいてソートしてパディングを削減し、
• model card に記載されている通り encode_document をバッチ処理で呼び出し、
• GPU のメモリ不足を検知した場合は自動的にバッチサイズを縮小し、
• Matryoshka representation を 256 次元にトリuncateして正規化し、
• float16 形式の Parquet シャードを原子的に書き出し、
• スループット、パッケージバージョン、ハードウェア仕様、ピーク VRAM 使用量、行数、出力チェックサムを記録します。

各完了したシャードには固有のマーカーが付与されるため、Job が再起動された場合でも検証済みの作業はスキップされます。これは大規模なコーパスを対象とする際に特に有用で、再試行時には既存の埋め込みを上書きするのではなく、処理を継続して行うことができます。

5,000 件の論文を対象としたパイロットテストでは、L4 GPU 上で Qwen Job が 1024 次元のベクトルを約 1 秒間に 75 件エンコードできました。同じ処理で 512 次元や 256 次元への確定生成も可能だったため、追加の推論コストをかけずに、ストレージと検索のトレードオフを比較検証することができました。

バケットがシステムをつなぐ接着剤

Storage Buckets は、Hugging Face Hub 上で提供される可変型の S3 互換オブジェクトストレージで、AI ワークロードに最適化されています。hf://buckets/... というパスを通じてアクセスでき、hf-mount を利用すれば、別途ストレージ統合を構築することなく、Jobs 内で読み書きマウントが可能です。

私たちが考えるバケットは、単にベクトルを格納する場所ではありません。それは、異なるライフサイクルを持つ 3 つのシステム間の境界線です。

• 生産環境のデータベースからソースレコードがエクスポートされ
• 一時的な Jobs がそれらを読み込んでベクトルを生成し
• インポーターが検索インデックスに反映させる前に結果を検証する

アーティファクトは、不変の実行プレフィックスの下で整理されます。

runs/<run-id>/
├── input/
│   ├── manifest.json
│   └── papers-*.jsonl
└── output/
    ├── manifest.json
    ├── embeddings-*.parquet
    └── embeddings-*.complete.json

バケット自体は意図的に可変ですが、不変性はアプリケーションレベルのルールとして実装されています。実行 ID は決して上書きされず、すべてのアーティファクトはマニフェストとチェックサムによって保護されます。

これにより、以下のような有用な特性が得られます。

• 再現性: データベース生成を、正確なコーパスのスナップショット、モデルの改訂履歴、および関連するアーティファクトセットまで遡って追跡できます。
• 安全な再試行: Jobs は、同じ実行プレフィックス内ですでに完了したシャードから再開可能です。

安価な実験が可能：複数のモデルや次元設定でも、1 つの検証済み入力スナップショットを共有して再利用できます。

段階的なロールアウト：生成データのインポートは即座に有効化されません。まずカバレッジを検証し、インデックスを構築してから公開します。

シンプルなロールバック：新しい生成物が安定性が確認されるまで、以前のバージョンとその関連アセットは引き続き利用可能です。

インポーターがスキーマ、チェックサム、次元、正規化、一意の論文 ID、現在のコンテンツハッシュを再検証し、すべての条件を満たした後にのみベクトルを PostgreSQL に読み込みます。その後、新しい生成物用の HNSW インデックスを構築し、対象となる現在の論文すべてがカバーされた時点で、原子操作によってアクティブな状態に切り替えます（HNSW は高速ベクトル検索を実現するグラフベースのアルゴリズムです）。

Inference Endpoints がリクエストパス上にセマンティック検索を実装

バッチ埋め込みは文書側の検索課題を解決します。一方、ユーザーからのクエリも、同じモデル契約に基づいてリクエスト時に埋め込む必要があります。

私たちは、Text Embeddings Inference (TEI) をバックエンドに持つ認証済み Inference Endpoint として固定されたモデルをデプロイしています。このエンドポイントはクエリテキストを受け取り、モデルの query プロンプトを用いて正規化された 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 インデックスが、正確な検索と比較して Recall@20 で 0.9955 を達成しました。また、p50 レイテンシは 1.31 ms、p95 は 2.21 ms でした。このテーブルとインデックスのストレージ使用量は、1024 次元版と比較して約 27% で済みながら、そのテストではほぼ同等の ANN（近似最近傍探索）精度を維持しています。

Endpoint は最大 1 つのレプリカで構成されており、アイドル時には スケール・トゥ・ゼロ します。これはコスト削減に役立つ機能です。つまり、利用がない間は課金されないためです。ただし、Endpoint が起動してトラフィックを処理するまでに時間がかかるため、「コールドスタート」は例外事象として扱うのではなく、アプリケーション設計の前提条件として組み込む必要があります。

そのため、クエリクライアントでは意図的に厳格な動作を実装しています。

• 1 秒間のプロダクションタイムアウト；
• ブロッキングしない並行処理制限；
• 応答次元、有限性、およびノルム値の検証；
• クエリと埋め込み生成をキーとした短期キャッシュ；
• 連続する失敗発生後のサーキットブレイカー。

ログには生のクエリテキストは含まれず、正規化された指紋のみが記録されます。

エンドポイントのスケールアップ、タイムアウト、 malformed ベクトルの返却、または並行処理リソースの枯渇が発生した場合、セマンティックブランチの処理は即座にスキップされます。この場合でもユーザーは待たされず、信頼性の低い依存関係に頼らず、レキシカル検索の結果を即座に受け取ることができます。

Inference Endpoints は非常に安定して動作し、主要な分析指標をすばやく確認できる優れたダッシュボードも備えています。

ハイブリッド検索は単独のブランチよりも強力

すべてのクエリに対して、レキシカルブランチでは重み付き PostgreSQL 全文検索を用いて最大 50 の候補を取得します。一方、セマンティックブランチでは pgvector から同様に最大 50 の候補を取得します。

これら 2 つのランキング結果は、重み付き逆順位融合（RRF）によって統合されます：

score(d)=∑r ∈ {lexical, semantic}wrk+rankr(d)

\text{score}(d) = \sum_{r \,\in\, \{\text{lexical},\, \text{semantic}\}} \frac{w_r}{k + \text{rank}_r(d)}

RRF（ReRank by Fusion）は、異なるスケールを持つ 2 つのシステムからのスコアではなくランクを組み合わせるため、シンプルかつ堅牢です。基本的には、ある論文がレキシカルブランチとセマンティックブランチの両方で高い順位に位置する場合、ハイブリッド検索でも上位に表示される可能性が高まります。

現在、私たちは等しいブランチ重みを使用しており、k=60（RRF アルゴリズムの超パラメータである「ランク定数」）を設定しています。

密度化検索は概念的なクエリに対する再現率を向上させます。一方、全文検索は正確な用語、識別子、および固有名詞に対して依然として優れた性能を発揮します。また、融合されたランキングの上に決定論的なアイデンティティ動作も維持しています：

• 完全一致するタイトルと arXiv ID は常に最上位に表示されます。
• method taxonomy（メソッド分類体系）は、「元の BERT 論文」のようなナビゲーション検索を認識します。
• タイトルが不完全な場合や、限定的なスペルミスがある場合は、保守的な trigram キャンディデートを使用します。
• 曖昧なファジーマッチの場合は、無理に結果を提示するのではなく、該当なしと判断します。

注意：ハイブリッド検索が常に最良の選択肢とは限りません。まずは安価で高速なベースラインとしてキーワード検索から始め、セマンティック検索やハイブリッド検索が検索品質の向上に有効であると確認されて初めて追加することをお勧めします。さらに、Qwen3-Reranker などのモデルを用いて、キーワード/セマンティック/ハイブリッド検索後に再ランク付けを行うことで、検索精度を向上させることも可能です。

One Endpoint, two update paths

大規模な初期コーパスは Jobs に埋め込まれていますが、Papers with Code は絶えず変化しています。新しい論文が到着し、抄録が修正され、arXiv の新バージョンが現在のものになります。

数行の変更に対して GPU Job を起動すると、不要な起動やオーケストレーションのオーバーヘッドが生じます。代わりに、毎時間の増分プロセスで欠落しているか内容が変更された論文を選択し、同じ TEI Endpoint に制限されたデルタを送信します。今回は document プロンプトを使用します。

各実行では最大 500 件の論文を 16 件ずつバッチ処理します。埋め込みを書き込む前に、ソース行がロックされ、そのコンテンツハッシュが再確認されます。推論中に論文が変更された場合、そのベクトルは破棄され、次の実行で取得されます。

これにより、以下のような役割分担が可能になります:

• Jobs は完全な再構築、新しいモデルの生成、大規模なバックフィル処理を担当します。

• Inference Endpoints は対照的なクエリの埋め込みと小規模な文書の増分更新を扱います。

• Buckets は大規模ビルドのアートファクトを保存し、これらのビルドを再開可能かつ監査可能にします。

毎時間のパスにより、アクティブなインデックスはライブカタログに近い状態を維持しつつ、オンラインエンドポイントを無制限のバッチプロセッサ化することはありません。

関連論文へのリンクはオンラインでほぼ無料になる

同じ文書埋め込みが、各論文ページでの関連論文推薦機能にも活用されています。

「SenseNova-U1」に関連する論文

ソースとなる論文にはすでに保存されたベクトルが含まれているため、関連論文の検索にはリクエスト時にモデルを呼び出す必要はありません。これはアクティブな生成に対する単一の最近傍探索クエリで完結します。一時的にベクトルが見つからない場合でも、アプリケーションは以前の arXiv 版を利用するか、既存のタスクベースや引用ベースのフォールバックから結果を補完できます。

私たちが学んだこと

1. スループット処理とレイテンシ敏感な処理を分離する

コーパス埋め込みとクエリ埋め込みには同じモデルを使用しますが、これらは異なるインフラストラクチャの課題です。Jobs はスループットとコストの上限管理に最適化されており、Inference Endpoints は可用性とリクエストのレイテンシ低減に最適化されています。

2. ストレージを計算と運用の明確な契約として定義する

Buckets（バケット）は、計算処理から本番環境への移行を明示的に管理します。チェックサム付きのアートファクトにより、データが生索引に組み込まれる前に検証可能な境界線が設けられます。

3. モデル名だけでなく、関連パラメータも固定する

リビジョン、次元数、プロンプト、正規化方法、入力フォーマッターはすべて検索結果に影響を与えます。これらをまとめて保存し、あらゆる場所で整合性を検証する必要があります。

4. コールドスタートを想定した設計にする

トラフィックが断続的な場合、スケール・トゥー・ゼロ（Scale-to-zero）は有効ですが、製品側ですぐに使えるフォールバック機能があることが前提条件です。ハイブリッド検索はこのフォールバックを自然に提供します。なぜなら、単独でも常に有用なlexical search（語彙検索）が存在するからです。

5. より小さなベクトルもシステム上のメリットになり得る

Matryoshka 埋め込み（embeddings）を使えば、品質、メモリ使用量、インデックスサイズ、レイテンシを一つのトレードオフとして評価できます。パイロット試験では、1024 次元と比較してストレージが大幅に削減される一方で、ANN（近似最近傍探索）の検索精度は 256 次元でも十分に維持されました。

6. アクティベーションは地味であるべきだ

新しいバージョンは現在のバージョンの横にインポートされ、独立してインデックス化され、完全性と最新性が確認された後、原子的にアクティベートされます。ロールバックは緊急時の再計算ではなく、設定変更として行われます。

ぜひ https://paperswithcode.co の検索機能や、https://paperswithcode.co/chat のチャットインターフェースでお試しください。フィードバックをお待ちしています！