CRINN: Aprendizaje por Refuerzo Contrastivo para la Búsqueda de Vecinos Más Cercanos 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

Resumen

Los algoritmos de búsqueda de vecinos más cercanos aproximados (ANNS, por sus siglas en inglés) se han vuelto cada vez más críticos para las aplicaciones recientes de IA, particularmente en la generación aumentada por recuperación (RAG) y en aplicaciones de modelos de lenguaje basados en agentes. En este artículo, presentamos CRINN, un nuevo paradigma para los algoritmos ANNS. CRINN aborda la optimización de ANNS como un problema de aprendizaje por refuerzo, donde la velocidad de ejecución sirve como señal de recompensa. Este enfoque permite la generación automática de implementaciones de ANNS progresivamente más rápidas, manteniendo las restricciones de precisión. Nuestra evaluación experimental demuestra la efectividad de CRINN en seis conjuntos de datos de referencia ampliamente utilizados para NNS. En comparación con los algoritmos ANNS de código abierto más avanzados, CRINN logra el mejor rendimiento en tres de ellos (GIST-960-Euclidiano, MNIST-784-Euclidiano y GloVe-25-angular), y empata en el primer lugar en dos de ellos (SIFT-128-Euclidiano y GloVe-25-angular). Las implicaciones del éxito de CRINN van más allá de la optimización de ANNS: valida que los modelos de lenguaje aumentados con aprendizaje por refuerzo pueden funcionar como una herramienta efectiva para automatizar optimizaciones algorítmicas sofisticadas que requieren conocimiento especializado y refinamiento manual intensivo en mano de obra. El código se puede encontrar en 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