index = nmslib.init(space='cosinesimil', method='sw-graph')
nmslib.addDataPointBatch(index, np.arange(model.wv.syn0.shape[0], dtype=np.int32), model.wv.syn0)
index.createIndex({}, print_progress=True)
start = time.time()
items = nmslib.knnQuery(index, 10, king_vec.tolist())
print(time.time() - start)
print([model.wv.index2word[idx] for idx in items])

Получаем:

    0.000545024871826
    [u'king', u'kings', u'queen', u'monarch', u'crown_prince', u'prince', u'sultan', u'ruler', u'princes', u'royal']

Рассмотрим развитие описанной выше идеи в алгоритме Hierarchical Navigable Small World (HNSW). Он во многом схож с NSW, однако теперь мы имеем дело с иерархией графов: на нулевом слое представлены все объекты, а по мере увеличения слоя — все меньшая и меньшая их подвыборка. При этом все объекты на слое $$ есть и на слое $$.

При поиске старт происходит со случайной вершины в графе верхнего слоя, там мы быстро находим близкие к запросу вершины (кандидаты) и возобновляем поиск с них на предыдущем слое.

SearchAtLayer (object q, Set[object] enterPoints, integer: M, ef, layer)
1  Set [object] visitedSet
2  priority_queue [object] candidates (closer - first), result (further - first)
3  candidates, visitedSet, result < enterPoints
7
4  repeat:
5     object c =candidates.top()
6     candidates.pop()
7     //check stop condition:
8     if d(c,q)>d(result.top(),q) do:
9        break
10    //update list of candidates:
11    for_each object e from c.friends(layer) do:
12       if e is not in visitedSet do:
13          add e to visitedSet
14          if d(e, q)< d(result.top(),q) or result.size()<ef do:
15             add e to candidates, result
16             if result.size()>ef do:
17                result.pop()
18 return best k elements from result

K-NNSearch (object query, integer: ef)
1 Set [object] tempRes, enterPoints=[enterpoint]
2 for i=maxLayer downto 1 do:
3    tempRes=SearchAtLayer (query, enterPoints, M, 1, i)
4    enterPoints =closest elements from tempRes
5 tempRes=SearchAtLayer (query, enterPoints, M, ef, 0)
6 return best K of tempRes

1.1. Pros & Cons

+ Алгоритм просто понять

+ Он показывает state-of-the-art результаты

+ Существует эффективная реализация в библиотеке nmslib

+ Небольшие дополнительные расходы памяти на хранение ребер графа

– Алгоритм не поддерживает сжатие векторных представлений, которое мы рассмотрим далее

index = nmslib.init(space='cosinesimil', method='hnsw')
nmslib.addDataPointBatch(index, np.arange(model.wv.syn0.shape[0], dtype=np.int32), model.wv.syn0)
index.createIndex({}, print_progress=True)
start = time.time()
items = nmslib.knnQuery(index, 10, king_vec.tolist())
print(time.time() - start)
print([model.wv.index2word[idx] for idx in items])

Получаем:

    0.000469923019409
    [u'king', u'kings', u'queen', u'monarch', u'crown_prince', u'prince', u'sultan', u'ruler', u'princes', u'Prince_Paras']

2. FAISS

В марте 2017 года Facebook представила свое решение для ANN — библиотеку FAISS. Она объединяет множество методов и алгоритмов. В алгоритме, который мы рассмотрим ниже, расстояния до групп векторов будут приближаться расстоянием до опорной точки рядом с ними. Так мы можем выяснить расстояния от запроса до небольшого количества опорных точек, а затем полным перебором посчитать расстояния до векторов, принадлежащих опорной точке, которая ближе остальных к запросу. Разберем этот алгоритм по частям.

2.1. ADC — asymmetric distance computation

Рассмотрим подробнее следующую идею: расстояния до групп векторов можно приблизить расстояниями до опорных точек рядом с ними. Опорные точки делят пространство на области. Для поиска опорных точек в FAISS используется широко известный алгоритм кластеризации k-means, векторам сопоставляются полученные центроиды.

(0.1, 0.2) > 1
(0.5, -0.2) > 2
(0.1, 0.1) > 1
(0.6, -0.1) > 2

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

Эффективно хранить и быстро получать векторы, принадлежащие центроиду, помогает простой трюк под названием inverted file.

2.2. IVF — inverted file

В IVF мы инвертируем присвоение. Теперь центроидам сопоставляются списки векторов.

1 > [(0.1, 0.2), (0.1, 0.1)]
2 > [(0.5, -0.2), (0.6, -0.1)]

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

2.3. PQ — product quantizer

Последняя составляющая, которой мы коснемся в статье, называется product quantizer. Она обеспечивает сжатие векторов с потерями и применяется, когда векторные представления не влезают в память. Предположим, что наши векторы имеют размерность 128 и мы хотим кодировать их 64 битами (всего 0,5 бита на компоненту), тогда нам придется заниматься квантованием с количеством центроидов, равным $$. Это нетривиальная задача $$, которая также требует огромной обучающей выборки.

Упростим задачу, разбив вектор на $$ частей $$, и по традиции найдем 256 центроидов для каждой из частей. То есть вектор можно переписать как набор индексов центроидов — например $$, а занимает это хозяйство $$ байт.

