CLOVER🍀

That was when it all began.

Qdrantのチュートリアルから、「検索品質を測定する(Measure retrieval quality)」を試す

これは、なにをしたくて書いたもの?

Qdrantのチュートリアルから、「検索品質を測定する(Measure retrieval quality)」を試してみたいと思います。

Measure retrieval quality - Qdrant

今回のチュートリアルの狙い

今回扱うチュートリアルは、こちらの「検索品質を測定する(Measure retrieval quality)」です。

Measure retrieval quality - Qdrant

どういうことをするのか?というのは、まずはこのページを読み進めて見てみようと思います。

まずは冒頭を読むと、このチュートリアルでは「セマンティック検索の品質をどのようにして測定するのか、そしてQdrantが使用する
ANNアルゴリズムのHNSWパラメーターをどのようにしてチューニングして最良の結果を得るのか、その方法を示す」ことを
行うようです。

In this tutorial, we will show how to measure the quality of the semantic retrieval and how to tune the parameters of the HNSW, the ANN algorithm used in Qdrant, to obtain the best results.

扱う品質のテーマとしては、埋め込みの品質と検索の品質の2つがあるようです。

まずは埋め込みの品質から。

Measure retrieval quality / Embeddings quality

埋め込みの品質は、通常はMTEB(Massive Text Embedding Benchmark)などのベンチマークで測定および比較されます。

MTEB Leaderboard - a Hugging Face Space by mteb

テキスト埋め込みのベンチマークMTEB(Massive Text Embedding Benchmark)って? - CLOVER🍀

評価プロセスは簡単で、人が作成したデータセットに基づく一連のクエリーと期待するドキュメントがあるので、評価プロセスでは
クエリーを取得してベクトル空間で最も類似したドキュメントを見つけて正解と比較します。

これで、最も類似したドキュメントの検索は、近似を行わない完全なkNN検索として実装されます。

finding the most similar documents is implemented as full kNN search, without any approximation.

その結果、ANNアルゴリズムの影響を受けることなく、埋め込みの品質を測定できるそうです。

少し馴染みがない言葉が出てきました。まずはkNN検索から。

k近傍法(ケイきんぼうほう、英: k-nearest neighbor algorithm, k-NN)は、入力との類似度が高い上位 k個の学習データで多数決/平均するアルゴリズムである。

k近傍法は以下の手順からなる:
1. 入力と全学習データとの類似度(距離)測定
2. 類似度上位 k 個の選出
3. 選出されたデータの多数決あるいは平均
すなわち「入力とよく似た k 個のデータで多数決/平均する」単純なアルゴリズムである

k近傍法 - Wikipedia

では、ANNの方は?

最近傍探索(英: Nearest neighbor search, NNS)は、距離空間における最も近い点を探す最適化問題の一種、あるいはその解法。近接探索(英: proximity search)、類似探索(英: similarity search)、最近点探索(英: closest point search)などとも呼ぶ。問題はすなわち、距離空間 M における点の集合 S があり、クエリ点 q ∈ M があるとき、S の中で q に最も近い点を探す、という問題である。

最近傍探索 - Wikipedia

最近傍探索の派生で、近似最近傍探索をANNと呼ぶみたいです。

近似最近傍探索 (英: approximate nearest neighbor, ANN)

つまり、こちらは近似になります、と。

ベクトル空間が大きくなると、計算量が指数関数的に増加する「次元の呪い」と呼ばれる事象への対策になっているみたいです。

次元の呪い - Wikipedia

つまり、kNN検索に比べてANNは精度を犠牲にして高速化することを狙っているわけですね。

次は、検索品質です。

Measure retrieval quality / Retrieval quality

Qdrantは純粋なkNN検索を実行せず、代わりにANNを使います。これが意味することは高速ではあるものの最適ではない結果が返されることが
あるということです。ただ、その近似の検索品質は測定することができるようです。

検索品質の指標にはいくつかの方法があるようです。

  • Precision@k#Precision_at_k)
    • 上位k件の検索結果内の関連ドキュメントの数に基づく指標
  • MRR(Mean reciprocal rank)
    • 検索結果内の最初の関連ドキュメントの位置が考慮される指標
  • DCG、NDCG
    • ドキュメントの関連性スコアに基づいた指標

