向量检索在大模型应用场景的技术和实践
1. 向量检索应用简介
向量是多维数学空间里的一个点,在各维度上的坐标的一串数字。这个点就是来源于真实世界的物体进行数字化之后在数学空间的投影。那么不同点之间有一个数学关系,就是距离,距离远近就代表两个物体的相似程度。
非结构化数据转换成向量的过程称为 embedding。通过深度学习的训练,可以将真实世界数字化后的离散特征提取出来,投影到数学空间上,成为一个数学意义上的向量,同时很神奇的保留着通过向量之间的距离表示语义相似度的能力,这就是 embedding 的效果。
在大语言模型出现之前(2020 年以前),向量检索这项技术就已经发展成熟。随着深度学习的技术,广泛应用于图片、音频、视频的搜索和推荐、人脸识别、语音识别等传统人工智能应用领域。
大模型的出现改变了人机交互方式,带来了人工智能技术的新革命,一下子火起来,进入了一个大模型的时代。当然现在处于初期阶段,在实际应用上还存在很多问题。
首先是知识能力不够强,大模型记忆比较有限,比如 6 亿参数,130 亿参数,就像人的脑细胞一样,能记住的东西有限。
其次大模型训练周期和成本都很高,所以没法跟上实时热点,比如我们说 chatgpt 只知道 2021 年前发生的事情,因为他只训练了 2021 年前的数据。此外大模型还有幻觉问题,回答会歪曲事实。
除了知识方面,大模型也很难保证隐私数据安全。比如我为大模型提供了一些关于自己的私人数据,当这个大模型为其他人提供服务的时候,这些隐私信息很有可能作为答案被说出来。最后大模型的推理成本也很高,回答一个问题的成本是很高的。
那么如何增强大模型的知识能力,同时也保护私有数据的安全性呢?我们可以看看下图右侧的这个例子。比如提问的时候附带上一些气象局的信息,大模型将会依据这些额外信息准确回答问题。
这种通过外置的数据资料和工具来增强大模型能力的技术体系,称为提示词工程。
从以上讨论可以看出,大模型应用离不开提示词工程。提示词工程怎么做呢?主要就是为大模型整理一个资料库,然后在访问流程上,先从海量资料库里找到最匹配的内容,拼接提示词来增强回答,本质上就是一个搜索引擎。
向量检索技术可以非常完美地支持这个场景。现在万物皆可 embedding,向量就是 AI 理解世界的通用形式,而向量检索的工程技术则连接 AI 和这些资料,为 AI 提供知识和记忆。
以下两张图是工程上向量技术来支持提示词工程的流程。
因此大模型应用现阶段是离不开向量检索技术的,向量数据库也会是每个大模型应用里必然要用到的组件,正如关系型数据库在 web 应用里一样。
2. 向量检索技术概述
向量相关的工程技术里最核心的当然是向量检索算法,即如何在海量向量里找到跟目标向量最相似的 K 个,又叫 topK。
最简单地当然是暴力算法, 将目标向量与集合里每个向量的距离都算一遍,然后排个序,取前面 K 个就行了。这样有个明显的问题,随着数据量增加,耗时会线性增长到无法接受的程度,比如我的数据日积月累,今天只要 10 毫秒,明天可能就需要 100 毫秒,后天可能就需要一秒,再过一个星期,可能就要分钟级别了。而向量检索速度直接影响到大模型回答问题的速度,这个就很难接受。
如何解决?最简单地就是通过分布式计算来解决。比如 1 亿向量,分成十份,放十台机器上并行计算,那相比用一台机器来说就快了十倍。而且数据即使增长的话,我也就增加机器就行,结果也是保证精确的。当然这样算力成本就会很高。
因此业界探索有没有可能放弃全局最优解,寻找局部最优解,来节省计算量。由此发展出了近似最近邻算法,业界有四种算法:哈希、树搜索、倒排、图搜索。哈希、树搜索、倒排这三种类似,都是通过某种分类的方式,将数据预分类,划分空间,对号入座来减少计算量;而图搜索则是一种比较新的思路。
我们考核算法的效果,主要看两个指标,一是性能,也就是查询的耗时以及能承受的 QPS;另一个就是召回率,代表查询准确度。我们将近似算法查出来的结果跟从全局视角看的结果的集合对比,重叠度就是召回率。比如暴力算法,那他召回率就是 100%,近似算法则有高有低,有的也可以逼近 99% 以上。
接下来我将介绍下具体的算法思路。
IVF 即倒排索引,倒排索引是搜索引擎的核心技术之一,指的从文档网页里提取关键词,来建立倒排检索结构。那么对向量数据来说,关键词是什么呢?真实世界的向量在空间分布上一般会扎堆分布,具备聚类特征,请看下图。
通过 k-means 算法将向量的聚类中心提取出来,那这个向量所处的聚类中心就是这个向量的关键词,用这个来建立倒排索引,就可以像搜索引擎一样,先命中聚类中心,再暴力搜索聚类中心下面的各个向量,这样相比全局搜索能过滤大量的数据。如果觉得找一个聚类中心不够准确,也可以多找几个,找的越多,结果越准。
通过聚类的思想能不能进一步做优化?答案是可以的,就是使用乘积量化技术,可以将一个 D 维浮点向量降维成一个 M 维的二进制向量,这样不仅仅是存储空间的压缩,也压缩了内存使用和计算量。
思路是首先将 D 维向量分解成 M 个子向量,然后用子向量训练,得到聚类中心,编成编号来代替子向量,这样就从一串浮点数,编成一个小整数,起到压缩效果。示例中 D=128,M=8,算下来压缩了 64 倍。
图算法是一种比较新的近似向量检索思路,它基于最小世界理论,就是说世界上任意两个人,都可以通过六跳联系起来。如果把向量看做人的话,按向量间的距离关系构造一个这样类似于真实世界的「小世界网络」,通过贪心算法,按距离建立联系一跳一跳逼近目标的向量定点。
随着小世界里的点不断插入,根据前面讲的挑选近邻点(熟人)构建边的思路,新插入的点会越来越局限于一个小圈子,这样的话跟一个很远的点建立联系就需要非常多的跳数了。因此怎么破圈呢?通过一些“机缘巧合”让新插入的点也能跟远处的点连起来是关键。
业界摸索出了 HNSW 算法,即采用类似链表查找算法里的跳表的思路。链表是一维的,图是二维的。我们来建立不同层级的图,往上指数递减定点数量来形成稀疏图,这样越稀疏的图自然就越能连接远方。
向量检索需要结合一套分布式系统设计来做工程落地。相关话题在数据类的分布式存储系统的工程理论都已经非常成熟,向量数据在存储层面也没有特别的地方,因此只要在存储系统上实现向量能力即可。
因此向量检索系统的架构,是一套向量数据存储系统加上向量检索能力,再附加上一些标量字段的存储和过滤能力。比如我可以给向量数据打上一些字符串标签用于特定场景的过滤。
考虑到向量存储体量不会特别大,即使百亿级的向量算到存储空间上也不是一个非常大的规模,因此主要还是考虑算力扩展性,追求小巧精悍和极致性能。因此采用下图右侧所示架构,同时着重关注单节点的列式数据存储格式的实现。
随着数据规模的增长,为了能够更方便地弹性扩展算力,这里的技术趋势也必然会往云原生和存算分离的方向发展,在遇到流量高峰时能快速拉起多副本的算力来承接。
3. 向量检索工程实践
百度的大模型场景跟业界通用的大模型应用场景类似,主要也是用向量检索技术来做增强知识和提示词工程。在百度的大模型场景里,向量检索技术会面对一个实际的大规模的工程挑战。
百度智能云在 2020 年就已经在 ElasticSearch 上集成向量检索能力了,并广泛应用于百度自己的业务线,有着充分的工程实践,并且能基于公有云的资源提供云原生的运维保障。如今在大模型场景下,需要针对新的挑战进行技术优化。
百度智能云是国内云厂商里最早提供托管 ES 服务的,经过这么多年的发展,持续迭代增强产品能力。
在 ES 中使用向量检索的功能非常简单,这里给大家做个演示。先用 ES 标准方式建表,指定一些向量索引相关的参数,通过 ES 的标准 bulk 方式批量导入数据,然后通过我们定义的跟 ES 风格接近的语法进行查询。
百度 ElasticSearch 的架构由管控平台和 BES 集群实例两部分组成。管控平台是全局层面来进行一些统一的集群管理、监控报警,以及执行定时创建删除 merge 索引的任务的平台。BES 集群实例则是一套构件在云服务器 BBC 和云磁盘 CDS 上面的 ElasticSearch 集群服务,前面通过 BLB 四层代理做节点负载均衡。磁盘上的数据可以通过策略定期下沉存储到对象存储 BOS 上,降低存储成本。
向量数据在 ES 的标准架构下按下图所示方式组织和访问。
我们选择自研插件,一是希望基于 C++ 实现来获得更贴近底层的极限性能,二是能改写存储格式等也是方便做到更极致的性能,三是能更灵活地控制检索逻辑,改写执行计划来实现比较复杂的查询。
这里核心的向量检索引擎部分,我们选择基于社区优秀的向量库实现做二次开发。我们对比了下 nmslib 和 faiss 各自在 ES 上的开箱性能(没做优化)。可以看出 hnsw 整体比较耗内存,其中 nmslib 的实现更胜一筹。因此我们在此基础上做进一步优化。
HNSW 构图过程中每插入一个点需要检索计算,插入大量的点也是一笔很大的计算开销,因此导入数据会很慢,导致前台阻塞。
这里改造成异步构建索引的机制,数据写入落盘之后就可以直接返回。然后后台通过 ES 的合并策略或者用户的定时或主动触发的方式来触发在后台构建 HNSW 索引,并且使用独立线程池构建,以免对前台查询请求造成影响。
这里同时也对 ES 的片段合并策略进行优化,在最终合并时统一构图,来降低逐个合并构图带来的中间计算开销。
有很多场景提出先按标量条件过滤数据再进行向量检索的需求,例如为向量数据打上标签,要保证检索出来的向量能匹配上标签。标准的 HNSW 实现无法做到这一点,如果先 HNSW 再做过滤呢,则无法保证返回的结果完整。因此我们改造了 HNSW 实现,让算法在遍历图挑选近邻的过程中,只考虑符合 filter 条件的定点。
实际测试的时候,发现性能和召回不是很理想。通过测试数据和研究一些资料发现,在过滤比例升高时,因为顶点被过滤导致能连通的路径变少,直接影响了 HNSW 算法的收敛速度,以及很容易走进死胡同导致召回率低。资料表明,过滤比达到 90% 以上时,性能和召回率会急剧下降。
因此这里我们选择改写执行计划,将标量过滤和暴力检索算法结合,可以取得令人满意的效果。
4. 总结展望
来看一张流传已久的图。就像人类从猴子进化到直立智人一样,半/非结构化数据的存储系统也经历了从最简单的 key-value 发展出各种更加复杂的结构的系统,最后长得越来越像传统数据库。
业界的向量数据库随着大模型的热度也纷纷冒出来,有这些专用的向量数据库,创业团队入局的商用的、开源的。一些传统开源的存储系统也都加上了向量能力。
结合自身业务场景,百度自研了 Puck/Tinker 向量检索算法,比业界开源向量库具备一定优势,曾获得 BigANN 大赛第一。