HNSW: строим маленькие миры для быстрого поиска

Предположим, у нас есть отсортированный связанный список, в котором мы хотим найти элемент. Единственный способ сделать это — пройтись по всему списку за \(O(n)\) операций. Если бы это был массив — можно было бы использовать бинарный поиск, но в списках у нас нет возможности обращаться по индексу.

Если мы хотим искать быстрее — можно воспользоваться списком с пропусками (Skip list). Это вероятностная структура данных, позволяющая находить элемент в отсортированном списке в среднем за \(O(log(n))\), используя \(O(n)\) памяти.

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

Для построения такой структуры необходимо пройтись по исходному списку и каждый элемент с вероятностью \(P\)(обычно — 0.5) поднять на следующий уровень. Затем перейти на второй уровень и повторить процедуру. В результате каждый последующий уровень будет меньше предыдущего в \(P\) раз. Таким образом, всего уровней \(O(log(n))\).

Вероятность того, что элемент окажется на уровне \(k\), равна \(P^k\)(геометрическое распределение). Таким образом, среднее число уровней, где есть произвольный элемент, равно \(1+P^1+P^2+…=\frac{1}{1-P}\). Это является константой, следовательно, асимптотика памяти остаётся равной \(O(n)\).

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

Поскольку всего уровней — \(O(log(n))\), и на каждом уровне мы делаем ограниченное число шагов вправо (в среднем — константное), итоговая сложность остаётся в среднем \(O(log(n))\).

Интуиция: представьте, что вы находитесь в Москве на станции «ВДНХ» и хотите добраться до центра Тимирязевского парка. Сначала вы доберётесь на метро достаточно близко до необходимой точки, потом перейдёте на менее разреженный уровень — уровень автобусов, и приблизитесь к парку, затем дойдёте пешком (самый не разреженный уровень). Также и в HNSW: при поиске точки мы сначала используем самый неточный и разреженный уровень, потом, поняв, что приблизиться на нём мы не сможем, переходим на более точный уровень, пока не доберёмся до самого точного из них.

Для начала давайте разберёмся, как происходит поиск в HNSW. Также, как и в skip list, начинаем с верхнего уровня, где entry-point может быть либо фиксированной, либо случайной.

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

Для нулевого уровня используется beam search — best-first поиск с ограничением числа одновременно хранящихся вершин. На каждом шаге выбирается наиболее близкая к запросу вершина из текущих кандидатов, рассматриваются её соседи, которые добавляются во множество кандидатов. После —отбор лучших по расстоянию. При этом в каждый момент времени хранится не более efSearch лучших кандидатов: при добавлении новых худшие отбрасываются. Поиск продолжается, пока среди кандидатов остаются вершины, потенциально более близкие к запросу, после чего из накопленных вершин выбираются ближайшие.

Рассмотрим beam search на конкретном примерн. Пусть efSearch=2:

Третий шаг: достаём H из Candidates. Никто из соседей не может улучшить W (так как они либо уже там, либо дальше). Аналогично и с G.

Последний шаг: поскольку кандидатов больше нет, выбирается ближайшая точка (H) из W.

При добавлении нового элемента в HNSW из экспоненциального распределения сэмплируется максимальный уровень для вершины — \(L\). Напомним, если точка на уровне \(L\) — она есть на всех уровнях ниже. На практике экспоненциальное распределение ограничивают сверху, чтобы число уровней не было слишком большим.

Следующий шаг — жадный поиск из entry point (такой же, как и при операции поиска элемента) на всех уровнях до L.

Оказавшись на уровне \(L\), мы вставляем элемент, то есть связываем точку с максимум M ближайшими точками, не нарушающими структуру графа. Для этого делаем beam search с параметром efConstruction, аналогичным efSearch из поиска, и находим кандидатов на связывание. Далее для каждого кандидата, начиная с ближайшего, пока не совершится максимум M вставок, добавляем связь, если выполняется следующее условие.

