CRINN: Aprendizado por Reforço Contrastivo para Busca de Vizinhos Mais Próximos Aproximados

CRINN: Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search

August 4, 2025

Autores: Xiaoya Li, Xiaofei Sun, Albert Wang, Chris Shum, Jiwei Li

cs.AI

Resumo

Algoritmos de busca por vizinhos mais próximos aproximados (ANNS) tornaram-se cada vez mais críticos para aplicações recentes de IA, particularmente em geração aumentada por recuperação (RAG) e aplicações de LLM baseadas em agentes. Neste artigo, apresentamos o CRINN, um novo paradigma para algoritmos ANNS. O CRINN trata a otimização de ANNS como um problema de aprendizado por reforço, onde a velocidade de execução serve como sinal de recompensa. Essa abordagem permite a geração automática de implementações de ANNS progressivamente mais rápidas, mantendo restrições de precisão. Nossa avaliação experimental demonstra a eficácia do CRINN em seis conjuntos de dados de referência amplamente utilizados para NNS. Quando comparado com algoritmos ANNS de código aberto de última geração, o CRINN alcança o melhor desempenho em três deles (GIST-960-Euclidiano, MNIST-784-Euclidiano e GloVe-25-angular), e empatou em primeiro lugar em dois deles (SIFT-128-Euclidiano e GloVe-25-angular). As implicações do sucesso do CRINN vão muito além da otimização de ANNS: ele valida que LLMs aumentados com aprendizado por reforço podem funcionar como uma ferramenta eficaz para automatizar otimizações algorítmicas sofisticadas que exigem conhecimento especializado e refinamento manual intensivo. O código pode ser encontrado em 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