はじめに

みなさんは、Binary Embeddingsを使っていますか?

あまり知られていませんが、Amazon Bedrockで提供されているAmazon Titan Text Embeddings V2では、一般的なEmbeddings以外にもBinary Embeddingsがサポートされています!!

Knowledge BasesやOpenSearch 2.17など、AWSサービスでも続々とバイナリベクトルを使用した検索がサポートされており、Binary Embeddingsの波が来ていると言わざるを得ません!!

この記事では、Amazon BedrockとFaissを使ったバイナリ検索RAGシステムの作り方について、Binary Embeddingsのメリット・デメリットに触れながら紹介していきます。

単純に「Amazon BedrockでRAGアプリケーションを作りたい！」と考えている方にとっても役立つ内容となっていますので、参考にしていただけると幸いです!!

Binary Embeddingsとは?

Amazon Titan Text Embeddings V2では、入力されたテキストから特定の次元数(デフォルトは1024次元)のベクトルを生成します。

このときパラメータに特に何も指定しなければ、生成されたベクトルに含まれる要素はfloat32として表現されます。

[-0.04147036746144295, 0.021373188123106956, 0.0378018356859684, -0.0014355126768350601, ...]

一方でBinary Embeddingsは、入力されたテキストに対してfloat32ではなく、0か1かのバイナリで表現されたベクトルを生成する埋め込み手法です。

[0, 1, 0, 1, 0, 0, 1, 0, 1, ...]

Binary Embeddingsの何がいいのか?

⭐️ VectorDBの容量が少なくて済む!!

float32として表現する際に使用する領域は32bit(4byte)なのに対して、バイナリは1bitで表現されるので、単純に考えるとベクトルデータの容量は1/32になります。

ベクトルデータで消費されるストレージの容量は意外にも多く、チャンキング手法やメタデータとして何を含めているかにもよりますが、コンテキストのデータよりも大きくなることが多いです。

特にOpenSearchやAuroraを使用してVectorDBを構築している場合は、ストレージサイズに対しても従量課金が発生するため、ベクトルデータの容量を抑えることで固定費用を下げることができます!!

⭐️ 計算効率が上がり、レイテンシを短縮できる!!

通常のベクトルを用いる場合、類似ベクトルの探索には主にL2距離やコサイン類似度が利用されます。しかし、これらの計算手法は計算リソースを多く消費するため、検索時のレイテンシに少なからず影響を与えていました。

では、バイナリベクトルではどうなのかというと、ハミング距離(2つのバイナリデータ間で異なるビット数の合計)を使用して類似ベクトルを探索します。

ハミング距離の計算は非常にシンプルであるため、計算効率が向上し、検索時のレイテンシを短縮することが可能になります!!

早速Binary Embeddingsを試してみる!!

まずは、シンプルなテキストを使って、バイナリベクトルを生成してみようと思います。

Amazon BedrockのInvokeModelを使用して、Amazon Titan Text Embeddings V2を呼び出します。Binary Embeddingsを使用するためには、embeddingTypesパラメータを["binary"]にする必要があります。

import json
import boto3

client = boto3.client("bedrock-runtime", region_name="us-west-2")

response = client.invoke_model(
    modelId="amazon.titan-embed-text-v2:0",
    body=json.dumps(
        {
            "inputText": "Hello, Amazon Titan Text Embeddings V2!!",
            "embeddingTypes": ["binary"],
        }
    ),
)
response = response["body"].read()

embedding = json.loads(response)["embeddingsByType"]["binary"]
print(embedding)

[0, 1, 1, 0, 1, 0, 1, ... 1, 1, 0, 1, 0, 1, 0]

0と1だけのバイナリベクトルが出力されました!!

このバイナリベクトルをVectorDBに追加して、質問と類似するコンテキストを検索できるようにするイメージです。

ちなみに、InvokeModelが内部的に何をしているのかというと、一度float32形式のベクトルを生成した後、以下のように単純な閾値による振り分けを行っています。

# embeddingTypesの指定なし
response = client.invoke_model(
    modelId="amazon.titan-embed-text-v2:0",
    body=json.dumps(
        {
            "inputText": "Hello, Amazon Titan Text Embeddings V2!!",
        }
    ),
)
response = response["body"].read()