ANNには関連性スコアやランキングに基づくものは適用されず、ベクトル検索におけるランキングはベクトル空間内のクエリーと
ドキュメント間の距離に依存するものの使用する距離関数(距離メトリクス)は同じであるため、距離は近似によって変化しません。

よってANNの品質測定では、Precision@kのような上位k件の検索結果内の関連ドキュメントの数等によって測定するものが意味のある
結果になります。
これは上位k件の検索結果内のドキュメントの数をkで割ったものとして計算されます。

ANNのテストの場合、kを固定した正確なkNN検索を正解として使用できます。これはANNが正確な検索とどれくらい近似しているかを表す
指標になります。

つまり、上位k件に対する純粋なkNNとANNを比較することで検索品質を測定しようということですね。

やりたいことはわかったので、実際に試してみましょう。

環境

今回の環境はこちら。Qdrantは172.17.0.2で動作しているものとします。

$ ./qdrant --version
qdrant 1.8.2

QdrantのWeb UIは0.1.22を使っています。

Qdrant Clientを使うPython環境。

$ python3 --version
Python 3.10.12


$ pip3 --version
pip 22.0.2 from /usr/lib/python3/dist-packages/pip (python 3.10)

Qdrantのチュートリアル「検索品質を測定する(Measure retrieval quality)」を試す

それでは、Qdrantのチュートリアル「検索品質を測定する(Measure retrieval quality)」をドキュメントに沿って進めていきます。

Measure retrieval quality - Qdrant

このチューと有りあるでは、Qdrant ClientとHugging FaceのDatasetsを使用します。

それぞれインストール。

$ pip3 install qdrant-client datasets

インストールしたライブラリーの一覧。

$ pip3 list
Package            Version
------------------ -----------
aiohttp            3.9.3
aiosignal          1.3.1
annotated-types    0.6.0
anyio              4.3.0
async-timeout      4.0.3
attrs              23.2.0
certifi            2024.2.2
charset-normalizer 3.3.2
datasets           2.18.0
dill               0.3.8
exceptiongroup     1.2.0
filelock           3.13.1
frozenlist         1.4.1
fsspec             2024.2.0
grpcio             1.62.1
grpcio-tools       1.62.1
h11                0.14.0
h2                 4.1.0
hpack              4.0.0
httpcore           1.0.4
httpx              0.27.0
huggingface-hub    0.21.4
hyperframe         6.0.1
idna               3.6
multidict          6.0.5
multiprocess       0.70.16
numpy              1.26.4
packaging          24.0
pandas             2.2.1
pip                22.0.2
portalocker        2.8.2
protobuf           4.25.3
pyarrow            15.0.1
pyarrow-hotfix     0.6
pydantic           2.6.4
pydantic_core      2.16.3
python-dateutil    2.9.0.post0
pytz               2024.1
PyYAML             6.0.1
qdrant-client      1.8.0
requests           2.31.0
setuptools         59.6.0
six                1.16.0
sniffio            1.3.1
tqdm               4.66.2
typing_extensions  4.10.0
tzdata             2024.1
urllib3            2.2.1
xxhash             3.4.1
yarl               1.9.4

まずはQdrantにデータを登録するプログラムを作成。

upload_data.py

from datasets import load_dataset
from qdrant_client import models, QdrantClient

dataset = load_dataset(
    "Qdrant/arxiv-titles-instructorxl-embeddings", split="train", streaming=True
)

dataset_iterator = iter(dataset)
traint_dataset = [next(dataset_iterator) for _ in range(60000)]
test_dataset = [next(dataset_iterator) for _ in range(1000)]


client = QdrantClient("http://172.17.0.2:6333", timeout=300)
client.create_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    vectors_config=models.VectorParams(
        size=768,
        distance=models.Distance.COSINE
    )
)

client.upload_points(
    collection_name="arxiv-titles-instructorxl-embeddings",
    points=[
        models.PointStruct(
            id=item["id"],
            vector=item["vector"],
            payload=item
        )
        for item in traint_dataset
    ]
)

while True:
    collection_info = client.get_collection(collection_name="arxiv-titles-instructorxl-embeddings")
    if collection_info.status == models.CollectionStatus.GREEN:
        # コレクションのステータスがgreenになると、インデックスの構築が完了
        # インデックスの構築中はyellow
        break

Qdrantが提供するベクトル入りのデータセットを使います。データは、Qdrantに登録する60000件と検索のテストで使う1000件に分けます。

