一文带你了解检索增强生成中的神兵利器 —— 近似近邻搜索

引言

随着大语言模型Chatgpt的横空出世，大语言模型（Large Language Model, LLM）频繁地出现在公众的视野中，成为了商业、娱乐、教育等领域讨论的热点。在LLM众多的出色能力中，其强大的检索能力（Information Retrieval）能力备受瞩目。大语言模型本身不联网，但却好像能回答互联网上能搜到的大部分问题，包括包括事情发生的具体时间、人物关系和前因后果等等。然而，LLM的记忆能力和检索能力也不是无限的。比如，LLM的幻觉（Hallucination）问题就是学术界和工业界目前致力于解决的问题 [1]。幻觉指的是即使在不确定答案的情况下，LLM不但不会承认无法回答，还会以自信的口吻凭空捏造出事实，通常可以以假乱真。为了解决这一现象，许多研究方向被提了出来，而检索增强生成（Retrieval-Augmented Generation, RAG）就是其中的一种方法。对于用户的提问，RAG首先生成信息检索请求，然后在数据库中寻找相关的信息，最后，结合相关信息和用户的提问向大语言模型进行提问（流程示意图见图1）。因为在数据库中寻找到的信息都是真实可靠的，大语言模型会根据提供的真实数据进行回答，减少其幻觉的可能。不仅如此，RAG的范式极大的扩展了大语言模型的应用场景，使得其可以实现大规模内容的记忆与整理。许多应用也由此催生出来，包括虚拟人设、文章理解/总结等。在RAG中，如何在大量的内容向量（数以万计）中找到与检索向量相匹配的内容直接决定了生成的质量和效率。能否在短时间内得到丰富翔实的内容对于最后回答的生成起到了近乎决定行性的作用。在本篇文章中，我们将介绍近似近邻搜索的概念，并介绍其中三种常见的方法。

图1: 检索增强生成

何为近似近邻搜索

K邻近算法 (K Nearest Neighbor Search)，顾名思义，会帮助我们在海量的内容向量中找到与检索向量最匹配的K个内容向量。然而，K近邻算法的算法复杂度是O(N)。当N非常大的时候，检索的速度将会不可接受。为了解决这一问题，近似K邻近算法 (Approximate K Nearest Neighbor Search，AKNNS) 被提了出来 [2, 3]。研究发现，当我们舍弃极小的准确性时，我们能极大的提高算法的效率。AKNNS中的“近似”强调的就是，我们无法保证得到的K个向量是最近的K个，但他们不会比最近的K个向量远太多。接下来，我们将介绍三种常见的AKNNS算法：量化 (Quantization)、 局部敏感哈希(Local Sensitive Hashing)和基于图的算法 (Graph-based algorithms)。

三种算法