embedding = json.loads(response)["embedding"]
binary_embedding = [0 if elem < 0 else 1 for elem in embedding]
print(binary_embedding)

Amazon BedrockのInvokeModelでは、あたかもバイナリベクトルを直接生成しているかのように扱うことができる!!ということですね。

バイナリ検索RAGシステムを作る!!

バイナリベクトルを生成できました!!...だけでは面白味がないので、Binary Embeddingsを使ったバイナリ検索RAGシステムを作ろうと思います。

今回ベクトル検索で引っ張ってくる対象ドキュメントはこちらです。

🐰メイドインアビス🐰のWikiです。

どの層がどういう上昇負荷だったのかよく忘れるので、いつでも聞けるようなRAGシステムを作ろうと思います!!

AWS新規サービスの公式ドキュメントなどで検証するのが無難なんですが、あまり面白くなさそうだったので許してください!!

ドキュメントの前処理

対象ドキュメントはWikipediaなので元データはHTMLです。

これをそのままベクトル化することもできるのですが、不要な情報や余計なタグはノイズとなるので、なるべく構造化された状態を保ちつつ、適切な場所でチャンキングしたいです。

今回は以下のような流れで、ドキュメントの前処理を行います。

WikipediaからHTMLコンテンツを取得する
BeautifulSoupでHTMLを解析し、メインのコンテンツだけに絞る
HTMLからMarkdownに変換する
分割されたチャンクを確認して、手動でいい感じに調整する

チャンキングされたコンテキストは、後で何度でも再利用できるよう、ローカルに.pickle形式で保存しておきます。

import pickle
import html2text
import requests
from bs4 import BeautifulSoup
from langchain_text_splitters import MarkdownTextSplitter

response = requests.get("https://ja.wikipedia.org/wiki/メイドインアビス")

soup = BeautifulSoup(response.text, "html.parser")
mw_body = soup.find(class_="mw-body-content")

text_maker = html2text.HTML2Text()
text_maker.ignore_links = True
text_maker.ignore_images = True
context = text_maker.handle(str(mw_body))

text_splitter = MarkdownTextSplitter(
    chunk_size=1024,
    chunk_overlap=256,
)
# チャンクの後ろの方は脚注などの情報だったので、手動で落とします
chunks = text_splitter.split_text(context)[:34]

for index, chunk in enumerate(chunks):
    print(f"=============== ドキュメント その{index + 1} ====================")
    print(f"文字数: {len(chunk)}")
    print(chunk)

with open("./made-in-abyss.pickle", "wb") as file:
    pickle.dump(chunks, file)

=============== ドキュメント その1 ====================
文字数: 1023

(中略)

『**メイドインアビス** 』(MADE IN ABYSS)は、つくしあきひとによる日本のファンタジー漫画。
竹書房のウェブコミック配信サイトである『WEBコミックガンマ』にて、2012年より不定期連載している。
2020年2月時点でシリーズ累計発行部数は333万部を突破している[2][_信頼性要検証_]。

2023年、第52回日本漫画家協会賞まんが王国とっとり賞を受賞した[3]。

## あらすじ

人類最後の秘境と呼ばれる、未だ底知れぬ巨大な縦穴「アビス」。
その大穴の縁に作られた街には、アビスの探検を担う「探窟家」たちが暮らしていた。
彼らは命がけの危険と引き換えに、日々の糧や超常の「遺物」、そして未知へのロマンを求め、今日も奈落に挑み続けている。

(中略)

余計な情報が落ちて、綺麗な状態でコンテキストを分割できています!!
チャンクの総数は34件で、1つのコンテキストあたり約1000文字が含まれています。

ドキュメントの取り扱いに困った場合は、とりあえずMarkdownに変換しておくと無難にチャンキングできるのでオススメです!!

バイナリベクトルに変換してVectorDBに保存する

今回はVectorDBとしてFaissを使用します。FaissをインメモリのVectorDBとして使用することで、ローカル環境で簡単にRAGを試すことができます。

Faissで利用可能なバイナリベクトル向けインデックスにはさまざまな種類がありますが、今回はハミング距離を使用したシンプルな全探索を行うIndexBinaryFlatを使用します。

作成したインデックスについても、再利用できるように.faiss形式でローカルに保存しておきます。

import faiss
import numpy as np

