company_banner

Методы приближенного поиска ближайших соседей


    Довольно часто программисты и специалисты из области data science сталкиваются с задачей поиска похожих профилей пользователей или подбора схожей музыки. Решения могут сводиться к преобразованию объектов в векторную форму и поиску ближайших.


    Мы тоже столкнулись с необходимостью поиска ближайших соседей в задаче распознавания лиц. Там мы формируем векторные представления лиц при помощи нейросети и ищем ближайшие векторы уже известных людей. Изначально для поиска мы выбрали Annoy, как хорошо известный и проверенный алгоритм, используемый в том числе в Spotify. Но быстро поняли, что с его аппетитами по памяти мы либо не вмещаемся в RAM, либо сильно теряем в точности. Это привело к небольшому исследованию. О результатах которого пойдет речь ниже.


    Чтобы разбавить теорию практикой, в статье будет немного кода, где мы ищем ближайших соседей для слов. Получим их векторные представления, используя популярный word2vec. Этот алгоритм выдает близкие векторы для семантически похожих слов. В word2vec CBOW векторные представления получаются как побочный продукт обучения небольшой нейросети, которая предсказывает слово по его окружению. Любопытно, что с векторами можно проворачивать арифметические операции наподобие king + (woman – man) = queen.


    word embeddings


    Посмотрим, как с этим работать в коде.


    model = gensim.models.KeyedVectors.load_word2vec_format('./GoogleNews-vectors-negative300.bin', binary=True)
    start = time.time()
    pprint.pprint(model.wv.most_similar(positive=['king']))
    print('time:', time.time() - start)
    print('king + (woman - man) = ', model.wv.most_similar(positive=['woman', 'king'], negative=['man'])[0])
    print('Japan + (Moscow - Russia) = ', model.wv.most_similar(positive=['Moscow', 'Japan'], negative=['Russia'])[0])

    Получаем:


        [(u'kings', 0.7138045430183411),
         (u'queen', 0.6510956883430481),
         (u'monarch', 0.6413194537162781),
         (u'crown_prince', 0.6204219460487366),
         (u'prince', 0.6159993410110474),
         (u'sultan', 0.5864823460578918),
         (u'ruler', 0.5797567367553711),
         (u'princes', 0.5646552443504333),
         (u'Prince_Paras', 0.5432944297790527),
         (u'throne', 0.5422105193138123)]
        time: 0.236690998077
        king + (woman - man) =  (u'queen', 0.7118192911148071)
        Japan + (Moscow - Russia) =  (u'Tokyo', 0.8696038722991943)

    В примере выше использовалась библиотека для работы с текстом gensim и word2vec модель (1,5 Гбайт) от Google, которая насчитывает 3 миллиона слов и коротких фраз. В выводе кода видно, что к королю близки королевы, монархи и принцы. Также мы убедились, что арифметика с векторами работает. Однако четверть секунды на один запрос — не очень привлекательно, а ведь в gensim сравнительно хорошая реализация bruteforce-поиска (с подсчетом расстояний до всех известных объектов). Далее мы рассмотрим методы, позволяющие сократить это время в сотни раз лишь с небольшими потерями в точности.


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


    nbrs = NearestNeighbors(algorithm='brute', metric='cosine')
    nbrs.fit(model.wv.syn0norm)
    king_vec = model.wv['king'][np.newaxis, :]
    # замерим скорость поиска сосдей к королю без разделения пространства и заодно выведем результат
    start = time.time()
    idxs = nbrs.kneighbors(king_vec, return_distance=False, n_neighbors=10)[0]
    print('full search time:', time.time() - start)
    print([model.wv.index2word[idx] for idx in idxs])
    
    # выберем 2 случайных вектора и получем коэффициенты задающие плоскость между ними
    vec1_idx = random.randint(0, model.wv.syn0norm.shape[0])
    vec2_idx = random.randint(0, model.wv.syn0norm.shape[0])
    plane = model.wv.syn0norm[vec1_idx] - model.wv.syn0norm[vec2_idx]
    
    # в результате следующего умножения матрица-вектор, мы получим вектор. 
    # Знаки элементов этого вектора указывают с какой стороны разделяющей плоскости оказалось слово
    scalar = model.wv.syn0norm.dot(np.transpose(plane))
    
    # определим с какой стороны плоскости запрос и подготовим бинарную маску для выборки векторов по ту же сторону плоскости
    if king_vec.dot(plane) > 0:
        mask = scalar > 0
    else:
        mask = scalar < 0
    
    print('elements in mask:', mask.sum())
    
    half_nbrs = NearestNeighbors(algorithm='brute', metric='cosine')
    half_nbrs.fit(model.wv.syn0norm[mask])
    half_index2word = list(compress(model.wv.index2word, mask))
    
    start = time.time()
    idxs = half_nbrs.kneighbors(king_vec, return_distance=False, n_neighbors=10)[0]
    print('half search time:', time.time() - start)
    print([half_index2word[idx] for idx in idxs])

    Получаем:


        full search time: 20.3163180351
        [u'king', u'kings', u'queen', u'monarch', u'crown_prince', u'prince', u'sultan', u'ruler', u'princes', u'Prince_Paras']
        elements in mask: 1961204
        half search time: 9.15824007988
        [u'king', u'kings', u'queen', u'monarch', u'crown_prince', u'prince', u'sultan', u'ruler', u'princes', u'Prince_Paras']

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



    1. HNSW


    В 2011 году был опубликован один из лучших алгоритмов поиска ближайших соседей Navigable Small World (NSW). В 2016-м появился его наследник Hierarchical Navigable Small World (HNSW).


    Начнем с родительского алгоритма NSW. Он основан на графе «мир тесен». Эти графы имеют любопытную и полезную нам особенность: пара вершин с большой вероятностью не смежна, но они достижимы за сравнительно небольшое число шагов ($\log{N}$ в среднем). Такие графы встречаются довольно часто. К примеру, нейронные сети мозга, группы в социальных сетях и семантическая сеть WordNet — это графы SW. В нашем случае вершинами являются векторы, а ребра соединяют их с ближайшими. В графе также представлены ребра, соединяющие вершины на большом расстоянии.


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



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


    Псевдокод
    K-NNSearch(object q, integer: m, k)
    1 TreeSet [object] tempRes, candidates, visitedSet, result
    2 for (i←0; i < m; i++) do:
    3    put random entry point in candidates
    4    tempRes←null
    5    repeat:
    6       get element c closest from candidates to q
    7       remove c from candidates
    8       //check stop condition:
    9       if c is further than k-th element from result
    10          than break repeat
    11    //update list of candidates:
    12    for every element e from friends of c do:
    13       if e is not in visitedSet than
    14          add e to visitedSet, candidates, tempRes
    15
    16    end repeat
    17    //aggregate the results:
    18    add objects from tempRes to result
    19 end for
    20 return best k elements from result

    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, однако теперь мы имеем дело с иерархией графов: на нулевом слое представлены все объекты, а по мере увеличения слоя — все меньшая и меньшая их подвыборка. При этом все объекты на слое $n+1$ есть и на слое $n$.



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


    Псевдокод
    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
    

    Векторы коллекции аппроксимируются своими центроидами $y \approx q_c(y)$, где $q_c:\mathbb{R}^d \rightarrow C_1 \subset \mathbb{R}^d$ и $C_1$ — множество центроидов. Тогда расстояние от запроса $x$ до $y$ может быть приближено $d(x, y) \approx d(x, q_c(y))$. Такой способ вычисления дистанции называют асимметричным. Простыми словами: мы разбили пространство на области и сказали, что расстояние от запроса до группы векторов, попавших в одну область, приблизительно равно расстоянию до центроида, образующего эту область.



    Эффективно хранить и быстро получать векторы, принадлежащие центроиду, помогает простой трюк под названием 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 бита на компоненту), тогда нам придется заниматься квантованием с количеством центроидов, равным $k=2^{64}$. Это нетривиальная задача $O(iknd)$, которая также требует огромной обучающей выборки.


    Упростим задачу, разбив вектор на $m$ частей $y_1, y_2, ..., y_m \in \mathbb{R}^{d/m}$, и по традиции найдем 256 центроидов для каждой из частей. То есть вектор можно переписать как набор индексов центроидов — например $(134, 20, 244, ..., 12)$, а занимает это хозяйство $m$ байт.


    Такой вид кодирования будем применять к остаточным векторам $r(y) = y - q_c(y)$, $r(y) \approx q_r(r(y))$, и тогда $y \approx q_c(y) + q_r(y - q_c(y))$


    2.4. Поиск


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



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


    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, содержащий миллион векторов.


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


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


    Method search time 1-R@1 index size index build time
    Flat-CPU 9.100 s 1.0000 512 MB 0 s
    nmslib (hnsw) 0.081 s 0.8195 512 + 796 MB 173 s
    IVF16384,Flat 0.538 s 0.8980 512 + 8 MB 240 s
    IVF16384,Flat (Titan X) 0.059 s 0.8145 512 + 8 MB 5 s
    Flat-GPU (Titan X) 0.753 s 0.9935 512 MB 0 s

    В таблице методы FAISS без сжатия, в частности IVF16384,Flat. Значит, используется IVFADC c 16 384 центроидами. Расходы памяти указаны для случая с миллионом векторов размерности 128 в float32. HNSW в пять раз быстрее при вычислениях на CPU, но на хранение ребер графа ($32*2*4*1000000/1024^2 = 244$) требуется больше памяти, чем на центроиды ($16384 * 128 * 4 / 1024^2 = 8$).


    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. Литература


    • +53
    • 12,3k
    • 4
    Mail.Ru Group 764,25
    Строим Интернет
    Поделиться публикацией
    Комментарии 4
    • +2
      Хороший обзор!
      P. P. S. Для себя мы выбрали HNSW.

      Я так понимаю HNSW у вас используется где-то в продакшне? Если да, то у меня есть куча вопросов :)
      А какая размерность вектора у вас используется?
      Какой размер данных (штук векторов) обычно в индексе и какие параметры M / ef выбрали для генерации индекса, какой ef для поиска?
      Данные изменяются? Если да, то как часто делаете перейстройку индекса и что делаете с теми данными, которые пока еще не попали в индекс, но по которым хочется искать?
      Как решали проблему удаления из индекса, или такой задачи нет/подходит решение перестроением индекса?
      Реализовали HNSW самостоятельно или взяли какую-то готовую реализацию?
      • +2
        Спасибо!
        До прода еще не успели доехать. Размерность вектора — 128. В индексе сотни миллионов, к проду будет 500-1000M, что может повлиять на настройки, а пока они такие: M=32, efConstruction=800, ef=200. Действительно, готовая реализация в NMSLIB не поддерживает обновления графа, хоть сам алгоритм и позволяет как удалять так и добавлять объекты. В нашей задаче нет необходимости в онлайн обновлениях. Потому свою реализацию делать не стали и взяли NMSLIB.
      • 0
        Картинка заглавная классная!
        • 0
          Вероятно, ручная реализация диаграмм Воронова и триангуляции Делоне.

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

        Самое читаемое