▐ 路径二：索引扁平化模型
1. 模型结构介绍
此模型将原本 TDM 模型中十余层的二叉树索引压缩到了三四层，第一层展开为数千节点，之后每一层按照十几的度展开。我们从第二层开始进行 beam search ，总共经过若干轮模型打分以及 TopK 的筛选，每次模型打分的数量在数万，如图所示。
相比于 TDM 模型，打分轮数减少了23倍，而每轮打分的广告数扩充了46倍。为了拿到更精准的打分结果，算法上在原来 TDM 的打分模型 DNN 的基础上引入了用户的序列特征与广告特征交叉进行 multi-head attention 的计算。这种结构在广告系统上用的相当广泛，如精排的 DIEN 模型中就有应用[3]。这里的挑战主要有两个，一是如何在TF里用表示树型索引的结构并在这种表示上高效的进行beam search所需的操作；二是高达数万的广告候选集的大小会在乘法效应的作用下影响所有的算子，如何管控它带来的计算资源（尤其是访存）的膨胀。
2. 树索引的设计
排他索引
由于广告的索引是一棵完全树的结构，同时从在线推理的角度看，它并不会变化，因此是一个 const。因此，我们设计了一个高效的完全数树（complete tree）结构的表示，节省空间的同时还能实现高效的子节点的查找。将一棵树的节点按层序遍历编号，然后从0号节点开始，在一个数组中依次记录下每个节点的子节点编号中最小的那个，直到叶子节点为止，最后再在数组末尾加上整棵树的节点个数。这样一来，对于一棵树的表示 {a0, a1, a2, a3 …}，整数区间[ai,ai+1) 就表示编号为i的节点的子节点编号。在这种表示下，查找一个节点的子节点的时间复杂度为O(1)。例子：下图中的树就能表示成：{1, 5, 8, 10, 11, 13}，节点1的子节点就是区间 [5, 8)={5, 6, 7}。
非排他索引：ragged tensor
上述数据结构只能表达树的结构，也就是排他的索引：每两个聚类之间不能存在交集。如果算法上放宽这条限制，索引结构就会变成图，上述的数据结构就无能为力了。这种情况下，我们可以使用TF原生的 ragged tensor 来表示这个索引，即第i行表示第i个节点的子节点序号。在这种表示下，对子节点的查找可以通过 gather+flat+unique 来进行。这样做的效率会低于上一节中的数据结构，但由于索引计算的占比不大，性能差异尚可以接受。
3. 访存优化
在优化的过程中，我们发现主要瓶颈都集中在显存带宽上。此模型对显存带宽的占用达到90%以上，触及天花板。为了减少访存带宽，我们运用了下面三种图优化的手段，通过数学上的等价变换，尽可能减少中间结果大小。
交换 broadcast 与 transpose
这样 elementwise multiply + transpose 的结构出现在 attention 中，这里的 elementwise multiply 包含一个隐式的 broadcast，使得之后的 transpose 所要移动的内存量增加了L倍。这里我们可以交换这两个op的位置，先做 transpose，这样可以避免隐式 broadcast 带来的内存膨胀的问题。
核函数近似与交换矩阵乘顺序
一般来说，对于 linear 或者 elementwise 的op，我们会比较容易做数学上的变换从而进行各方面优化。注意到 attention 中存在 softmax(AB) 的结构，它是一个核函数，可以表示成内积形式。因此我们将原本的 softmax(AB)*C 近似替换成了f(A)*f(B)C[4]。由于这里A矩阵较大的，而B和C矩阵较小，我们可以根据矩阵乘的结合律按照 f(A)(f(B)*C) 的顺序计算，从而保证中间结果也维持在了较小的规模。
拆分 tile+concat+matmul
DNN 第一层的 matmul 是整个模型中最大的op，它是输入是由两个两部 concat 而来的。第一部分的 batchsize 是1，而第二部分的 batchsize 高达数万，而在 concat 前会将前者复制多份(tile操作)，从而使其 batchsize 与后者保持一致。观察到 tile，concat，matmul 都是线性操作，因此可以将这个过程进行线性变换，将两个部分分开计算后再合并，这样就避免了第一部分的复制操作，节省了大量的计算和内存资源。
4. Beam Search 宽度的调节
由于访存的消耗都与打分量成正比，最直接有效进行优化的方式就是控制打分量，也就是 beam search 的宽度。最终选出的广告只有数百个，而目前 beam search 的宽度在数万的级别，相当于只选出了前几十分之一。考虑到缩小打分量可能会影响召回的效果，所以我们考察了召回率随之变化的情况，如下表。可见即使将宽度缩小一半（表上从15W缩减到7.5W），召回率也不会相差太多。
在我们的架构设计中，beam search 宽度可以由上游请求中的参数决定（作为模型输入传进来），这样业务可以实时对效果和水位进行调整 tradeoff。这相当于提供了无需替换模型就进行降级的能力。
▐ 总结
本文详细介绍了我们是如何解决召回新模型上线过程中的工程问题的。这离不开与算法同学的通力合作。其中很多问题都需要工程和算法的协同优化：比如 TopK 的筛选对区分度的需求需要量化算法来保证；再比如索引扁平化模型中 att 结构的选择和需要从效果和计算资源两方面进行考量的。
我们认为相比于解决上述问题的具体方法，如何得到这些方法的思路更为重要。还是拿TopK和索引扁平化模型来举例子，TopK 的优化中通过预筛选来降低问题规模的思路，和全库模型中通过数学上的等价变换来进行计算优化的思路，在许多优化问题上都能应用。希望本文的读者能从这些思路中有所收获。
引用
[1] Zhu, Han, et al. “Learning tree-based deep model for recommender systems.” Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2018.
[2] Johnson, Jeff, Matthijs Douze, and Hervé Jégou. “Billion-scale similarity search with gpus.” IEEE Transactions on Big Data (2019).
[3] Zhou, Guorui, et al. “Deep interest network for click-through rate prediction.” Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining. 2018.
[4] Choromanski, Krzysztof, et al. “Rethinking attention with performers.” arXiv preprint arXiv:2009.14794 (2020).
