【CTR】DR：字节深度检索召回模型

本文介绍的内容为字节跳动 2020 年的工作——《Deep Retrieval: An End-to-End Learnable Structure Model for Large-Scale Recommendations》，一个用于大规模召回的端到端模型。

工业界推荐系统大多分为两个部分：召回和排序。我们知道由于召回阶段的面对的候选集数量非常庞大，所以其核心问题在于如何高效地在接近线性时间的消耗内获取最相关的召回列表。目前大部分的做法是多路召回，其中也会有利用 Embedding 进行最大内积搜索优化算法（Maximum Inner Product Search，MIPS）来获取召回列表，但这种情况会带来一部分的精度损失。

为此，字节的同学提出了一个端到端的模型架构——Deep Retrieval（以下简称 DR）。DR 将所有的候选集编码到离散的隐式空间中，随其他网络参数一起学习。模型上线时则采用 beam search 的方式获取最相关的候选集。最终实验证明，DR 模型时间复杂度接近线性，同时可以取得与暴力枚举相当的结果。

1.背景

目前工业界基于向量內积检索的召回算法有两大缺点：

• 首先，Embdding 向量学习的目标和最大內积搜索算法的结构其实并不是完全契合；
• 其次，依赖 user 和 item 的 Embedding 向量的內积来表示相似度，这种方式限制了模型的表达能力。

为了克服这两大缺点，阿里提出了基于树的检索算法 TDM/JTM。TDM/JTM 将索引建模成一棵树结构，候选集的每个 Item 为树中的叶子结点，并将模型的参数学习和树结构的参数学习结合起来训练。这种很好的提高了检索的精度，但是基于树的检索算法也有很明显的问题：

• 首先，树结构本身很难学习，而且树结构的部分叶子结点可以会因为稀疏数据而导致学习不充分；
• 其次，候选集只属于一个叶子节点，不符合常理。同时这也限制了模型只能从一个角度来刻画候选集，影响了模型的表达。

为此，作者提出了一种端到端的模型训练框架“深度检索” DR，该模型使用 D*K 维的矩阵来作为索引结构（如下图所示）。每个 Item 都需要走 D 步，每一步有 K 种选择。走到最后会有 K^D 可能的路径（也可以叫做编码），每一条路径可以代表一个包含多个 Item。每个 Item 可以被这样编码一次或多次，所以 Item 也可以属于多条路径。

DR 模型有两个优势：

• 训练方便：Item 的路径可以和神经网络参数一起使用 EM 类型的算法进行联合训练；
• 表达能力强：Item 和 Embedding 之间是多对多的编码范式，故其可以表达更复杂的关系。

2.模型架构

本节详细介绍下 DR 模型的细节。先从单路的目标函数说起，再延伸到多路目标函数，最后简单介绍下 Beam Search 算法。

2.1 目标函数

前文说到，模型包含 D 层，且每一层包含 K 个节点。每层的输出为 K 个节点的概率分布，以上绿色路径为例，我们来看一下前向传播的过程。

另 x 为 User，y 为 Item，给定训练样本 (x,y) 表示 User 和 Item 的交互（点击、收藏等）。设 为所有 Items 的 Label，则 Item 到 Path 的映射为 。

第 1 层我们输入的 User Embedding：Emb(x)，经过 Softmax 后输出 ，其中 为参数。

第 2～D 层我们将 User Embedding：Emb(x) 与前面所有层的输出进行拼接，经过 Softmax 后得到 ；

最后根据链式法则我们可以得到：

其中 。

给定 N 个训练样本，我们用最大似然估计得到：

2.2 多路经目标函数

在以前的树模型中每个 Item 只属于一类，这并不符合基本认知（比如说玫瑰属于花类但也属于礼物），同时也限制了模型的表达能力。为此 DR 模型打破了该限制：每个 Item 可以属于多类。在模型里表现为每个 Item 可以有多个路径：。

假设候选 Item 有 J 条路径，那么相应的训练目标变化为：

2.3 束搜索

对于 MLE 算法训练的模型，Beam search 只在预测的时候需要。假设每一层有三个节点(a,b,c)，Beam Size 为 2，Beam search 搜索过程大致如下：

• 在第一层时，选择概率最大的 2 个，假设为 a b，那么当前的 2 个路径就是 a 和 b；
• 到第二层时，我们将当前路径 a 和 c 分别与该层的所有节点进行组合，得到新的 6 个组合 aa ab ac ca cb cc，计算每个组合的得分并选择得分最高 2 个序列，作为新的当前序列，假如为 aa cb。
• 后面会不断重复这个过程，直到遇到结束符或者达到最大长度为止。最终输出得分最高的 2 个路径。

预测过程中，我们为每个 User Embedding 进行搜索从而获取多条候选路径，并将路径下包含的 Items 合并起来作为候选集。算法伪代码如下：

在搜索过程中，模型在每一层只保留最好的 B 个节点，所以到最后一层后会获得 B 个路径。故时间复杂度。

3.训练

介绍完框架和目标函数后，我们来讨论下模型如何训练。

与常见的使用梯度下降算法进行训练不同的是，我们采用 EM 算法进行训练。主要原因在于我们目标的优化目标同时涉及路径和 Item 之间的映射关系，属于离散型关系，而梯度下降是针对连续型关系的。

3.1 EM 算法

给定一个 user-item 的 pair-wise ，Item 的 paths 有 ，参数为 。

根据所有可能的映射关系 ，我们来最大化目标函数。由于路径的搜索空间过大 ，我们没法遍历所有的路径。所以我们只考虑 Beam Search 给出的路径，并且令其他路径为 0。但此时又会碰到 log(0) 的问题，为此我们使用目标函数的上界来作为目标函数：

其中， 为指示函数。

在 E 步骤，实际过程中无法计算出所有可能路径的分数，因此我们主要通过 Beam Search 获取一个较高的分数；然后在 M 步骤，我们最大化目标函数。

算法伪代码如下：

3.2 过拟合

为了丰富路径的表达，防止出现某一条路径对于所有输入都可取得最高值，我们需要增加过拟合项。即对于包含 Item 数目较大的路径进行惩罚。

其中， 为惩罚因子， 为 Item 的数量，f 是一个递增凸函数。作者使用 。

惩罚项只应用于 M 步，而不是在连续参数的线性回归训练中使用，（因为 M 步会分配 Item）。

3.3 联合训练

作者在实验过程中发现加入一个 Softmax 分类模型和 DR 模型进行联合训练可以大大提高性能。作者推测这是由于项目的路径在开始是随机分配的，所以增大了优化的难度。而通过与一个容易训练的 Softmax 模型共享输入，能够加速模型的优化。我们最终的优化目标为：

在使用算法 3 进行 BeamSearch 来检索一组候选项之后，我们使用 Softmax 函数重新对这些候选项进行排序，从而获得最终的最优候选项。

4.实验

本节将介绍下 DR 模型与树模型 TDM、JTM 在不同数据集中的效果对比。

下面两张图我们可以看到，DR 模型的效果大大提升，并且和暴力枚举的效果相当：

我们再来看一下参数敏感性：

• 每一层的 K 控制着模型的表达能力；
• 路径的数量 J 控制着 Item 的多样性；
• Beam Size B 控制着被召回的候选路径的个数；
• 控制惩罚项。

5.结论

总结：本文提出了一个用于大型推荐系统的端到端可学习的召回模型——DR（Deep Retrieval）。DR 使用 EM 算法和联合训练来优化模型的参数。实验表明，DR 模型的效果远远高于其他模型，并且效果接近暴力枚举。