只改2行代码,RAG效率暴涨30%,多种任务适用,可扩展至百亿级数据规模应用
只需修改两行代码,RAG向量检索效率暴涨30%!
不仅适用于文搜文”、“图搜图”、“文搜图”、“推荐系统召回”多种任务;而且具备良好扩展性,适合十亿、百亿级别大规模应用。
浙江大学高云君、柯翔宇团队联手向量检索领域大佬傅聪,开源新方法PSP(Proximity graph with Spherical Pathway),突破RAG两大难题。
简单来说,主流向量检索方法都是基于欧几里得距离设计,主要看“谁离你最近”;但有时AI其实更需要比较“语义相关性”,也就是最大内积、看谁最相似。
以往的内积检索办法,不能像欧式距离检索方法那样满足数学上的三角关系,所以很多老方法失效。
PSP发现,只要进行微小改动,老图结构也能找到最大内积最优解。
而且PSP还设置了提前停止策略,能判断检索是否应该提前结束,避免浪费算力,让搜索更快。
AI产品背后的技术核心
向量检索,是支撑起明星AI产品的核心技术组件。它不仅大大拓宽传统语义检索(关键词检索)的边界,和大模型的配合更是浑然天成。
如何发挥这项技术的真正潜力,让向量模型和向量数据库的组合真正跑出效果,关键在于——选对“度量空间”。
尽管基于图的向量检索算法,如HNSW、NSG等,因其优秀的检索速度备受青睐,但往往被忽视的是,它们都是面向欧式空间设计的向量检索算法。
“度量错配”在很多场景下是毁灭性的,很多适合用“最大内积”检索的向量数据,搭配欧式向量算法,往往会出现“检索结果和query语义无关”的问题。
回看最大内积检索领域,其实还没有出现类似HNSW、NSG这样现象级的检索算法。之前的很多工作往往只在某些数据集上面表现良好,但换了数据集,效果就会剧烈退化。
破局关键:仅需修改2行代码,实现全局最大内积解
研究团队通过理论探索发现,在最大内积检索领域的研究泾渭分明地分成两种范式:
一是把最大内积转换为最小欧式距离,进而可以用HNSW、NSG来解决。但这种转化往往会伴随着信息损失或者拓扑空间的非线性转换,而这些问题会对搜索效果带来不同程度的负面影响。
二是不进行空间转化,直接在内积空间进行检索。这样做的好处是避免了信息损失或空间扭曲,但相对应的痛点是,缺少有效手段对无效检索空间进行裁剪,进而难以达到更好的检索速度。
为什么在内积空间直接做检索这么难呢?
最核心原因在于内积空间并不是一个严格意义上的“度量空间”。从数学上来说,一个空间可以称之为“度量空间”,需要满足诸多条件,典型地,我们最常接触的欧式空间就是一个度量空间。而作为一个“残缺空间”,内积空间缺少的最重要的属性就是“三角不等式”。
根据NSG论文的理论部分,HNSW、NSG、SSG等state-of-the-art的向量检索算法之所以能如此高效,就是因为他们都利用了三角不等式对索引结构(图结构)进行了高效的裁剪。
而以内积作为距离度量,构建的三角形,不满足我们耳熟能详的那句口诀“三角形中任意两边之和大于第三边,而任意两边之差小于第三边”。正是这一属性的缺失,阻碍了最大内积检索算法进一步发展。
PSP研究团队对这一问题进行了深入研究,从理论上证明了一件事情:对任意搜索请求,即Query点q,它在一个为欧式距离设计的图索引结构上,可以通过简单的贪心算法找到全局最优的最大内积解。
基于图的向量检索算法都利用贪心算法进行检索:当我们从随机点开始在图上游走时,NSG这类算法会从路径上的点的邻居中,寻找一个距离目标“最近”的邻居进行跳转,这样从邻居的邻居逐步跳转到全局最优解。
而这种贪心算法曾经隐含的理论要求的是,如果构建图用的是欧式距离表达“远和近”,那么贪婪游走也需要用欧式距离来定义远和近。
而PSP团队的研究成果意义在于,如果构建图用的是欧式距离,在贪婪游走的时候可以用内积来定义远和近,最终到达的终点就是全局最优的最大内积解!
因此,研究团队可以通过仅修改检索(贪婪游走)算法中的两行代码,就实现将一个现有的欧式算法向最大内积的适配:
△
优化:合理引导搜索行为规避冗余计算
PSP研究团队发现,最大内积检索的过程中,会存在大量冗余计算,而这些冗余是可以通过合理引导搜索行为来规避的。
最大内积中的搜索行为与欧式空间中的搜索行为有极大差异,如下图所示:
左图中,绿色方框(query)的最近欧式近邻是红色三角,但它的最大内积近邻是橙色方块。因此,在搜索query的最近欧式邻居的时候,游走行为会很快在三角形附近停止,但搜索他的最大内积邻居会继续走到“外围”橙色方块附近。
从更宏观的角度看,研究团队发现,最大内积检索的解空间往往在数据集“外围”(不同于欧式距离最近邻,可以存在于数据空间的任意位置)。因此,最大内积的搜索行为往往服从一种“由内而外,再外围扩展”的模式(如上图右图)。
针对这种特性,PSP会设计针对性的策略,让图上搜索的起始点就尽量分布在距离“答案”更近的区域。
同时,冗余不仅仅发生在搜索过程的前段,也非常多地集中在搜索过程的后段。
如上图,PSP研究团队发现,在图索引上搜索到精确解的“最少步数”因Query而异,呈现明显的长尾分布(图a),而他们也通过大量实验挖掘出四类“特征”帮助我们判断搜索应该在什么时候停下来(图b)。这四类特征可以在搜索过程中以非常低的成本被计算和记录,实现自适应的“早停”策略。
具体来说,可以在数据库中随机采样一部分点作为query,通过对它们进行搜索来收集最优停止步数前后的数据构成可分类的样本,再用这个样本去训练一颗决策树,就可以辅助搜索过程判断停止条件:
如上图,研究团队通过对决策树剪枝,可以让整棵决策树保留较小的高度。选择决策树作为分类器,可以有效拟合少量样本,并直接翻译为if-else语句嵌入搜索代码中,实现高效的“停止判断”。
性能实测:稳定、高效、可扩展性
研究团队为了充分测试PSP算法的效果,在8个大规模、高维度的数据集上进行了充分测试。从维度看,DBpedia100K和DBpedia1M分别高达1536和3072维,用OpenAI text-embedding-3-large模型抽取;从数量看,最大的数据集Commerce100M包含1亿数据库点。
比较向量检索算法,往往关注相同召回率下的检索速度,即Query-Per-Second(QPS)。从上图中可看出,PSP相对于现有state-of-the-art的方法有着稳定、明显的提升。在MNIST数据上,甚至超过第二名4倍之多。
值得注意的是,baseline的方法里,往往有一些会在图中“缺席”。这是因为它们性能远差于其它方法,而很难和其它方法画到同一张图中。比如ip-HNSW在MNIST数据集中缺席;ScaNN在Laion10M和Commerce100M上缺席等等,这突出了PSP的表现稳定性。
另外,所使用的数据集包含了“文搜文”“图搜图”“文搜图”“推荐系统召回”等诸多数据模态,体现出PSP强大的泛化性。
除了比较检索性能,另外一个考察向量检索算法的应用价值的重要维度是scalability。好的检索需要远低于线性增长的时间复杂度(time complexity)。
上图可以看出,PSP在Top-1近邻上表现出log(N)速率增长的时间复杂度。而在Top-K检索上表现出接近log(N)的复杂度。这体现出PSP优秀的可扩展性,即在十亿乃至百亿级别的数据上进行高效检索的潜力。
论文链接: https://arxiv.org/pdf/2503.06882
Github链接:https://github.com/ZJU-DAILY/PSP
本文来自微信公众号“量子位”,作者:PSP团队,36氪经授权发布。