DiskANN: дисковое решение ANNS с высокой полнотой и высоким QPS на наборе данных миллиардного масштаба
«DiskANN: Fast Accurate Billion-point Nearest Neighbor Search on a Single Node» — это статья, опубликованная на NeurIPS в 2019 году. В статье представлен современный метод построения индекса и поиска по набору данных миллиардного масштаба с использованием одной машины всего с 64 ГБ оперативной памяти и достаточно большим SSD. Более того, он удовлетворяет трём требованиям ANNS (Approximate Nearest Neighbor Search) для крупномасштабного набора данных: высокая полнота, низкая задержка и высокая плотность (число узлов на одной машине). Этот метод строит графовый индекс на наборе данных миллиардного масштаба SIFT-1B с использованием одной машины с 64 ГБ оперативной памяти и 16-ядерным CPU, достигая 5000 QPS (запросов в секунду) при recall@1 более 95 % и средней задержке ниже 3 мс.
Авторы
Сухас Джаярам Субраманья: Бывший сотрудник Microsoft India Research Institute, докторант CMU. Основные исследовательские интересы — высокопроизводительные вычисления и алгоритмы машинного обучения для крупномасштабных данных.
Devvrit: ассистент-исследователь магистратуры в The University of Texas at Austin. Его исследовательские интересы — теоретическая информатика, машинное обучение и глубокое обучение.
Рохан Кадекоди: докторант University of Texas. Его направление исследований — системы и хранение данных, в основном включая энергонезависимое хранение, файловые системы и kV-хранилища.
Равишанкар Кришасвами: главный исследователь Microsoft Indian Research Institute. Доктор CMU. Направление исследований — аппроксимационные алгоритмы на основе графов и кластеризации.
Харша Вардхан Симхадри: главный исследователь Microsoft Indian Research Institute. Доктор CMU. Ранее он изучал параллельные алгоритмы и системы времени выполнения. Сейчас его основная работа — разработка новых алгоритмов и написание моделей программирования.
Мотивация
Большинство основных алгоритмов ANNS идут на определённые компромиссы между производительностью построения индекса, производительностью поиска и полнотой. Графовые алгоритмы, такие как HNSW и NSG, в настоящее время являются передовыми методами с точки зрения производительности поиска и полноты. Поскольку метод графового индексирования, размещённого в памяти, занимает слишком много памяти, индексировать и выполнять поиск по крупномасштабному набору данных на одной машине с ограниченными ресурсами памяти относительно сложно.
Многим приложениям требуются быстрые ответы ANNS на основе евклидова расстояния для наборов данных миллиардного масштаба. Ниже приведены два основных решения:
Инвертированный индекс + квантование: кластеризовать набор данных на M разделов и сжать набор данных с использованием схем квантования, таких как PQ (Product Quantization). Это решение даёт низкую полноту из-за потери точности, вызванной сжатием данных. Увеличение topk помогает повысить полноту, при этом QPS будет соответственно снижаться.
Разделение и индексирование: разделить набор данных на несколько непересекающихся шардов и построить индекс в памяти для каждого шарда. Когда поступают запросы, поиск будет выполняться по индексам каждого шарда, а результаты будут возвращаться после объединения. Это решение вызывает чрезмерное расширение масштаба набора данных, и поэтому требуется больше машин из-за ограничения ресурсов памяти на одной машине, что приводит к низкому QPS.
Оба упомянутых выше решения ограничены ограничением памяти одной машины. В этой статье предлагается проект механизма индексирования, размещённого на SSD, для решения этой проблемы. Задача индексирования, размещённого на SSD, состоит в уменьшении числа случайных обращений к диску и числа запросов на доступ к диску.
Вклад
В этой статье представлена схема ANNS, размещённая на SSD, под названием DiskANN, которая может эффективно поддерживать поиск по крупномасштабным наборам данных. Эта схема основана на графовом алгоритме, представленном в этой статье: Vamana. Вклад этой статьи включает:
DiskANN может индексировать и выполнять поиск по набору данных миллиардного масштаба с более чем 100 измерениями на одной машине с 64 ГБ RAM, обеспечивая recall@1 более 95% с задержками менее 5 миллисекунд.
Новый графовый алгоритм под названием Vamana с меньшим радиусом поиска, чем у NSG и HNSW, был предложен для минимизации количества обращений к диску.
Vamana может работать в памяти, и его производительность не ниже, чем у NSG и HNSW.
Меньшие индексы Vamana, построенные на перекрывающихся разделах большого набора данных, можно объединить в один граф без потери связности.
Vamana можно комбинировать со схемами квантования, такими как PQ. Структура графа и исходные данные хранятся на диске, а сжатые данные — в памяти.
Vamana
Этот алгоритм похож на идею NSG[2][4] (тем, кто не понимает NSG, см. Reference [2], а если вы не хотите читать статьи, можете обратиться к Reference [4]). Их главное отличие заключается в стратегии отсечения. Точнее, в стратегию отсечения NSG был добавлен переключатель alpha. Основная идея стратегии отсечения NSG заключается в том, чтобы выбор соседей целевой точки был как можно более разнообразным. Если новый сосед находится ближе к соседу целевой точки, чем к самой целевой точке, нам не нужно добавлять эту точку в множество соседних точек. Другими словами, для каждого соседа целевой точки не может быть других соседних точек в пределах окружающего радиуса dist (целевая точка, соседняя точка). Эта стратегия отсечения эффективно контролирует исходящую степень графа и является относительно радикальной. Она уменьшает объем памяти, занимаемый индексом, повышает скорость поиска, но также снижает точность поиска. Стратегия отсечения Vamana позволяет свободно контролировать масштаб отсечения с помощью параметра alpha. Принцип работы заключается в умножении dist (соседняя точка, кандидатная точка) в условии отсечения на параметр alpha (не меньше 1). Только когда dist (целевая точка, определенная кандидатная точка) больше увеличенного эталонного расстояния, применяется стратегия отсечения, что повышает допустимость взаимного исключения между соседями целевой точки.
Процесс индексирования Vamana относительно прост:
Инициализировать случайный граф;
Вычислить начальную точку, которая похожа на навигационную точку NSG. Сначала находится глобальный центроид, а затем точка, ближайшая к глобальному центроиду, выбирается в качестве навигационной точки. Разница между Vamana и NSG заключается в том, что на вход NSG уже подается граф ближайших соседей, поэтому пользователи могут просто выполнить приближенный поиск ближайшего соседа для точки центроида непосредственно на исходном графе соседей. Однако Vamana инициализирует случайный граф ближайших соседей, поэтому пользователи не могут напрямую выполнять приближенный поиск на случайном графе. Им нужно выполнить глобальное сравнение, чтобы получить навигационную точку в качестве начальной точки последующих итераций. Цель этой точки — минимизировать средний радиус поиска;
Выполнить Approximate Nearest Neighbor Search для каждой точки на основе инициализированного случайного графа соседей и начальной точки поиска, определенной на шаге 2, сделать все точки на пути поиска кандидатными множествами соседей и выполнить стратегию отсечения ребер, используя alpha = 1. Аналогично NSG, выбор множества точек на пути поиска, начинающемся от навигационной точки, в качестве кандидатного множества соседей добавит некоторые длинные ребра и эффективно уменьшит радиус поиска.
Настроить alpha > 1 (в статье рекомендуется 1.2) и повторить шаг 3. Поскольку шаг 3 основан на случайном графе ближайших соседей, после первой итерации граф имеет низкое качество. Поэтому необходима еще одна итерация для улучшения качества графа, что очень важно для показателя recall.
В этой статье сравниваются три графовых индекса, то есть Vamana, NSG и HNSW. С точки зрения производительности индексирования и запросов Vamana и NSG относительно близки, и оба немного превосходят HNSW. Данные см. в разделе Experiment ниже.
Рисунок 1.
Чтобы визуализировать процесс построения индекса Vamana, в статье представлен граф, в котором 200 двумерных точек используются для имитации двух раундов итерации. В первой строке используется alpha = 1 для обрезки рёбер. Можно видеть, что стратегия обрезки относительно радикальна, и большое количество рёбер обрезается. После увеличения значения alpha и ослабления условий обрезки множество рёбер, очевидно, добавляется обратно. В итоговом графе добавлено довольно много длинных рёбер. Это может эффективно уменьшить радиус поиска.
DiskANN
Персональный компьютер всего с 64GB памяти не смог бы вместить даже миллиард единиц исходных данных, не говоря уже об индексе, построенном на них. Здесь возникают две проблемы: 1. Как индексировать такой крупномасштабный набор данных при ограниченных ресурсах памяти? 2. Как вычислять расстояние при поиске, если исходные данные не могут быть загружены в память?
В статье предложены следующие решения:
Для первой проблемы: сначала разделить данные на k кластеров с помощью k-means, а затем распределить каждую точку в ближайшие i кластеров. Как правило, для числа i достаточно 2. Построить индекс Vamana на основе памяти для каждого кластера и в итоге объединить k индексов Vamana в один.
Для второй проблемы: строить индекс по исходным векторам и запрашивать сжатые векторы. Построение индексов по исходному вектору обеспечивает качество графа, тогда как сжатый вектор может быть загружен в память для грубого поиска. Хотя поиск со сжатыми векторами может вызвать потерю точности, общее направление будет правильным, если качество графа достаточно высоко. Итоговый результат расстояния будет вычислен с использованием исходного вектора.
Структура хранения индекса DiskANN похожа на структуру обычных графовых индексов. Набор соседей каждой точки и исходные векторные данные хранятся вместе. Это позволяет лучше использовать локальность данных.
Как упоминалось ранее, если данные индекса хранятся на SSD, количество обращений к диску и запросов чтения и записи на диск должно быть максимально сокращено, чтобы обеспечить низкую задержку поиска. Поэтому DiskANN предлагает две стратегии оптимизации:
Кэширование горячих точек: кэшировать в памяти все точки в пределах C переходов от начальной точки. Значение C лучше задавать в пределах от 3 до 4.
Лучевой поиск: проще говоря, это предварительная загрузка информации о соседях. При поиске точки p соседнюю точку p необходимо загрузить с диска, если её нет в памяти. Поскольку небольшое количество операций случайного доступа к SSD занимает примерно столько же времени, сколько операция доступа к одному сектору SSD, информацию о соседях W непосещённых точек можно загружать за один раз. W нельзя задавать слишком большим или слишком маленьким. Большое W будет тратить вычислительные ресурсы и пропускную способность SSD впустую, тогда как маленькое увеличит задержку поиска.
Experiment
Эксперимент состоит из трёх групп:
Сравнение между индексами на основе памяти: Vamana VS. NSG VS. HNSW
Наборы данных: SIFT1M (128 измерений), GIST1M (960 измерений), DEEP1M (96 измерений) и набор данных 1M, случайно выбранный из DEEP1B.
Параметры индекса (все наборы данных используют один и тот же набор параметров):
HNSW：M = 128, efc = 512.
Vamana: R = 70, L = 75, alpha = 1.2.
NSG: R = 60, L = 70, C= 500.
Параметры поиска в статье не приведены, что может соответствовать параметрам индексирования. Что касается выбора параметров, параметры NSG, упомянутые в статье, основаны на параметрах, перечисленных в репозитории GitHub NSG, чтобы выбрать группу с лучшей производительностью. Vamana и NSG относительно близки, поэтому параметры также заданы близкими. Однако причина выбора параметров HNSW не указана. Мы считаем, что параметр M у HNSW задан относительно большим. Это может привести к менее убедительному сравнению между графовыми индексами, если их исходящие степени не заданы на одном уровне.
При указанных выше параметрах индексирования время индексирования Vamana, HNSW и NSG составляет 129 с, 219 с и 480 с соответственно. Время индексирования NSG включает время построения исходного графа соседей с помощью EFANN [3].
Кривая Recall-QPS:
Рисунок 2.
Из рисунка 3 видно, что Vamana демонстрирует отличную производительность на трех наборах данных, аналогичную NSG и немного лучшую, чем HNSW.
Сравнение радиуса поиска:
Из рисунка 2.c видно, что Vamana имеет самый короткий средний путь поиска при одинаковом уровне recall по сравнению с NSG и HNSW.
Сравнение между индексом, построенным за один раз, и большим объединенным индексом
Набор данных: SIFT1B
Параметры индекса, построенного за один раз: L = 50, R = 128, alpha = 1.2. После работы в течение 2 дней на машине с 1800G DDR3 пиковое потребление памяти составляет около 1100 G, а средняя исходящая степень — 113,9.
Процедура индексирования на основе объединения:
Обучить 40 кластеров на наборе данных с использованием kmeans;
Каждая точка распределяется в 2 ближайших кластера;
Построить индекс Vamana с L = 50, R = 64 и alpha = 1.2 для каждого кластера;
Объединить индексы каждого кластера.
Этот индекс сформировал индекс размером 384GB со средней исходящей степенью 92,1. Этот индекс работал 5 дней на машине с 64GB DDR4.
Результаты сравнения приведены ниже (рисунок 2a):
Рисунок 3.
В заключение:
Индекс, построенный за один раз, значительно лучше, чем индекс на основе объединения;
Индекс на основе объединения также превосходен;
Схема индексирования на основе объединения также применима к набору данных DEEP1B (рисунок 2b).
Дисковый индекс: DiskANN VS. FAISS VS. IVF-OADC+G+P
IVFOADC+G+P — это алгоритм, предложенный в источнике [5].
В этой статье DiskANN сравнивается только с IVFOADC+G+P, поскольку в источнике [5] доказано, что IVFOADC+G+P лучше, чем FAISS. Кроме того, FAISS требует ресурсов GPU, которые поддерживаются не всеми платформами.
IVF-OADC+G+P, по-видимому, представляет собой комбинацию HNSW и IVF-PQ. Он определяет кластеры с использованием HNSW и выполняет поиск, добавляя некоторые стратегии отсечения к целевому кластеру.
Результат приведен на рисунке 2a. 16 и 32 на рисунке — это размер кодовой книги. Набор данных — SIFT1B, квантованный с помощью OPQ.
Детали реализации кода
Исходный код DiskANN открыт на https://github.com/microsoft/DiskANN
В январе 2021 года исходный код дискового решения был открыт.
Ниже в основном описываются процесс индексирования и процесс поиска.
Построение индекса
Для построения индекса используется 8 параметров:
data_type: варианты включают float/int8/uint8.
data_file.bin: Исходный бинарный файл данных. Первые два целых числа в файле соответственно представляют общее количество n векторов набора данных и размерность вектора dim. Последние n * dim * sizeof(data_type) байт — это непрерывные векторные данные.
index_prefix_path: Префикс пути выходного файла. После построения индекса будет сгенерировано несколько файлов, связанных с индексом. Этот параметр является общим префиксом каталога, в котором они хранятся.
R: Максимальная исходящая степень глобального индекса.
L: Параметр L индекса Vamana, верхняя граница размера множества кандидатов.
B: Порог памяти при запросе. Он управляет размером кодовой книги PQ, в GB.
M: Порог памяти при построении индекса. Он определяет размер фрагмента, в GB.
T: Количество потоков.
Процесс индексирования (входная функция: aux_utils.cpp::build_disk_index):
Сгенерировать различные имена выходных файлов согласно index_prefix_path.
Проверка параметров.
Прочитать метаданные data_file.bin, чтобы получить n и dim. Определить число подпространств m кодовой книги PQ согласно B и n.
generate_pq_pivots: Выбрать центральную точку обучающего набора PQ с использованием коэффициента выборки p = 1500000/n равномерно, чтобы обучить PQ глобально.
generate_pq_data_from_pivots: Сгенерировать глобальную кодовую книгу PQ и сохранить центральную точку и кодовую книгу отдельно.
build_merged_vamana_index: нарезать исходный набор данных, построить индексы Vamana по сегментам и, наконец, объединить индексы в один.
partition_with_ram_budget: Определить количество фрагментов k согласно параметру M. Выполнить выборку набора данных с помощью kmeans, распределяя каждую точку в два ближайших кластера. Фрагментировать набор данных, при этом каждый фрагмент создает два файла: файл данных и файл ID. Файл ID и файл данных соответствуют друг другу, и каждый ID в файле ID соответствует вектору в файле данных. ID получаются путем нумерации каждого вектора исходных данных от 0 до n-1. ID относительно важен и связан с объединением.
Глобально равномерно выбрать обучающий набор с коэффициентом выборки 1500000 / n;
Инициализировать num_parts = 3. Итерировать с 3:
- Выполнить num_parts-means++ на обучающем наборе на шаге i;
- Использовать коэффициент выборки 0.01, чтобы равномерно глобально выбрать тестовый набор, и разделить тестовый набор на ближайшие 2 кластера;
- Подсчитать количество точек в каждом кластере и разделить его на коэффициент выборки, чтобы оценить количество точек в каждом кластере;
- Оценить память, требуемую крупнейшим кластером на шаге 3, согласно размеру индекса Vamana; если она не превышает параметр M, перейти к шагу iii, иначе num_parts ++ и вернуться к шагу 2;
Разделить исходный набор данных на num_parts групп файлов, каждая группа файлов включает фрагментированные файлы данных и файлы ID, соответствующие фрагментированным данным.
Создать индексы Vamana отдельно для всех срезов на шаге a и сохранить их на диск;
merge_shards: объединить num_parts shard Vamana в глобальный индекс:
Прочитать файл ID фрагментов num_parts в idmap. Этот idmap эквивалентен созданию прямого отображения fragment->id;
Создать обратное отображение id-> fragments согласно idmap и узнать, в каких двух фрагментах находится каждый вектор;
Использовать reader с кешем 1GB, чтобы открыть индексы Vamana срезов num_parts, и использовать writer с кешем 1GB, чтобы открыть выходной файл, готовый к объединению;
Поместить num_parts навигационных точек индекса Vamana в файл центральных точек, который будет использоваться при поиске;
Начать объединение по ID от меньшего к большему, по очереди считывать набор соседних точек каждого исходного вектора в каждом фрагменте согласно обратному отображению, удалять дубликаты, перемешивать, усекать и записывать в выходной файл. Поскольку нарезка изначально была глобально упорядочена, и теперь объединение также идет по порядку, ID в итоговом сброшенном индексе и ID исходных данных соответствуют один-к-одному.
Удалить временные файлы, включая файлы фрагментов, индексы фрагментов и файлы ID фрагментов.
7.create_disk_layout: Глобальный индекс, созданный на шаге 6, имеет только компактную таблицу смежности. Этот шаг предназначен для выравнивания индекса. Таблица смежности и исходные данные хранятся вместе. При поиске загружать таблицу смежности и считывать исходный вектор вместе для точного вычисления расстояния. Также существует понятие SECTOR, размер по умолчанию равен 4096. Каждый SECTOR содержит только 4096 / node_size элементов векторной информации. node_size = размер одного вектора + размер таблицы смежности одного узла.
8.Наконец, выполнить глобальную равномерную выборку 150000 / n, сохранить ее и использовать для warmup при поиске.
Search
Существует 10 параметров поиска:
index_type: варианты включают Float/int8/uint8, аналогично первому параметру data_type при построении индекса.
index_prefix_path: см. параметр индекса index_prefix_path.
num_nodes_to_cache: количество горячих точек кеша.
num_threads: количество потоков поиска.
beamwidth: верхний предел количества точек предварительной загрузки. Система определяет, если установлено 0.
query_file.bin: файл набора запросов.
truthset.bin: файл набора результатов, "null" означает, что набор результатов не предоставлен, программа вычисляет его самостоятельно;
K: topk;
result_output_prefix: путь для сохранения результатов поиска;
L*: Список параметров поиска. Можно добавить несколько значений. Для каждого L при поиске с разными L будет предоставлена статистическая информация.
Процесс поиска:
Загрузить связанные данные: загрузить набор запросов, данные центральных точек PQ, данные codebook, начальную точку поиска и другие данные, а также прочитать метаданные индекса.
Использовать набор данных, сэмплированный во время индексирования, чтобы выполнить cached_beam_search, подсчитать количество обращений к каждой точке и загрузить в кеш num_nodes_to_cache точек с наибольшей частотой обращений.
По умолчанию выполняется операция WARMUP. Как и на шаге 2, этот набор сэмплированных данных также используется для выполнения cached_beam_search.
В соответствии с заданным количеством параметров L каждый L будет снова обработан с помощью cached_beam_search с набором запросов, и будут выведены такие статистические данные, как recall rate и QPS. Процесс warmup и статистика hotspot data не учитываются во времени запроса.
О cached_beam_search:
Найти ближайшего к точке запроса кандидата из начальной точки-кандидата. Здесь используется расстояние PQ, а начальная точка добавляется в очередь поиска.
Начать поиск:
Из очереди поиска берется не более beam_width + 2 непосещенных точек. Если эти точки находятся в кеше, добавить их в очередь попаданий кеша. Если попадания нет, добавить их в очередь промахов. Убедиться, что размер очереди промахов не превышает beam_width.
Отправить асинхронные запросы доступа к диску для точек в очереди промахов.
Для точек, попавших в кеш, использовать исходные данные и данные запроса для вычисления точного расстояния, добавить в очередь результатов, а затем использовать PQ для вычисления расстояния до соседних точек, которые еще не были посещены, перед добавлением в очередь поиска. Длина очереди поиска ограничивается параметрами.
Обработать точки промаха кеша на шаге a, аналогично шагу c.
Когда очередь поиска пуста, поиск завершается, и возвращается topk из очереди результатов.
Резюме
Хотя это довольно объемная работа, в целом она превосходна. Идеи статьи и кода ясны: разделить несколько перекрывающихся buckets с помощью k-means, затем разделить buckets для построения индекса map и в итоге объединить индексы, что является относительно новой идеей. Что касается memory-based graph index Vamana, по сути это случайно инициализированная версия NSG, которая может управлять гранулярностью обрезки. При запросе она максимально использует cache + pipeline, скрывает часть времени io и повышает QPS. Однако, согласно статье, даже если характеристики машины не являются выдающимися, время обучения занимает до 5 дней, а удобство использования относительно низкое. Оптимизация обучения определенно необходима в будущем. С точки зрения кода качество относительно высокое, и его можно напрямую использовать в production environment.
Ссылки
[Cong Fu, Chao Xiang, Changxu Wang, and Deng Cai. Быстрый приближенный поиск ближайшего соседа с navigating spreading-out graphs. PVLDB, 12(5):461 – 474, 2019. doi: 10.14778/3303753.3303754.] (http://www.vldb.org/pvldb/vol12/p461-fu.pdf)
Cong Fu and Deng Cai. GitHub - ZJULearning/efanna: fast library for ANN search and KNN graph construction.
Search Engine For AI：Промышленное решение для поиска в высокоразмерных данных
Читать далее
Why We Built Vector Lakebase: Rethinking Unstructured Data Architecture for AI
Vector Lakebase: a unified, lake-native data foundation for AI workloads — and an answer to what happens after vector databases succeed.
Zilliz Cloud On-Demand Compute: Pay Only for What You Use
The customer case behind Zilliz Cloud On-Demand: how a $10K vector search bill came down to under $500, and the engineering changes that made it possible.
Why Teams Are Migrating from Weaviate to Zilliz Cloud — and How to Do It Seamlessly
Explore how Milvus scales for large datasets and complex queries with advanced features, and discover how to migrate from Weaviate to Zilliz Cloud.