パララックスにおける文献検索:集合代数、全コーパス再現、およびJevスクリーニング
動機
AI論文はかつてない速さで出版されており、文献検索は本当に難しい課題になっています。汎用エージェントやコーディングエージェントを文献検索(関連研究の探索や、ある分野の全体像の把握)に使うと、いつも2つの限界にぶつかりました:
- 汎用エージェントはブラックボックスな検索APIを通じてオープンウェブを検索します。通常は上位k件の結果しか得られないため、網羅的な検索はできず、「この分野には論文が何本あるか」といった質問にも答えられません。
- 検索APIの結果はノイズが多いです。インターネット全体を対象としているため、データのクリーニングと整理に余分な労力がかかります。
このため、エージェントは論文検索で早々に打ち切ってしまいがちで、そうでなければ遅くて高コストになります。今年の初めに、自分用のパイプラインを vibe coding で作りました:arXivメタデータのクロール -> 埋め込み -> キーワードまたは意味検索によるリコール -> LLMによるスクリーニング。それでもかなり高コストで、スクリーニングも遅く、使い心地の良いものではありませんでした。
9月中旬にJevが登場し、リコール後の論文のスクリーニングと分類に非常によく合っていると気づきました。そこでJevの上で文献インフラを再構築し始めました。このGIFでスクリーニングの速さが伝わるでしょう:
この記事の残りでは、検索システムの設計と、本番環境で動かすために必要だったいくつかの最適化について説明します。
検索構造
しばらくレコメンダーシステムに携わっていたものの、エージェント型検索についてはほとんど経験がなく、その頃の習慣をいくつか持ち込んでいました。最初に思いついたのはファネル型でした:リコール、粗いランキング、精細なスクリーニング。しかし論文検索はフィード推薦とはかなり異なることが分かりました:
- 拡張が必要:論文は互いに引用し合っており、初期の論文の引用をたどるのは、さらなる研究を見つける非常によくある方法です。
- 補完が必要:人間が検索するにせよエージェントが検索するにせよ、見落としが発生します。最初のパスが終わったら、不適切なものが混入していないか、あるいは以前に気づいたものが漏れていないかを確認するレビューのステップが必要です。
ファネルだけでは明らかに不十分です。複数のリコールチャネルのマージと重複排除、そして偽陽性の一括除去の必要性を加えると、集合がこの問題にとって良いデータ構造であることが分かりました。そこで検索レイヤー全体を集合代数のインターフェースを中心に書き直しました。
エージェントが扱う基本単位は論文の集合です(通常の性質を備えています:well-defined、要素は互いに異なる、順序なし)。和集合、積集合、差集合をサポートし、さらにいくつかの実用的なオブジェクトと演算子があります:
- 仮想ルート:コーパス全体や、ある日のarXivアナウンス内の論文など、出発点を表す仮想オブジェクトです。仮想ルートは大量のメンバーを永続化してストレージを圧迫するのを避けます。
- ルールフィルタ:arXivカテゴリ、時間枠、リソース種別、埋め込み類似度、キーワードマッチなどで論文をフィルタします。
- AIスクリーニング:Jevを使って、エージェントがユーザーの意図に基づいて書いた基準に照らして各論文を判定します。
- AI分類:Jevのマルチラベル確率インターフェースを使って論文を分類します。論文は複数のカテゴリに属することができるため、カテゴリは互いに排他的ではありません。
検索システム全体はノードグラフとして可視化され、フロントエンドではReactFlowで構築されています。ノードグラフはあらゆる種類の現実的なニーズに対応できる十分な柔軟性があります。
典型的なワークフローは上の図のようなものです。いくつかのキーワード検索とセマンティック検索のチャネル、それらを統合して重複排除し大きな候補集合にする union ノード、Jev を呼び出す AI スクリーニングノード、そして論文をトピックごとに分類する AI 分類ノードが続きます。
ほとんどの場合、エージェントはこの典型的なワークフローに従います。しかし、より大規模で複雑なリクエストでは、レビューの途中で不足を見つけ、それを埋めるための新しい集合を構築し、「最終」集合ができるまで反復してから分類に進むことがあります。その結果として得られるワークフローは、回路図のように複雑になることがあります。
汎用エージェントやコーディングエージェントは通常、結果をテキストファイルに、あるいはスクリプトで処理する XML/JSON に記録します。集合をインデックス付きのデータベースに保持する方がはるかにスケールし、効率も格段に良くなります。この種の大規模処理を決定論的なコードに任せることは、モデルのハルシネーションや、コンテキストが圧縮されたときに詳細が失われることに対する防御にもなります。エージェントが保持するのは集合へのハンドルだけであり、メンバーレベルの操作はすべてデータベース内で行われます。
もう一つの実用的な問題は検索の深さです。キーワードマッチはヒットするかしないかのどちらかですが、セマンティック検索はランキングであり、300位でも関連がある一方で3000位はすでに的外れかもしれません。事前に K を選ぶのは困難です。そこで集合を作成する前に、エージェントはまず総件数、少数のサンプル、分布だけを返す安価なプローブを実行します。セマンティック検索のプローブにはランク対類似度の曲線も含まれ、エージェントはその曲線の肘の位置に基づいてどこで切るかを決められます。
初期のバージョンでは、プローブと集合の作成は2つの別々の計算でした。同じ上位10000件のセマンティッククエリに対して、プローブに2分、集合の作成はすべてを再計算するためさらに37分かかりました。現在はプローブの完全なランキングが保存され、集合の作成はそこから上位 N 件を取るだけで、再計算はありません。
検索パフォーマンス
上記のアイデアはごく自然なものです。ComfyUI のようなローコードのノードグラフツールはすでに広く知られており、文献検索について似たようなことを考えた人は多いでしょう。このようなシステムを本番運用する上で難しいのは、データ抽象化ではありません。大規模環境でのパフォーマンスであり、速度、コスト、品質のバランスです。
Parallax は現在、OAI インターフェースを通じて2020年以降の arXiv cs.* カテゴリのメタデータ更新を収集しており、約82万件の論文があります。メタデータとそのインデックスは約7.8 GBです。セマンティック検索には Perplexity の pplx-embed-v1-0.6b を使用し、各アブストラクトを1024次元のベクトルに埋め込みます。float32 では1ベクトルあたり4 KBで、82万件では理論上約3.3 GB、halfvec(半精度)では各2 KBで合計約1.65 GBです。ただしこれは生データにすぎません。PostgreSQL + pgvector では、正確さと高速さを両立するためにベクトルを複数の形式で保存しており、インデックスを含めた合計は約17 GBになります。
キーワード検索
PostgreSQL でファジーマッチングを行う通常の方法はトライグラムインデックスです。テキストを3文字の断片に分割し、転置インデックスを構築し、その断片を使って候補を見つけ、その後テーブルに戻って元のテキストを確認します。しかし、一般的なフレーズではパフォーマンスが悪くなります。約69万件の論文のコーパスで、アブストラクトから motion generation を検索すると、トライグラムインデックスは使われましたが34秒かかり、そのうち32秒は IO 待ちでした。mot、ion、gen のような断片はあまりに一般的すぎます。インデックスは25,000件の候補を生成しましたが、実際に一致したのは776件だけで、数万回のテーブルルックアップは無駄になりました。
私たちの修正は、隣接する単語の断片のインデックスを追加することでした。タイトルやアブストラクト内の隣接する単語の各ペアについて、最初の単語の最後の3文字と2番目の単語の最初の3文字(さらに5文字版)を取り出し、それらを結合して断片とします。例えばmotion generationはion:genとotion:generを生成します。これらはGIN配列インデックスに格納されます。特定の2つの単語が隣り合って現れることは、3文字の断片単体よりもはるかに珍しいため、候補集合は大幅に縮小されます。単語が十分に長い場合は5文字版を優先します。なぜならion:genのような組み合わせは依然としてかなり一般的だからです。
断片は必要条件としてのみ使用されます。最終的なチェックは依然として元のフィールドごとのILIKEをソーステキストに対して行うため、結果セットは旧実装と同一であり、返される件数も正確です。変更後、同じ単一フレーズのクエリは約30 msで完了し、エージェントがよく実行する7フレーズのクエリは23秒から0.09秒になり、同じスナップショット上で旧実装と新実装の論文IDが1件ずつ一致しました。
セマンティックリコール
ベクトル検索には常に速度と精度のトレードオフが伴います。正確検索はすべての候補ベクトルを走査します。HNSWのような近似最近傍(ANN)インデックスははるかに高速ですが結果を取りこぼし、Kが大きくなるほど取りこぼしも増えます。私たちは最終的に2つの経路を維持し、それらを自動的に選択するようにしました。
近似経路はパーティショニング + 半精度 + 再ランキングを使用します:
- 1024次元のfloat32 HNSWインデックスは4.4 GBです。キャッシュされていない場合、1回のクエリで数百MBのインデックスページを読み込み、インデックスがあっても20秒以上かかります。そのためインデックスはhalfvecを使用してサイズを半分にし、候補の検索にのみ使用されます。
- 必要以上の候補を数倍取得し、元のfloat32ベクトルで再スコアリングと並べ替えを行って半精度による誤差を補正します。
- ベクトルは年ごとにパーティション化され、各パーティションは独自のHNSWグラフを持ちます。ユーザーが時間範囲を指定すると、無関係な年はスキップされます。
主要なANNパラメータは検索幅ef_searchと候補プール倍率です。本番環境からの32件の実セマンティッククエリで設定を比較し、正確なtop-Kを正解として使用しました:
| ef_search / プール倍率 | Recall @ K=300 | K=1000 | K=2000 | レイテンシ |
| 1000 / 50 | 100% | タイムアウト | タイムアウト | 67–74 s |
| 400 / 3(現在のデフォルト) | 99.7% | 98.0% | 96.5% | 主に1.6–7 s |
| 200 / 2 | 99.0% | 95.8% | 94.8% | 1.4–6 s |
| 100 / 1 | 97.0% | 94.0% | 93.2% | 3–4 s |
取りこぼされた論文はほぼすべてK位付近に位置し、類似度はK番目の結果より数千分の一低いだけなので、実際の結果への影響はリコール数値が示唆するよりも小さいです。
広い時間範囲、大きなK、またはANNがタイムアウトした場合には、正確経路に切り替えます。コーパス全体に対して単一の正確なリコールSQLを実行すると実行時間が予測できず、中断された場合にチェックポイントがないため最初からやり直す必要があります。代わりに候補をバッチに分割し、各バッチのtop-Kを計算してから、マージして全体のtop-Kを取得します。
精度は一切失われません。ある論文が全体の上位K件に入っているなら、コーパス全体でそれより高いスコアを取る論文はK件未満であり、したがって同じバッチ内でそれより高いスコアを取る論文もK件未満です。その論文は必ず自身のバッチの上位K件に入り、マージの過程で脱落することはありません。すべてのバッチが同じスコアリングと同じ順序付け(同点は論文IDで解決)を使う限り、バッチ処理の結果は全体を一度にソートした場合とまったく同じになります。
バッチ処理には多少の時間オーバーヘッドが加わりますが、信頼性とリソース制御の両方が向上します。各バッチは完了時に保存されるため、再起動したプロセスは停止したところから続行できます。タイムアウトしたバッチは2つのより小さなバッチに分割されます。また、大きなジョブはバッチ間でリソースを解放するため、他のユーザーの小さなクエリをブロックしません。本番データベースで受け入れテストを実施しました。全コーパスから81万件の候補、K=5000で、407バッチ、約5分で完了しました。同じスナップショットに対するワンショットのfloat32完全一致クエリと比較して、5000件の結果すべてがID、順序、スコアで一致しました。
スクリーニングと分類
最初に使っていたDeepSeek-Flashのスクリーニングと比較して、9月下旬にJevへ切り替えた後、新しいパイプラインは本番サーバー上で36,564件の論文を約5分でスクリーニングしました。これはLLM版より10倍以上速く、コストは5分の1未満です。
なぜLLMの方が高コストで遅いのでしょうか。私たちのパイプラインは各判定の監査情報(確率または理由)を収集するからです。LLMはtrueかfalseを出力するだけでなく、判定と理由を含むJSONを出力します。CoT推論を有効にすると推論の長さが予測できないため、レイテンシも予測できません。CoTを無効にすると、バッチ処理で不正なJSONが頻繁に生成され、バッチ全体を破棄して再試行するか、失敗した項目を新しいバッチに集める必要があります。
一方、Jevのインターフェースはコンテキスト(状態)と一連の質問を受け取り、較正済みの確率を直接返します。スクリーニングと分類はどちらもJevのNoulプリミティブ上に構築されています。スクリーニングでは、1件の論文とエージェントが記述したinclude/
マルチラベル分類の入力形式
マルチラベル分類では、各論文が十数個のカテゴリ質問に答える必要があります。元の形式では、すべての質問で論文のタイトルとアブストラクトを繰り返していたため、多くのトークンを浪費していました。改善策は、論文を共有状態に置き、すべてのカテゴリ質問がそれを参照するようにすることでした。さらに一歩進めて、5件の論文を状態に置き、5件すべてについて全カテゴリを1回のリクエストで問い合わせることはできないでしょうか。12カテゴリを持つ120件の実論文でこれらの選択肢を比較しました:
| 形式 | 平均時間 | 入力トークン |
| すべての質問で論文を繰り返す | 57.9 s | 617K |
| 共有状態、1リクエスト1論文 | 35.8 s | 210K |
| 共有状態、1リクエスト5論文 | 14.6 s | 114K |
速度とコストだけを見れば、1リクエストにつき5本の論文を処理する方法が最適だ。しかし、同じグループ分けを2回実行しても結果が変わったのはわずか1本だったのに対し、120本の論文をシャッフルして新しいグループに分けると、約20本のラベルが変わった。例えば、アブストラクトにグラフニューラルネットワークと明記されている論文は、単独で判定すると「グラフ学習」カテゴリで0.95だったが、他の4本とグループ化すると0.78になり、ヒット閾値を下回った。
つまり、グループ内の他の論文が、モデルによるこの論文の判定に影響を与えるのであり、これは分類としては受け入れられない。そこでデフォルトでは、各リクエストに1本の論文のみを含め、そのリクエスト内で12カテゴリすべてを判定するようにした。それでも元の形式と比べて入力トークンは約66%削減され、論文間の干渉もない。
バッチ処理と並行性
Jevのリクエスト1回は数百ミリ秒しかかからないが、数万本の論文を確実に処理するには、モデルの外側で多くの作業が必要になる:
- 適応的パッキング:スクリーニングのバッチサイズは固定ではない。Jevのコンテキストウィンドウと、現在のバッチ内のメタデータのトークン長から計算される。最初の実装では論文を1本ずつ試し、そのたびにバッチ全体のトークン数を再見積もりしていたため、バッチが大きくなるほど遅くなった。収まる最大のプレフィックスを二分探索で求める方式に切り替えたことで、パッキングのオーバーヘッドは大幅に削減され、送信されるリクエストはまったく同じままである。
- レート制限:アカウントには1秒あたりのリクエスト数と1秒あたりのトークン数のクォータがある。GCRAアルゴリズムを用いて理論上の送信時刻で進めるようにしたため、ロック待ちやスレッドスケジューリングがクォータを消費することがなくなった。レート制限がボトルネックとなる場合、実スループットはクォータの73%から98%に向上した。
- 公平性:すべてのジョブが1つのアカウントクォータを共有しており、キューイングの際には処理中のリクエストが少ないユーザーが優先されるため、小さいジョブが大きいジョブに圧迫されることがない。
プロファイリングにより、しばらくの間、スループットを制限していたのはネットワークでもモデルでもなく、自社のワーカープロセスがCPUを使い切っていたことだと判明した。各HTTPリクエストには約120のSQL文が含まれており、コストの予約、送信前チェック、決済、結果の永続化をカバーしていた。これらはすべて、すべての呼び出しを復旧可能にし、二重課金を決して起こさず、問題発生時に追跡可能にするためのものだった。クォータ関連のSQLを1リクエストあたり約36文から約15文に削減し、ほぼ同時に完了する結果を1つのコミットにまとめることで、さらなるスループット向上が可能になった。
入力トークンで測ると、1本の論文のスクリーニングには300トークン強かかるため、1万本の論文をスクリーニングするコストはおよそ1元(約15セント)である。スクリーニングがこれほど安価になると、エージェントの動作も変わった。以前はスクリーニング前に候補セットを可能な限り絞り込んでいたが、今では広くリコールし、関連する可能性のあるものをすべて取り込み、Jevに各論文を判定させることができる。
結果
「LLMエージェントのメモリ機構」に関する検索では、エージェントはまずユーザーのヒントに基づいてコーパスに対して安価なプローブを実行し、論文のおおよその分布を把握した。その後、6つのキーワードと5つのセマンティックなリコールチャネルを作成し、それらが統合されて3万件以上の候補になった。エージェントはJevのための基準、すなわちどのような論文を含め、何を除外するかを書き、バックグラウンドでの並行バッチ判定に提出した。
別のタスクである「マルチモーダル理解と統合マルチモーダルモデルに関する8月の研究レビュー」では、エージェントは3つのキーワードおよびセマンティックリコールチャネルから466件の候補を取得し、スクリーニングによって428件の論文をトピックに関連するものとして保持しました。その後、エージェントはセット内のタイトルを読み、カテゴリリストの草案を作成し、Jev分類ノードを呼び出しました。結果は、ベンチマークとデータセット、エンボディドAI、安全性とプライバシー、統合理解と生成を含む10個のカテゴリでした。
ユーザーは、ベンチマークカテゴリにはまだ100件以上の論文があり、さらに分割できると指摘したため、エージェントはそれらのエントリを調査し、ベンチマークサブセットを11個のサブカテゴリに分割しました。
分類中、Jevは各論文に対して各カテゴリの基準への包含確率を付与します。これらの確率はすべてデータベースに保存され、ユーザーとエージェントが確認できるようになっています。
なぜParallaxなのか?
なぜParallaxを構築するのか?再帰的自己改善が本当に到来し、AI論文がその価値を完全に失うまでは、arXiv論文はフロンティアラボの外にいる研究者、学生、開発者にとって、この分野に関する一次情報源であり続けます。しかし、出版される論文がますます増え、それらが集合的にますます多くの洞察を含むようになるにつれて、人々が論文の増加に匹敵するペースでこの分野を理解することが難しくなっています。
私は関連するすべての論文を見つけ、分野全体、あるいは自分が関心を持つニッチな領域で、空間と時間を超えて何が起きたのかを見たいのです。それを行う唯一の方法は、全文コーパスに対する高速なリコールスキャンです。Google SearchやarXiv SearchのトップK結果ではそこに到達できません。「すべてを見つける」という需要は、GitHub上の無数のawesome論文リストから明らかです。
1年半前、まだゲーム業界で働いていたとき、ユーザーリサーチチームによる、新しいジャンルがどのように形成されるかを観察しているという講演に参加しました。Steamタグはプレイヤーによって作成され、レビューとタグを継続的に追跡することで、あるフレーズが散発的な言及からほとんどのプレイヤーが使うものへと変わり、やがて新しい「ジャンル」になる様子を見ることができます。「ソウルライク」と「メトロイドヴァニア」はそのように始まりました。もともとプレイヤーは「このゲームはDark Soulsと同じくらい難しく、同じように物語を語る」あるいは「これはMetroidとCastlevaniaのようにプレイできる」という意味で使っていました。
arXiv論文の爆発はAI研究を高頻度で急速に動く時代に突入させ、Jevはこの分野がどのように変化しているかを観察する、はるかに安価で柔軟な方法を私に示してくれました。
Parallaxは公開されています。https://parallax.ink でサインアップして試してみてください。