▐ 背景

阿里妈妈展示广告召回大多采用 Tree-based Deep Model(以下简称TDM)模型，它通过对候选广告的聚类，构造了深达十余层的二叉树索引，并使用 beam search 在此索引上进行检索[1]。由于线上服务的 latency 约束及当时的硬件算力限制使我们不能直接对整个候选集（百万甚至千万量级）进行模型打分。

随着近几年硬件（GPU & ASIC）的发展，我们在模型上的计算能力有了很大提升。于是，算法同学从扩大打分量和提升模型复杂度的两个方向对目前的模型进行了升级。第一个方向是进行全库打分，即抛弃索引结构，用相对简单的双塔模型一次性给所有广告打分；第二个方向诞生了索引扁平化项目，保留树结构索引的同时将其扁平化，减少 beam search 的深度并增加宽度，同时引入 attention 计算，增加了打分模型的复杂度。

这些给工程优化提出了很高的要求，也带来很大挑战。首先，模型大小和计算量的暴增，使得打分的广告数成倍增长，模型的复杂度也大大增加，例如索引扁平化模型中，单次打分消耗的浮点计算次数（FLOPs）就要比原 TDM 模型增长约十余倍。其次，模型中引入了一些非典型的非数值计算（比如TopK和beam search），这些计算如果只使用TensorFlow原生算子会带来性能热点，使 latency 过长（数十上百毫秒）无法满足线上服务的要求。

本文主要分享我们是如何通过系统和算法的协同设计，逐一解决这两个模型在上线过程中的工程难点，为广告召回业务迭代打开新的空间。

▐ 路径一： 全库模型

1. 模型结构介绍

由于每一次用户 page view 模型计算需要处理所有的广告，对于在线系统来说打分量过于巨大（百万至千万量级），只能采用相对简单的模型结构，这里我们选用了经典的双塔模型结构：用户特征和广告特征分别经过DNN得到特征向量，用这两个向量的内积表示匹配程度，并由TopK从中选出最大的那部分。由于广告特征都是固定的，所以广告所在那一路的 DNN 可以在离线提前算完，线上模型直接拿到所有广告的向量作为常量模型参数。如下图所示。

此模型最大的性能挑战来自于 TopK。由于规模达到了数百万，如使用 TF 原生 TopK 算子，其实现并不适合我们的数据特点，latency 将高达数十 ms，不满足线上服务的要求。第二个挑战来自于内积，它是全模型计算量和访存量最大的部分，对于它的处理也相当重要。

2. 对于 TopK 的优化

总体思路

首先，我们调研了 TF 的原生 TopK 算子(即tf.math.topk())的 GPU 实现。它是由若干 RadixSort 组成，其时间复杂度也的确与输入规模n成线性，如下表所示。

规模	10W	20W	50W	100W	200W	500W	1000W
latency(ms)	0.926	1.825	4.755	9.628	19.198	47.821	95.301

观察到在我们问题中 k(500or1000) 是远远小于 n(100W-1000W) 的，可以设定一个阈值，并通过筛选的方式过滤掉大部分数据来减少计算量。例如，我们可以选取千分之一的分位数过滤，这样就相当于将问题规模缩减了1000倍。此筛选操作主要是一个 elementwise 比较加很小规模数据的搬运，在GPU上并行会非常快（在1000W的规模上也不会超过1ms），而这个分位数可以通过采样+排序来进行粗略的估计。虽然这一过程会带来一定的 overhead，但相比于其带来的收益，此 overhead 可以说微乎其微。

筛选算法

对于数据采样的数量，既不能太多，也不能太少。太多会给排序带来很大的 overhead；太少则分位数粒度太粗且不准。通常，我们选取的采样数会在满足分位精度条件的基础上尽可能小，同时要求采样数大于等于k，原因如下：

我们选作阈值的分位数是估计出来的，并不是准确值，所以可能存在经过此阈值筛选出的个数实际上小于k，从而导致后续 TopK 出错。当这种 badcase 出现时，我们的做法是不断选取下一个分位数（即经排序的样本的下一个）再进行筛选，直到筛选出的个数大于等于k。为保证此循环能够正确的结束，我们需要保证样本个数s大于等于k。这样，只要当阈值取到第k个之后，就能保证一定有足够的数被筛选出来(即前k个样本)。

同样，为了降低这一 badcase 出现的概率，我们可以适当放宽此阈值。在实践中，我们的做法是将阈值的分位数放宽到了 k/n 的两倍，即选取第 2_ceil(k/n_s)+1 大的样本作为阈值。根据 order statistics，第k个的样本的概率分布满足如下公式：

下图显示了均匀分布上的1000个样本，1倍分位阈值(第二大样本)和2倍分位阈值(第三大样本)筛选出来的个数的分布情况。左图为概率密度函数 PDF，右图为累积分布函数 CDF。横坐标是筛出的比例，即筛选出的个数除以n。可计算出在 k=1000，n=1000W 的情况下，即横坐标为 k/n=0.0001 时，此方法能将badcase发生概率降低30倍。

与原生TopK性能对比

