埋め込みベクトルと類似度検索の理論とPython実装

埋め込みベクトル(Embedding)は、テキストや画像などのデータを高次元のベクトル空間に変換したものです。意味的に類似したデータは、ベクトル空間上で近くに配置されます。この性質を利用することで、セマンティック検索やレコメンデーションなど、様々なアプリケーションが可能になります。

本記事では、埋め込みベクトルの概念、類似度計算の数学、そして実践的な類似度検索の実装について解説します。

本記事の内容

  • 埋め込みベクトルの概念と直感
  • 類似度指標(コサイン類似度、ユークリッド距離など)
  • テキスト埋め込みの取得方法
  • セマンティック検索の実装
  • 近似最近傍探索(ANN)
  • ベクトルデータベースの活用

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

埋め込みベクトルとは

概念

埋め込みベクトル(Embedding)は、離散的なデータ(単語、文、画像など)を連続的な高次元ベクトル空間に射影したものです。

$$ \text{Embedding}: \mathcal{X} \to \mathbb{R}^d $$

ここで $\mathcal{X}$ は入力空間(例: 単語の集合)、$d$ は埋め込み次元です。

重要な性質

1. 意味的類似性の保存

意味的に類似したデータは、ベクトル空間上で近くに配置されます。

$$ \text{sim}(x_1, x_2) \text{ が高い} \Leftrightarrow \|\bm{e}_1 – \bm{e}_2\| \text{ が小さい} $$

2. 算術演算による関係の捉え方

有名な例として、Word2Vecで発見された以下の関係があります:

$$ \bm{e}_{\text{king}} – \bm{e}_{\text{man}} + \bm{e}_{\text{woman}} \approx \bm{e}_{\text{queen}} $$

埋め込みの可視化

import numpy as np
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE


def visualize_embeddings(embeddings, labels, title="Embedding Space"):
    """
    高次元埋め込みを2Dに可視化

    Args:
        embeddings: (n_samples, embed_dim) の配列
        labels: 各サンプルのラベル
        title: グラフタイトル
    """
    # t-SNEで次元削減
    tsne = TSNE(n_components=2, random_state=42, perplexity=30)
    embeddings_2d = tsne.fit_transform(embeddings)

    plt.figure(figsize=(12, 8))
    unique_labels = list(set(labels))
    colors = plt.cm.Set3(np.linspace(0, 1, len(unique_labels)))

    for label, color in zip(unique_labels, colors):
        mask = [l == label for l in labels]
        plt.scatter(
            embeddings_2d[mask, 0],
            embeddings_2d[mask, 1],
            c=[color],
            label=label,
            alpha=0.7
        )

    plt.legend()
    plt.title(title)
    plt.xlabel('t-SNE 1')
    plt.ylabel('t-SNE 2')
    plt.tight_layout()
    plt.show()

類似度指標

コサイン類似度

2つのベクトル間の角度の余弦を計算します。最も広く使われる類似度指標です。

$$ \text{cos\_sim}(\bm{a}, \bm{b}) = \frac{\bm{a} \cdot \bm{b}}{\|\bm{a}\| \|\bm{b}\|} = \frac{\sum_{i=1}^{d} a_i b_i}{\sqrt{\sum_{i=1}^{d} a_i^2} \sqrt{\sum_{i=1}^{d} b_i^2}} $$

性質: – 範囲: $[-1, 1]$ – 1: 完全に同じ方向 – 0: 直交(無関係) – -1: 完全に反対方向

利点: – スケール不変(ベクトルの大きさに依存しない) – 高次元でも有効

ユークリッド距離

ベクトル間の直線距離を計算します。

$$ d_{\text{eucl}}(\bm{a}, \bm{b}) = \|\bm{a} – \bm{b}\|_2 = \sqrt{\sum_{i=1}^{d} (a_i – b_i)^2} $$

性質: – 範囲: $[0, \infty)$ – 0に近いほど類似

内積(ドット積)

単純な内積も類似度として使用できます。

$$ \text{dot}(\bm{a}, \bm{b}) = \bm{a} \cdot \bm{b} = \sum_{i=1}^{d} a_i b_i $$

