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

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 分析训练性能。 Profiler 是开源 TensorBoard Profiler 的托管式实现。

通过在训练作业中传递 –profiler 实参,您可以启用 TensorFlow 回调,以便为每个 epoch 分析指定数量的批次。该配置文件会捕获来自主机 CPU 以及设备 GPU 或 TPU 硬件的轨迹。轨迹提供有关训练作业资源消耗的信息。为避免出现内存不足错误,我们建议您先将分析时长设置为 2 到 10 个训练步数,然后根据需要增加。

如需了解如何将 Profiler 与 Vertex AI Training 和 Vertex AI TensorBoard 搭配使用,请参阅分析模型训练性能。 如需了解调试最佳实践,请参阅优化 GPU 性能。 如需了解如何优化性能,请参阅使用 Profiler 优化 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 | 资深开发技术推广工程师