Пять способов индексации векторов для RAG, протестированных на миллионе эмбеддингов Википедии

Система RAG существует или погибает в зависимости от способности к поиску, причем поиск становится всё сложнее по мере увеличения объёма корпуса данных. В этом обзоре с помощью реального эксперимента сравниваются пять способов индексации векторов, поэтому выбор метода основывается на цифрах, а не на привычке.

Что такое индексация векторов и почему она важна

Индексация векторов — это способ организации большого набора векторов так, чтобы можно было быстро находить похожие из них, не сравнивая запрос со всеми векторами. В реляционных базах данных существуют аналогичные методы — бакетирование и порядок Z, которые сокращают объём данных, сканируемых при выполнении запроса. Базы данных векторов используют специализированные структуры для уменьшения объёма пространства поиска в том же духе.

Точный поиск с низкой задержкой является основой любой масштабируемой системы RAG. Крупные сервисы уже зависят от этого: Amazon выполняет поиск в каталогах, содержащих миллионы товаров, Netflix — среди тысяч фильмов, адаптированных под глобальную аудиторию, а X — среди миллиардов постов и аккаунтов. Объем данных сейчас достигает терабайт в день, поэтому масштабирование не может быть вопросом, решаемым позже. Оно должно быть частью архитектуры с самого начала проектирования.

Каждый из приведенных ниже методов оценивается по трём критериям:

• Задержка: скорость, с которой индекс возвращает топ-K результатов.
• Точность поиска: степень совпадения результатов с результатами точного поиска.

Сложность технического обслуживания: степень сложности поддержания актуальности индекса. Исходные данные приложения типа RAG редко остаются неизменными, и их может потребоваться обновлять каждые 24 часа, поэтому индекс должен позволять осуществлять обновления в течение этого времени без значительных технических усилий.

Настройка теста

В эксперименте использовались первые миллион строки набора данных CohereLabs/wikipedia-2023-11-embed-multilingual-v3 — объем, достаточный для выявления различий между методами. Запросы эмбеддинговались с помощью модели mxbai-embed-large, которая генерирует 1024-мерные векторы, совместимые с эмбеддингами данного набора данных. Индексы создавались в PostgreSQL с использованием pgvector.

Метод 1: плоский поиск

Поиск с прямым сравнением, также называемый методом грубой силы, является самым простым вариантом. Вычисляется степень сходства между запросом и каждым вектором, после чего сохраняются наилучшие результаты.

• Задержка: самая высокая среди всех методов, поскольку сравниваются все векторы. Время ростет линейно с количеством векторов.
• Точность: наивысшая, так как поиск происходит точно и возвращаются истинно ближайшие векторы.
• Техническое обслуживание: не требуется, поскольку индекс отсутствует.

Средняя задержка в ходе тестирования составила 6,25 секунды. Крупный ритейлер немедленно отклонит такой вариант, как и небольшая стартап-компания.

Метод 2: IVFFlat

Inverted File Flat — это метод приближенного поиска ближайшего соседа. Он использует алгоритм K-Means для разделения векторного пространства на кластеры. Вместо сканирования всего пространства он находит центроиды, наиболее близкие к запросу, и ищет только внутри этих кластеров, применяя точные расчеты расстояний внутри них.

Алгоритм работает следующим образом:

Query → closest centroids → search selected clusters → Top-K

В pgvector для управления этим процессом используются два параметра. Количество кластеров, называемых списками, определяется заранее; в данном тесте было использовано 1 000 списков. Во время выполнения запроса второй параметр — пробные элементы — определяет, сколько из этих списков будут просмотрены. При значении пробных элементов 10 сканируются десять списков, центроиды которых находятся ближе всего к запросу.

Время отклика было приемлемым. При наличии миллиона векторов и 1 000 кластеров запросы обрабатывались не более чем за секунду.

Точность оказалась низкой по двум причинам. Во-первых, существует проблема границ: близость измеряется относительно центроидов, а не относительно самих векторов, поэтому некоторые из самых близких векторов остаются незамеченными, если они принадлежат кластеру, центроид которого случайно находится дальше от запроса. Во-вторых, параметр пробных элементов ограничивает поиск фиксированным количеством кластеров, поэтому релевантные векторы в непросмотренных кластерах никогда не учитываются.