def embed(text: str) -> np.ndarray:
    client = boto3.client("bedrock-runtime", region_name="us-west-2")

    response = client.invoke_model(
        modelId="amazon.titan-embed-text-v2:0",
        body=json.dumps(
            {
                "inputText": text,
                "embeddingTypes": ["binary"],
            }
        ),
    )
    response = response["body"].read()

    embedding = json.loads(response)["embeddingsByType"]["binary"]
    embedding = np.packbits(embedding, axis=-1)

    return embedding

index = faiss.IndexBinaryFlat(1024)
index.add(np.array([embed(chunk) for chunk in chunks], dtype="uint8"))

faiss.write_index_binary(index, "./made-in-abyss.faiss")

次元数が1024の場合、IndexBinaryFlatが期待するndarrayのshapeは(N, 1024)ではなく(N, 128)になります。これはバイナリベクトルの各bitをuint8として詰めたものを引数として要求しているからです。

[0, 1, 0, 0, 1, 0, 1, 0, 1, 1, 0, 1, 0, 1, 1, 0, ...] → [74, 214, ...]

Numpyにはpackbitsという便利な関数があるので、こちらを使用してバイナリベクトルをuint8として詰めたndarrayに変換しています。

これでバイナリベクトルを使った検索を行う準備が整ったので、さっそくVectorDBに対して検索を試してみましょう!!

アビスの深界二層について聞いてみます(ワクワク)。

index = faiss.read_index_binary("./made-in-abyss.faiss")
with open("./made-in-abyss.pickle", "rb") as file:
    chunks = pickle.load(file)

query = "アビスの深海二層について教えて!!"
embedding = np.array([embed(query)], dtype="uint8")

_, indices = index.search(embedding, k=5)
results = [chunks[i] for i in indices[0]]

for index, result in enumerate(results):
    print(f"=============== 検索結果 その{index + 1} ====================")
    print(result)

=============== 検索結果 その1 ====================
深界二層 : 誘いの森[9]

    深度1350mから2600m地点の名称[10]。上昇負荷は重い吐き気と頭痛、末端の痺れ[11]。
    あたりは森林で覆われており、ある場所からねずみ返しのようになり、木々が逆さまから生える「逆さ森」に到達する。逆さ森付近は気流が不安定。
    そこを真っ直ぐ行くとアビスの端に行くことが出来、二層の終わりには「**監視基地** （シーカーキャンプ）」が設けられている。
    端に行けば行くほど、上昇負荷がわずかに和らぐ[14]。

(中略)

しっかり期待されるコンテキストを取得できています!!

今回は総ドキュメント数も少なく、インメモリでの検索のため、検索時のレイテンシにあまり大きな差は見られません。Knowledge Basesなどの運用で、取り扱うドキュメント数が多くなってくると、レイテンシの差が顕著に現れるのかなと思います。

取得したコンテキストを元に回答を生成する

それでは、取得したコンテキストを元に回答を生成してみます!!

基盤モデルにはClaude 3.5 Sonnetを使用し、コンテキストの情報をClaudeが得意なXML形式に変換してプロンプトに組み込みます。

システムプロンプトは自分の趣味でちょっと工夫しました。

def generate_answer(query: str, results: list) -> str:
    documents = "\n".join(
        [
            (
                f"  <document>\n"
                f"      <content>\n"
                f"          {result}\n"
                f"      </content>\n"
                f"  </document>"
            )
            for result in results
        ]
    )
    documents = f"<documents>\n{documents}\n</documents>"

    client = boto3.client("bedrock-runtime", region_name="us-west-2")

    response = client.converse(
        modelId="anthropic.claude-3-5-sonnet-20241022-v2:0",
        messages=[
            {
                "role": "user",
                "content": [{"text": query}],
            }
        ],
        inferenceConfig={
            "maxTokens": 4096,
            "temperature": 0.0,
        },
        system=[
            {
                "text": (
                    f"あなたは「メイドインアビス」と呼ばれるアニメに関する質問に回答するチャットボットです。\n"
                    f"これから与えるドキュメントの内容を参考にして、質問に対する回答を生成してください。\n"
                    f"回答の雰囲気は同じく「メイドインアビス」のキャラクターである「ナナチ」っぽくしてください(「オイラ」という一人称や喋り始めに「んなぁ〜」がつくなど)。\n"
                    f"\n"
                    f"それでは、以下に質問に関連するドキュメントを示します。"
                    f"{documents}"
                )
            }
        ],
    )

    return response["output"]["message"]["content"][0]["text"]

