在之前的博客中，我们概述了在 Google Cloud 上实施推荐系统的三种方法，包括 (1) 使用Recommendations AI 的完全托管解决方案，(2) 来自BigQuery ML的矩阵分解，以及 (3) 使用双塔的自定义深度检索技术编码器和Vertex AI 匹配引擎。在此博客中，我们深入探讨选项 (3)，并演示如何通过使用 Vertex AI 从头开始实施端到端候选检索工作流来构建播放列表推荐系统。具体来说，我们将涵盖：
- 检索建模的演变以及为什么双塔编码器在深度检索任务中很受欢迎
- 使用Spotify 百万播放列表数据集(MPD)构建播放列表延续用例
- 使用TensorFlow Recommenders (TFRS) 库开发自定义双塔编码器
- 使用Vertex AI 匹配引擎在近似最近邻 (ANN) 索引中提供候选嵌入
- 所有相关代码都可以在这个GitHub 存储库中找到。
背景
为了满足低延迟服务需求，大规模推荐系统通常作为多阶段系统部署到生产环境中。第一阶段（候选检索）的目标是筛选大型（>100M 元素）候选项目语料库并检索相关项目子集（~数百）用于下游排名和过滤任务。为了优化这个检索任务，我们考虑了两个核心目标：
- 在模型训练期间，找到将所有知识编译成嵌入的最佳方法
query, candidate。 - 在模型服务期间，以足够快的速度检索相关项目以满足延迟要求
图 1：多阶段推荐系统的概念组件；本博客的重点是第一阶段，候选人检索。
双塔架构在检索任务中很受欢迎，因为它们捕获查询和候选实体的语义，并将它们映射到共享的嵌入空间，以便语义相似的实体更紧密地聚集在一起。这意味着，如果我们计算给定查询的向量嵌入，我们可以在嵌入空间中搜索最接近（最相似）的候选者。由于这些基于神经网络的检索模型利用了元数据、上下文和特征交互，因此它们可以生成信息丰富的嵌入，并提供针对各种业务目标进行调整的灵活性。
图 2：双塔编码器模型是一种特定类型的基于嵌入的搜索，其中一个深度神经网络塔生成查询嵌入，第二个塔计算候选嵌入。计算两个嵌入向量之间的点积可以确定候选与查询的接近程度（相似程度）。资料来源：宣布 ScaNN：高效向量相似性搜索。
虽然这些功能有助于实现有用的query, candidate嵌入，但我们仍然需要解决检索延迟要求。为此，双塔架构提供了另一个优势：能够分离查询和候选项目的推理。这种解耦意味着可以预先计算所有候选项嵌入，从而将服务计算减少为 (1) 将查询转换为嵌入向量，以及 (2) 搜索相似向量（在预先计算的候选项中）。
随着候选数据集扩展到数百万（或数十亿）个向量，相似性搜索通常成为模型服务的计算瓶颈。放宽搜索以近似计算距离可以显着改善延迟，但我们需要尽量减少对搜索准确性（即相关性、召回率）的负面影响。
在论文Accelerating Large-Scale Inference with Anisotropic Vector Quantization中，Google 研究人员使用一种新颖的压缩算法解决了这种速度与准确性的权衡问题，与以前的最先进方法相比，该算法提高了检索的相关性和速度。在谷歌，这项技术被广泛采用，以支持跨搜索、YouTube、广告、智能镜头等的深度检索用例。虽然它在开源库 ( ScaNN ) 中可用，但实施、调整和扩展仍然具有挑战性。为了帮助团队在没有运营开销的情况下利用这项技术，谷歌云提供这些功能（以及更多）作为Vertex AI 匹配引擎的托管服务。
这篇文章的目的是演示如何使用 Vertex AI 实施这些深度检索技术，并讨论团队需要评估其用例的决策和权衡。
图 3：在 Vertex AI 上进行双塔训练和部署的参考架构。
用于深度检索的双塔
为了更好地理解双塔架构的好处，让我们回顾一下候选人检索中的三个关键建模里程碑。
检索建模的演变
传统的信息检索系统严重依赖于基于标记的匹配，其中使用 n-gram 的倒排索引检索候选对象。这些系统是可解释的、易于维护的（例如，没有训练数据），并且能够实现高精度。然而，他们通常召回率很低（即，难以找到给定查询的所有相关候选词），因为他们寻找与关键词完全匹配的候选词。虽然它们仍然用于特定的搜索用例，但今天的许多检索任务要么采用基于嵌入的技术，要么被基于嵌入的技术所取代。
图 4：基于标记的匹配通过匹配在查询和候选项目中找到的关键词来选择候选项目。
基于分解的检索引入了一个简单的基于嵌入的模型，该模型通过捕获对之间的相似性并将它们映射到共享嵌入空间来提供更好的泛化。这种协同过滤query, candidate技术的主要好处之一是嵌入是从隐式查询候选交互中自动学习的。从根本上说，这些模型分解了完整的查询-候选交互（共现）矩阵，以生成查询和候选的更小、更密集的嵌入表示，其中这些嵌入向量的乘积是交互矩阵的良好近似。这个想法是，通过将整个矩阵压缩成 k 维，模型学习了前 k 个潜在因素来描述query, candidate对关于建模任务。
图 5：基于分解的模型将查询候选交互矩阵分解为两个捕获查询候选交互的低阶矩阵的乘积。
最新的检索建模范例，通常称为神经深度检索 (NDR)，产生相同的嵌入表示，但使用深度学习来创建它们。NDR 模型（如双塔编码器）通过处理具有连续网络层的输入特征来应用深度学习，以学习数据的分层表示。实际上，这会产生一个充当信息蒸馏管道的神经网络，其中重复转换原始的多模态特征，从而放大有用的信息并过滤掉不相关的信息。这导致能够学习非线性关系和更复杂的特征交互的高度表达模型。
图 6：像双塔编码器这样的 NDR 架构在概念上类似于因式分解模型。两者都是基于嵌入的检索技术，计算查询和候选的低维向量表示，其中这两个向量之间的相似性是通过计算它们的点积来确定的。
在双塔架构中，每个塔都是一个神经网络，处理查询或候选输入特征以生成这些特征的嵌入表示。因为嵌入表示只是长度相同的向量，所以我们可以计算这两个向量之间的点积以确定它们的接近程度。query, candidate这意味着嵌入空间的方向由训练示例中每对的点积决定。
最佳服务的解耦推理
除了提高表现力和泛化能力之外，这种架构还提供了优化服务的机会。由于每个塔仅使用其各自的输入特征来生成向量，因此经过训练的塔可以单独运行。用于检索的塔的解耦推理意味着当我们在野外遇到它的对时，我们可以预先计算我们想要找到的东西。这也意味着我们可以不同地优化每个推理任务：
- 使用经过训练的候选塔运行批量预测作业，为所有候选塔预计算嵌入向量，附加 NVIDIA GPU 以加速计算
- 将预先计算的候选嵌入压缩到针对低延迟检索优化的 ANN 索引；将索引部署到服务端点
- 将经过训练的查询塔部署到端点以实时将查询转换为嵌入，连接 NVIDIA GPU 以加速计算
训练双塔模型并使用 ANN 索引为它们提供服务不同于训练和提供传统机器学习 (ML) 模型。为了阐明这一点，让我们回顾一下实施该技术的关键步骤。
图 7：在 Vertex AI 上进行双塔训练和部署的参考架构。
- 离线训练组合模型（双塔）；每个塔都为不同的任务单独保存
- 将查询塔上传到 Vertex AI 模型注册表并部署到在线端点
- 将候选塔上传到 Vertex AI 模型注册表
- 请求候选塔预测每个候选轨道的嵌入，将嵌入保存在 JSON 文件中
- 从嵌入 JSON 创建 ANN 服务索引，部署到在线索引端点
- 用户应用程序使用播放列表数据调用 endpoint.predict()，模型返回表示该播放列表的嵌入向量
- 使用播放列表嵌入向量搜索N个最近的邻居（候选曲目）
- 匹配引擎返回 N 个最近邻居的产品 ID
问题框架
在此示例中，我们使用 MPD 构建一个推荐用例，播放列表延续，其中为给定播放列表（查询）推荐候选曲目。该数据集是公开可用的，并为该演示提供了几个好处：
- 包括难以复制的实体（例如，播放列表、曲目、艺术家）之间的真实关系
- 大到足以复制生产中可能发生的可伸缩性问题
- 各种特征表示和数据类型（例如，播放列表和曲目 ID、原始文本、数字、日期时间）；能够使用来自Spotify Web Developer API的额外元数据丰富数据集
- 团队可以通过收听检索到的候选曲目来分析建模决策的影响（例如，为您自己的 Spotify 播放列表生成推荐）
训练实例
为推荐系统创建训练示例是一项非常重要的任务。与任何 ML 用例一样，训练数据应该准确地代表我们试图解决的潜在问题。如果不这样做，可能会导致模型性能不佳，并对用户体验造成意想不到的后果。YouTube 推荐的深度神经网络论文中的一个这样的教训强调，与“观看时间”等功能相比，严重依赖“点击率”等功能可能会导致推荐点击诱饵（即视频用户很少完成）更好地捕捉用户的参与度。
训练示例应该表示数据中的语义匹配。对于播放列表延续，我们可以将语义匹配视为将播放列表（即一组曲目、元数据等）与足够相似的曲目配对，以保持用户参与他们的收听会话。我们的训练示例的结构如何影响这一点？
- 训练数据来自正
query, candidate对 - 在训练期间，我们通过它们各自的塔向前传播查询和候选特征以产生两个向量表示，我们从中计算表示它们相似性的点积
- 在训练之后和服务之前，调用候选塔来预测（预计算）所有候选项目的嵌入
- 在服务时，模型处理给定播放列表的特征并生成向量嵌入
- 播放列表的向量嵌入用于搜索以在预先计算的候选索引中找到最相似的向量
- 候选和播放列表向量在嵌入空间中的位置以及它们之间的距离由训练示例中反映的语义关系定义
最后一点很重要。similar因为嵌入空间的质量决定了检索的成功与否，创建该嵌入空间的模型需要从最能说明给定播放列表和要检索的曲目之间关系的训练示例中学习。
这种高度依赖于配对数据选择的相似性概念突出了准备描述语义匹配的特征的重要性。成对训练的模型playlist title, track title将不同于成对训练的模型来定位候选轨道aggregated playlist audio features, track audio features。
从概念上讲，由对组成的训练示例playlist title, track title将创建一个嵌入空间，其中属于相同或相似标题（例如，beach vibes和beach tunes）的播放列表的所有曲目将比属于不同播放列表标题（例如，beach vibesvs workout tunes）的曲目靠得更近；由对组成的示例aggregated playlist audio features, track audio features将创建一个嵌入空间，其中属于具有相似音频配置文件（例如，live recordings of instrumental jams和high energy instrumentals）的播放列表的所有曲目将比属于具有不同音频配置文件（例如，live recordings of instrumental jamsvs acoustic tracks with lots of lyrics）的播放列表的曲目更靠近。
这些例子的直觉是，当我们以描述曲目如何出现在某些播放列表上的格式构建丰富的曲目播放列表特征时，我们可以将此数据提供给双塔模型，该模型学习父播放列表和播放列表之间的所有利基关系。儿童轨道。现代深度检索系统通常会考虑用户配置文件、历史参与和上下文。虽然我们在此示例中没有用户和上下文数据，但可以轻松地将它们添加到查询塔中。
使用 TFRS 实施深度检索
在使用 TFRS 构建检索模型时，两座塔是通过模型子类化实现的。每个塔都是单独构建的，可调用以处理输入特征值，将它们传递给特征层，然后连接结果。这意味着塔只是生成一个连接的向量（即查询或候选的表示；无论塔代表什么）。首先，我们定义塔的基本结构并将其实现为子类化的 Keras 模型：
class Playlist_Tower(tf.keras.Model): ''' produced embedding represents the features of a Playlist known at query time ''' def __init__(self, layer_sizes, vocab_dict): super().__init__() # TODO: build sequential model for each feature here def call(self, data): ''' defines what happens when the model is called ''' all_embs = tf.concat( [ # TODO: concatenate output of all features defined above ], axis=1) # pass output to dense/cross layers if self._cross_layer is not None: cross_embs = self._cross_layer(all_embs) return self.dense_layers(cross_embs) else: return self.dense_layers(all_embs)
我们通过为塔正在处理的每个特征创建 Keras 顺序模型来进一步定义子类塔：
# Feature: pl_name_src self.pl_name_src_text_embedding = tf.keras.Sequential( [ tf.keras.layers.TextVectorization( vocabulary=vocab_dict['pl_name_src'], ngrams=2, name="pl_name_src_textvectorizor" ), tf.keras.layers.Embedding( input_dim=MAX_TOKENS, output_dim=EMBEDDING_DIM, name="pl_name_src_emb_layer", mask_zero=False ), tf.keras.layers.GlobalAveragePooling1D(name="pl_name_src_1d"), ], name="pl_name_src_text_embedding" )
因为播放列表中表示的特征STRUCT是序列特征（列表），我们需要重塑嵌入层输出并使用 2D 池化（与应用于非序列特征的 1D 池化相反）：
# Feature: artist_genres_pl self.artist_genres_pl_embedding = tf.keras.Sequential( [ tf.keras.layers.TextVectorization( ngrams=2, vocabulary=vocab_dict['artist_genres_pl'], name="artist_genres_pl_textvectorizor" ), tf.keras.layers.Embedding( input_dim=MAX_TOKENS, output_dim=EMBED_DIM, name="artist_genres_pl_emb_layer", mask_zero=False ), tf.keras.layers.Reshape([-1, MAX_PL_LENGTH, EMBED_DIM]), tf.keras.layers.GlobalAveragePooling2D(name="artist_genres_pl_2d"), ], name="artist_genres_pl_emb_model" )
两座塔都建成后，我们使用 TFRS 基础模型类 ( tfrs.models.Model ) 来简化组合模型的构建。我们将每个塔都包含在类中__init__并定义compute_loss方法：
class TheTwoTowers(tfrs.models.Model): def __init__(self, layer_sizes, vocab_dict, parsed_candidate_dataset): super().__init__() self.query_tower = Playlist_Tower(layer_sizes, vocab_dict) self.candidate_tower = Candidate_Track_Tower(layer_sizes, vocab_dict) self.task = tfrs.tasks.Retrieval( metrics=tfrs.metrics.FactorizedTopK( candidates=parsed_candidate_dataset.batch(128).map( self.candidate_tower, num_parallel_calls=tf.data.AUTOTUNE ).prefetch(tf.data.AUTOTUNE) ) ) def compute_loss(self, data, training=False): query_embeddings = self.query_tower(data) candidate_embeddings = self.candidate_tower(data) return self.task( query_embeddings, candidate_embeddings, compute_metrics=not training, candidate_ids=data['track_uri_can'], compute_batch_metrics=True )
密集和交叉层
我们可以通过在连接的嵌入层之后添加密集层来增加每个塔的深度。由于这将强调学习连续的特征表示层，因此可以提高我们模型的表达能力。
同样，我们可以在嵌入层之后添加深层和交叉层，以更好地建模特征交互。在与建模隐式特征交互的深层相结合之前，跨层对显式特征交互进行建模。这些参数通常会带来更好的性能，但会显着增加模型的计算复杂性。我们建议评估不同的深层和跨层实现（例如，并行与堆栈）。有关详细信息，请参阅 TFRS深度和交叉网络指南。
特征工程
由于基于分解的模型提供了一种纯粹的协同过滤方法，使用 NDR 架构的高级特征处理使我们能够将其扩展到还包含基于内容的过滤的各个方面。通过包含描述播放列表和曲目的附加功能，我们为 NDR 模型提供了学习有关对的语义概念的机会playlist, track。包含标签特征（即关于候选轨道的特征）的能力也意味着我们训练的候选塔可以为训练期间未观察到的候选轨道计算嵌入向量（即冷启动）。从概念上讲，我们可以认为这样一个新的候选曲目嵌入编译了从具有相同或相似特征值的候选曲目中学习到的所有基于内容和协同过滤的信息。
有了这种添加多模态特征的灵活性，我们只需要处理它们以生成具有相同维度的嵌入向量，这样它们就可以连接起来并馈送到后续的深层和交叉层。这意味着如果我们使用预先训练的嵌入作为输入特征，我们会将它们传递到连接层（参见图 8）。
图 8：从输入到串联输出的特征处理图示。文本特征是通过 n-grams 生成的。n-gram 的整数索引被传递到嵌入层。散列产生不超过 1,000,000 的唯一整数；传递给嵌入层的值。如果使用预训练嵌入，则这些嵌入无需转换即可通过塔并与其他嵌入表示连接。
哈希与 StringLookup() 层
当需要快速性能时，通常建议使用散列法，并且优于字符串查找，因为它不需要查找表。为散列层设置合适的 bin 大小至关重要。当唯一值多于散列箱时，值开始被放入相同的箱中，这会对我们的建议产生负面影响。这通常称为散列冲突，可以在构建模型时通过为唯一值分配足够的 bin 来避免。有关更多详细信息，请参阅将分类特征转换为嵌入。
TextVectorization() 层
文本特征的关键是了解使用TextVectorization层创建额外的 NLP 特征是否有帮助。如果从文本特征派生的附加上下文很少，那么模型训练的成本可能不值得。该层需要根据源数据集进行调整，这意味着该层需要扫描训练数据以为前 N 个 n-gram（由 设置）创建查找词典max_tokens。
图 9：指导特征工程策略的决策树。
使用匹配引擎进行高效检索
到目前为止，我们已经讨论了如何将查询和候选映射到共享嵌入空间。现在让我们讨论如何最好地使用这个共享的嵌入空间来提供高效的服务。
回想一下在服务时，我们将使用经过训练的查询塔来计算查询（播放列表）的嵌入，并在最近邻搜索中使用该嵌入向量来寻找最相似的候选（曲目）嵌入。而且，由于候选数据集可以增长到数百万或数十亿个向量，因此这种最近邻搜索通常成为低延迟推理的计算瓶颈。
许多最先进的技术通过压缩候选向量来解决计算瓶颈，使得 ANN 计算可以在穷举搜索所需时间的一小部分内执行。Google Research 提出的新颖压缩算法修改了这些技术，以优化最近邻搜索精度。他们提出的技术的细节在此处描述，但从根本上说，他们的方法寻求压缩候选向量，以便保留向量之间的原始距离。与以前的解决方案相比，这导致向量及其最近邻居的相对排名更准确，即，它最大限度地减少了扭曲我们的模型从训练数据中学习的向量相似性。
完全托管的矢量数据库和 ANN 服务
匹配引擎是一种托管解决方案，利用这些技术进行高效的矢量相似性搜索。它为客户提供高度可扩展的矢量数据库和 ANN 服务，同时减轻开发和维护类似解决方案（例如开源ScaNN库）的运营开销。它包含多项可简化生产部署的功能，包括：
- 大规模：支持具有多达 10 亿个嵌入向量的大型嵌入数据集
- 增量更新：根据向量的数量，完整的索引重建可能需要数小时。通过增量更新，客户可以在不构建新索引的情况下进行小的更改（有关更多详细信息，请参阅更新和重建活动索引）
- 动态重建：当索引增长超出其原始配置时，匹配引擎会定期重新组织索引和服务结构以确保最佳性能
- 自动缩放：底层基础设施是自动缩放的，以确保在规模上保持一致的性能
- 过滤和多样性：每个向量包含多个限制和拥挤标签的能力。在查询推理时，使用布尔谓词来过滤和多样化检索到的候选对象（有关详细信息，请参阅过滤向量匹配）
创建 ANN 索引时，Matching Engine 使用Tree-AH策略来构建候选索引的分布式实现。它结合了两种算法：
- 用于分层组织嵌入空间的分布式搜索树。这棵树的每一层都是下一层节点的聚类，其中最终叶层是我们的候选嵌入向量的聚类
- 非对称哈希 (AH) 用于快速点积近似算法，用于对查询向量和搜索树节点之间的相似性进行评分
图 10：分区候选向量数据集的概念表示。在查询推理期间，对所有分区质心进行评分。在与查询向量最相似的质心中，对所有候选向量进行评分。对评分的候选向量进行聚合和重新评分，返回前 N 个候选向量。
该策略将我们的嵌入向量分成多个分区，其中每个分区由它包含的向量的质心表示。这些分区质心的集合形成一个较小的数据集，汇总了较大的分布式矢量数据集。在推理时，匹配引擎对所有分区质心进行评分，然后对质心与查询向量最相似的分区内的向量进行评分。
结论
在这篇博客中，我们深入了解了使用TensorFlow Recommenders和Vertex AI Matching Engine的候选检索工作流程的关键组件。我们仔细研究了双塔架构的基本概念，探索了查询和候选实体的语义，并讨论了训练示例的结构等因素如何影响候选检索的成功。
在后续博文中，我们将演示如何使用 Vertex AI 和其他 Google Cloud 服务来大规模实施这些技术。我们将展示如何利用 BigQuery 和 Dataflow 构建训练示例并将它们转换为TFRecords以进行模型训练。我们将概述如何构建用于使用 Vertex AI 训练服务训练双塔模型的 Python 应用程序。我们将详细介绍操作训练有素的塔的步骤。
