Fork me on GitHub

阿里妈妈 Dolphin 分布式向量召回技术详解

三贤、陌奈、赫琦 阿里妈妈技术 稿

前言

随着深度学习技术发展,万物皆可用向量表示,向量召回计算已经成为很多算法场景所需的必备能力,其广泛应用在搜索、推荐和广告等业务场景中。阿里妈妈工程平台智能分析引擎团队为了更好地支持智能物料推荐广告场景,针对其吞吐规模大、要求延迟低、查询条件灵活等需求,在 Dolphin 引擎 (面向阿里妈妈广告营销场景的智能分析引擎,兼容 Greenplum)中实现了支持分布式高性能的向量召回技术。该技术主要是将 Faiss 封装成 **gpdb-faiss-vector **插件,部署在 Dolphin/Greenplum 上,基于此,我们搭建了在离线向量召回系统。其已经服务阿里妈妈多个业务场景,包括达摩盘人群推荐、直通车关键词推荐、投放计划成效预估和广告工具推荐等。

目前该方案已经开源,本文将结合业界方案,对 gpdb-faiss-vector 的业务背景、核心功能做简要介绍,欢迎试用体验与交流讨论。

**GitHub地址:**https://github.com/AlibabaIncubator/gpdb-faiss-vector

业界方案

向量召回技术主要是求解 KNN 问题或 RNN 问题,KNN(K-Nearest Neighbor)是指在向量空间中找到距离查询点最近的 K 个点,RNN(Radius Nearest Neighbor)则是指找到查询点距离阈值内的点。工业界在应用向量召回技术时,通常要处理的向量数量巨大,准确求解的计算成本很高。于是人们常常将其简化为 ANN(Approximate Nearest Neighbor)问题,只求近似解。

针对ANN问题,业界也提出了多种解法,可分为三大类,适用场景不同:

  • 空间划分类 :如基于欧式空间、采用多维二叉树结构的 KD-Tree,在检索时可以快速定位到相关的子集,从而减少了要扫描的向量个数;
  • 空间编码和转换类 :如将高维向量降维映射到低纬空间的局部敏感哈希,可以大大减少扫描的计算量;
  • 邻居图类 :如基于邻居的邻居也可能是邻居思想的邻居图的 HNSW,预先建立关系图,减少要扫描的向量个数,加快检索时的距离收敛速度。

随着技术的发展,目前业界涌现出一批优秀的向量召回方案,各有千秋,可大致分为三类:

  • 开源向量召回库 :如 Facebook 的 Faiss,微软的 SPTAG 等。其通常内部包含多种检索算法,种类齐全,业界有名;但其只实现了向量召回的核心功能,通常是作为库嵌入到程序中提供服务,并未对其他周边如存储、管控等提供支持,仅支持单机使用,还需要进行产品化改造,使用成本高。
  • 分布式向量召回引擎 :Milvus,京东的 Vearch,微信的 SimSvr。其通常功能强大,配套设施齐全,除提供向量召回的功能外,还支持弹性伸缩故障转移等;但需要单独部署,向量数据需要在向量存储数据库和检索引擎之间进行大量迁移,另外其通常使用自定义的 API,不支持基于 SQL 的操作。
  • 数据库向量召回插件 :PASE,pgvector。其提供 SQL 接口,用户使用成本低,且代码与数据库深度结合,充分利用其机制,可以很自然的实现混合查询(标签查询+向量召回);但不支持批量检索,缺乏分布式能力,另外支持的检索算法较少(多是基于KD-Tree,HNSW),无法为每种场景提供最佳选择。

