💬 NLP

CRINN: Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search

This paper introduces CRINN, a novel paradigm that leverages contrastive reinforcement learning to automatically generate high-performance approximate nearest neighbor search algorithms by optimizing for execution speed while maintaining accuracy, achieving state-of-the-art results on multiple benchmarks and demonstrating the potential of LLMs for automating complex algorithmic optimization.

Original authors: Xiaoya Li, Albert Wang, Guoyin Wang, Chris Shum, Jiwei Li

Published 2026-07-14

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Xiaoya Li, Albert Wang, Guoyin Wang, Chris Shum, Jiwei Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a massive library with billions of books, and you need to find the one that sounds most like a specific sentence you just typed. In the world of computers, this is called "Approximate Nearest Neighbor Search" (ANNS). It's the engine behind cool tech like AI chatbots that remember your past conversations or robots that find similar images. But here's the catch: finding that perfect book in a library of billions is slow. If you check every single book, it takes forever. So, computer scientists have built "smart shortcuts" to find the right book quickly, trading a tiny bit of accuracy for huge speed gains.

For years, making these shortcuts faster has been a human-only job. It's like tuning a race car engine: an expert mechanic has to listen to the engine, guess which part is dragging, tweak a screw here, change a gear there, and then test it again. This takes deep knowledge and a lot of patience.

Enter CRINN, a new system that asks a big question: Can we teach an AI to tune the engine itself?

The Robot Mechanic with a Superpower

CRINN is a "reinforcement learning" framework. Think of it as a robot mechanic that doesn't just guess; it learns by playing a game. The game is simple: write a piece of code to search the library, run it, and see how fast it is.

The Reward: If the code finds the books fast, the robot gets a high score.
The Lesson: The robot compares its new code against its old code. It doesn't just look at the score; it looks at why one was faster. Did it skip a step? Did it organize the books better?
The Loop: The robot uses this comparison to write an even faster version next time. It keeps doing this, getting smarter and faster with every try.

The paper shows that this robot mechanic is incredibly good at its job. When tested on six different types of "libraries" (datasets ranging from images of digits to word embeddings), CRINN managed to build search systems that were the fastest in the world for three of them (GIST-960, MNIST-784, and GloVe-25) and tied for first place on two others (SIFT-128 and GloVe-25).

How It Learned to Win

The secret sauce is called Contrastive Reinforcement Learning. Imagine the robot is given two versions of a code snippet: one that runs in 1 second and one that runs in 2 seconds. The robot is asked to explain why the first one is faster. It might notice, "Ah, the fast one checks three books at once, while the slow one checks them one by one." By learning these patterns, the robot starts writing code that naturally includes those speed-boosting tricks.

The researchers didn't just let the robot guess randomly. They gave it a specific goal: find the best balance between speed and accuracy. They measured performance using Queries Per Second (QPS)—how many searches the system can do in one second—while keeping the accuracy (called recall) high. For example, on the MNIST dataset (handwritten digits), CRINN was 85.25% faster than the previous best system at a very high accuracy level (0.999 recall). That's a massive jump!

The "One Size Fits All" Surprise

Here is where it gets really interesting. The robot was trained using only one type of library: the SIFT-128 dataset, which uses a specific way of measuring distance called "Euclidean distance" (like measuring a straight line on a map). You might think, "If I train a mechanic on a sedan, they won't know how to fix a motorcycle."

But the paper found that CRINN's training on Euclidean distance actually worked surprisingly well on "Angular distance" datasets (like the GloVe word embeddings), which are mathematically different. It suggests the robot learned general principles of speed that apply even when the rules of the game change slightly. However, the paper notes a clear limit: when tested on the NYTimes-256 dataset, CRINN actually performed 82.85% worse than the best existing system. This tells us that while the robot is a genius, it's not magic; it still struggles with certain types of data that are very different from what it learned on.

The Step-by-Step Upgrade

The researchers didn't just throw the robot at the whole problem at once. They broke the search process into three stages, like upgrading a car in phases:

Building the Map (Graph Construction): This is where the library is organized. CRINN improved this stage the most, boosting speed by an average of 22.11%. It learned to use "adaptive search," meaning it spends more effort finding the book only when it's really needed, and "multi-entry points," allowing it to start looking from several doors at once.
Searching the Map (Search): Once the map is built, you have to find the book. CRINN improved this by 18.30%, using tricks like "batch processing" (checking groups of books together) and "early termination" (stopping the search as soon as it's sure it found the best match).
Polishing the Result (Refinement): The final step to make sure the answer is perfect. This gave a smaller boost of 9.69%, because there was less room for improvement by this stage.

What This Means

The paper argues that CRINN proves a powerful idea: AI can automate complex, expert-level tasks. For decades, making these search algorithms faster required human experts with years of training. Now, a system that learns by comparing its own code can do it automatically.

The authors are careful to say this isn't a solved problem for every single dataset (as seen with the NYTimes result), but it is a major step forward. They suggest that as AI applications like "Retrieval-Augmented Generation" (where AI uses outside knowledge to answer questions) become more popular, tools like CRINN will be essential to keep everything running fast without needing a human to tweak every single setting.

In short, CRINN is a robot that learned to be a master mechanic for data search, proving that with the right training, machines can not only write code but also figure out how to make it run faster than humans ever could.

Technical Summary: CRINN

Problem Statement

Approximate Nearest-Neighbor Search (ANNS) is a critical component for modern AI applications, particularly Retrieval-Augmented Generation (RAG) and agent-based Large Language Model (LLM) systems, which require rapid retrieval from massive vector databases. While existing open-source ANNS solutions (e.g., FAISS, DiskANN, ParlayANN) rely on robust backbones like HNSW (Hierarchical Navigable Small World), their optimization remains a largely manual process. This process requires deep expertise in computer architecture, parallel programming, and algorithmic properties to identify bottlenecks, tune parameters (such as graph degree and search beam width), and design data structures for memory locality. The paper posits that this manual refinement is labor-intensive and raises the question of whether LLMs, augmented with reinforcement learning (RL), can automate the generation and testing of algorithmic improvements to discover optimization patterns that human engineers might overlook.

Methodology

The authors propose CRINN (Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search), a framework that treats ANNS optimization as a reinforcement learning problem where execution speed serves as the reward signal.

Core Framework

CRINN utilizes a contrastive RL model to optimize ANNS algorithms module by module (specifically targeting Graph Construction, Search, and Refinement within an HNSW backbone). The process involves:

Contrastive Prompting: The LLM is provided with a prompt containing previous code implementations paired with their execution scores. The model is instructed to perform a comparative analysis to identify why certain implementations outperform others, distinguishing between effective and ineffective strategies.
Iterative Generation: Based on this analysis, the LLM synthesizes new code variants.
Reward Signal: The generated code is executed, and the reward is calculated based on performance. To address the difficulty of comparing ANNS algorithms due to the trade-off between Queries Per Second (QPS) and Recall, the authors employ a specific evaluation strategy:
- They vary the search parameter $ef$ to generate a set of (QPS, recall) points.
- They filter points within the recall range of $[0.85, 0.95]$ , a "sweet spot" where most algorithms have sufficient data and performance differences are meaningful.
- The reward is defined as the Area Under the Curve (AUC) formed by these filtered points, providing a single scalar value for the RL objective.
Training Algorithm: The system employs Group Relative Policy Optimization (GRPO). For each input prompt, the model generates a group of code completions. Rewards are normalized within the group, and the policy is updated to maximize the likelihood of generating higher-reward code while maintaining a constraint against deviating too far from a reference policy (KL divergence).

Implementation Details

Baseline: The framework uses GLASS as the initial starting point, a graph-based ANNS library, and sequentially optimizes its constituent modules.
Training Data: The RL model is trained exclusively on the SIFT-128-Euclidean dataset to maintain a single generalizable codebase.
Output Structure: The LLM is constrained to output a structured response containing a "Performance Analysis," "Algorithm Design," and "Code Implementation."

Key Contributions

Automated Optimization Paradigm: CRINN introduces a novel paradigm where ANNS optimization is automated via RL-augmented LLMs, removing the need for manual, expertise-driven tuning.
Contrastive Learning Mechanism: By integrating exemplar code and execution metrics directly into the prompt, the model learns to reason about performance-critical patterns through direct comparison rather than trial-and-error alone.
Reward Engineering: The paper addresses the challenge of defining a scalar reward for ANNS by utilizing the AUC of the QPS-recall curve within a specific, practical recall range ( $[0.85, 0.95]$ ), enabling fair comparison across discrete parameter settings.
Progressive Optimization: The system demonstrates that sequential optimization of distinct modules (Graph Construction, Search, Refinement) yields cumulative performance gains.

Experimental Results

The authors evaluated CRINN on six widely-used benchmark datasets: SIFT-128, GIST-960, MNIST-784, GloVe-100, GloVe-25, and NYTimes-256.

Performance: CRINN achieved best-in-class performance on three datasets (GIST-960-Euclidean, MNIST-784-Euclidean, and GloVe-25-angular) and tied for first place on two others (SIFT-128-Euclidean and GloVe-25-angular).
Quantitative Gains:
- On MNIST-784, CRINN showed up to an 85.25% improvement in QPS at 0.999 recall compared to the best baseline (ParlayANN).
- On GIST-960, it achieved a 72.68% improvement at 0.99 recall.
- On SIFT-128 (the training dataset), improvements ranged from +3.25% to +25.57% across recall levels.
Limitations: CRINN underperformed the best baseline on NYTimes-256 (an angular distance dataset) by 82.85% at 0.9 recall. The authors attribute this to the model being trained solely on Euclidean distance data, suggesting that optimization patterns for angular similarity may differ significantly.
Module Contributions: The Graph Construction module yielded the most significant individual gains (averaging +22.11%), followed by Search (+18.30%) and Refinement (+9.69%).

Significance and Claims

The paper claims that the success of CRINN extends beyond the specific optimization of ANNS algorithms. It validates that LLMs augmented with reinforcement learning can function as effective tools for automating sophisticated algorithmic optimizations that traditionally demand specialized domain knowledge and labor-intensive manual refinement.

The authors argue that as the demand for efficient vector search grows with the proliferation of RAG and agent-based LLM applications, frameworks like CRINN will become increasingly valuable for maintaining competitive performance across evolving hardware architectures and application requirements. The work demonstrates that automated systems can not only match but, in specific contexts, exceed the performance of state-of-the-art human-engineered solutions.