dataset = load_dataset(
    "Qdrant/arxiv-titles-instructorxl-embeddings", split="train", streaming=True
)

dataset_iterator = iter(dataset)
traint_dataset = [next(dataset_iterator) for _ in range(60000)]
test_dataset = [next(dataset_iterator) for _ in range(1000)]

もっとも、このプログラムでは検索のテストは行わないのですが。

QdrantClientを作成して、コレクションも作ります。

client = QdrantClient("http://172.17.0.2:6333", timeout=300)
client.create_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    vectors_config=models.VectorParams(
        size=768,
        distance=models.Distance.COSINE
    )
)

QdrantClientのタイムアウトはデフォルトが5秒で、自分の環境だと時々タイムアウトしていたので長めに設定しておきました。

https://github.com/qdrant/qdrant-client/blob/v1.8.0/qdrant_client/qdrant_client.py#L52-L54

コレクションの距離メトリクスはコサイン類似度ですね。

データのアップロード。

client.upload_points(
    collection_name="arxiv-titles-instructorxl-embeddings",
    points=[
        models.PointStruct(
            id=item["id"],
            vector=item["vector"],
            payload=item
        )
        for item in traint_dataset
    ]
)

最後に、登録したデータを含むインデックスの構築が完了するまで待ちます。インデックスの構築が完了しないと検索が実行できないからです。

while True:
    collection_info = client.get_collection(collection_name="arxiv-titles-instructorxl-embeddings")
    if collection_info.status == models.CollectionStatus.GREEN:
        # コレクションのステータスがgreenになると、インデックスの構築が完了
        # インデックスの構築中はyellow
        break

チュートリアルの内容をそのまま書いたのですが、リソースを使いすぎるので途中にスリープを入れるべきでした…。

実行。

$ python3 upload_data.py

完了までにはそこそこ時間がかかりますが、無事データの登録が完了です。

points_countとvectors_countは同じなのに対して、indexed_vectors_countの値が違うのがちょっと気になるところですが、ここでは
気にせずにおきます。

次は、検索品質の評価プログラムです。

measure_retrieval_quality.py

import time
from datasets import load_dataset
from qdrant_client import models, QdrantClient

dataset = load_dataset(
    "Qdrant/arxiv-titles-instructorxl-embeddings", split="train", streaming=True
)

dataset_iterator = iter(dataset)
traint_dataset = [next(dataset_iterator) for _ in range(60000)]
test_dataset = [next(dataset_iterator) for _ in range(1000)]

client = QdrantClient("http://172.17.0.2:6333")

def avg_precision_at_k(k):
    precisions = []

    for item in test_dataset:
        ann_result = client.search(
            collection_name="arxiv-titles-instructorxl-embeddings",
            query_vector=item["vector"],
            limit=k,
        )

        knn_result = client.search(
            collection_name="arxiv-titles-instructorxl-embeddings",
            query_vector=item["vector"],
            limit=k,
            search_params=models.SearchParams(
                exact=True
            )
        )

        ann_ids = set(item.id for item in ann_result)
        knn_ids = set(item.id for item in knn_result)

        precision = len(ann_ids.intersection(knn_ids)) / k
        precisions.append(precision)

    return sum(precisions) / len(precisions)

start_time = time.perf_counter()

avg_precision_at_k_result = avg_precision_at_k(k=5)

elapsed_time = time.perf_counter() - start_time

print(f"avg(precision@5) = {avg_precision_at_k_result}, elapsed time = {elapsed_time:.3f} sec")

データセットとQdrantClientの作成は先ほどとほぼ同じです。タイムアウトは伸ばす必要がないので、デフォルト値のままにしておきます。

ポイントになるのはこちらの関数ですね。

def avg_precision_at_k(k):
    precisions = []

    for item in test_dataset:
        ann_result = client.search(
            collection_name="arxiv-titles-instructorxl-embeddings",
            query_vector=item["vector"],
            limit=k,
        )

        knn_result = client.search(
            collection_name="arxiv-titles-instructorxl-embeddings",
            query_vector=item["vector"],
            limit=k,
            search_params=models.SearchParams(
                exact=True
            )
        )

        ann_ids = set(item.id for item in ann_result)
        knn_ids = set(item.id for item in knn_result)

        precision = len(ann_ids.intersection(knn_ids)) / k
        precisions.append(precision)

    return sum(precisions) / len(precisions)