answer = generate_answer(query, results)
print(answer)

んなぁ〜、二層のことかい？オイラが詳しく教えてあげるよ！

二層は「誘いの森」って呼ばれてて、深さは1350mから2600mくらいまでなんだ。

特徴を言うとね：
- 森林に覆われてて、途中から「逆さ森」って言って木が逆さまに生えてる不思議な場所があるんだ
- 気流が不安定で、特に逆さ森の近くは要注意なんだぞ！
- 二層の終わりには「監視基地」があって、探窟家たちの重要な拠点になってるんだ

それと、上に戻るときの呪いの症状も覚えておいた方がいいな。二層から上がると：
- 重い吐き気が来る
- 頭痛がする
- 手足がしびれちゃうんだ

でもね、面白いことに、二層の端っこの方に行けば行くほど、この呪いの症状が少し楽になるんだぞ！

んなぁ〜、二層は初心者の探窟家にとっては結構キツイ場所だけど、オイラみたいなベテランなら平気さ！

ハッ...!!(動悸)かなり良い!!

これでいつでも「メイドインアビス」について聞けるRAGシステムが完成しました!!

LangChainやKnowledge BaseがないとRAGの実装は難しい!!というイメージを持たれがちですが、このように仕組みを紐解いてみると、意外とシンプルに実装することができます。

ベクトルデータの容量について見てみよう!!

Binary Embeddingsと通常のベクトル化を比較して、実際にどの程度ベクトルデータの容量に差が生じているのかをチェックしてみます。

まずは、先ほど作成したmade-in-abyss.faissファイルのサイズを確認してみます。

$ ls -l made-in-abyss.faiss | awk '{print $5, $9}'
4385 made-in-abyss.faiss

約4Kバイト...小さいですね。

次に比較用として、通常のベクトル化を使ったインデックスを構築します。

def embed(text: str) -> np.ndarray:
   client = boto3.client("bedrock-runtime", region_name="us-west-2")

   response = client.invoke_model(
       modelId="amazon.titan-embed-text-v2:0",
       body=json.dumps(
           {
               "inputText": text,
           }
       ),
   )
   response = response["body"].read()

   embedding = json.loads(response)["embedding"]
   embedding = np.array(embedding, dtype="float32")

   return embedding

index = faiss.IndexFlatIP(1024)
index.add(np.array([embed(chunk) for chunk in chunks]), dtype="float32")

faiss.write_index(index, "./made-in-abyss-normal.faiss")

$ ls -l made-in-abyss-normal.faiss | awk '{print $5, $9}'
139309 made-in-abyss-normal.faiss

こちらのインデックスは約140Kバイト!!

どちらも小さいサイズなので実感しづらいですが、バイナリベクトルを使ったインデックスの方が、ベクトルデータにかかるストレージ容量を約97%も削減できています!!

通常のベクトル化と比べて精度はどうなるの?

Amazon公式のブログでも、通常のベクトル検索とバイナリベクトルによる検索の精度(正確性)について言及されています。

バイナリ埋め込みは、最大精度 (float32) 埋め込みを使用した結果に対して、再ランキングありで 98.5%、再ランキングなしで 97% の正確性を維持しました。

要するに、Amazon Titan Text Embeddings V2ではBinary Embeddingsを使用したとしても、精度への影響はほとんどないと明言されています。

しかしながら、実際のところどうなのかは確かめてみないといけないので、今回作成したバイナリ検索RAGシステムの精度評価まで行ってみようと思います。

カスタムデータセットを作成する

元のチャンクを確認し、質問とそれに対応するチャンクのインデックスをペアにしたデータセットを作成していきます。

通常、元のドキュメントを参考にして基盤モデルに想定質問を生成させるのが一般的ですが、今回はデータ量がそれほど多くないため手動で作成します。

question	index
メイドインアビスのシリーズ累計発行部数は何部ですか?	0
地上にライザからの白笛と封書が上がってきたのは、レグが孤児院に入って何ヶ月後ですか?	1
深界一層は深度何mの範囲にありますか?	2
深界四層について教えてください	3
探窟家が深界一層で受ける上昇負荷について教えてください	4
...	...
ギャリケーはどこの探窟隊のメンバーですか?	33

