Azure AI Search で最新のベクトル量子化機能を使って適応的検索システムを構築

Microsoft にとって最も急速に成長しているサービスの 1 つである Azure AI Search は 10 月に最新のベクトル検索機能を複数追加しました。

- **量子化のための MRL サポート**で Azure OpenAI の MRL トレーニング済み Embeddings モデルの次元要件を下げることが可能に
- `@search.score` にハイブリッド検索結果のサブスコアを表示
- ハイブリッド検索のターゲットフィルターをベクトルクエリのみにする
- テキスト分割スキルにトークンチャンクが追加
- 顧客管理キー暗号化 (CMK) のポータル サポート

など AI 統合、ベクトル検索周りの機能アップデートが続いています。

# 量子化のための MRL サポート
OpenAI の `text-embedding-3-large` と `text-embedding-3-small` では、[MRL(Matryoshka Representation Learning)](https://arxiv.org/abs/2205.13147)という技術を使用しています。この手法では、ベクトルの先頭部分に向けて情報密度が高まるようにモデルが学習されているため、一部の次元だけを利用しても重要な情報を保持できます。これにより、用途に応じてフルサイズの 3,072 次元のベクトルだけでなく、256 次元や 512 次元のサブベクトルも柔軟に活用でき、性能を大幅に損なうことなく効率的な検索や分析が可能になります。

今回インデックスにフルサイズのベクトルを保存しておくだけで、Azure AI Search がユーザーが指定した任意の次元数に削減してベクトルインデックスを構築してくれるようになりました。それだけでなくベクトル量子化機能と組み合わせることでさらなる削減が可能です。もちろん削減による精度の低下をカバーするためのリランク機能も使用できます。

## 要件
- `text-embedding-3-small`、`text-embedding-3-large` (テキストコンテンツのみ)。
- `Edm.Half` タイプまたは `Edm.Single` の新しいベクトル フィールド(既存のフィールドに MRL 圧縮を追加することはできません)。
- HNSW アルゴリズム(このプレビューでは網羅的 KNN はサポートされていません)。
- スカラーまたはバイナリ量子化を設定します。バイナリ量子化推奨。
- API バージョンは `2024-09-01-preview` を使用します。

## MRL をサポートするベクトル検索構成の例
`vectorSearch.compressions` に `truncationDimension` パラメーターが追加されました。ここに切り詰める次元を指定します。組み込みの `binaryQuantization` と `scalarQuantization` の両方で利用できます。**開発者はこれまで通りフルサイズのベクトルを保存するだでよく、設定を切り替えるだけで柔軟に実装できます。**（変更する場合インデックスの再構築が必要）

```json
    "profiles": [
      {
      ...
        "compression": "vector-bq-compressor-with-mrl"
      }
    ],
    "compressions": [
      {
        "name": "vector-bq-compressor-with-mrl",
        "kind": "binaryQuantization",
        "rerankWithOriginalVectors": true,
        "defaultOversampling": 4,
        "truncationDimension": 1024,
        "scalarQuantizationParameters": null
      },
      {
        "name": "vector-sq-compressor-with-mrl",
        "kind": "scalarQuantization",
        "scalarQuantizationParameters": {
            "quantizedDataType": "int8"
        },
        "rerankWithOriginalVectors": true,
        "defaultOversampling": 4,
        "truncationDimension": 1024
      }
    ]
```

※ compression にはどちらか一方の量子化機能をセットできる

## 精度評価とベクトルインデックスサイズの比較
私がいつも使っている MIRACL 日本語データセット 8,066 件を使用してそれぞれのベクトル量子化モードで精度評価を行いました。**切り詰めた次元ごとの精度とベクトルインデックスの縮小効果について注目してください。**

### 条件
- `truncationDimension`: 1024, 512, 256 次元に自動切り詰め
- `BQ`: 組み込みバイナリ量子化、bin、サイズ: 1bit
- `SQ`: 組み込みスカラー量子化、int8、サイズ: 8bit
- `rerank`: オリジナルベクトルによるリランク `rerankWithOriginalVectors`
- `defaultOversampling`: デフォルトの `4` を使用
- `top_k`: 上位 5 件を取得

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/674344/d2c496e8-a3c5-65fa-2326-fc1a8c1226ec.png)

