实现两塔检索以进行大规模候选生成

Last reviewed 2025-01-16 UTC

本文档提供了一个参考架构,介绍了如何使用 Vertex AI 实现端到端双塔候选生成工作流。双塔模型框架是一种强大的检索技术,适用于个性化用例,因为它可以学习两个不同实体(例如网页查询和候选项)之间的语义相似性。

本文档适用于正在开发具有低延迟服务要求的大规模推荐应用的技术从业者,例如数据科学家和机器学习工程师。如需详细了解构建双塔模型的建模技术、问题建模和数据准备,请参阅使用 TensorFlow Recommender 和 Vector Search 扩缩深度检索

架构

下图显示了用于训练双塔模型的架构,并分别部署每个塔,以执行不同的部署和服务任务:

用于训练双塔模型并单独部署每个塔的架构。

图中显示的架构包括以下组件:

  • 训练数据:训练文件存储在 Cloud Storage 中。
  • 双塔训练:使用 Vertex AI Training 服务对组合双塔模型进行离线训练;每个塔都会单独保存,并用于不同的任务。
  • 已注册的查询塔和候选塔:塔训练完毕后,每个塔都会单独上传到 Vertex AI Model Registry。
  • 已部署的查询塔:已注册的查询塔部署到 Vertex AI 在线端点。
  • 批量预测嵌入:已注册的候选塔在批量预测作业中用于预计算所有可用候选项的嵌入表示法。
  • 嵌入 JSON:系统会将预测的嵌入保存到 Cloud Storage 中的 JSON 文件。
  • ANN 索引:Vertex AI Vector Search 用于创建配置为近似最近邻 (ANN) 搜索的服务索引。
  • 已部署的索引:ANN 索引部署到 Vertex AI Vector Search 索引端点。

使用的产品

此参考架构使用以下 Google Cloud 产品:

  • Vertex AI Training:全托管式训练服务,可让您实现大规模模型训练。
  • Vector Search:向量相似度匹配服务,可让您存储、编入索引和搜索语义相似或相关的数据。
  • Vertex AI Model Registry:一个中央代码库,您可以在其中管理机器学习模型的生命周期。
  • Cloud Storage:适用于各种数据类型的费用低廉且不受限制的对象存储。数据可从 Google Cloud内部和外部访问,并且跨位置进行复制以实现冗余。

使用场景

为了满足低延迟服务要求,大规模推荐系统通常会作为两阶段系统或有时作为多阶段系统部署到生产环境。第一阶段(候选项生成)的目标是过滤大量候选项,并检索数百个项中的相关子集,以便进行下游过滤和排名任务。如需优化此检索任务,请考虑以下两个核心目标:

  1. 在模型训练期间,学习要解决的问题或任务的最佳表示法,并将此表示法编译为 <query, candidate> 嵌入。
  2. 在模型服务期间,检索相关项的速度足够快,能够满足延迟时间要求。

下图显示了两阶段推荐系统的概念组件:

两阶段推荐系统的概念组件。

在该图中,候选集生成会过滤数百万个候选项。然后,系统会对生成的数百个候选项进行排名和过滤,以返回数十个推荐项。

本文档中的参考架构会训练基于双塔的检索模型。在该架构中,每个塔都是一个神经网络,用于处理查询或候选商品特征,然后生成这些特征的嵌入表示法。每个塔都是单独部署的,因为每个塔在生产环境中都将用于不同的任务:

  • 候选塔:候选塔用于预计算所有候选项的嵌入。预先计算的嵌入会部署到针对低延迟检索进行了优化的 Vertex AI Vector Search 索引端点。
  • 已部署的塔:在线上投放期间,已部署的查询塔会将原始用户查询转换为嵌入表示法。然后,系统会使用嵌入表示法在已部署的索引中查找类似的项嵌入。

