广告深度学习计算：召回算法和工程协同优化的若干经验

▐  背景

阿里妈妈展示广告召回大多采用 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，但相比于其带来的收益，此 over