ベースラインを `3072@nocomp@5` とし、`text-embedding-3-large` のフルサイズ 3,072 次元で構築したベクトルインデックスの Recall は `0.813` であり、ベクトルインデックスサイズは `94.96 MB` です。ここから組み込みのスカラー量子化、バイナリ量子化、MRL 対応のベクトル切り詰め、オリジナルベクトルを使用したリランク機能の組み合わせを評価しました。

## ベクトルインデックスサイズの削減効果
それぞれのベクトル量子化モードごとの実際のベクトルインデックスサイズの比較表です。
| Mode | Vector Type | Vector Size | Vector Index Size |
|:-----------:|:------------:|:------------:|:------------:|
| nocomp(Baseline)       | float32        | 32bit         | 94.96 MB         |
| SQ     | int8      | 8bit       | 24.18 MB       |
| BQ     | bin      | 1bit       | 3.38 MB       |
| BQ-MRL-1024     | bin      | 1bit       | 1.41 MB       |
| BQ-MRL-512     | bin      | 1bit       | 0.94 MB       |
| BQ-MRL-256     | bin      | 1bit       | 0.68 MB       |

圧縮率は埋め込み次元と切り捨てに依存します。たとえば、3,072 次元をバイナリ量子化すると **28 倍の圧縮率**が、3,072 次元を 1,024 次元に切り捨てた `text-embedding-3-large` を使用すると、バイナリ量子化で **83 倍の圧縮率**が得られます。精度とサイズはトレードオフの関係にありますので要件を満たすところまで検証してください。Microsoft としては、元の次元の 1/2 または 1/3 を使用することを推奨しています。

## クエリーパフォーマンス比較
ベースラインと比較してだいたい 60% 高速化されました。ベクトル量子化機能の違いによる差はあまりありませんでした。リランクを行うと数ミリ秒増加します。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/674344/620823fc-fbfa-6c47-2137-d550ebb364a3.png)

ベクトルはループの度にランダムで生成し、それぞれ 100 回繰り返しました。

## リランクによる効果と適応的検索システム
MRL 論文における「適応的検索システム」を使用することで、検索システムは単一の固定次元の Embeddings を使用するのではなく、複数の次元にわたって段階的に情報を取り出せる仕組みが提供されます​。

検索第一段階で低次元の Embeddings を使用して候補を絞り込み、その後より高次元の Embeddings を用いて候補のリランキングを行います。例: 256 次元の埋め込みで 200 件の候補を絞り、その後 3,072 次元の Embeddings で再評価。

適応的検索システムによる効果としては以下のようなものが挙げられます。
1. 計算効率の向上
従来の検索システムと比べ、高速化を実現します。これは、高次元 Embeddings を必要な部分にのみ使うことで、全体の計算負荷を抑えるためです。

1. 柔軟な検索戦略
検索システムは、異なるデータセットや検索対象に応じて、Embeddings 次元を動的に調整できます。これにより、少ない計算資源で最大限の精度を達成します。

| Mode | Vector Type | Vector Size | Vector Index Size | **Recall**
|:-----------:|:------------:|:------------:|:------------:|:------------:|
| nocomp(Baseline)| float32| 32bit| 94.96 MB|0.813
| 256@BQ-MRL@5-rerank| bin| 1bit| 0.68 MB|0.803

