CRINN：通过强化学习自动优化近似近邻算法

24/08/2025

三个要点
✔️ CRINN 利用强化学习自动优化 ANNS 代码，无需手动调整
✔️ 考虑到 QPS 和可重复性之间的权衡，速度是提高性能的奖励信号
✔️ 在六个基准上进行了评估，在 MNIST 和 GIST 上实现了显著的速度改进。

CRINN: Contrastive Reinforcement Learning for Approximate Nearest Neighbor Search
written by Xiaoya Li, Xiaofei Sun, Albert Wang, Chris Shum, Jiwei Li
(Submitted on 4 Aug 2025 (v1), last revised 20 Aug 2025 (this version, v2))
Comments: Preprint Version
Subjects: Machine Learning (cs.LG); Artificial Intelligence (cs.AI); Computation and Language (cs.CL); Databases (cs.DB)

本文所使用的图片要么来自论文、介绍性幻灯片，要么是参考这些图片制作的。

概述

本文的重点是优化高维向量空间中的近似近邻搜索（ANNS）。

近似近邻搜索是一种以略微牺牲搜索精度为代价来显著提高搜索速度的技术，最近已成为检索增强生成（RAG）和基于代理的 LLM 应用基础技术的组成部分。
传统的优化方法是由人工专家进行剖析、分析缓存缺失、调整数据结构并反复手动调整参数。
然而，这种方法既专业又耗费人力，而且在跟上硬件和应用环境的发展方面存在局限性。

因此，作者提出了一种新的优化框架--CRINN，它结合了 LLM 和强化学习。
CRINN 将代码执行速度视为一种奖励，并通过基于对比学习的强化学习自动生成高效的 ANNS 代码。
这样就能在不依赖人工调整的情况下生成连续改进的实现，从而在检索性能方面实现新的突破。

建议的方法

CRINN 将 ANNS 视为一个具有强化学习功能的优化问题，并将强化学习与对比学习相结合，以提高性能。
具体来说，该设计将现有的实现代码及其执行速度纳入提示，并允许 LLM 对某种实现更快的原因进行对比分析。

这样，模型就能学习到提高速度的模式，并生成新的、更好的代码。
生成的代码会被执行，并根据速度和可重复性给予奖励。

奖励用于执行基于组相对策略优化（GRPO）的强化学习，并按顺序更新模型。
奖励设计还侧重于每秒查询次数（QPS）和召回率之间的权衡，召回率范围[0.85,0.95]内的曲线下面积被用作标量奖励。
此外，以现有的名为 GLASS 的 ANNS 库为初始基础，对每个模块进行了顺序优化：图构建、探索和完善。

这种结构化方法将传统的专家协调自动化，并能开发出高效的搜索算法。

实验

实验测试了 CRINN 在 SIFT-128、GIST-960、MNIST-784、GloVe-25、GloVe-100 和 NYTimes-256 六个基准数据集上的性能。

我们选择了具有代表性的开源 ANNS 实现（如 ParlayANN、GLASS、NNDescent、PyNNDescent、Vearch 和 Voyager）进行比较。只使用 SIFT-128（欧氏距离）进行训练，然后根据其他数据集评估泛化性能。
结果，CRINN 在 MNIST-784 和 GIST-960 数据集上的处理速度提高了 85%，图构建模块的改进尤为显著。

另一方面，在一些数据集（如 NYTimes-256）上也观察到了性能下降，这表明优化可能会受到距离规模和数据特征的限制。
此外，增量模块优化的有效性也得到了证明，证实了底层 GLASS 持续改进的潜力。

总之，CRINN 集速度和准确性于一身，优于现有方法，为使用强化学习进行代码优化提供了一个新方向。