CRINN: Контрастное обучение с подкреплением для поиска приближенных ближайших соседей

CRINN: Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search

August 4, 2025

Авторы: Xiaoya Li, Xiaofei Sun, Albert Wang, Chris Shum, Jiwei Li

cs.AI

Аннотация

Алгоритмы поиска приближенных ближайших соседей (ANNS) становятся все более важными для современных приложений искусственного интеллекта, особенно в задачах генерации с использованием поиска (RAG) и в приложениях на основе агентов с использованием крупных языковых моделей (LLM). В данной статье мы представляем CRINN — новую парадигму для алгоритмов ANNS. CRINN рассматривает оптимизацию ANNS как задачу обучения с подкреплением, где скорость выполнения служит сигналом вознаграждения. Этот подход позволяет автоматически генерировать все более быстрые реализации ANNS при соблюдении ограничений на точность. Наши экспериментальные результаты демонстрируют эффективность CRINN на шести широко используемых наборах данных для тестирования NNS. По сравнению с современными открытыми алгоритмами ANNS, CRINN показывает наилучшую производительность на трех из них (GIST-960-Euclidean, MNIST-784-Euclidean и GloVe-25-angular) и делит первое место на двух других (SIFT-128-Euclidean и GloVe-25-angular). Успех CRINN имеет значение, выходящее за рамки оптимизации ANNS: он подтверждает, что LLM, усиленные обучением с подкреплением, могут служить эффективным инструментом для автоматизации сложных алгоритмических оптимизаций, требующих специализированных знаний и трудоемкой ручной доработки. Код доступен по адресу: https://github.com/deepreinforce-ai/CRINN.

English

Approximate nearest-neighbor search (ANNS) algorithms have become increasingly critical for recent AI applications, particularly in retrieval-augmented generation (RAG) and agent-based LLM applications. In this paper, we present CRINN, a new paradigm for ANNS algorithms. CRINN treats ANNS optimization as a reinforcement learning problem where execution speed serves as the reward signal. This approach enables the automatic generation of progressively faster ANNS implementations while maintaining accuracy constraints. Our experimental evaluation demonstrates CRINN's effectiveness across six widely-used NNS benchmark datasets. When compared against state-of-the-art open-source ANNS algorithms, CRINN achieves best performance on three of them (GIST-960-Euclidean, MNIST-784-Euclidean, and GloVe-25-angular), and tied for first place on two of them (SIFT-128-Euclidean and GloVe-25-angular). The implications of CRINN's success reach well beyond ANNS optimization: It validates that LLMs augmented with reinforcement learning can function as an effective tool for automating sophisticated algorithmic optimizations that demand specialized knowledge and labor-intensive manual refinement.Code can be found at https://github.com/deepreinforce-ai/CRINN