Увеличение количества зондов с 10 до 500 привело к увеличению задержки, хотя в этом тесте это было лишь незначительно. Однако невозможно заранее определить, позволят ли 10, 100, 500 или 900 зондов достичь наилучшего уровня восстановления данных, поскольку настоящие близкие элементы могут находиться в кластере, который был пропущен. Использование всех 1 000 зондов сводит на нет всю цель, так как это превращает процесс поиска обратно в метод грубой силы.

Эксплуатационные характеристики также были плохими. Создание индекса для более чем миллиона векторов заняло примерно 2,37 минуты, что считается приемлемым. Проблемы возникают при обновлении: индексы, основанные на центроидах, могут потребовать пересчета или повторной индексации каждый раз, когда поступает новая партия данных. При миллиардах векторов такие технические затраты могут оказаться неоправданными.

Для данного случая использования соотношение между задержкой и точностью было просто неприемлемым. Следующий вопрос заключается в том, может ли другой метод снизить задержку и при этом обеспечить лучшее восстановление данных.

Метод 3: HNSW

Хорошей аналогией может служить почтальон, ищущий дом. Имея адрес, он не ищет на каждой улице — он движется по главным автострадам в нужный регион, где связей мало, но они покрывают большие расстояния. Затем он переходит на основные дороги, где сеть более плотная. Наконец он идет по местным улицам, где сеть наиболее плотная, пока не доберется до конкретной двери.

HNSW, сокращение от hierarchical navigable small world, работает по тому же принципу. Он начинает с редких, дальнобойных связей на верхних уровнях и переходит к всё более плотным связям на нижних уровнях, пока не найдет ближайшие векторы. Если иерархия дорог хорошо организована, поиск направляется прямо к цели, а качественная структура обеспечивает высокую точность поиска.

Время отклика снова было приемлемым — не более одной секунды при обработке миллиона векторов.

Точность была высокой, но с оговорками, связанными с параметрами сборки. Индекс создается с использованием предложения следующей структуры:

    CREATE INDEX wikipedia_embeddings_hnsw_idx
    ON wikipedia_embeddings
    USING hnsw (emb vector_cosine_ops)
    WITH (
        m = 16,
        ef_construction = 64
    );

Это утверждение позволяет создать индекс на основе корпуса. Важны три параметра:

• m — это количество соседей, к которым может быть связан каждый вектор в графе. Он определяет связность и плотность графа.
• ef_construction — это количество потенциальных соседей, учитываемых при создании индекса. Он контролирует объем пространства поиска во время формирования индекса.
• ef_search ограничивает пространство поиска во время выполнения запроса, так как определяет, сколько кандидатов будут рассматриваться при ответе.

Как граф выбирает свои m связей? Для миллиона векторов у каждого из них есть 999,999 потенциальных соседей. Механизм создания индекса не выбирает их случайным образом. Он оценивает кандидатов по расстоянию и с использованием эвристики выбора соседей, сохраняя небольшой набор, примерно из 16 элементов, которые полезны для навигации по графу.

Почему это обеспечивает более высокую точность?

1. Небольшое количество полезных соседей для каждого вектора образует «автомагистрали» графа.
2. Именно эти связи выбираются для того, чтобы граф оставался связным и удобным для навигации.
3. Во время поиска алгоритм шаг за шагом следует по этим связям к векторам, более похожим на запрос.
4. Поскольку граф предоставляет эффективные пути к релевантной области, существует больше шансов найти истинных ближайших соседей.
5. Таким образом, тщательная структуризация и навигация повышают точность, избегая одновременно полного сравнения.

Короче говоря, граф создается целенаправленно, а не случайным образом. Значения m и ef_construction определяют способ его построения, что в свою очередь влияет на уровень воспроизведения результатов.

Техническое обслуживание — здесь HNSW показывает свои слабые стороны. Создание индекса для более чем миллиона векторов заняло удивительные 37 минут. При обновлении, если первоначальный набор данных велик и хорошо отражает общее распределение, новые данные можно продолжать добавлять без пересоздания индекса. Однако по мере увеличения объема новых данных по сравнению с исходным набором распределение может измениться, граф становится менее оптимальным, и в конечном итоге возникает необходимость в пересоздании. Это происходит реже, чем при использовании метода IVF, но все же случается.

Благодаря принципиальной структуре графа этот метод является лучшим вариантом по сравнению с IVF, несмотря на дополнительные технические усилия. Баланс между задержкой и точностью результатов является приемлемым, хотя и не выдающимся.

Метод 4: BM25 плюс поиск векторов

