Re-ranking の理論と実装 — Cross-Encoder と ColBERT で検索精度を引き上げる

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)

class BiEncoderSim:
    """Bi-Encoderのシミュレーション。"""

    def __init__(self, dim=128):
        self.dim = dim

    def encode_query(self, query_tokens):
        """クエリをトークン埋め込みの平均で1ベクトルに集約する。"""
        return np.mean(query_tokens, axis=0)

    def encode_doc(self, doc_tokens):
        """文書をトークン埋め込みの平均で1ベクトルに集約する。"""
        return np.mean(doc_tokens, axis=0)

    def score(self, query_tokens, doc_tokens):
        q_vec = self.encode_query(query_tokens)
        d_vec = self.encode_doc(doc_tokens)
        # コサイン類似度
        return np.dot(q_vec, d_vec) / (np.linalg.norm(q_vec) * np.linalg.norm(d_vec) + 1e-10)


class CrossEncoderSim:
    """Cross-Encoderのシミュレーション。
    トークン間の全ペア相互作用を考慮する。"""

    def __init__(self, dim=128):
        self.dim = dim

    def score(self, query_tokens, doc_tokens):
        # 全トークン間の類似度行列
        sim_matrix = query_tokens @ doc_tokens.T
        norms_q = np.linalg.norm(query_tokens, axis=1, keepdims=True)
        norms_d = np.linalg.norm(doc_tokens, axis=1, keepdims=True)
        sim_matrix = sim_matrix / (norms_q @ norms_d.T + 1e-10)

        # 全ペアの平均（相互注意の近似）
        return np.mean(sim_matrix)


class ColBERTSim:
    """ColBERT（MaxSim）のシミュレーション。"""

    def __init__(self, dim=128):
        self.dim = dim

    def score(self, query_tokens, doc_tokens):
        # 類似度行列
        sim_matrix = query_tokens @ doc_tokens.T
        norms_q = np.linalg.norm(query_tokens, axis=1, keepdims=True)
        norms_d = np.linalg.norm(doc_tokens, axis=1, keepdims=True)
        sim_matrix = sim_matrix / (norms_q @ norms_d.T + 1e-10)

        # MaxSim: クエリの各トークンに対する最大類似度の和
        max_sims = np.max(sim_matrix, axis=1)  # (L_q,)
        return np.sum(max_sims)


def generate_synthetic_data(n_docs=200, n_relevant=20, dim=128, q_len=5, d_len=30):
    """関連文書と非関連文書を含む合成データを生成する。"""
    # クエリ
    query_tokens = np.random.randn(q_len, dim) * 0.5

    docs = []
    labels = []

    for i in range(n_docs):
        if i < n_relevant:
            # 関連文書: クエリトークンに類似したトークンを含む
            doc = np.random.randn(d_len, dim) * 0.5
            # ランダムな位置にクエリに類似するトークンを埋め込む
            for j in range(min(q_len, d_len)):
                pos = np.random.randint(0, d_len)
                doc[pos] = query_tokens[j] + np.random.randn(dim) * 0.1
            docs.append(doc)
            labels.append(1)
        else:
            # 非関連文書
            doc = np.random.randn(d_len, dim) * 0.5
            docs.append(doc)
            labels.append(0)

    return query_tokens, docs, np.array(labels)


# データ生成
query, docs, labels = generate_synthetic_data(n_docs=200, n_relevant=20)

# 各手法でスコアリング
bi_enc = BiEncoderSim()
cross_enc = CrossEncoderSim()
colbert = ColBERTSim()

scores_bi = [bi_enc.score(query, doc) for doc in docs]
scores_cross = [cross_enc.score(query, doc) for doc in docs]
scores_colbert = [colbert.score(query, doc) for doc in docs]

# Recall@k の計算
def recall_at_k(scores, labels, k):
    top_k_indices = np.argsort(scores)[::-1][:k]
    relevant_in_top_k = np.sum(labels[top_k_indices])
    total_relevant = np.sum(labels)
    return relevant_in_top_k / total_relevant

k_values = [5, 10, 20, 50, 100]
recall_bi = [recall_at_k(scores_bi, labels, k) for k in k_values]
recall_cross = [recall_at_k(scores_cross, labels, k) for k in k_values]
recall_colbert = [recall_at_k(scores_colbert, labels, k) for k in k_values]

# 可視化
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# 左: スコア分布
for ax_idx, (scores, name, color) in enumerate([
    (scores_bi, 'Bi-Encoder', '#00d4ff'),
    (scores_cross, 'Cross-Encoder', '#ffa726'),
    (scores_colbert, 'ColBERT (MaxSim)', '#66bb6a'),
]):
    relevant_scores = [s for s, l in zip(scores, labels) if l == 1]
    irrelevant_scores = [s for s, l in zip(scores, labels) if l == 0]

    axes[0].hist(relevant_scores, bins=20, alpha=0.3, color=color, label=f'{name} (relevant)')

axes[0].set_xlabel('Score')
axes[0].set_ylabel('Count')
axes[0].set_title('Score Distribution (Relevant Documents)')
axes[0].legend(fontsize=8)
axes[0].grid(True, alpha=0.3)

# 中央: Recall@k
axes[1].plot(k_values, recall_bi, 'o-', label='Bi-Encoder', color='#00d4ff', markersize=8)
axes[1].plot(k_values, recall_cross, 's-', label='Cross-Encoder', color='#ffa726', markersize=8)
axes[1].plot(k_values, recall_colbert, '^-', label='ColBERT', color='#66bb6a', markersize=8)
axes[1].set_xlabel('k')
axes[1].set_ylabel('Recall@k')
axes[1].set_title('Recall@k Comparison')
axes[1].legend()
axes[1].grid(True, alpha=0.3)

# 右: 二段階パイプラインのシミュレーション
# Stage 1: Bi-Encoderで上位N1件を取得
# Stage 2: Cross-Encoderで再ランキング
n1_values = [10, 20, 30, 50, 100, 200]
rerank_recalls = []

for n1 in n1_values:
    # Stage 1
    top_n1_indices = np.argsort(scores_bi)[::-1][:n1]

    # Stage 2: Cross-Encoderで再スコアリング
    rerank_scores = [(idx, cross_enc.score(query, docs[idx])) for idx in top_n1_indices]
    rerank_scores.sort(key=lambda x: x[1], reverse=True)

    # Top-10のRecall
    top_10_indices = [idx for idx, _ in rerank_scores[:10]]
    recall = sum(labels[idx] for idx in top_10_indices) / min(20, 10)
    rerank_recalls.append(recall)

axes[2].plot(n1_values, rerank_recalls, 'o-', color='#ab47bc', markersize=8,
             label='Bi-Encoder + Cross-Encoder rerank')
axes[2].axhline(y=recall_at_k(scores_bi, labels, 10), color='#00d4ff',
                linestyle='--', alpha=0.7, label='Bi-Encoder only @10')
axes[2].axhline(y=recall_at_k(scores_cross, labels, 10), color='#ffa726',
                linestyle='--', alpha=0.7, label='Cross-Encoder only @10')
axes[2].set_xlabel('Stage 1 Candidates (N1)')
axes[2].set_ylabel('Recall@10 after Reranking')
axes[2].set_title('Two-Stage Pipeline: Effect of N1')
axes[2].legend(fontsize=8)
axes[2].grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig('reranking_comparison.png', dpi=150, bbox_inches='tight')
plt.show()