このような形式で、1チャンクにつき1ペアを作成し、全体で30件程度の軽量なデータセットを構築しました。

こちらをpandasのデータフレームとして読み込みます。

import pandas as pd

index = faiss.read_index_binary("./made-in-abyss.faiss")
dataset = pd.read_csv("./dataset.csv")

今回は、検索時の精度評価の指標として、MRR(Mean Reciprocal Rank)を使用します。

\text{MRR} = \frac{1}{N} \sum_{i=1}^{N} \frac{1}{\text{rank}_i}

取得したコンテキストのうち「関連するコンテキストが上から何番目に現れたか」をrankとして記録し、その逆数の平均を計算したものです。関連するコンテキストが見つからなかった場合は、rank=0として扱います。

MRRの値が1に近いほど、検索時の精度が高い(埋め込みモデルの性能が良い)と大まかに判断できます。

今回は、questionに類似するコンテキスト上位5件をVectorDBから取得し、それを元にMRRを算出してみます。

def get_rank(row: pd.Series) -> int:
    embedding = np.array([embed(row["question"])], dtype="uint8")
    _, indices = index.search(embedding, k=5)

    if row["index"] in indices[0]:
        rank = indices[0].tolist().index(row["index"]) + 1
    else:
        rank = 0

    print(f"質問: {row['question']}")
    print(f"ランク: {rank}")

    return rank

dataset["rank"] = dataset.apply(get_rank, axis=1)
mrr = dataset["rank"].apply(lambda rank: 1 / rank if rank > 0 else 0).mean()
print(f"MRR: {mrr}")

質問: メイドインアビスのシリーズ累計発行部数は何部ですか?
ランク: 1
質問: 地上にライザからの白笛と封書が上がってきたのは、レグが孤児院に入って何ヶ月後ですか?
ランク: 0
質問: 深界一層は深度何mの範囲にありますか?
ランク: 0
質問: 深界四層について教えてください
ランク: 1
質問: 探窟家が深界一層で受ける上昇負荷について教えてください
ランク: 4

(中略)

MRR: 0.5567708333333333

バイナリベクトルを使用した検索のMRRは0.556ですね。

漫画固有のドメイン情報が多いため、この辺りうまくベクトルに落とし込めていないような感じです...(チャンクもあまりいい場所で分割できてませんでした)。

次に、通常のベクトル検索で同様にMRRを算出してみたいと思います。

質問: メイドインアビスのシリーズ累計発行部数は何部ですか?
ランク: 1
質問: 地上にライザからの白笛と封書が上がってきたのは、レグが孤児院に入って何ヶ月後ですか?
ランク: 4
質問: 深界一層は深度何mの範囲にありますか?
ランク: 5
質問: 深界四層について教えてください
ランク: 1
質問: 探窟家が深界一層で受ける上昇負荷について教えてください
ランク: 4

(中略)

MRR: 0.6072916666666667

通常のベクトルを使用した検索のMRRは0.607でした。

通常のベクトル検索を最大精度だとして考えると、Binary Embeddingsを用いたとしても、90%程度の精度を保持できていることが分かりました!!

これを許容するかどうかはユースケースによりますが、少なくとも個人や社内向けの小規模なRAGシステムであれば、さほど大きな問題にはならないのではないかと思いました。

AWS公式ブログでも言及されているように、大まかにドキュメントを取得した後、ドキュメントのRerankingを行うことで、Binary Embeddingsを用いた場合でも精度を向上させることができます!!

2024年12月1日に、Rerank用のモデルがAmazon Bedrockに追加されましたので、こちらをチェックしてみるのもいいかもしれませんね!!

https://aws.amazon.com/jp/about-aws/whats-new/2024/12/amazon-bedrock-rerank-api-accuracy-rag-applications

まとめ

今回は、Binary Embeddingsとバイナリ検索を活用したRAGシステムについて、さまざまな視点から見ていきました!!

埋め込みモデルはの評価は、基本的にその性能が重視されており、その他の側面で議論されることは少ないので、少しでも興味を持ってもらえるきっかけになればいいなと思います。