Иногда очевидный ответ является правильным. Пока внимание сосредотачивается на всё более сложных подходах, уже тридцать лет существует простой и эффективный вариант. BM25 — это алгоритм лексического ранжирования, который оценивает документы по частоте вхождения терминов запроса, редкости этих терминов в всем корпусе и длине документа.

Механизм работы:

1. Документы разбиваются на термины, а обратный индекс соотносит каждый термин с документами, в которых он встречается.
2. Частота термина показывает, как часто он встречается в конкретном документе.
3. Обратная частота документа показывает, насколько редок или распространён термин во всей коллекции.
4. Термин, часто встречающийся в одном документе, но редко в целом, получает более высокий приоритет.
5. BM25 сочетает эти показатели с длиной документа для оценки релевантности.

Во время поиска термины запроса ищутся в обратном индексе, оцениваются по соответствующим документам, и возвращаются документы с наивысшими оценками.

Сама формула не выглядит красиво, но идея проста: найти документы, наиболее релевантные словам в запросе.

После получения этих документов объем пространства поиска существенно сокращается. Плоский векторный поиск в этом небольшом наборе затем находит наилучшие результаты. Это и есть гибридный способ поиска.

Одно предостережение: на первом этапе ищутся точные ключевые слова, поэтому ошибка пользователя в написании может негативно сказаться на результатах. Шаг нормализации с использованием ИИ, который исправляет опечатки и варианты написания перед началом поиска, легко решает эту проблему.

Задержка. Векторная часть отвечает в ожидаемые миллисекунды, но лексическая часть в данной конфигурации работала примерно 12–14 секунд. Кэширование может решить эту проблему, хотя первый запрос по-прежнему будет медленным. Здесь стоит отметить кэширование, поскольку оно особенно эффективно для BM25: точные документы находятся путем сопоставления ключевых слов. Для других методов индексации кэширование должно быть семантическим, что требует наличия порога сходства и может немного снизить точность поиска.

Техническое обслуживание. Не требуется, поскольку векторный индекс не создается.

На практике для большинства ИИ-продуктов с количеством документов до примерно 50 000 гибридный метод поиска с использованием BM25 остается оптимальным вариантом, обеспечивая хороший баланс между задержкой и точностью в таком масштабе.

Метод 5: DiskANN и граф Vamana

Этот метод более абстрактен, как и сами эмбеддинги. Процесс состоит из восьми шагов:

1. Начните с большого набора векторов, которые требуют эффективного поиска.
2. Создайте граф, связав каждый вектор с тщательно выбранным набором соседей. Это и есть граф Вамана.
3. Храните граф на диске, а не в памяти, чтобы можно было обрабатывать наборы данных, значительно превышающие объем RAM.
4. При поступлении запроса начните с точки входа в граф.
5. Следуйте связям между векторами, пошагово двигаясь к областям, близким к запросу.
6. По мере продвижения поиска загружайте необходимые узлы с SSD.
7. После исследования достаточной части графа соберите небольшой набор перспективных кандидатов.
8. Сравните этих кандидатов более точно и верните K наиболее близких вариантов.

Ключевой момент заключается в том, что граф навигации хранится на SSD, а не в RAM, поэтому ограничение по размеру менее значимо, чем затраты памяти.

Сначала выполните фильтрацию метаданных

Одна техника заслуживает особого упоминания и должна рассматриваться в первую очередь среди всех вышеупомянутых методов. Фильтрация метаданных основана на том же принципе, что и группировка в реляционных базах данных: сужается набор кандидатов с помощью таких атрибутов, как категория, дата или владелец, ещё до начала векторного поиска.

Какой метод выбрать

• Метод грубой силы: небольшие и средние наборы данных, до примерно 20 000 документов. Максимальная точность при приемлемой задержке.
• BM25 в сочетании с векторным поиском: средние наборы данных, примерно от 20 000 до 100 000 документов. Хороший баланс между лексическим и семантическим поиском.
• HNSW: крупные наборы данных с миллионами векторов или более. Хорошие показатели задержки и точности, но с более высокой сложностью настройки.
• DiskANN и Vamana: когда оперативная память становится дорогой. SSD обеспечивает очень большую емкость, поэтому ограничением является не размер, а сама память.
• IVFFlat: большие наборы данных, состоящие примерно из от одного до десяти миллионов векторов, при условии хорошей кластеризации данных. В противном случае снижается точность поиска.