なぜ全件検索は高コストなのか?: 詳細解説とJava実装例(RAG)

1. 概要

RAG (Retrieval-Augmented Generation) パイプラインにおいて、retrieval ステップは通常、質問のベクトルを何十万ものpassageのベクトルと照合するプロセスです。最もシンプルな方法——各ベクトルを順序立てて比較する——は常に正確な結果をもたらしますが、処理コストがデータサイズに応じて線形に増加します。本記事では、なぜ全件比較がコストが高いのか、Inverted-File Index (IVF) がクラスタリングを通じてコストを削減する仕組み、そして近似結果を受け入れるために支払う代価について説明します。この記事を読み終えたら、JavaでシンプルなフラットインデックスとIVFインデックスを自分で書くことができ、速度とrecallの間のトレードオフを測定でき、いつapproximate searchを使うべき(あるいは使うべきでない)かが理解できるようになります。

2. 問題: フラットインデックスとO(n)の計算量

2.1. フラットサーチの仕組み

フラットインデックス(またはbrute-force search)は、クエリを全てのベクトルとデータセット内で比較し、スコア(通常はcosine similarityまたはdot product)が最も高いベクトルを選び出します。中間的なフィルタリングステップはありません。

public record Passage(String id, float[] vector) {}

public class FlatVectorIndex {

    private final List<Passage> passages;

    public FlatVectorIndex(List<Passage> passages) {
        this.passages = passages;
    }

    public List<Passage> search(float[] query, int topK) {
        return passages.stream()
                .sorted(Comparator.comparingDouble(p -> -cosine(query, p.vector())))
                .limit(topK)
                .toList();
    }

    private double cosine(float[] a, float[] b) {
        double dot = 0, na = 0, nb = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            na += a[i] * a[i];
            nb += b[i] * b[i];
        }
        return dot / (Math.sqrt(na) * Math.sqrt(nb));
    }
}

searchを呼び出すたびに、stream()内の暗黙的なループが全てのpassagesに対してcosineを計算する必要があります。データセットにn個のベクトルがあり、各ベクトルがd次元である場合、各クエリのコストはO(n × d)です。passageの数が2倍になると、比較操作も2倍になります——これが「O(n)」の意味です。

注記: フラットサーチは「exact」な結果を提供します——つまり、選択した距離公式に従って完全に正確です。しかし、「exact」であることは、passageが意味論的に正しい答えを含んでいるという意味ではありません。幾何学的な正確性と意味論的な正確性は異なるものです。

2.2. データサイズに応じたコスト増加の測定

public static void main(String[] args) {
    for (int n : List.of(10_000, 100_000)) {
        var passages = TestData.generatePassages(n, 128, 42);
        var index = new FlatVectorIndex(passages);
        float[] query = TestData.randomVector(128, 7);

        long start = System.nanoTime();
        index.search(query, 5);
        long elapsedMs = (System.nanoTime() - start) / 1_000_000;

        System.out.printf("n=%d -> %d ms%n", n, elapsedMs);
    }
}

開発マシン上での実行例(実際の数値はハードウェアに依存しますが、傾向は変わりません):

n=10000  -> 18 ms
n=100000 -> 176 ms

データが10倍になると、実行時間も約10倍になります——O(n)の予測通りです。100万passageを使用し、毎秒複数のクエリを処理する場合、このコストは急速に蓄積されます。これが、フラットインデックスがproductionのRAGシステムでスケールしない理由です。

3. 解決策: Inverted-File Index (IVF)

3.1. クラスタリングの考え方

IVFは全てのベクトルと比較しないことでコストを削減します。代わりに:

1. インデックスを構築する際、クラスタリング(通常はk-means)を実行して、互いに近いベクトルのグループを代表する「centroid」を見つけます。
2. データセット内の各ベクトルは、最も近いcentroidを持つグループに割り当てられます。
3. 検索時には、centroidがクエリに最も近い複数のグループのみを選択し、それらのグループ内のベクトルのみと比較します。

100万のpassageが1000個のグループに分割されていると想像してください。10個のグループのみをスキャンすれば、100万個全てではなく約10,000個のpassageと比較します——コストが約100倍削減されます。

3.2. nprobeパラメータ

スキャンされるグループの数をnprobeと呼びます。より多くのグループをスキャンすると、recall(正しい結果を見つける確率)が向上する傾向がありますが、コストも増加します。これは、productionで直接調整するパラメータです。

public class IvfIndex {

    private final List<float[]> centroids;
    private final Map<Integer, List<Passage>> buckets = new HashMap<>();

    public IvfIndex(List<Passage> passages, List<float[]> centroids) {
        this.centroids = centroids;
        for (Passage p : passages) {
            int cluster = nearestCentroidIndex(p.vector());
            buckets.computeIfAbsent(cluster, k -> new ArrayList<>()).add(p);
        }
    }

    public List<Passage> search(float[] query, int topK, int nprobe) {
        List<Integer> clusters = nearestCentroidIndices(query, nprobe);
        return clusters.stream()
                .flatMap(c -> buckets.getOrDefault(c, List.of()).stream())
                .sorted(Comparator.comparingDouble(p -> -cosine(query, p.vector())))
                .limit(topK)
                .toList();
    }

    private int nearestCentroidIndex(float[] v) {
        return nearestCentroidIndices(v, 1).get(0);
    }

    private List<Integer> nearestCentroidIndices(float[] v, int k) {
        return IntStream.range(0, centroids.size())
                .boxed()
                .sorted(Comparator.comparingDouble(i -> -cosine(v, centroids.get(i))))
                .limit(k)
                .toList();
    }

