IR reading 2024春で紹介しようと考えていたSparseモデルの論文

背景

IR reading 2024春で発表予定だった論文の紹介。

sigir.jp

都合により発表できなくなったのでここで供養する。 数式等細かいところまでまとめる時間なかったので読んでたときのメモ+αくらいの内容。

どっちかにしようかなと思っていたけど、OneSparseのほうは発表があるっぽい。

Improved Learned Sparse Retrieval with Corpus-Specific Vocabularies(ECIR2024) Puxuan Yu, Antonio Mallia

arxiv.org

Sparseな学習済検索システムのefficiency と effectivenessを改善するために、コーパス固有のVocabularyを活用するという研究。

Corpus-Specific Vocabularies (CSV)というのは、ターゲットとするコーパスを用いてBERTで事前学習したものを指す。 この論文だとBERTの標準的な語彙サイズである3万語よりも多い、10万語,30万語という大きいサイズを検討している。

このCSVをつかうことで effectiveness と efficiency を改善できるらしい。

effectivenessの改善

そもそも標準的なBERTは事前学習に、BooksCorpus(800Mワード)やEnglish Wikipedia(2,500Mワード)のデータセットを用いられるが、これが検索対象のコーパスと語彙が一致するわけではない。 一方で、検索対象のコーパスを使って事前学習を行うことでVocabularyに含まれる語彙が検索対象のコーパスに含まれるもののみになる(この論文でいうCorpus-Specific Vocabularies)。

検索クエリで使われる語彙に対しては、標準的なBERTのVocabulariesよりも検索対象のコーパスで事前学習したCorpus-Specific Vocabulariesのほうがカバレッジが高いと考えられる。

ちなみに、Abstractでは12% qualityが上がったと書いてあったが、どの実験の話をしているのかわからなかった。 一番近いのはこの実験結果

efficiencyの改善

CSVを使うとeffectivenessが向上するのは予想通りな感じはするが、efficiencyも改善するらしい。

語彙数が増えるほどクエリのトークン数が減り、また、pasageの長さ、クエリのポスティングリスト、MRT(mean retrieval time)が減少する。らしい。

またlatencyも向上する。 これはuniCOILとSPLADEで原因が異なるよう。

以下は私の理解なので正しいか怪しいところ。

  • uniCOILの場合、CSVを用いると、転置インデックス内のポスティングリストの長さが短縮されるため、クエリ処理に必要な時間が短縮され、レイテンシが短縮される。
  • SPLADEの場合、effectivenessとefficiencyのバランスをとるためにおこなれている?SPLADEモデルで使用されるFLOPSスパース正則化の効果が語彙サイズが大きくなるほど高まるからなのだろうか。 そうだとすると、CSVがというより語彙サイズの問題のようにも思える。

デメリット

一見よさそうに見えるがデメリットはもちろんある。

  1. 事前学習にコストかかる
  2. 大規模なコーパスが必要

CSVは独自のコーパスを使ってBERTの事前学習を行うため、単純に時間とコストがかかる。 2に関して、ベースラインを作るために、30kの語彙を持つBERT-largeで事前学習する実験をしたがMSMARCOコーパスだと小さすぎたとのことで、事前学習にはそれなりに大きなコーパスが必要になるよう。