双塔架构非常适合许多检索任务,因为双塔架构可以捕获查询和候选实体的语义关系,并将它们映射到共享的嵌入空间。将实体映射到共享嵌入空间后,语义上相似的实体会更紧密地聚集在一起。因此,如果您计算给定查询的向量嵌入,则可以搜索嵌入空间,以查找最接近(最相似)的候选项。这种架构的主要优势在于能够将查询和候选表示的推理分离。这种分离的主要优势有两个:

  • 您无需重新训练新的商品词汇,即可投放新的(新鲜)商品。通过将任意一组商品特征馈送给候选商品塔,您可以计算任意一组候选商品的商品嵌入,即使是训练期间未见过的候选商品也是如此。执行此计算有助于解决冷启动问题。
    • 候选塔可以支持任意一组候选项,包括尚未与推荐系统互动的项。之所以能够提供此支持,是因为双塔架构会处理每个 <query, candidate> 对的丰富内容和元数据特征。这种处理方式可让系统根据已知项来描述未知项。
  • 您可以通过预计算所有候选项嵌入来优化检索推理。这些预先计算的嵌入可以编入索引并部署到针对低延迟检索进行了优化的服务基础设施。
    • 借助塔的协同学习,您可以用查询来描述项,反之亦然。如果您有某个对的一半(例如查询),并且需要查找另一个对应的项,则可以提前预计算方程的一半。借助预计算,您可以尽快做出其余决策。

设计考虑事项

本部分提供的指导可帮助您在 Google Cloud 中开发满足安全性和性能需求的候选生成架构。本部分中的指导并非详尽无遗。根据您的具体要求,您可以选择考虑其他设计因素和权衡因素。

安全

Vertex AI Vector Search 支持公共端点和虚拟私有云 (VPC) 端点部署。如果您想使用 VPC 网络,请先按照设置 VPC 网络对等互连连接中的说明操作。如果向量搜索索引部署在 VPC 边界内,用户必须从同一 VPC 网络中访问关联的资源。例如,如果您是在 Vertex AI Workbench 中进行开发,则需要在与已部署的索引端点位于同一 VPC 网络中创建 Workbench 实例。同样,任何预计要创建端点或将索引部署到端点的流水线都应在同一 VPC 网络中运行。

性能优化

本部分介绍在使用此参考架构在 Google Cloud 中设计满足工作负载性能要求的拓扑时应考虑的因素。

分析训练作业

如需优化数据输入流水线和整体训练图,我们建议您使用 Cloud Profiler 分析训练性能。性能分析器是开源 TensorBoard 性能分析器的托管式实现。

通过在训练作业中传递 –profiler 参数,您可以让 TensorFlow 回调针对每个数据集对固定数量的批处理进行性能分析。该配置文件会捕获来自主机 CPU 和设备 GPU 或 TPU 硬件的轨迹。轨迹会提供有关训练作业的资源消耗情况的信息。为避免发生内存不足的错误,我们建议您先将配置文件时长设置为 2 到 10 个训练步骤,然后根据需要进行增加。

如需了解如何将 Profiler 与 Vertex AI Training 和 Vertex AI TensorBoard 搭配使用,请参阅分析模型训练性能。如需了解调试最佳实践,请参阅优化 GPU 性能。 如需了解如何优化性能,请参阅使用性能分析器优化 TensorFlow 性能

充分利用加速器

当您附加 NVIDIA GPU 或 Cloud TPU 等训练加速器时,请务必充分利用它们。充分利用训练加速器是费用管理方面的最佳实践,因为加速器是架构中最昂贵的组件。充分利用训练加速器也是提高作业效率的最佳实践,因为没有空闲时间会减少总体资源消耗。

为了让加速器保持充分利用,您通常需要执行几次迭代,找到瓶颈、优化瓶颈,然后重复这些步骤,直到加速器设备利用率达到可接受的水平。由于此用例的许多数据集太大而无法放入内存,因此存储空间、主机虚拟机和加速器之间通常会出现瓶颈。