量化 (Quantization）

量化的主要思路在于去掉空间中的冗余信息，从而提高空间利用率、节省内存。在近似邻近算法中，量化算法将原本数据集中的N个数据点，量化到S个中心点。给定一个检索向量，算法只需要在S个中心点之间找到最近的一个中心点，并返还该中心点的所有数据点。然而，中心点离检索向量的距离最近不代表中心点的数据集离检索向量最近。比如，当检索向量离两个中心点的距离差不多近时，检索向量会离两个中心点所包含的许多内容向量都差不多远，但是量化算法会优先选择其中一个中心点所包含的向量。为了提高匹配的准确性，算法可以选择L个离检索向量最近的中心点，然后再在这L个中心点所包含的所有数据点中进行距离的排序，选择最近的K个数据点。当L增大时，算法的准确性提升，但是因为需要排序的数据点的数量增多，效率则会降低。所以需要选择一个合适的L。

我们简单的计算一下空间复杂度和时间复杂度。首先，时间复杂度为 O(S)，因为我们只需要在S个中心点中找到离检索向量最近的L个中心点。当S比N小很多的情况下，这相比较O(N)是一个很大的速度提升。我们接下来计算空间复杂度。如果向量空间的维度是D，空间复杂度为O(D S + N \log S)。第一部分用于存储中心点的坐标，第二部分用于存储每个数据点对应的中心点。当D和S都偏大的时候，第一项可能是不小的空间成本。为了节省内存、提高空间效率，乘积量化的方法被提了出来。

乘积量化 (Product Quantization, PQ) [4]

乘积量化将向量分解成多个较短的向量，并对每个短向量在其相对应的子向量空间进行量化。为什么要这么做呢？我们计算一下空间复杂度和空间复杂度就知道了。如果我们将向量分为m部分，假设可以整除，那么每个向量的维度则为D / m（如图2）。因为我们将空间分成了m部分，为了保证中心点的数量为S不变，每个子向量空间我们只需要S^{\frac{1}{m}}个中心点。所以，空间复杂度为O(D/m * S^{\frac 1m} * m + m N \log S^{\frac 1m})。化简后可得O(DS^{\frac 1m} + N \log S)。原本量化算法的空间复杂度对于中心点的数量线性增长，乘积量化的空间复杂度对于中心点的数量却是次线性的。当m较大时，S^{\frac 1m}随S的增长速度十分缓慢 。时间复杂度为O(m S^{\frac 1m})。如果S>>m的话，这比单纯的量化算法好。乘积量化在实际应用中也表现的十分优异 [4]。

图 2: 乘积量化向量分解

局部敏感哈希 (Local Sensitive Hashing, LSH) [5]

图 3: 局部敏感哈希

另外一种近似近邻搜索算法是局部敏感哈希。与正常的哈希算法不同，局部敏感哈希不希望减少哈希碰撞 (Collisions) 的数量。相反，哈希碰撞是被希望看到的，因为局部敏感哈希的目的在于让两个相似的内容拥有相同的哈希码。如果这个假设是真的，那么在匹配向量的过程中，我们只需要在众多向量中，选择哈希码与检索向量哈希码相同的向量进行匹配，大大提升匹配效率。如何让哈希函数具有局部敏感性吗呢？位采样（bit sampling）是最早提出的局部敏感哈希函数之一 [6]。位采样通过对内容的二进制编码进行采样，当两个向量在二进制表示中随机决定的特定位置有相同的数值时，说明他们可能拥有类似的信息。通常需要对信息编码中的多个位置进行采样，以提高准确性。然而，基于编码的哈希函数并不能很好地将向量的语义信息考虑进去，随着神经网络的普及，越来越多人开始用神经网络作为局部敏感哈希函数 [7]。

基于图的算法 (Graph-based algorithms)

最后介绍的方法是基于图的近似近邻搜索。基于图的近似近邻算法在实验和应用中的表现尤为突出，因而得到了学界的极大关注 [8]。我们这次介绍一种比较经典的图算法：分级可导航小世界 (Hierarchical Navigable Small World) [9]。首先，我们介绍可导航小世界。可导航小世界是一个非指向图（见图3）。为了寻找离检索向量最近的向量，从一个预先选好的起点 (entry point) 出发，在该点的所有邻近点 (neighbors) 中，选择离检索向量最近的一个向量作为下一个目的地，并重复这一过程，直到当前点是局部最优点，即所有邻近点都离检索向量更远。如何构建非指向图是非常关键的一步，如果节点数量过少、连接距离过长，算法很可能找到离检索节点较远的向量。相反，如果节点数量多、连接距离短，算法的效率则十分地低效，需要花很长时间才能停止。可分级的导航小世界在一定程度上解决了这个问题。

图 4：可导航小世界 (Navigable Small World)

分级可导航小世界包含分为多个不同等级的可导航小世界，等级越高的小世界的连接距离越长，同时节点的数量越少，对应的，等级越低的小世界连接距离越短、节点数量越多，并且等级低的小世界节点包括等级高的小世界的节点。搜索算法从等级最高的小世界开始，执行可导航小世界的算法。运行结束后，降级到更低一级的小世界，并在低一级的小世界中重复刚刚的流程，直至降低到最低级的小世界并且停止（见图4）。分级的好处在于，在高级的小世界中，可以进行较为粗略的搜索，将搜索点移至检索向量附近，通过降级的方式，进行更加精细化的搜索，实现效率和准确率的平衡。如果对于图算法有进一步的兴趣，可以参见[8]中新提出的算法。

图 5：分级可导航小世界

结语

在这篇文章中，我们描述了LLM存在的幻觉问题并指出了检索信息增强的生成方法 (RAG) 是一个非常有希望的解决方案。通过对大量数据的检索，RAG利用召回的数据进行信息增强，再通过大语言模型生成更为可靠且丰富的回答。在检索过程中，为了在大量的信息中快速地找到相关信息，近似近邻搜索的方法被提了出来。近似近邻搜索通过舍弃少量的准确性，极大地提升了检索的效率。我们讨论了三种常见的近似近邻搜索算法，分别是：量化、局部敏感哈希和基于图的算法。这三种算法在各种工业场景中展示出优越的表现，在保证准确率的同时大大地提升了效率。由Meta开发的向量相似搜索库FAISS就集成了这三种方法和他们的变种 [11]。然而，在一些新兴的应用场景，如RAG中，如何改良这种方法以适用于RAG独特的工作方式仍然是现在的研究热点。

参考资料

[1] L. Huang et al., "A survey on hallucination in large language models: Principles, taxonomy, challenges, and open questions," arXiv preprint arXiv:2311.05232, 2023.

[2] M. Wang et al., "A comprehensive survey and experimental comparison of graph-based approximate nearest neighbor search," arXiv preprint arXiv:2101.12631, 2021.

[3] A. Andoni, P. Indyk, and I. Razenshteyn, "Approximate nearest neighbor search in high dimensions," in Proceedings of the International Congress of Mathematicians: Rio de Janeiro, 2018.

[4] H. Jegou, M. Douze, and C. Schmid, "Product quantization for nearest neighbor search," IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 33, no. 1, pp. 117-128, 2010.

[5] O. Jafari et al., "A survey on locality sensitive hashing algorithms and their applications," arXiv preprint arXiv:2102.08942, 2021.

[6] P. Indyk and R. Motwani, "Approximate nearest neighbors: towards removing the curse of dimensionality," in Proceedings of the Thirtieth Annual ACM Symposium on Theory of Computing, 1998.

[7] R. Liu et al., "Can LSH (locality-sensitive hashing) be replaced by neural network?," Soft Computing, vol. 28, no. 2, pp. 1041-1053, 2024.

[8] M. Wang et al., "A comprehensive survey and experimental comparison of graph-based approximate nearest neighbor search," arXiv preprint arXiv:2101.12631, 2021.

[9] Yu A. Malkov and D. A. Yashunin, "Efficient and robust approximate nearest neighbor search using hierarchical navigable small world graphs," IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 42, no. 4, pp. 824-836, 2018.

[10] J. Briggs, "Faiss: The Missing Manual," Pinecone, [Online]. Available: https://www.pinecone.io/learn/series/faiss/. [Accessed: 21 Feb. 2024].

[11] Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-scale similarity search with GPUs. IEEE Trans. Big Data 7, 3 (2019), 535–547.

图源：

图 1: https://www.linkedin.com/pulse/what-retrieval-augmented-generation-grow-right/

图 2: https://www.pinecone.io/learn/series/faiss/product-quantization/

图 3: https://www.pinecone.io/learn/series/faiss/locality-sensitive-hashing/

图 4,5: https://www.pinecone.io/learn/series/faiss/hnsw/