字节最新复杂召回模型，提出深度检索DR框架解决超大规模推荐系统中的匹配问题

导读：今天分享一下ByteDance最新公开的一篇关于复杂深度召回模型的论文“深度检索”，使用EM算法学习图路径模型，取得了与暴力算法相当的效果，推荐工业界实战干货论文，值得细读。

论文：Deep Retrieval: An End-to-End Learnable Structure Model for Large-Scale Recommendations

地址：https://arxiv.org/abs/2007.07203

摘要

工业界大规模推荐系统中一个非常核心的问题是如何高效地在接近线性时间的消耗内获取最相关的候选集。之前大家一般的做法是分为两步：首先训练一个內积模型；然后使用最大內积搜索优化算法（MIPS）来获得候选集，即便可能损失一部分的检索精度。在本文中，我们提出了一个端到端的模型框架深度检索DR。DR将所有的候选集编码到离散的隐式空间中，随着其他的网络参数一起学习。模型训练完成后的在线服务阶段，则通过beam search的方式获取最相关的候选集。实验证明了DR可以接近线性的计算复杂度，同时取得了与暴力算法相当的效果。

背景

尽管在工业界大规模推荐系统中，基于向量內积检索的召回算法获取了广泛的应用，但是它有着两个方面不足：首先，表征向量学习的目标和最大內积搜索算法的结构其实并不是完美契合的；其次，依赖于用户和item的embedding向量的內积来表示相似度限制了模型的能力。

为了打破基于向量內积检索模型的限制，阿里提出了基于树的检索算法TDM/ JTM。它们将索引建模成为一棵树结构，候选集的每个item则是树中的叶子节点。并且将模型参数的学习和树结构的学习完美结合起来提升检索的精度。但是基于树的检索算法也有着明显的问题：

1. 树结构本身很难学习，而且树结构对结果较为敏感，叶子节点的数据可能是稀疏的学习不充分；
2. 候选的item只能分配到一个叶子节点当中去，不符合常理。这其实也限制了模型从多个角度来刻画候选集item的表达。

本文中我们提出了一种端到端的模型训练框架“深度检索”DR，使用一个D x K维的矩阵来作为索引结构。模型预测需要走D步，每一步有K种选择。走到最后会有K^D可能的路径（也可以叫做编码），每一条路径可以代表一个包含多个item的集合，每个item反过来也可以同时在多条路径的集合里面。

DR框架的优势主要有如下的两个方面：首先在训练阶段，item的路径可以和NN网络参数一起使用EM类型的算法联合训练得到；其次从模型能力上，多对多的编码范式也允许DR框架可以学习表示user和item之间更复杂的关系。

DR框架模型

详细介绍一下DR框架的细节。首先，在给定模型参数θ的前提下构建用户x_i选择路径c_i的概率并且给出训练的目标；然后，介绍一下多路径机制来允许模型从多个角度捕获item的属性；最后在预测阶段，引入了一种beam search的算法基于user embedding来选择候选路径。

路径构建方式

整个建模思路是包含D层，每一层包含K个节点而且每一层是一个MLP的网络结构（当然这里也可以使用其他譬如RNN等的网络结构）。每一层输出的就是在K个节点上的概率分布。以下作图的绿色路径为例，根据链式法则可以得到它的概率为：

也就是说，c是一条路径的概率如下式所示。

训练目标

给定训练集，训练的目标函数就是最大化其似然函数：

前面有提到过TDM/JTM的主要问题是每个item只能属于一个叶子节点，DR框架打破了这个限制：每个候选item可以属于多个路径下的集合里面。假设候选item可以存在与J条路径的集合当中，那么相应的训练目标变化为：

在线的Beam search检索

在线预测阶段，给定user embedding作为输入，我们使用Beam search算法来获取多条候选路径并将他们下面包含的item合并起来作为最后的候选集。走到每一层始终只保留TopB个路径，直到走完最后一层，然后将TopB个路径对应的item召回，时间复杂度O(DKB logB)。

DR模型训练

训练步骤

前面介绍了DR框架的建模思路以及优化目标，优化目标上如果单纯只考虑网络参数的优化是连续的所以可以通过梯度下降的方式来进行优化，但是优化目标同时涉及了路径和item集合之间的映射关系，因为是离散的所以不能通过梯度优化器进行优化。因此，我们使用了EM类的算法来联合地优化映射关系以及网络参数。下面主要介绍EM算法的细节以及防止过拟合的惩罚项。

由于总的路径的可能性空间很大K^D量级太大，如果通过beam search算法记录Top的路径，剩下的路径置为0的话会导致目标函数出现log(0)而无法计算的情况，因此使用目标函数的上界来代替目标函数：

在E步骤，由于实际操作中计算出所有可能路径的打分不现实，因此我们主要通过beam search算法获取一个分数更高的路径子集；然后在M步骤，我们简单的最大化在这个路径子集上的目标函数。完整的EM算法步骤如下图所示：

防止过拟合的惩罚项

举例来说如果某个路径对于任意的输入都是最高值，那么在M步时item会全部被分配到该路径下面。因此为了防止过拟合，我们对于包含item数目较大的路径进行惩罚，惩罚项只应用在M步骤上：

加入联合训练

从我们的实验中发现，DR框架的训练如果加入softmax分类模型的联合训练可以明显提升效果。我们猜测主要的原因是路径和item集合在最开始是随机分配的，导致了优化训练中的困难。通过共享softmax分类模型的输入，DR框架在优化方向上可以受到一些积极的影响。因此我们最终模型的优化目标是二者的结合：

也就是在DR路径模型召回item候选集后，我们再使用softmax分类模型对这些候选集进行排序返回最终top的候选集合。

实验结果

实验结果不用多说好就完事了，主要的对比baseline是ICF、YoutubeDNN、TDM都是推荐工业界应用广泛的方案。值得一提的是，论文宣称DR框架可以几乎达到暴力穷举算法相当的效果。

参考

1. Deep Retrieval: An End-to-End Learnable Structure Model for Large-Scale Recommendations