下图展示了机器学习训练输入流水线的概念阶段:

机器学习训练输入流水线的概念阶段。

在该示意图中,系统会从存储空间读取数据并进行预处理。数据预处理完毕后,会发送到设备。如需优化性能,请先确定整体性能是受主机 CPU 还是加速器设备(GPU 或 TPU)所限。设备负责加速训练循环,而主机负责向设备提供训练数据并接收设备的结果。以下部分介绍了如何通过提高输入流水线性能和设备性能来解决瓶颈问题。

提升输入流水线性能
  • 从存储空间读取数据:如需提高数据读取速度,请尝试使用缓存、预提取顺序访问模式和并行 I/O。
  • 预处理数据:如需改进数据预处理,请为数据提取和转换配置并行处理,并调整数据输入流水线中的 interleave 转换
  • 将数据发送到设备:为了缩短总作业时间,请将数据从主机并行传输到多台设备。
提升设备性能
  • 增大小批次大小。小批量是指每台设备在训练循环的一次迭代中使用的训练样本数量。通过增加微批次大小,您可以提高操作之间的并行性并改进数据重用。不过,小批次必须能够与训练程序的其余部分一起放入内存。如果您过度增加小批次大小,可能会遇到内存不足错误和模型偏差。
  • 将用户定义的函数向量化。通常,数据转换可以表示为用户定义的函数,该函数描述了如何转换输入数据集的每个元素。如需对此函数进行向量化处理,您可以一次对一批输入应用转换操作,而不是一次转换一个元素。任何用户定义的函数都有与调度和执行相关的开销。转换一批输入时,您需要为每批输入支付一次开销,而不是为每个数据集元素支付一次开销。
先纵向扩容再横向扩容

为训练作业配置计算资源时,我们建议您先纵向扩容,然后再横向扩容。这意味着,您应先选择体积更大、性能更强的设备,然后再使用多台性能较弱的设备。我们建议您按以下方式扩容:

  1. 单个工作器 + 单个设备
  2. 单个工作器 + 功能更强大的设备
  3. 单个工作器 + 多台设备
  4. 分布式训练

如需评估 ANN 搜索的优势,您可以衡量给定查询的延迟时间和召回率。为了帮助调整索引,Vertex AI Vector Search 提供了创建暴力破解索引的功能。暴力破解索引会执行穷尽搜索(延迟时间较长),以查找给定查询向量的真实最近邻。暴力索引不适合用于生产环境,但在索引调整期间计算召回率时,它可以提供一个良好的基准。

为了评估召回率与延迟时间之间的关系,您可以将预计算的候选嵌入部署到一个配置为 ANN 搜索的索引,并部署到另一个配置为暴力搜索的索引。蛮力索引会返回绝对最近邻,但通常比 ANN 搜索需要更长时间。您可能愿意牺牲一些检索召回率来缩短检索延迟时间,但应评估这种权衡。影响召回率和延迟时间的其他特征包括:

  • 建模参数:许多建模决策会影响嵌入空间,而嵌入空间最终会成为服务索引。比较针对使用浅层检索模型和深层检索模型构建的索引检索到的候选项。
  • 维度:维度是最终由模型决定的另一个方面。ANN 索引的维度必须与查询和候选塔向量的维度一致。
  • 标记重合和过滤标记:标记可以提供强大的功能,以针对不同的生产用例量身定制结果。最佳实践是了解标记如何影响检索到的候选项并影响性能。
  • ANN 计数:增加此值会提高召回率,但可能会按比例增加延迟时间。
  • 要搜索的叶节点的百分比:要搜索的叶节点的百分比是评估检索率与延迟时间权衡的最关键选项。增加此值会提高召回率,但也可能会按比例增加延迟时间。

后续步骤

如需查看更多参考架构、图表和最佳实践,请浏览云架构中心

贡献者

作者:

其他贡献者:Kaz Sato | 开发者宣传主管