データセットから抜き出した1000件のテストデータを使い、ANNと正確なkNN検索の結果を比較します。
つまり、検索としては2000回実行します。

Qdrantの検索をふつうに行うとANNなのですが、検索時のパラメーターでexact=Trueを指定すると正確なkNN検索として実行します。

            search_params=models.SearchParams(
                exact=True
            )

そして、ANNとkNN検索の結果のidの集合を作成して

        ann_ids = set(item.id for item in ann_result)
        knn_ids = set(item.id for item in knn_result)

集合の差(ANN - kNN)を取って、その差集合の長さをkで割ります。その値はリストとして持ちます。

        precision = len(ann_ids.intersection(knn_ids)) / k
        precisions.append(precision)

最後に、算出した値の平均値を算出します。

    return sum(precisions) / len(precisions)

つまり、これはPrecision@kの実装ですね。

これをk=5として呼び出します。

start_time = time.perf_counter()

avg_precision_at_k_result = avg_precision_at_k(k=5)

elapsed_time = time.perf_counter() - start_time

print(f"avg(precision@5) = {avg_precision_at_k_result}, elapsed time = {elapsed_time:.3f} sec")

この後、検索精度の変更を行うのでドキュメントの例から少し変えて、実行時間も計測するようにしておきました。

実行してみます。

$ python3 measure_retrieval_quality.py

結果。

avg(precision@5) = 0.9947999999999995, elapsed time = 93.955 sec

これでも十分検索品質が高いと言えそうですね。

実行時間についてはなんとも言えませんが、実行を繰り返すと速くなっていきました。キャッシュでしょうね…。

avg(precision@5) = 0.9947999999999995, elapsed time = 70.645 sec
avg(precision@5) = 0.9947999999999995, elapsed time = 52.928 sec

ちなみにデータセットのロードなどの方が圧倒的に時間がかかるので、スクリプトの実行はこれくらいでは全然終わりません。

それでもこれ以上に精度が必要で、かつ遅延が許容できる場合はコレクションのHNSWパラメーターを変更することで精度を
上げられるようです。

まず、現在のHNSWパラメーターは以下のようにm=16、ef_construct=200になっています。

HNSWは階層グラフになっていて、mはそのノードごとのエッジの数です。mの値を大きくするほど検索の精度は高くなりますが、
より多くのディスクスペースを消費します。ef_constructはインデックス構築中に考慮する近傍の数で、値を大きくするほど検索の精度は
高くなりますが、インデックスの構築に時間がかかるようになります。

これを変更するプログラムを作成。

modify_hnsw_config.py

from qdrant_client import models, QdrantClient

client = QdrantClient("http://172.17.0.2:6333")

client.update_collection(
    collection_name="arxiv-titles-instructorxl-embeddings",
    hnsw_config=models.HnswConfigDiff(
        m=32,  # ノードごとのエッジの数をデフォルトの16から32へ変更
        ef_construct=200,  # インデックス構築中に考慮する近傍の数をデフォルトの100から200へ変更
    )
)

while True:
    collection_info = client.get_collection(collection_name="arxiv-titles-instructorxl-embeddings")
    if collection_info.status == models.CollectionStatus.GREEN:
        break

実行。

$ python3 modify_hnsw_config.py

確かにパラメーターが変更されました。

検索品質を測定しなおしてみます。

$ python3 measure_retrieval_quality.py

結果。

avg(precision@5) = 1.0, elapsed time = 110.615 sec

確かに精度が上がりました。というか、1.0になりましたが…。処理時間も長くなりましたね。

何回か実行すると、速度が変わりのも同じですね。

avg(precision@5) = 1.0, elapsed time = 71.224 sec

何回かコレクションの設定変更 → 検索を繰り返していると以下のような結果を出すこともあったので、1.0というのはたまたま(?)
ということでしょうけど。

avg(precision@5) = 0.9981999999999999, elapsed time = 58.315 sec

これで、今回のチュートリアルで確認したいことは見れたと思います。

おわりに

Qdrantのチュートリアルから、「検索品質を測定する(Measure retrieval quality)」を試してみました。

率直に、おもしろかったです。検索品質を測定とはどうするのか?と思って見ていましたが、ANNとkNN検索の差を見て評価するというのが
なるほどな、と思いました。

また、そもそもkNN検索やANNも知らなかったので勉強になりました。

ベクトル検索を使うにあたり、このあたりの使い方や考え方は押さえておきたいですね。