注意: ベクトルが正規化されている場合、内積はコサイン類似度と等価です。

$$ \text{if } \|\bm{a}\| = \|\bm{b}\| = 1: \quad \bm{a} \cdot \bm{b} = \text{cos\_sim}(\bm{a}, \bm{b}) $$

マンハッタン距離(L1距離)

$$ d_{\text{manh}}(\bm{a}, \bm{b}) = \|\bm{a} – \bm{b}\|_1 = \sum_{i=1}^{d} |a_i – b_i| $$

Pythonでの実装

import numpy as np
from numpy.linalg import norm


def cosine_similarity(a, b):
    """コサイン類似度"""
    return np.dot(a, b) / (norm(a) * norm(b))


def euclidean_distance(a, b):
    """ユークリッド距離"""
    return norm(a - b)


def dot_product(a, b):
    """内積"""
    return np.dot(a, b)


def manhattan_distance(a, b):
    """マンハッタン距離"""
    return np.sum(np.abs(a - b))


# バッチ処理版(効率的)
def batch_cosine_similarity(query, documents):
    """
    1つのクエリと複数のドキュメント間のコサイン類似度

    Args:
        query: (embed_dim,) クエリベクトル
        documents: (n_docs, embed_dim) ドキュメントベクトル
    Returns:
        similarities: (n_docs,) 類似度スコア
    """
    query_norm = query / norm(query)
    doc_norms = documents / norm(documents, axis=1, keepdims=True)
    return np.dot(doc_norms, query_norm)

テキスト埋め込みの取得

Sentence Transformers

from sentence_transformers import SentenceTransformer


# モデルの読み込み
model = SentenceTransformer('all-MiniLM-L6-v2')

# テキストの埋め込み
sentences = [
    "機械学習は人工知能の一分野です。",
    "深層学習はニューラルネットワークを使います。",
    "今日は天気が良いです。"
]

embeddings = model.encode(sentences)
print(f"Embedding shape: {embeddings.shape}")  # (3, 384)

# 類似度の計算
from sentence_transformers import util
cos_sim = util.cos_sim(embeddings[0], embeddings[1])
print(f"Similarity between 1 and 2: {cos_sim.item():.4f}")

OpenAI Embeddings

import openai


def get_openai_embedding(text, model="text-embedding-3-small"):
    """OpenAI APIで埋め込みを取得"""
    response = openai.embeddings.create(
        input=text,
        model=model
    )
    return response.data[0].embedding


# 使用例
text = "機械学習の基礎を学ぶ"
embedding = get_openai_embedding(text)
print(f"Embedding dimension: {len(embedding)}")  # 1536

カスタム埋め込み(BERT)

import torch
from transformers import AutoTokenizer, AutoModel


class BERTEmbedder:
    """BERTによるテキスト埋め込み"""

    def __init__(self, model_name='bert-base-uncased'):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.model.eval()

    def mean_pooling(self, model_output, attention_mask):
        """トークン埋め込みの平均プーリング"""
        token_embeddings = model_output.last_hidden_state
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(
            token_embeddings.size()
        ).float()
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
        return sum_embeddings / sum_mask

    def encode(self, texts, batch_size=32):
        """テキストリストを埋め込みに変換"""
        all_embeddings = []

        for i in range(0, len(texts), batch_size):
            batch_texts = texts[i:i + batch_size]
            encoded = self.tokenizer(
                batch_texts,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors='pt'
            )

            with torch.no_grad():
                outputs = self.model(**encoded)

            embeddings = self.mean_pooling(outputs, encoded['attention_mask'])
            # L2正規化
            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
            all_embeddings.append(embeddings)

        return torch.cat(all_embeddings, dim=0).numpy()


# 使用例
embedder = BERTEmbedder()
texts = ["Hello world", "Machine learning is great"]
embeddings = embedder.encode(texts)
print(f"Shape: {embeddings.shape}")  # (2, 768)

セマンティック検索の実装

基本的な実装

import numpy as np
from sentence_transformers import SentenceTransformer