Faiss(https://github.com/facebookresearch/faiss)是用于高效相似性检索和密集向量聚类的库。Greenplum(https://github.com/greenplum-db/gpdb)是基于 PostgreSQL 用于分析的大规模并行处理架构的分布式数据库,Dolphin 是面向阿里妈妈广告营销场景的智能分析引擎,具有计算存储分离、快速 scale-up 等特点,同时也兼容Greenplum。我们通过以插件的形式将 Faiss 集成到 Dolphin 引擎中,以使其可对外提供向量召回的能力,基于 gpdb-faiss-vector 插件与 Dolphin 引擎,形成了一套向量召回技术方案。该方案具有 SQL 化、分布式、索引类型丰富和支持在离线计算等特点,让向量召回的的使用体验、性能和扩展性大幅提升,可满足绝大多数场景。

业务背景

该框架已在阿里妈妈多个业务场景落地,既有离线查询场景也有在线查询场景。

  • 离线查询场景 :从 MaxCompute 上拉取特征向量集并建立索引。然后批量拉取查询向量进行向量召回,并将结果批量写入到 MaxCompute,期望高吞吐。
  • 在线查询场景 :预先建立向量索引,通常是交互式的传入单个向量来检索,然后系统将结果直接返回,期望低延迟。

系统架构

传统的向量召回从使用角度看都是调用接口,算法服务代码需要兼容客户端和接口API,开发和学习成本高,对于使用向量召回应用的算法同学来说,使用 SQL 学习成本最低,所以我们选择在数据库上插入向量计算能力。

得益于 PostgreSQL 灵活的插件机制,我们在 Dolphin 引擎的每个节点上部署了gpdb-faiss-vector 插件,使其具备了向量召回能力。插件内嵌缓存,可在数据库会话级别缓存 Faiss index,以便加速召回过程。

向量召回执行流程如下:

  • Master 节点在接收到向量召回 SQL 后,将任务分发到各个 Segment 节点去执行检索操作。

  • Segment 节点上的 PostgreSQL 去查询 Faiss缓存中是否有期望的 Faiss index 内存对象。如果没有,PostgreSQL 会从数据库中读取相关的 Faiss index 字节序列并将其反序列化成内存对象存入到缓存中。

  • 然后 PostgreSQL 会调用 Faiss index 内存实例对象进行检索,并将检索结果转换成 PostgreSQL 数据格式。最后 Master 节点将各个 Segment 节点检索得到的局部 TopK 个向量结果通过多路归并合并得到全局 TopK 个向量结果。

    图片

核心优势

该方案相对于现有业界方案有两大优势:一是操作简便,功能强大,可以全程使用 SQL 操作,和利用数据库的各种功能;二是高性能,可以分布式并行批量交互式检索。

SQL化算子

我们为向量召回的各个环节均提供了相应的 SQL 函数,可以使整个流程在数据库内部完成。库内处理(In-database processing)相较于传统方法不仅有显著的性能改进,还提供了更多功能和更好的使用体验。

图片

整个插件的 SQL 函数 API 在与 Faiss API 保持基本一致的同时还稍作扩充,既降低了 Faiss 使用者的上手成本,又提供了更多的自由度。整个 SQL 函数使用环节如下:

  1. create :根据用户参数创建相应的 Faiss index。
  2. train :训练 Faiss index,有些 Faiss index 需要通过训练阶段分析向量的分布。
  3. add :向 Faiss index 中添加待检索的向量。可以根据需求为每个向量指定特定的 ID,而不是默认的顺序递增序号。
  4. set_runtime_parameters :可调整 Faiss index 的运行期参数。
  5. search/range_search :检索并返回距离位于topk或者阈值内的邻居向量们 ID 和距离。该环节可以使用 Faiss缓存加速以达到低延迟的效果,还可以一次传入多个查询向量,进行批量并行查询,提高吞吐。
  6. topk_merge (可选):在分布式查询的情况下,不同 Segment 节点运行上述步骤得到局部 TopK 结果,在 Master 节点上进行合并,求得全局的 TopK 结果。

针对向量召回流程,我们提供了完备的功能支持,使整个检索流程全程可以通过 SQL 的形式组织起来,极大地方便了算法及工程同学使用。

灵活可配置

用户通常依据不同数据维度和分布选择不同算法或算法组合,进行相关的参数调整,根据具体场景需求实现精度和性能之间的平衡。

为了更灵活地满足用户需求,在 create 阶段,我们采用基于 Faiss 的 index_factory 工厂模式函数,可以解析字符串来生成 Faiss index,支持 Faiss 各种 index 类型(如 IVF、HNSW)和各种向量量化压缩(如 SQ、PQ)。这使得用户可以通过 SQL 创建任意复杂结构的 Faiss index,甚至可以是嵌套的。同时还支持指定预处理模块,倒排模块、编码模块。方便用户根据自身业务特点等灵活选用合适的检索方案。以 IDMap、HNSW32、Flat 为例,含义分别为:启用 ID 映射、使用顶点连接数是32(用户指定的构造期参数)的 HNSW 索引、索引中原样存储向量数据。

另外,在 set_runtime_parameters 阶段,可以为不同的 Faiss index 配置不同的运行期参数。基于 Faiss 的ParameterSpace 对象可以配置 Faiss index 的运行期参数的特性,SQL 函数在接收到传入的配置项字符串后,会调用 ParameterSpace 读取分析,以进行相应的修改。在使用 HNSW 的实践中,我们发现召回率不及预期,经过排查发现是 efSearch 参数设置过小导致的。为此我们只需要简单的修改下运行期参数即可,并不需要重新创建 Faiss index。

至此,我们提供了对 Faiss index 各类型的全面支持,以及配置各种参数的能力,用户可以根据业务需求进行各种选型试验。

其他功能

使用数据库搭建向量召回系统可以大量利用数据库自身的各种特性,轻松实现各种高级用法,而不需要重新进行大量基础设施建设。比如:

  • 原子更新 :在使用 Faiss index 检索到结果后,通常还需要与之相应的向量特征数据集,才能产生正确的召回结果。使用数据库的事务,可以轻松实现 Faiss index 与数据集的同步更新,省去了用户进行数据同步与对齐的烦恼。
  • 多索引 :一份向量特征数据集,可能涉及到业务上不同的应用场景,用户可以轻松建立多种 Faiss index,以应对不同的检索需求。
  • 联合检索 :很多业务场景,需要进行“有条件的向量召回”,即向量既要满足标签检索条件,又要满足相似性检索的要求,如电商中宝贝类目下的相似性检索。借助数据库自身的能力,预先将向量特征数据按照标签聚合建立Faiss index,检索时先进行标签过滤,然后再进行加载相关 Faiss index 来进行向量召回。

高性能

业务场景需要大量的向量召回操作,为此,我们从设计之初就着重考虑性能。

PostgreSQL 是最先进的单机数据库,具有强大的性能,但是在大量数据面前还是有些吃力。而 Dolphin 是分布式并行架构,因此我们可以将数据打散分片到不同节点,然后利用强大的并行能力,并行读取数据,并行检索,极大地降低了整体耗时。

另外可以把业务向量数据按照一定规律分布存放在 Dolphin 引擎上,并按照相同的分布规律创建 Faiss index,这样在创建索引时每个节点只需读写本地数据,省去大量网络开销。

批量查询优化

数据库传统的插件通常是面向点查询场景,并不适用批量查询的场景。而插件整合的 Faiss 支持批量查询,其底层可以利用多线程进行并行查询,同时还会用 OpenBLAS 中的矩阵运算取代向量计算,使用 CPU 高阶扩展指令集SIMD 等来加速关键环节,满足高吞吐的需求。

交互式查询优化

交互式查询场景通常是算法实时推荐场景,对响应时间的要求很高,需要我们着重降低延迟。

Dolphin 提供了 bytea 数据类型用于存储二进制数据。我们将 Faiss 序列化后的二进制数据以 bytea 的形式存储在Dolphin 中。在进行向量查询时,Dolphin 先是读取出 Faiss 序列化数据,然后将其反序列化成 Faiss 内存对象,最后根据检索向量在Faiss内存对象中进行检索。

在实践中,我们发现数据库读取 Faiss index 字节序列和将其反序列化为内存对象为主要耗时环节。为此,我们引入了可以缓存 Faiss index 内存对象的 缓存,以期减少乃至取消这两环节的耗时。

为了让缓存区别不同的 Faiss index,可以给每个 Faiss index 独一无二的标识符,如 Faiss index 字节序列的 MD5 值。在调用检索函数的时候,我们传入 Faiss index 字节序列的同时还传入 Faiss index字节序列标识符,流程如下:

  1. 检索函数先是拿标识符去查询 Faiss缓存中是否存在对应的 Faiss index;
  2. 如果不存在,则读取 Faiss index 字节序列,并反序列化成 Faiss index 内存对象,加入到 Faiss 缓存中;
  3. 如果存在,则直接使用 Faiss 缓存中的 Faiss index 内存对象去查询。

图片

引入缓存后,在命中缓存的情况下,系统可以提供近似在线服务的低延迟,典型查询场景总耗时由 100+ms 降至 2ms 左右。

为了防止缓存占用过多的内存影响到数据库的正常运行,我们需要限定缓存占用内存总大小,而 Faiss index 不是固定大小的,无法像传统缓存那样通过设置缓存条目数量来限制内存使用量。好在缓存提供了为每个缓存条目配置权重或费用(charge)的功能,我们只需将其设置成 Faiss index 字节序列的大小即可。

输出优化

实践中发现,在 Dolphin 中将组合类型展开成多列时,特别是多行时,性能损耗特别明显。因此,函数的输出格式多采用数组形式,一个检索向量对应的 K 个结果放在同一行里,极大地降低了输出结果的行数,减轻了数据库的负担。

但这样也带来了问题,就是分布式查询时无法利用数据库 ORDER BY 得到全局的 TopK 结果。为此,我们专门提供了 topk_merge 函数,接收各个 Segment 节点上检索得到的局部 TopK 结果,根据 search 结果内距离有序的特点,在 Master 节点上进行多路归并,求得全局的 TopK 结果。

未来规划

我们在 Dolphin 内实现了分布式、SQL 化和在离线一体的向量计算技术框架,使用成本显著降低,用户基于该方案可以快速开发上线向量召回业务,但在实践中发现还存在一些问题:

  • 当前采用的索引缓存策略在一定程度上解决了交互式查询的高延迟问题,但为实现该策略我们要为 Dolphin 的每个会话维护一块独立缓存,这无疑增加了内存压力。
  • 此外缓存策略不可避免地带来预热问题,即索引未命中的首次查询时间大大增加,造成交互式向量召回查询性能的波动。

为解决这两个问题,我们计划将向量召回计算能力从 Dolphin 中分离出来,在集群的每个节点上只保留一份独立的 Faiss search 服务,Dolphin 的每个 Segment 节点通过请求该服务即可实现向量召回。这样,内存中仅需缓存一份索引数据,也无需针对新会话重新加载数据,从而可以提供更加稳定、高效的向量召回计算服务。为进一步提升计算性能并扩大使用场景,未来我们将继续对 Dolphin 的向量召回计算引擎进行优化。也欢迎感兴趣的同学试用体验和我们交流。

关于我们

阿里妈妈工程平台智能分析引擎团队负责开发面向百万广告主的数据产品和计算引擎,支持万亿级数据毫秒级交互式人群圈选,关于人群的各种高阶细分洞察分析等。同时计算引擎还集成SQL语法的ML、AI能力,支持大规模广告物料智能推荐场景。


本文地址:https://www.6aiq.com/article/1661483821846
本文版权归作者和AIQ共有,欢迎转载,但未经作者同意必须保留此段声明,且在文章页面明显位置给出