Azure AI Search のベクトル最適化パフォーマンスを調査

Azure AI Search は柔軟なスケーリングに対応したフルマネージドのエンタープライズグレード・ベクトルデータベースです。

私はこれまで、Azure AI Studio 上のモデルカタログからデプロイした　Cohere-embed-v3-multilingual モデルを使用したベクトル最適化の機能やベクトル検索の精度評価を行ってきました。今回はこれまで本格的に実施していなかった検索パフォーマンスのテストをしていきたいと思います。

検索クエリーと測定方法

検索クエリーの送信には REST API を使用します。これはレスポンスヘッダーの中に、Azure AI Search が検索に要した実処理時間 elapsed-time が格納されているからです。単にクライアントにタイマーをセットして検索処理の時間を計測しようとしても、クライアント マシンの構成やネットワークパフォーマンス、クライアント処理の実装方法など、多くの要因によって正しく計測するのが難しいです。正確な elapsed-time が分かれば、計算によってネットワークレイテンシのみを分離して分析することも可能になります。

    api_version = "2024-05-01-preview"

    # HTTPヘッダーを設定
    headers = {
        "Content-Type": "application/json",
        "api-key": api_key
    }

    # クエリを設定
    query = {
        "search": "",
        "vectorQueries": [  
                {
                    "vector": vector,
                    "k": 3,
                    "fields": vector_field,
                    "kind": "vector",
                    "exhaustive": False
                }
            ]
    }

    # クエリURLを設定
    url = f"{endpoint}/indexes/{index_name}/docs/search?api-version={api_version}"

    start = time.perf_counter()
    response = requests.post(url, headers=headers, json=query)
    end = time.perf_counter()
    measured_time = (end - start) * 1_000 # ミリ秒に変換
    
    server_time = float(response.headers['elapsed-time'])  # サーバー処理時間
    elapsed_time = response.elapsed.total_seconds() * 1_000  # HTTP 時間
    overhead_time = measured_time - elapsed_time # オーバーヘッド時間+コンテンツの受信時間
    
    print("server_time(ms)", server_time)
    print("elapsed_time(ms)", elapsed_time)
    print("overhead_time(ms)", overhead_time)

今回はサーバー処理時間のみを分析しますが、上記のコードのように HTTP 通信時間や Python オーバーヘッド+コンテンツの受信時間を算出することもできます。

response.elapsed.total_seconds() は正確にはリクエストの最初のバイトを送信してからレスポンスヘッダの解析が終了するまでにかかる時間を返します。コンテンツの受信時間は含みませんので、measured_time - elapsed_time とすることでコンテンツの受信時間が得られますが、Python の実行にかかるオーバーヘッドも含まれます。

検索サービス

• Plan: Standard 1
• Region: Japan East
• Replica: 1
• Partition: 1
• Deploy Date: 2024/04/03~2024/05/17
  検索サービスをいつデプロイしたかによって、パフォーマンスが大きく異なります。これは、容量の増強とともにインフラスペックの増強が行われたためです。

ベクトル検索

ベクトルはループの度にランダムで生成し、100 回繰り返します。float32 は実際に格納されている値の範囲でだいたいに絞って生成しています。

# 生成する値の範囲
low = -0.12219238
high = 0.10455322

float32_vector = (np.random.rand(1024) * (high - low) + low).astype(np.float32).tolist() 
int8_vector = np.random.randint(low=-128, high=128, size=1024, dtype=np.int8).tolist()
uint8_vector = np.random.randint(low=0, high=256, size=1024, dtype=np.uint8).tolist()
bin_vector =  np.random.randint(low=0, high=256, size=128, dtype=np.uint8).tolist()

データセットはこれまでの精度評価で使用してきた MIRACL 日本語データセット 8,066 件をそのまま使用します。

平均検索速度（ミリ秒）

結果は、
float32>complessions(float32)>complessions(int8)>uint8>int8>bin
のようになりました。

バイナリが最も早い

バイナリはこの実験データの中で最も高速な平均 7.4 ms という結果になりました。精度をトレードオフにして速度を求めたい場合、バイナリ化することにより一桁ミリ秒のオーダーが狙えるかもしれません。

float32 vs int8

Cohere 社によると float32 から int8 への量子化で約 30% のスピードアップが実現できるとしていますが、今回の実験では 34.8% のスピードアップとなりました。

int8-embeddingsを使用することで、99.99%の検索品質を保ちながら、4倍のメモリ節約と約30%の検索スピードアップを実現します。私たちの意見では、これは素晴らしいソリューションであり、ほとんどのデプロイメントにお勧めします。
https://cohere.com/blog/int8-binary-embeddings#using-int8-embeddings

uint8 と int8 の差は 1ms 程度 で uint8 が高速でした。

complession モード

Azure AI Search にはスカラー量子化という機能が搭載されており、float32 のベクトルを自動的に int8 に量子化します。これによりベクトルストアは約 1/4 に削減できます。速度については検索クエリーのベクトル型によって大きく変わります。float32 vs comp(float32) で 4.3% の高速化でした。スカラー量子化が有効化されている場合、クエリーベクトルを int8 にしても検索できます。この場合、float32 vs comp(int8) で 28.2% の高速化となります。

exhaustive KNN(完全な KNN) モード

Azure AI Search は ANN のアルゴリズムに HNSW を使用しており、精度をトレードオフとして検索速度を向上させることができます。もし高い Recall を求める場合は、速度を犠牲にしてベクトル空間全体をスキャンして類似度を計算することができます。これを exhaustive KNN(完全な KNN) と呼び、exhaustive=True とすることにより有効化できます。

vectorSearch: algorithms を hnsw に設定していたとしても exhaustive を利用できます。ただし、その逆はできません。

測定結果はすべてのベクトル型において、処理時間が増加しています。（15〜45%↑）

100 回の計測における処理時間の比較

今回 100 回分の平均を取っていますが、時系列的に見たときに速度は一定でないことを理解しておく必要があります。あと最初のクエリーだけ時間がかかるのが気になるところです🤔

1 億件のベクトルインデックスを検索

それでは、1 億件のベクトルを Azure AI Search に挿入して検索したとき、どのようなパフォーマンスとなるのでしょうか。

検索サービス

• Plan: Standard 2
• Region: East US
• Replica: 1
• Partition: 1
• Deploy Date: 2024/05/17~

データセット

使用する大規模データセットは Microsoft が公開している SPACEV1B を利用します。SpaceV (Bing Web ベクトル検索) からリリースしたデータセットです。これは、Microsoft SpaceV Superior モデルによってエンコードされた 10 億オーバーのドキュメント ベクトルと 29,000 を超えるクエリ ベクトルで構成されています。今回は 100 次元の int8 を 1つのインデックスに 132,772,505 件挿入したところで評価します。

平均検索速度（ミリ秒）

• int8: 22.29 ms
• int8-exhaustive: 13980 ms

1 億件までいくと明らかに HNSW が高速であることが分かります。Azure AI Search では検索単位のスケールアップを行わない素の S2 プランでも 1 億件を約 22 ms でベクトル検索できます。

パフォーマンス監視

Azure AI Search のリソースログは Azure Monitor の Log Analytics で参照することができます。検索サービスのポータル ページで、[診断設定] を使用してログ記録を有効にした後、[ログ] を選択して Log Analytics に対する kusto クエリを発行できます。

注意

本検証結果はすべて個人的な条件設定で行なったもので公式なものではありません。あくまで参考程度に捉えてください。