class SemanticSearch:
    """セマンティック検索エンジン"""

    def __init__(self, model_name='all-MiniLM-L6-v2'):
        self.model = SentenceTransformer(model_name)
        self.documents = []
        self.embeddings = None

    def add_documents(self, documents):
        """ドキュメントを追加"""
        self.documents.extend(documents)
        new_embeddings = self.model.encode(documents)

        if self.embeddings is None:
            self.embeddings = new_embeddings
        else:
            self.embeddings = np.vstack([self.embeddings, new_embeddings])

    def search(self, query, top_k=5):
        """
        クエリに類似したドキュメントを検索

        Args:
            query: 検索クエリ
            top_k: 返す結果の数
        Returns:
            list of (document, score) tuples
        """
        query_embedding = self.model.encode([query])[0]

        # コサイン類似度を計算
        similarities = np.dot(self.embeddings, query_embedding)

        # 上位k件を取得
        top_indices = np.argsort(similarities)[::-1][:top_k]

        results = [
            (self.documents[i], similarities[i])
            for i in top_indices
        ]
        return results


# 使用例
search_engine = SemanticSearch()

documents = [
    "機械学習は人工知能の一分野で、データからパターンを学習します。",
    "深層学習はニューラルネットワークを多層にしたものです。",
    "自然言語処理は人間の言語をコンピュータで処理する技術です。",
    "強化学習は試行錯誤を通じて最適な行動を学習します。",
    "コンピュータビジョンは画像や動画を理解する分野です。"
]

search_engine.add_documents(documents)

query = "AIが画像を認識する方法"
results = search_engine.search(query, top_k=3)

print(f"Query: {query}\n")
for doc, score in results:
    print(f"Score: {score:.4f}")
    print(f"  {doc}\n")

近似最近傍探索(ANN)

大規模なデータセットでは、全てのベクトルとの類似度を計算する(ブルートフォース)のは非効率です。近似最近傍探索(Approximate Nearest Neighbor)を使用します。

FAISS

Facebookが開発した高速なベクトル検索ライブラリです。

import faiss
import numpy as np


class FAISSIndex:
    """FAISSを使った高速ベクトル検索"""

    def __init__(self, dimension, index_type='flat'):
        """
        Args:
            dimension: 埋め込み次元
            index_type: 'flat' (正確), 'ivf' (高速), 'hnsw' (高速)
        """
        self.dimension = dimension
        if index_type == 'flat':
            # 正確だが遅い
            self.index = faiss.IndexFlatIP(dimension)  # 内積
        elif index_type == 'ivf':
            # IVF (Inverted File) インデックス
            quantizer = faiss.IndexFlatIP(dimension)
            self.index = faiss.IndexIVFFlat(quantizer, dimension, 100)
            self.needs_training = True
        elif index_type == 'hnsw':
            # HNSW (Hierarchical Navigable Small World)
            self.index = faiss.IndexHNSWFlat(dimension, 32)
        else:
            raise ValueError(f"Unknown index type: {index_type}")

        self.index_type = index_type
        self.needs_training = index_type == 'ivf'

    def add(self, embeddings):
        """埋め込みをインデックスに追加"""
        # 正規化(コサイン類似度のため)
        faiss.normalize_L2(embeddings)

        if self.needs_training and not self.index.is_trained:
            self.index.train(embeddings)

        self.index.add(embeddings)

    def search(self, query, k=5):
        """
        上位k件の類似ベクトルを検索

        Returns:
            distances: (n_queries, k) 類似度スコア
            indices: (n_queries, k) インデックス
        """
        faiss.normalize_L2(query)
        distances, indices = self.index.search(query, k)
        return distances, indices


# 使用例
dimension = 384
n_vectors = 100000

# ダミーデータ
vectors = np.random.randn(n_vectors, dimension).astype('float32')
query = np.random.randn(1, dimension).astype('float32')

# インデックス作成
index = FAISSIndex(dimension, index_type='flat')
index.add(vectors)

# 検索
distances, indices = index.search(query, k=5)
print(f"Top 5 indices: {indices[0]}")
print(f"Top 5 distances: {distances[0]}")

