CRINN: Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search

2025-08-04

AI Summary

This work addresses the critical challenge in retrieval-augmented generation (RAG) and large language model (LLM) agents: balancing speed and accuracy in approximate nearest neighbor search (ANNS). We propose the first automated ANNS optimization framework based on contrastive reinforcement learning. By modeling algorithm generation as a sequential decision-making process and using query latency as the reward signal—while enforcing precision constraints—the framework autonomously evolves efficient index structures and search policies. Crucially, we introduce contrastive learning to enhance policy stability, enabling end-to-end, human-prior-free algorithm design. Evaluated on six standard benchmarks, our method achieves state-of-the-art performance: ranking first on three datasets and jointly first on two others. It significantly improves retrieval efficiency and cross-dataset generalization, establishing a novel paradigm for AI-driven, automated algorithm discovery.

Technology Category

• Search and Optimization: Learning to Search
• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Agentic search
• Graph Algorithms and Modeling for the Web: Querying, indexing, and retrieval in Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking

Abstract

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

Problem

• Optimizes ANNS algorithms using reinforcement learning for speed
• Automates generation of faster ANNS implementations with accuracy
• Enhances retrieval performance in AI applications like RAG

Innovation

• Uses reinforcement learning for ANNS optimization
• Rewards execution speed while maintaining accuracy
• Automates generation of faster ANNS implementations