DiskANN - дисковое решение ANNS с высоким показателем Recall и QPS на миллиардных массивах данных
Ченгминг Ли, инженер по исследованиям и разработкам компании Zilliz, окончил Юго-Восточный университет со степенью магистра в области компьютерных наук. В настоящее время он занимается проблемами ANNS на высокоразмерных данных, включая решения на основе графов и квантования.
"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-ядерным процессором, достигая 5000 QPS (запросов в секунду) при более чем 95 % recall@1 и средней задержке менее 3 мс.
Авторы
Сухас Джаярам Субраманья: бывший сотрудник индийского исследовательского института Microsoft, докторант CMU. Основные исследовательские интересы - высокопроизводительные вычисления и алгоритмы машинного обучения для крупномасштабных данных.
Девврит: Аспирант-исследователь в Техасском университете в Остине. Его научные интересы - теоретическая информатика, машинное обучение и глубокое обучение.
Рохан Кадекоди: Докторант Техасского университета. Направление его исследований - системы и системы хранения данных, в основном включающие постоянные хранилища, файловые системы и kV-хранилища.
Равишанкар Кришасвами (Ravishankar Krishaswamy): Главный научный сотрудник индийского исследовательского института Microsoft. Доктор CMU. Направление исследований - алгоритмы аппроксимации на основе графов и кластеризации.
Харша Вардхан Симхадри: главный научный сотрудник индийского исследовательского института Microsoft. Доктор CMU. В прошлом изучал параллельные алгоритмы и системы времени выполнения. Сейчас его основная работа - разработка новых алгоритмов и написание моделей программирования.
Мотивы
Большинство основных алгоритмов ANNS идут на компромисс между производительностью построения индекса, производительностью поиска и запоминанием. Графовые алгоритмы, такие как HNSW и NSG, являются современными методами с точки зрения производительности поиска и запоминания. Поскольку метод индексирования на основе графов занимает слишком много памяти, относительно сложно индексировать и искать в больших массивах данных на одной машине с ограниченными ресурсами памяти.
Многие приложения требуют быстрой реакции ANNS на основе евклидова расстояния на миллиардных массивах данных. Ниже приведены два основных решения:
- Инвертированный индекс + квантование: кластеризация набора данных на M разделов и сжатие набора данных с помощью схем квантования, таких как PQ (Product Quantization). Это решение дает низкий показатель recall из-за потери точности, вызванной сжатием данных. Увеличение topk помогает улучшить отзыв, но при этом QPS соответственно падает.
- Разделить и проиндексировать: разделить набор данных на несколько разрозненных осколков и построить индекс в памяти для каждого осколка. При поступлении запросов поиск будет выполняться по индексам каждого осколка, а результаты будут возвращаться после объединения. Такое решение приводит к чрезмерному увеличению масштаба набора данных, а значит, требуется больше машин из-за ограничения ресурсов памяти на одной машине, что приводит к низкому QPS.
Оба вышеупомянутых решения ограничены объемом памяти одной машины. В данной статье предлагается разработать механизм индексирования на SSD для решения этой проблемы. Задача индексирования на SSD состоит в том, чтобы уменьшить количество случайных обращений к диску и количество запросов на обращение к диску.
Материалы
В данной статье представлена схема SSD-резидентной ANNS под названием DiskANN, которая может эффективно поддерживать поиск в больших массивах данных. Эта схема основана на графовом алгоритме, представленном в данной работе: Vamana. Вклад данной работы заключается в следующем:
- DiskANN может индексировать и искать в миллиардном наборе данных с более чем 100 измерениями на одной машине с 64 ГБ оперативной памяти, обеспечивая более 95 % recall@1 при задержках менее 5 миллисекунд.
- Для минимизации количества обращений к диску был предложен новый алгоритм на основе графов под названием Vamana с меньшим радиусом поиска, чем у NSG и HNSW.
- Vamana может работать в памяти, и его производительность не ниже, чем у NSG и HNSW.
- Меньшие индексы Vamana, построенные на пересекающихся разделах большого набора данных, могут быть объединены в один граф без потери связности.
- Vamana можно комбинировать со схемами квантования, такими как PQ. Структура графа и исходные данные хранятся на диске, а сжатые данные - в памяти.
Vamana
Этот алгоритм похож на идею NSG[2][4] (для тех, кто не понимает NSG, обратитесь к ссылке [2], а если вы не хотите читать статьи, можете обратиться к ссылке [4]). Их основное отличие заключается в стратегии обрезки. Если быть точным, то в стратегию обрезки NSG добавлен переключатель alpha. Основная идея стратегии обрезки NSG заключается в том, чтобы выбор соседей целевой точки был как можно более разнообразным. Если новый сосед находится ближе к соседу целевой точки, чем целевая точка, нам не нужно добавлять эту точку в набор соседних точек. Другими словами, для каждого соседа целевой точки не может быть других соседних точек в радиусе окружения dist (целевая точка, соседняя точка). Эта стратегия обрезки эффективно контролирует степень out-degree графа и является относительно радикальной. Она уменьшает объем памяти, занимаемый индексом, повышает скорость поиска, но при этом снижает точность поиска. Стратегия обрезки Вамана заключается в свободном управлении масштабом обрезки с помощью параметра alpha. Принцип работы заключается в умножении dist (соседней точки, точки-кандидата) в условии обрезки на параметр alpha (не менее 1). Только когда dist (целевая точка, точка-кандидат) больше увеличенного опорного расстояния, принимается стратегия обрезки, увеличивающая допустимость взаимного исключения между соседями целевой точки.
Процесс индексирования в Vamana относительно прост:
- Инициализируем случайный граф;
- Вычисление начальной точки, которая аналогична навигационной точке NSG. Сначала находим глобальный центроид, а затем находим точку, ближайшую к глобальному центроиду, в качестве навигационной точки. Разница между Vamana и NSG заключается в том, что входные данные NSG уже представляют собой граф ближайших соседей, поэтому пользователи могут просто выполнить приблизительный поиск ближайших соседей в точке центроида непосредственно на исходном графе соседей. Однако Vamana инициализирует случайный граф ближайших соседей, поэтому пользователи не могут проводить приблизительный поиск непосредственно на случайном графе. Необходимо провести глобальное сравнение, чтобы получить навигационную точку в качестве начальной точки последующих итераций. Цель этой точки - минимизировать средний радиус поиска;
- Выполнить приближенный поиск ближайших соседей для каждой точки на основе инициализированного случайного графа соседей и начальной точки поиска, определенной на шаге 2, сделать все точки на пути поиска наборами соседей-кандидатов и выполнить стратегию обрезания краев, используя альфа = 1. Как и в NSG, выбор набора точек на пути поиска, начиная с навигационной точки, в качестве набора соседей-кандидатов приведет к увеличению некоторых длинных ребер и эффективному уменьшению радиуса поиска.
- Настройте альфа > 1 (в статье рекомендуется 1,2) и повторите шаг 3. Если шаг 3 основан на случайном графе ближайших соседей, то после первой итерации граф оказывается некачественным. Поэтому требуется еще одна итерация для улучшения качества графа, что очень важно для коэффициента запоминания.
В данной статье сравниваются три графовых индекса - Vamana, NSG и HNSW. С точки зрения производительности индексирования и запросов Vamana и NSG относительно близки, и оба немного превосходят HNSW. Данные см. в разделе "Эксперимент" ниже.
2.png
Для визуализации процесса построения индекса Vamana в статье приводится график, на котором 200 двумерных точек используются для моделирования двух раундов итерации. В первом ряду для обрезки краев используется альфа = 1. Видно, что стратегия обрезки относительно радикальна, и обрезается большое количество ребер. После увеличения значения alpha и ослабления условий обрезки многие ребра, очевидно, добавляются обратно. В итоговый граф добавляется довольно много длинных ребер. Это позволяет эффективно уменьшить радиус поиска.
DiskANN
Персональный компьютер с объемом памяти всего 64 Гб не сможет вместить даже миллиард необработанных данных, не говоря уже о построенном на них индексе. Перед нами стоят две задачи: 1. Как проиндексировать столь масштабный набор данных при ограниченных ресурсах памяти? 2. Как рассчитать расстояние при поиске, если исходные данные не могут быть загружены в память?
В статье предложены следующие решения:
- Для первой задачи: сначала разделить данные на k кластеров с помощью k-средних, а затем распределить каждую точку по ближайшим i кластерам. Обычно для числа i достаточно 2. Построить индекс Вамана на основе памяти для каждого кластера и, наконец, объединить k индексов Вамана в один.
- Для второй задачи: построение индекса на оригинальных векторах и запрос сжатых векторов. Построение индексов на оригинальном векторе обеспечивает качество графа, а сжатый вектор может быть загружен в память для грубого поиска. Хотя поиск по сжатым векторам может привести к потере точности, общее направление будет верным, если качество графа достаточно высокое. Окончательный результат расстояния будет вычислен с использованием исходного вектора.
Расположение индексов DiskANN аналогично общим индексам графов. Набор соседей каждой точки и исходные векторные данные хранятся вместе. Это позволяет лучше использовать локальность данных.
Как упоминалось ранее, если индексные данные хранятся на твердотельном накопителе, то для обеспечения низкой задержки поиска необходимо максимально сократить количество обращений к диску, а также количество запросов на чтение и запись. Поэтому DiskANN предлагает две стратегии оптимизации:
- Cache hotspot: кэшировать все точки в пределах C прыжков от начальной точки в памяти. Значение C лучше установить в пределах от 3 до 4.
- Лучевой поиск: Проще говоря, это предварительная загрузка информации о соседях. При поиске точки p соседнюю точку p нужно загрузить с диска, если ее нет в памяти. Поскольку операция произвольного доступа к SSD занимает примерно столько же времени, сколько операция доступа к одному сектору SSD, за один раз можно загрузить информацию о соседях W точек, к которым нет доступа. Значение W не может быть слишком большим или маленьким. Большое значение W приведет к трате вычислительных ресурсов и пропускной способности SSD, а малое - к увеличению задержки поиска.
Эксперимент
Эксперимент состоит из трех групп:
Сравнение индексов на основе памяти: Vamana VS. NSG VS. HNSW
Наборы данных: SIFT1M (128 измерений), GIST1M (960 измерений), DEEP1M (96 измерений) и 1M набор данных, случайно выбранный из DEEP1B.
Параметры индекса (во всех наборах данных используется один и тот же набор параметров):
HNSW:M = 128, efc = 512.
Вамана: R = 70, L = 75, альфа = 1,2.
NSG: R = 60, L = 70, C = 500.
Параметры поиска в статье не приводятся, что может соответствовать параметрам индексации. Для выбора параметров NSG, упомянутые в статье, основаны на параметрах, перечисленных в репозитории GitHub для NSG, чтобы выбрать группу с лучшей производительностью. Vamana и NSG относительно близки, поэтому и параметры задаются близкие. Однако причина выбора параметров HNSW не указана. Мы считаем, что параметр M в HNSW задан относительно большим. Это может привести к менее убедительному сравнению между индексами на основе графов, если их out-degrees не установлены на одном уровне.
При указанных выше параметрах индексирования время индексирования Vamana, HNSW и NSG составляет 129 с, 219 с и 480 с соответственно. Время индексации NSG включает время построения начального графа соседей с помощью EFANN [3].
Кривая Recall-QPS:
3.png
Из рисунка 3 видно, что Vamana демонстрирует отличную производительность на трех наборах данных, аналогичную NSG и немного лучшую, чем HNSW.
Сравнение радиуса поиска:
Из рисунка 2.c видно, что Vamana имеет самый короткий средний путь поиска при одинаковом коэффициенте запоминания по сравнению с NSG и HNSW.
Сравнение между единовременно созданным индексом и большим объединенным индексом
Набор данных: SIFT1B
Параметры однократно построенного индекса: L = 50, R = 128, alpha = 1,2. После работы в течение 2 дней на машине с 1800 Гб DDR3 пиковый объем памяти составляет около 1100 Гб, а средний показатель out-degree - 113,9.
Процедура индексирования на основе слияния:
- Обучить 40 кластеров на наборе данных с помощью kmeans;
- Каждая точка распределяется по 2 ближайшим кластерам;
- Построить индекс Вамана с L = 50, R = 64 и альфа = 1,2 для каждого кластера;
- Слияние индексов каждого кластера.
Этот индекс сгенерировал индекс объемом 384 ГБ со средним значением out-of-degree 92,1. Этот индекс работал в течение 5 дней на машине с 64 ГБ DDR4.
Результаты сравнения выглядят следующим образом (рис. 2a):
4.png
Выводы:
- Однократно построенный индекс значительно лучше, чем индекс на основе слияния;
- Индекс на основе слияния также превосходен;
- Схема индексирования на основе слияния также применима к набору данных 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 года исходный код решения DiskANN был выложен в открытый доступ.
Далее в основном представлены процесс индексирования и процесс поиска.
Построение индекса
Существует 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 в ГБ.
M: Порог памяти при построении индекса. Определяет размер фрагмента, в ГБ.
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-файле соответствует вектору в файле данных. Идентификаторы получаются путем нумерации каждого вектора исходных данных от 0 до n-1. Идентификатор относительно важен и связан с объединением.
- Глобально равномерно дискретизируйте обучающее множество с частотой дискретизации 1500000 / n;
- Инициализируем num_parts = 3. Итерация от 3:
- Выполните num_parts-means++ на обучающем множестве на шаге i;
- Используйте частоту дискретизации 0,01 для равномерной глобальной выборки тестового набора и разделите его на 2 ближайших кластера;
- Подсчитайте количество точек в каждом кластере и разделите его на частоту выборки, чтобы оценить количество точек в каждом кластере;
- Оценить объем памяти, требуемый самым большим кластером на шаге 3 в соответствии с размером индекса Вамана, если он не превышает параметр M, перейти к шагу iii, в противном случае num_parts ++ вернуться к шагу 2;
- Разделите исходный набор данных на файлы группы num_parts, каждая группа файлов включает файлы фрагментированных данных и ID-файлы, соответствующие фрагментированным данным.
- Создайте индексы Vamana отдельно для всех фрагментов на шаге a и сохраните их на диск;
- merge_shards: объединить num_parts шардов Vamana в глобальный индекс:
- Считайте ID-файл фрагментов num_parts в idmap. Эта idmap эквивалентна созданию прямого отображения фрагмент->id;
- Установите обратное отображение id->фрагменты в соответствии с idmap и узнайте, в каких двух фрагментах находится каждый вектор;
- Используйте считыватель с кэшем 1 ГБ для открытия индексов Vamana фрагмента num_parts, и используйте записывающее устройство с кэшем 1 ГБ для открытия выходного файла, готового к слиянию;
- Поместите num_parts навигационных точек индекса Vamana в файл центральных точек, который будет использоваться при поиске;
- Начать объединение по ID от малого к большому, прочитать набор соседних точек каждого исходного вектора в каждом фрагменте по очереди в соответствии с обратным отображением, дедуплицировать, перетасовать, усечь и записать в выходной файл. Поскольку нарезка изначально была глобально упорядочена, теперь объединение также упорядочено, поэтому ID в итоговом индексе и ID исходных данных соответствуют один к одному.
- Удалите временные файлы, включая файлы фрагментов, индексы фрагментов и файлы идентификаторов фрагментов.
Создайте_дисковую_разметку: Глобальный индекс, созданный на шаге 6, имеет только компактную таблицу смежности. Этот шаг предназначен для выравнивания индекса. Таблица смежности и исходные данные хранятся вместе. При поиске загружайте таблицу смежности и считывайте исходный вектор вместе для точного расчета расстояния. Существует также концепция SECTOR, размер которого по умолчанию составляет 4096. Каждый SECTOR содержит только 4096 / node_size фрагментов векторной информации. node_size = размер одного вектора + размер таблицы смежности одного узла.
Наконец, сделайте глобальную равномерную выборку размером 150000 / n, сохраните ее и используйте для разминки при поиске.
Поиск
Имеется 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, данных кодовой книги, начальной точки поиска и других данных, а также чтение метаданных индекса.
- Используя набор данных, отобранных во время индексирования, выполните cached_beam_search, подсчитайте время доступа к каждой точке и загрузите в кэш точки num_nodes_to_cache с наибольшей частотой доступа.
- По умолчанию выполняется операция WARMUP. Как и на шаге 2, этот набор данных также используется для выполнения cached_beam_search.
- В соответствии с количеством заданных параметров L, для каждого L будет снова выполнен cached_beam_search с набором запросов, и будет выведена статистика, такая как recall rate и QPS. Процесс прогрева и получения статистики горячих точек не учитывается во времени запроса.
О cached_beam_search:
- Поиск ближайшего кандидата к точке запроса от начальной точки кандидата. При этом используется расстояние PQ, а начальная точка добавляется в очередь поиска.
- Начать поиск:
- В очереди поиска есть не более чем beam_width + 2 непосещенных точки. Если эти точки есть в кэше, добавьте их в очередь попадания в кэш. Если они не попали в кэш, добавьте их в очередь промахов. Убедитесь, что размер очереди промахов не превышает beam_width.
- Отправьте асинхронные запросы на доступ к диску точкам в очереди промахов.
- Для точек, попавших в кэш, используйте исходные данные и данные запроса для расчета точного расстояния, добавьте их в очередь результатов, а затем используйте PQ для расчета расстояния до соседних точек, которые не были посещены перед добавлением в очередь поиска. Длина очереди поиска ограничивается параметрами.
- Обработка кэшированных пропущенных точек на шаге a аналогична шагу c.
- Когда очередь поиска пуста, поиск завершается, и возвращается очередь результатов topk.
Подведение итогов
Несмотря на то что работа получилась довольно объемной, в целом она превосходна. Идеи статьи и кода понятны: разделить несколько пересекающихся ведер с помощью k-means, затем разделить ведра для построения картографического индекса и, наконец, объединить индексы, что является относительно новой идеей. Что касается графового индекса Vamana, основанного на памяти, то это, по сути, случайно инициализированная версия NSG, которая может управлять гранулярностью обрезки. При запросах он полностью использует кэш + конвейер, покрывает часть времени ио и улучшает QPS. Однако, согласно статье, даже если состояние машины не является экстраординарным, время обучения занимает до 5 дней, а удобство использования относительно низкое. Оптимизация обучения, безусловно, необходима в будущем. С точки зрения кода, его качество относительно высокое и может быть непосредственно использовано в производственных условиях.
Ссылки
- Suhas Jayaram Subramanya, Fnu Devvrit, Harsha Vardhan Simhadri, Ravishankar Krishnawamy, Rohan Kadekodi. DiskANN: быстрый точный поиск ближайших соседей по миллиарду точек на одном узле. NeurIPS 2019.
- [Cong Fu, Chao Xiang, Changxu Wang, and Deng Cai. Быстрый приближенный поиск ближайших соседей с помощью навигационных расходящихся графов. 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: быстрая библиотека для поиска ANN и построения KNN-графов.
- Поисковая система для искусственного интеллекта:高维数据检索工业级解决方案
Try Managed Milvus for Free
Zilliz Cloud is hassle-free, powered by Milvus and 10x faster.
Get StartedLike the article? Spread the word