▐ 背景
阿里妈妈展示广告召回大多选用 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 Jgou. “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、推荐体系和比照学习等机器学习范畴前沿，日常还会共享我的创业心得和人生感悟。想进一步沟通的同学也能够经过大众号加我的微信，和我一同讨论技术问题，谢谢！