Diversity heuristic: \(∀s∈S:d(q,c)<d(s,c)\), где q — вставляемая точка, с — текущий кандидат, а S — множество выбранных соседей. То есть мы добавляем точку, только если она ближе к искомой, чем к выбранным соседям. Это обеспечивает разнообразие соседей, предотвращая избыточные связи между близкими друг к другу точками и улучшая связность, «ориентируемость» по графу. Без этого условия соседи часто лежали бы в одном плотном направлении, а за счёт эвристики обеспечивается геоометрически разнообразное покрытие.

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

Данная процедура производится на всех уровнях ниже \(L\).

Построение HNSW заключается в последовательном добавлении всех элементов (обычно в случайном порядке) в структуру с использованием описанной процедуры вставки. Также существует способ построения (Batch construction), который позволяет строить граф с лучшей структурой и улучшает качество поиска, но требует знания всего датасета заранее и дорогого ресурса по времени построения.

Давайте вернёмся к гиперпараметрам и обсудим их влияние на качество и память.

Для начала разберёмся, как мы вообще оцениваем качество ANN-поиска. В отличие от точного поиска, мы «жертвуем» точностью ради скорости, поэтому основной критерий — насколько хорошо алгоритм приближает истинных ближайших соседей.

Чаще всего используют метрику recall@k:

\(Recall@k=\frac{| \text{true neighbors} \cap \text{found neighbors} |}{k}\)

То есть мы сравниваем набор из \(k\) найденных HNSW-соседей с истинными \(k\) ближайшими соседями и смотрим долю совпадений.

Помимо качества нас также интересует latency — время ответа на запрос и использование памяти. Таким образом, HNSW существует в треугольнике компромиссов: качество ↔ скорость ↔ память.

Самый главный параметр — \(M\)— улучшает связность графа и увеличивает recall@k, при этом замедляет поиск и вставку, так как нам нужно просмотреть больше соседей, и увеличивает использование памяти. Обычно лежит в пределах от 8 до 64.

Параметр efSearch управляет шириной поиска (beam width) во время запроса и увеличивает recall@k, поскольку мы рассматриваем больше соседей и увеличиваем шанс найти настоящего соседа. При этом считаем больше расстояний и замедляем поиск. Обычно в ~2-4 раза больше M.

Параметр efConstruction управляет качеством построения графа. Его увеличение приводит к улучшению структуры графа при построении и увеличивает recall. При этом он замедляет вставку. Обычно лежит в пределах от 20 до 200.

Интуиция:

• \(M\) — ёмкость графа (количество связей для хранения). Обычно лежит от 8 до 64.
• \(efConstruction\) — насколько хорошо мы заполняем эту ёмкость. Обычно от 100 до 200.
• \(efSearch\) — насколько тщательно мы исследуем граф при поиске.

Главная магия HNSW в том, что он превращает пространство в small-world граф, где большинство связей — локальные: лежат между близкими точками. Но есть дальние прыжки, которые сильно сокращают путь. Если в обычном KNN-графе мы движемся только локально (как в BFS), следовательно, долго, то в HNSW мы «пропрыгиваем» большие расстояния на высоких уровнях и уточняем результат на нижних. А за счёт использования экспоненциального распределения мы обеспечиваем то, что на высоких уровнях лежит сильно меньше точек, чем на низких.

Также благодаря diversity heuristic граф имеет связи в разных направлениях, что делает жадный поиск эффективным: при поиске появляется возможность выходить за пределы локально плотных областей и быстрее находить пути к удалённым частям пространства. Без этой эвристики граф деградировал бы в набор сильно кластеризованных подграфов с преобладанием локальных связей, из-за чего поиск часто «застревал» бы в локально плотных регионах и хуже исследовал глобальную структуру данных.

Наивная реализация заключается в хранении структуры такого вида:

class Node:
    vector
    neighbors

Но этот вариант слабоват тем, что он не обеспечивает последовательный доступ к памяти: объекты разбросаны, переходы по указателям приводят к частым cache miss и замедляют поиск. Поэтому на практике используют хранение в виде плотных массивов (contiguous memory), что улучшает локальность данных и позволяет эффективно использовать кэш и SIMD-инструкции.