Такой вид кодирования будем применять к остаточным векторам $$, $$, и тогда $$

2.4. Поиск

А теперь соберем все это в одной схеме.

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

import faiss

index = faiss.index_factory(model.wv.syn0norm.shape[1], 'IVF16384,Flat')
index.verbose = True
train = model.wv.syn0norm[np.random.binomial(1, 1./3, size=model.wv.syn0norm.shape[0]).astype(bool)]
index.train(train)
index.add(model.wv.syn0norm)
index.nprobe = 100
start = time.time()
distances, items = index.search(king_vec, 10)
print(time.time() - start)
print([model.wv.index2word[idx] for idx in items[0]])

Получаем:

    0.0048999786377
    [u'king', u'kings', u'queen', u'monarch', u'crown_prince', u'prince', u'sultan', u'ruler', u'princes', u'Prince_Paras']

2.5. Pros & Cons

+ Поддержка сжатия

+ Малые накладные расходы на хранение центроидов

+ Возможность вычислений на GPU*

– В пять раз медленнее HNSW на CPU

* Мы не смогли быстро завести GPU-реализацию из коробки и решили не тратить время на это.

3. ANNOY

Идея деления пространства плоскостями хорошо реализована в Annoy. Алгоритм прекрасно описан автором в блоге, рекомендую прочесть, если хотите разобраться в деталях. Я попробую коротко изложить суть. В алгоритме мы рекурсивно делим пространство плоскостями, образуя бинарное дерево. В каждом узле дерева хранится вектор, задающий текущую плоскость. При поиске мы начинаем с корня и выбираем дочернюю ноду на основе положения запроса относительно плоскости. Так мы спускаемся к листовым элементам дерева, в которых хранятся векторы, оказавшиеся по одну сторону группы плоскостей (это небольшой кусочек пространства), они с высокой вероятностью окажутся искомыми ближайшими соседями. Посмотрим на достоинства и недостатки Annoy по сравнению с другими алгоритмами.

3.1. Pros & Cons

+ Алгоритм просто понять

– Он требует много памяти

– Проигрывает в скорости работы

4. Сравнение алгоритмов

У каждого из алгоритмов есть набор параметров, будь то максимальное количество друзей для вершины (в NSW, HNSW) или количество центроидов (в FAISS). Эти параметры влияют на объем потребляемой памяти, качество и скорость поиска. Автор Annoy реализовал тесты для группы ANN алгоритмов в репозитории ann-benchmarks на разных параметрах. В них оценивается точность поиска десяти ближайших соседей в датасетах, полученных при помощи алгоритмов GloVe и SIFT.

GloVe — это еще один способ получить векторные представления слов, он превосходит word2vec по всем показателям при обучении на корпусе одного размера. Датасет составлен из 1,2 миллиона векторных представлений слов, обученных на 2 миллиардах твитов. SIFT — старый алгоритм получения ключевых точек изображения и их векторных представлений, устойчивых к трансформациям. Он использовался для распознавания объектов, и важной частью этого распознавания был поиск похожих векторных представлений. Есть несколько вариаций датасетов, нас интересует SIFT 1M, содержащий миллион векторов.

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

Видно, что HNSW уверенно лидирует. Однако на графиках нет FAISS. Facebook самостоятельно сравнил HNSW и FAISS в разных конфигурациях, результаты приведены в таблице.

В таблице методы FAISS без сжатия, в частности IVF16384,Flat. Значит, используется IVFADC c 16 384 центроидами. Расходы памяти указаны для случая с миллионом векторов размерности 128 в float32. HNSW в пять раз быстрее при вычислениях на CPU, но на хранение ребер графа ($$) требуется больше памяти, чем на центроиды ($$).

5. Заключение

Мы рассмотрели ряд алгоритмов, применяемых для быстрого поиска ближайших соседей. Annoy проиграл и по памяти, и по скорости работы, но идея хороша и может помочь в решении смежных задач. Например, удобный для чистки датасета алгоритм поиска аномалий isolation forest очень похож по своей задумке. FAISS — отличное решение при ограничениях по памяти, с ним вполне можно уложить миллиард векторных представлений в 60 Гбайт RAM, используя IVF16384, PQ64. Однако если память не узкое место, то стоит выбрать HNSW.

P. S. Самыми интересными оказались публикации об алгоритмах в FAISS. Там, к примеру, можно прочесть об оптимизации под GPU, об улучшенном методе квантования (Optimized Product Quantization) и о более хитром способе построении индекса (inverted multi-index).

P. P. S. Для себя мы выбрали HNSW.

6. Литература

• Approximate Nearest Neighbor Search Small World Approach
• Approximate nearest neighbor algorithm based on navigable small world graphs
• Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs
• Product quantization for nearest neighbor search
• SEARCHING IN ONE BILLION VECTORS: RE-RANK WITH SOURCE CODING
• GloVe: Global Vectors for Word Representation
• Object Recognition from Local Scale-Invariant Features
• Isolation Forest
• Billion-scale similarity search with GPUs
• Optimized Product Quantization
• The inverted multi-index