OneSparse: A Unified System for Multi-index Vector Search (WWW '24) Chen , 他著者多数

https://www.microsoft.com/en-us/research/uploads/prodnew/2024/04/OneSparse-A-Unified-System-for-Multi-index-Vector-Search.pdf

タイトル通り複数のvector searchのindexを統合するシステムOneSparseの提案。 ここでいう複数のvector searchのindex(Multi-index Vector Search)はsparseとdenseの二つのモデルを別のindexして検索するようなシステムを指す。

従来のシステムでは、例えばsparseとdenseの二つのモデルを別のindexしていた場合、それぞれのindexからtopKをそれぞれ検索した後に統合する仕組になっている。

この仕組はいくつか課題があって、

  • それぞれのindexで適切なkの設定が難しい
  • 最初にそれぞれスコアを行うため、結果的に最終的に使わない結果も計算したり低品質な結果が混じったりする

Onesparseはこのあたりの課題を解決しつつ、検索精度の維持しつつ検索速度を向上させる。

Onesparseの手法

OneSparseでは、従来のそれぞれのtopkの計算を待ってから結合していくのではなく、そもそも計算中に不要なドキュメントを排除していく手法をとる。

そのためのアイディアとして、dense vectorをSPANNアルゴリズムクラスタリングする。各クラスタは属するvectorとidのペアを持ち、このクラスタをポスティングリスト(SPANNポスティングリスト)として扱うことで、 従来のポスティングリストの走査時の性質をそのまま使えるようにしている。

この性質を使ってfast multi-way merge algorithmが提案されている。 このアルゴリズムは、転置インデックスが転置リストの走査中にインデックスの積集合/和集合を実行できるという特徴に基づいて処理を効率化するもので、すべてのindexに含まれるドキュメントのみをスコア計算する。

従来の方法だとすべてのインデックスが TopK検索を完了するまで待つが、複数のインデックスを同時にトラバースすることでそもそも両方のindexに含まれないドキュメントはskipしてそもそもスコア計算をしないことで計算量をかなり節約できる。 これによって、ANN 距離と BM25 スコア計算を実行する前に、平均して 99% 以上の低品質の候補をフィルタリングできるらしい。

さらに計算量を減らすアイディアとして、SPANNポスティングリストの圧縮も提案されている。 属するvectorの重心をそのクラスタのembeddingにすることでより効率的にスコア計算が可能になり、検索時はクエリがどのクラスタに近いかだけを計算するようにする。

したがってOneSparseでは、すべてのindexに含まれるdocumentのみに対してスコア計算を行い、さらにANN距離の計算はそのポスティングリストひとつに対してのみ実施することで、パフォーマンスの改善を実現している。

結果として、Bing Web 検索で品質を担保して、レイテンシーが5倍になったらしい。

感想

論文ではElasticserachでの実現方法も記載されており、実現しやすそうな気もしたが、以下の二つが気になっている。

  1. 検索時のクエリとクラスタとの距離計算

論文中では、最初にin-memoryのSPTAGインデックスを利用して、いくつかの最も近いクラスタを特定する。とあるが、in-memoryのSPTAGインデックス構築をどうやるのか、どれくらいのスペックが必要なのかという点が実用を考えると気になるポイントだった。

  1. fast multi-way merge algorithmの最悪ケースの精度

全てのindexに含まれるドキュメントのみスコア計算するが、もしドキュメントがかぶらなかったらかなり低いランクのドキュメントばかりになるのでは?と思ったが、そのあたりどう対処するのかよくわからなかった。

クラスタ構成のQdrantでsnapshotを作成/復元する

クラスタ構成のQdrantでsnapshotを復元する際のめも

前提

[重要] Qdrantはクラスタ構成であること。

面倒なことにQdrantは1ノードの構成なのかクラスタ構成なのかでsnapshotの復元方法が異なる。

環境

github.com

手順

collection作成~point追加

これを使ってamazon_reviewというcollectionを作成 HuggingFaceのデータセットからamazonのレビューを取得して5000件ほどembeddingして追加。

github.com

data downloadとかembeddingに少し時間かかるので、自分で適当なcollectionとpointを作成したほうが早い。

qdrant.tech

snapshot作成

POST collections/amazon_review/snapshots

デフォルトでsnapshotのディレクトリは./snapshotsになっている。

コンテナの中に入って確認したらこんな感じ。

root@3396583d4622:/qdrant# ls -lh snapshots/amazon_review/
total 73M
-rw-r--r-- 1 root root 73M Aug 26 07:52 amazon_review-8556042130106098-2023-08-26-07-52-28.snapshot

snapshotが作成されていることが確認できたので、collectionを削除。

DELETE collections/amazon_review

snapshot復元

snapshotの復元方法は二つあるがクラスタ構成の場合はAPIを使って復元する方法一択になる。

まずcollectionを作成する。

PUT collections/amazon_review
{
    "vectors": {
      "size": 768,
      "distance": "Cosine"
    }
}

次にsnapshotを復元する。

PUT /collections/amazon_review/snapshots/recover
{
  "location": "http://qdrant_primary:6333/collections/amazon_review/snapshots/amazon_review-8556042130106098-2023-08-26-07-32-20.snapshot"
}

ポイントはsnapshot復元前にcollectionを作成しておくこと。そして作成時にvectorsのsizeはsnapshotのものに揃えておくこと。 ここが間違っているとエラーになる。

collectionを作成せずにrecoverすると以下のエラーが返ってくる。

{
  "error": "Wrong input: Failed to download snapshot from http://qdrant_primary:6333/collections/amazon_review/snapshots/amazon_review-8556042130106098-2023-08-26-07-52-28.snapshot: status - 404 Not Found"
}

download snapshot でエラーが発生してstatus 404なのでlocationが間違っているように見えるが、単純にcollectionがないだけである。

そしてややこしいことに、本当にsnapshotが存在しない場合も

PUT /collections/amazon_review/snapshots/recover
{
  "location": "http://qdrant_primary:6333/collections/amazon_review/snapshots/hogehogehoge.snapshot"
}

同様のエラーが返ってくる...

{
  "error": "Wrong input: Failed to download snapshot from http://qdrant_primary:6333/collections/amazon_review/snapshots/hogehogehoge.snapshot: status - 404 Not Found"
}

collectionがないことに気づかず無駄な時間を使ったのでこのエラーが出た場合は注意.

参考

qdrant.tech

docker composeでqdrantのクラスタ構成を構築する

ベクトル検索エンジンのqdrantをdocker-composeを使ってクラスタ構築した際のメモ。

結論

composeファイルはこんな感じ。

version: "3.9"
services:
  qdrant_primary:
    image: "qdrant/qdrant:v1.4.1"
    ports:
      - "6333:6333"
    environment:
      QDRANT__CLUSTER__ENABLED: "true"
    command: [ "./qdrant", "--uri", "http://qdrant_primary:6335" ]
  qdrant_secondary:
    image: "qdrant/qdrant:v1.4.1"
    environment:
      QDRANT__CLUSTER__ENABLED: "true"
    command:
      [
        "./qdrant",
        "--bootstrap",
        "http://qdrant_primary:6335"
      ]

説明

環境変数を使う場合は、QDRANT__CLUSTER__ENABLED をTrueにする。

qdrantでクラスタ構築する場合は、config.ymlか環境変数クラスタモードを指定する必要がある。

今回はお手軽に環境変数で指定する。

ノード間での内部通信ができるようにする

qdrantではデフォルトで内部通信にport 6335が使われる。 クラスタ管理する際はすべてノードがこのportで通信できるようにしておく必要がある。

またクラスタ内の全ノードと連携するために、二つ目以降のノードはいずれかの一つのノードを知っておく必要がある。なので、最初のノードが自身のURLを提供するようにしておく。

このあたりの設定はCLIでserver起動する際の引数で指定する。

github.com

今回使うのは --uri--bootstrapの二つ。

--uriでノードのURIを指定する。このURIは他のノードが到達できるようにしておく必要がある。

--bootsrapで--uriで指定された既存のノードのURIを指定する。

qdrant.tech

問題

100%再現していないけど、docker imageのv1.4.1を使うと稀にqdrant_secondaryのほうでleader is not established within 10 secsというエラー(panic)が発生して、qdrant_secondaryがクラスタに追加されなかった。

他にも、collectionを作成しようとするとqdrant_primaryのほうでleader is not establishedが発生して安定しない時期があった。 v1.4.0を使うと発生しなくなり、その後v1.4.1に戻してからは発生していない。

自分の環境の問題な気がするけど問題の切り分けができてない。

参考

qdrant.tech

www.youtube.com