我们通过TF原生op搭建了一个子图，它不依赖任何 custom op 的实现，并且能在 GPU 上高效的运行。为了匹配 TF 原生的 TopK 算子的语义，我们通过 tf.while() 实现了对高维输入的支持。

下图显示了 batchsize=8 的不同尺寸下，原生 TF 实现和我们的优化算法的性能对比。我们的优化算法相比原生算计性能有了成倍的提升，且在越大的尺寸上优势越明显。

相较于现在常规使用的向量检索框架（例如 Faiss），这种基于 TF 的做法相对灵活，在百万千万的数量级上也有不错的性能。在一千万的规模上，我们的 TopK 算法能达到显存带宽理论上限的10%左右，并且随着问题规模的扩大这一指标能继续提升（Faiss中的 WarpSelection 算法能做到16%，但k的大小受到GPU寄存器数量的限制[2]）。当然，对于 GPU 上类似 WarpSelection 的 state-of-art 的 TopK 算法，我们也能很容易的迁移进TF里来。

3. 对于内积的优化

Batching

当 batchsize 大于1的时候，这里的内积表现为矩阵乘。TF (v1.15)原本用的 cublasSgemmEx 接口并不能在 fp16 类型的 GEMM 上启用 TensorCore，修改成 cublasGemmEx 可以解决这一问题。在 m=8 时，获得约 20%~40% 不等的性能提升，实测 latency 见下表。

size m_n_k, k=64	m=1, n=100W	m=8, n=100W	m=1, n=200W	m=8, n=200W	m=1, n=500W	m=8, n=500W	m=1, n=1000W	m=8, n=1000W
FP16	1.48	2.92	2.96	4.40	5.76	8.21	9.26	10.64
TensorCore	1.48	1.64	2.96	3.31	5.20	5.55	8.56	10.61

观察到在 TensorCore 的加持下，当 batch 从1到8增大时，latency 几乎没怎么变，也就是增大 batchsize 可以显著提升内积这部分的计算效率。不过考虑到线上 latency 的约束，batchsize 还是需要控制在合理范围之内，实践中，我们就将其控制在8以下。

我们测试了 batching 带来的收益。对于一百万左右规模的模型，在线上可以容许的 latency 范围内，按 batchsize=8 做 batching 可以将有效服务容量提升约3倍。

关于 INT8 量化

内积部分的访存在整个模型计算中的占比达到60%以上，不做 batching 时，更是高达90%。为进一步减少访存，我们尝试将这个 GEMM 用 INT8 量化。这里有一个问题需要注意，在 cublasLt 接口中，INT8 GEMM 的输入矩阵需要遵循特殊的内存排列格式，A矩阵和C矩阵遵循 CUBLASLT_ORDER_COL32 格式而B矩阵需要遵循 CUBLASLT_ORDER_COL4_4R2_8C 格式，因此需要再计算前后进行 transform 操作。

幸运的是，B矩阵（即广告向量的矩阵），是一个常量，它的 transform 过程可以被 constant fold，从而避免很多访存开销。对与C矩阵，可以将它的 transform 与一个定制的近似 TopK 的算子进行融合，降低 transform 的开销。而A矩阵的 transform 是无法避免的。

需要注意的时，上文中所描述的利用分位数筛选的 TopK 算法对数据的区分度是有要求的。如果数据的区分度不够大，这个筛选算法就不能有效地约减问题规模，会拖慢后面寻找真实 TopK 的过程的效率。而内积部分使用 INT8 量化的话有可能会导致这个问题，需要对量化算法进行调整和校准。

关于 cache 的利用

另一个节省访存的做法就是充分利用 cache。在内积计算的过程中，大部分的访存都集中在对广告向量的读取上，如果我们能让这个 constant 常驻高速缓存的话，就可以大幅减少内存访问。

GPU 的缓存太小，不太适合进行缓存常驻，而目前涌现的众多人工智能ASIC就拥有相对来说大得多的缓存，例如平头哥的含光800芯片，更适合这样的优化。同时这些芯片在矩阵乘等密集计算上的算力相比 GPU 也毫不逊色。但由于这些芯片支持的算子可能不那么全面，某些定制的量化算法和 TopK 之类的计算就需要被 offload 到 CPU 端进行，此过程中会造成 PCIE 传输的 overhead。

▐ 路径二：索引扁平化模型

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），召回率也不会相差太多。

总打分量	1W	1.5W	2W	2.5W	7.5W	15W	全库检索
召回率	0.382	0.464	0.510	0.521	0.541	0.545	0.580

在我们的架构设计中，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).

机器学习/对比学习算法交流群

已建立机器学习/对比学习算法交流群！想要进交流群学习的同学，可以直接加我的微信号：duibai996。

加的时候备注一下：昵称+学校/公司。群里聚集了很多学术界和工业界大佬，欢迎一起交流算法心得，日常还可以唠嗑~

最后欢迎大家关注我的微信公众号： 对白的算法屋（duibainotes），跟踪NLP、推荐系统和对比学习等机器学习领域前沿，日常还会分享我的创业心得和人生感悟。想进一步交流的同学也可以通过公众号加我的微信，和我一同探讨技术问题，谢谢！