    private double cosine(float[] a, float[] b) {
        double dot = 0, na = 0, nb = 0;
        for (int i = 0; i < a.length; i++) {
            dot += a[i] * b[i];
            na += a[i] * a[i];
            nb += b[i] * b[i];
        }
        return dot / (Math.sqrt(na) * Math.sqrt(nb));
    }
}

nearestCentroidIndicesはクエリに最も近いk個のcentroidを見つけます——ここがnprobeがsearch(query, topK, nprobe)を呼び出すときに活躍する場所です。インデックスの構築(constructor)は1回実行されますが、検索は毎回実行されます——これがIVFがクエリ時にはるかに高速である理由です。

注記: 上記の例では、事前に用意されたcentroidを使用してコードを短くしています(データからサンプルされた)。実際には、高品質なクラスタを得るために真のk-means(centroidを更新する複数のイテレーション)を実行する必要があります——FAISS、Milvus、pgvectorなどのライブラリがこの作業を既に行っています。

3.3. pgvectorを使った実装例

PostgreSQLと拡張機能pgvectorを使用している場合、IVFとnprobeの概念がSQL内で直接現れます:

CREATE INDEX ON passages USING ivfflat (embedding vector_cosine_ops)
WITH (lists = 1000);

SET ivfflat.probes = 10;

listsはインデックス構築時のクラスタ数に相当し、probesはクエリ時のnprobeに相当します。k-meansを自分で書く必要はありません——トレードオフを理解して適切な数値を選択するだけです。

4. トレードオフ: recall vs. 速度

IVFはapproximate nearest neighbor (ANN)の一形態です。絶対的な精度と引き換えに速度を手に入れます。問題は:クラスタが数学的な距離から形成され、人間が文書を分類する方法に従った「きれいな」主題の境界からは形成されていないということです。

RAGの具体例:「フライト遅延時の払い戻し」に関する条項は、「出張費用ポリシー」のクラスタよりも保険文書のベクトルクラスタに近い場合があります——ただし、主題的には、従業員は出張文書で検索することになります。nprobeが「出張費用ポリシー」クラスタのみをスキャンする場合、正しい答えがデータに存在していても、完全に見落とされる可能性があります。

4.1. nprobeに応じたrecallの低下テスト

class IvfIndexTest {

    @Test
    void nprobeThapCoTheLamMatKetQuaDung() {
        var passages = TestData.generatePassages(2000, 64, 42);
        var centroids = TestData.sampleCentroids(passages, 40, 42);
        var ivf = new IvfIndex(passages, centroids);
        var flat = new FlatVectorIndex(passages);

        float[] query = passages.get(1337).vector();

        var exact = flat.search(query, 5);
        var approxFewProbes = ivf.search(query, 5, 1);
        var approxManyProbes = ivf.search(query, 5, 40);

        assertThat(approxManyProbes).containsExactlyInAnyOrderElementsOf(exact);
        assertThat(approxFewProbes).isNotEqualTo(approxManyProbes);
    }
}

approxManyProbesは全40個のクラスタをスキャンするため、結果はフラットサーチと完全に一致します——これはIVFにロジックエラーがないことを確認するためのベースラインです。approxFewProbesは1つのクラスタのみをスキャンし、最後のアサーションは結果が正確な場合と異なることを確認します——nprobeが低すぎる場合のrecall喪失のリスクを示しています。

exact            = [p1337, p0892, p0456, p1890, p0023]
approxFewProbes  = [p1337, p0456, p0901, p0234]
approxManyProbes = [p1337, p0892, p0456, p1890, p0023]

nprobeが低いと、インデックスは単に少ない結果を返すだけでなく、「クラスタの外れ」に位置する重要なpassageを見落とす可能性があります——上の保険/出張費用ポリシーの例と同じです。

5. フラットサーチとIVFの使い分け

• フラットサーチを使用する場合: データセットが小さい場合(数万passage以下)、またはaudit/compliance目的で絶対的な精度が必要な場合、またはプロトタイプ段階にあり、さらにパラメータを調整したくない場合。
• IVF(またはHNSW)を使用する場合: データセットが大きい場合(数十万から数百万passage)、クエリレイテンシが重大な要件である場合(チャットボット、リアルタイム検索)、そして速度と引き換えにrecallが100%未満であることを受け入れる場合。

注記: スピードベンチマークが良いという理由だけでnprobeを過度に低く設定しないでください。RAGシステムの実際の質問セットでrecallを測定してください——重要なコンテキストが欠けた回答は、数十ms遅いクエリより悪いです。

productionの一般的なパターン:IVF(またはHNSW)を使用して許容可能なrecallで迅速にtop-100 candidatesを取得し、その後rerank ステップ(cross-encoderまたは小さな100 candidate セット上のフラットサーチ)を使用して失った精度を補うことです。このアプローチはANNの速度を保ちながら、最終結果の品質を向上させます。

6. 結論

フラットサーチは完全に正確ですが、passageの数に応じてコストが線形に増加します——RAGが数百万の文書を提供する必要がある場合、スケールしません。IVFはクラスタリングとnprobeパラメータでこの問題を解決し、recallの一部を速度と引き換えにします。代償は、ベクトルクラスタが人間が理解する主題の境界と一致しないため、重要なpassageが「迷い出た」スキャンされないクラスタに位置することがあるということです。このトレードオフを明確に理解することで、正しいインデックスを選択し、正しくパラメータを調整でき、精度を補うためのrerank ステップが必要な時期が分かります。