**ベクトルインデックスサイズは `1/83`、精度は元の `99%` に抑えられるという驚異的な効果**を体験しました。なぜこんなことができるかというと、`rerankWithOriginalVectors`によるリランクと `defaultOversampling` によるオーバーサンプリングを行っているからなんですね。[以前紹介](https://qiita.com/nohanaga/items/d0922a812a41af2c7f64#%E3%83%99%E3%82%AF%E3%83%88%E3%83%AB%E6%A4%9C%E7%B4%A2%E6%9C%80%E9%81%A9%E5%8C%96%E6%A9%9F%E8%83%BD)した通り、Azure AI Search は**非圧縮のオリジナルベクトルのデータを内部に保管**しているので、そのデータを使用してリランクしているんです。これぞまさに適応型検索と言えます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/674344/f7370690-03b5-e9ef-4427-5d5c547568c4.png)

まずは 256 次元に切り詰めたバイナリ量子化ベクトルで検索（256@BQ-MRL@5: 0.704）して、その後オリジナルベクトルでリランクします（256@BQ-MRL@5-rerank: 0.803）。

`defaultOversampling` パラメーターはリランキングのために元の非圧縮ベクトルをどれだけ取ってくるかを計算（k に乗ずる）するパラメータです。今回は `k = 5`、`defaultOversampling = 4` の場合 `k * defaultOversampling = 20` で 20 件のドキュメントを内部で取得しました。


つまり、**Azure AI Search では適応的検索システムを自動的に利用**できるということです。

# ベクトルを自分で切り詰めることは可能？
OpenAI の `text-embedding-3-large` と `text-embedding-3-small` は `dimensions` パラメータを使用すれば任意の次元が得られます。ふと思ったんですが、これと単純にフルサイズの Embeddings を取得した後、先頭から n 次元で切り詰めたものと同一になるのでしょうか？これができないと Azure AI Search は `truncationDimension` の度に Embeddings API にアクセスすることになってしまいます。一応 [OpenAI](https://openai.com/index/new-embedding-models-and-api-updates/) は以下のように言ってます。

>新しい Embeddings モデルは両方とも、開発者が Embeddings の使用に伴うパフォーマンスとコストをトレードオフできる手法(MRL)を使用してトレーニングされました。具体的には、開発者は API パラメータを渡すことで、Embeddings の概念表現プロパティを失うことなく、Embeddings を短縮 (**つまり、シーケンスの末尾からいくつかの数字を削除**) できます。

```python
def generate_embeddings(text): 
    return client.embeddings.create(input = [text], model= model).data[0].embedding

def generate_embeddings_with_dimensions(text, dimensions): 
    return client.embeddings.create(input = [text], model= model, dimensions=dimensions).data[0].embedding
```

```python
vector3072 = generate_embeddings("源実朝って何した人？")
print(vector3072[0:10])
[0.025496920570731163, -0.005837578326463699, -0.00656503951177001, -0.028669128194451332, -0.03868662565946579, -0.03043411672115326, 0.01302274875342846, -0.009063451550900936, 0.020440468564629555, -0.054857730865478516]

vector1024 = generate_embeddings_with_dimensions("源実朝って何した人？", 1024)
print(vector1024[0:10])
[0.03459785878658295, -0.007921258918941021, -0.00890838261693716, -0.038902364671230316, -0.05249553546309471, -0.04129735007882118, 0.017671123147010803, -0.012298583984375, 0.027736544609069824, -0.07443880289793015]
```

そのまま比較すると違いますね。元のベクトルを単位ベクトルに変換します。

```python
vector = np.array(vector3072)

# 先頭から 1024 次元に切り詰める
truncated_vector = vector[:1024]

# ノルム（ベクトルの長さ）で正規化
normalized_vector = truncated_vector / np.linalg.norm(truncated_vector)

print(normalized_vector[0:10])

[ 0.03459785915541529 -0.00792125904724126 -0.00890838216120472
 -0.03890236299825162 -0.05249553262918173 -0.04129735122000008
  0.01767112330046206 -0.01229858403292695  0.02773654373311119
 -0.07443879510093626]
```

```python
0.03459785878658295
0.03459785915541529

-0.007921258918941021
-0.00792125904724126
```
ほぼ一致しましたね。おそらく内部でもこのように実装しているのでしょう🤔




# 参考
https://techcommunity.microsoft.com/t5/ai-azure-ai-services-blog/azure-ai-search-october-updates-nearly-100x-compression-with/ba-p/4265447