Annoy

Spotifyが開発したANNライブラリです。

from annoy import AnnoyIndex


def build_annoy_index(embeddings, num_trees=10):
    """Annoyインデックスを構築"""
    dimension = embeddings.shape[1]
    index = AnnoyIndex(dimension, 'angular')  # コサイン類似度

    for i, embedding in enumerate(embeddings):
        index.add_item(i, embedding)

    index.build(num_trees)
    return index


def search_annoy(index, query, k=5):
    """Annoyで検索"""
    indices, distances = index.get_nns_by_vector(
        query, k, include_distances=True
    )
    # angularは距離なので、類似度に変換
    similarities = [1 - d**2 / 2 for d in distances]
    return indices, similarities

ベクトルデータベース

主要なベクトルデータベース

DB 特徴 用途
Pinecone フルマネージド、スケーラブル プロダクション
Weaviate オープンソース、GraphQL 中規模
Milvus 高性能、分散対応 大規模
Chroma シンプル、組み込み型 開発・小規模
Qdrant Rust製、高速 中〜大規模

Chromaの使用例

import chromadb
from chromadb.utils import embedding_functions


# クライアント作成
client = chromadb.Client()

# 埋め込み関数
embedding_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="all-MiniLM-L6-v2"
)

# コレクション作成
collection = client.create_collection(
    name="documents",
    embedding_function=embedding_fn
)

# ドキュメント追加
documents = [
    "機械学習は人工知能の一分野です。",
    "深層学習はニューラルネットワークを使います。",
    "自然言語処理はテキストを扱います。"
]

collection.add(
    documents=documents,
    ids=[f"doc_{i}" for i in range(len(documents))]
)

# 検索
results = collection.query(
    query_texts=["AIの技術について"],
    n_results=2
)

print(results)

RAG(Retrieval-Augmented Generation)

埋め込みベクトルと類似度検索は、RAGの中核を担います。

class SimpleRAG:
    """シンプルなRAG実装"""

    def __init__(self, embedding_model, llm_client):
        self.embedding_model = embedding_model
        self.llm_client = llm_client
        self.documents = []
        self.embeddings = None

    def add_documents(self, documents):
        """ドキュメントを追加"""
        self.documents.extend(documents)
        new_embeddings = self.embedding_model.encode(documents)
        if self.embeddings is None:
            self.embeddings = new_embeddings
        else:
            self.embeddings = np.vstack([self.embeddings, new_embeddings])

    def retrieve(self, query, top_k=3):
        """関連ドキュメントを検索"""
        query_embedding = self.embedding_model.encode([query])[0]
        similarities = np.dot(self.embeddings, query_embedding)
        top_indices = np.argsort(similarities)[::-1][:top_k]
        return [self.documents[i] for i in top_indices]

    def generate(self, query):
        """RAGで回答を生成"""
        # 関連ドキュメントを検索
        relevant_docs = self.retrieve(query)

        # プロンプトを構築
        context = "\n".join(relevant_docs)
        prompt = f"""以下の情報を参考に、質問に答えてください。

情報:
{context}

質問: {query}

回答:"""

        # LLMで生成
        response = self.llm_client.generate(prompt)
        return response

まとめ

本記事では、埋め込みベクトルと類似度検索について解説しました。

  • 埋め込みベクトル: テキストを高次元ベクトル空間に射影し、意味的類似性を捉える
  • 類似度指標: コサイン類似度が最も広く使われる。ユークリッド距離、内積も用途により選択
  • テキスト埋め込み: Sentence Transformers、OpenAI API、BERTなどで取得可能
  • 近似最近傍探索: FAISS、Annoyなどで大規模データを高速検索
  • ベクトルデータベース: Pinecone、Chroma、Milvusなどでスケーラブルに管理
  • RAG: 検索と生成を組み合わせた強力なアーキテクチャ

埋め込みベクトルと類似度検索は、セマンティック検索、レコメンデーション、RAGなど、現代のAIアプリケーションの基盤となる重要な技術です。

次のステップとして、以下の記事も参考にしてください。