Press "Enter" to skip to content

ICLR 2021 | 快手联合多所高校提出: 一体化模型压缩框架(UMEC), 加速推荐场景落地

本站内容均来自兴趣收集,如不慎侵害的您的相关权益,请留言告知,我们将尽快删除.谢谢.

模型压缩能够 帮助提升模型在服务器端或者手机端的推理效率,快手联合多所高校通过借鉴模型压缩的想法创造性地提出了 一体化的模型压缩框架(UMEC) ,本文将从研究背景、方法介绍、实验结果 等 多方面对这 一框架展开系统描述。

 

 

论文链接:

 

https://openreview.net/forum?id=BM—bH_RSh  

 

代码链接:

 

https://github.com/VITA-Group/UMEC   

 

模型压缩最早诞生于图像视觉领域,目的是为了提升模型在服务器端或者手机端的推理效率。常见的需要压缩的模型包括CNN、Bert模型等。推荐的模型虽然与上述模型相去甚远,但是也有提升模型,效率也非常迫切,比如降低embedding的维度,减少feature的数量等等,推理效率的提升往往都能带来收入上或者用户时长上的直接提升。

 

快手 联合 德州农工大学、罗切斯特大学和德州奥斯汀大学 的研究者们,借鉴模型压缩的想法针对推荐场景提出一种 一体化的模型压缩框架(UMEC) ,将推荐系统里开源的benchmark模型压缩了三倍且不损失模型精度,多个压缩尺度下均达到当前最优的模型压缩效果,目前相关的技术已经在快手场景落地。

 

| 背景和动机

 

在当今信息过剩的时代背景下,如何构建有效的推荐系统来给用户推荐有用的信息,并为用户定制化内容就显得尤为重要。而推荐系统的核心在于推荐模型,在如今机器学习火热的背景下,推荐模型通常是以神经网络的形式发挥作用。在实际场景中,推荐模型的性能,包括预测的准确度和耗费的资源(时间、能量等),会显着影响到推荐的效果以及用户的实时体验。 如何在保持推荐效果的情况下,减少其资源的消耗,并且提高系统整体的推荐效率,成为亟待解决的痛点问题 。

 

通常来说,一个有效的推荐模型由两模块构成: 特征嵌套(feature embedding) 和 预测模型(prediction model) 。前者是对稀疏及稠密特征的一个预处理模块,目的是初步提取信息,并生成适合下游任务使用的数据格式;后者为一个排序神经网络,执行推荐任务(排序或分类)。下图是一个典型推荐模型的结构 DLRM[1] ,特征嵌套和预测模型两个模块占用了主要资源,均有很多优化空间。已有的一些研究工作通常将这两模块的优化看作独立问题,进行独立的优化处理,比如针对特征嵌套模块的特征压缩或者特征选择方法,或是针对预测模型模块的模型压缩方法。 这样的分开优化也能取得一定的效果,但并未考虑到两模块之间的关系,未能从整体上去联合优化两个模块,达到更好的压缩优化结果 。

 

 

图1: 推荐模型实例 (DLRM[1]),原始稀疏和稠密特征在经由特征嵌套模块 (灰色)处理后,被拉长并连接为一维向量作为预测模型(绿色)的输入,第i组特征的维度为gi 。预测模型为一个全连接神经网络,执行二分类任务,其输入为按上述方式处理好的输入特征(紫色)。

 

基于上述不足,本文提出了如下问题:在限制计算存储资源,且不损害模型推荐效果的情况下, 如何对于特征嵌套(图中灰色模块)和预测模型(图中绿色模块)两个模块进行联合压缩优化,从而达到资源的最优配置? 对这个问题,本文第一次尝试做出系统性解答,并取得了极佳的优化结果。

 

| 方法介绍

 

针对上述提到的推荐模型的压缩问题,本文 独创性地提出了UMEC框架 :我们将特征嵌套模块和预测模型模块看作一个整体,在有整体资源约束的情况下,将特征压缩和模型压缩两个任务进行 联合优化 。这样框架求解出的推荐模型,不仅能满足有效降低实际资源消耗的目标,同时也能满足预测模型精度保持不变的需求。

 

1) 带资源约束的联合优化框架 – UMEC

 

下面来简单介绍一下我们的算法框架。上述的优化问题可被数学建模成如下形式:

 

 

该式分别申明了目标函数、对整体计算资源的上界约束、以及对特征模块和预测模块的结构化稀疏程度的下界约束。具体来说:

 

1. 目标函数 :即原始推荐模型的训练损失函数。

 

2. 消耗资源的约束 :推荐模型实际消耗的资源应小于给定的资源预算 (budget),实际消耗资源是跟模型权重的结构化稀疏度
s
s 相关联的,一般来说,结构化稀疏度越大,模型越小,所消耗资源越少。本文以模型计算量Flops作为使用的资源来指导模型的压缩。在实际应用时,本框架也可以自由替换为其它的资源,如latency, energy等等。

 

3. 结构化稀疏程度的约束 :对于预测模型,我们想要第
L
L 层删除的神经元的个数不少于
s^{(l)}
s(l) 个;对于第一层,我们想要删除的特征的个数不少于
s^{(1)}
s(1) 个。

 

2)优化算法

 

上述优化问题包含的难点有:含有非连续不可微的不等式约束条件;含有非凸的不等式约束条件等等。我们使用了一些优化技巧来解决该问题。

 

用拉格朗日乘子法来重构问题

 

 

我们采用拉格朗日乘子法,通过引入对偶变量

y,z
y , z ,将原问题改写为一个minimax问题。

\left \lceil \cdot \right \rceil
⌈ ⋅ ⌉ 代表向上取整函数。Sparsity Loss 的两项soft regularization 用来等价代替 (3c) (3d) 两式。Resource Loss用来等价代替 (3b) 。其中权重项

y,z
y , z 称为对偶变量,会在minimax优化中被自动调整。该minimax问题的最优解即为原问题的最优解。

 

用ADMM求解minimax问题

 

我们采用ADMM算法来求解这个优化问题,其中训练权重

\mathcal{W}
W 可以用proximal-SGD求解[2]. 稀疏度

s
s 的求解依赖于 Sparsity Loss 和 资源Resource, 这两项均含有不可微项

\left \lceil s \right \rceil
⌈s⌉ , 我们采取了Straight Through Estimator (STE) [3]方法来处理这个可微问题。对于对偶变量

y,z
y , z , 可以直接用梯度上升算法求解。详细的推导可见论文的附录部分。算法的完整版本如下:

 

 

| 实验结果

 

下图展示了UMEC框架算法在开源广告推荐模型DLRM[1]上的压缩效果,并与其它SOTA等baseline压缩算法方法进行了对比。图中所示的 ECC[5] 是一种基于模型推理能耗来指导模型压缩的优化方法,我们的UMEC在多个压缩尺度上超越了这个SOTA方法。

 

 

图2: 在Criteo AI Labs Ad Kaggle dataset开源数据集上的关于特征选择和模型压缩两个任务联合优化的结果对比。

 

每条曲线展示了模型压缩率和模型预测准确度之间的trade-off曲线。可以看到,我们的算法在保证了同样的压缩率下,预测准确率持续优于其它的压缩算法。我们在压缩掉65%的计算量的情况下,依然能够保持和原模型同样的准确率,而其它的模型无法做到这一点。同时,其余的方法随着压缩率的增长,有较大幅度的下跌。

 

 

图3: 在Criteo AI Labs Ad Kaggle dataset开源数据集上的特征选择任务。

 

每条曲线展示了模型压缩率和模型预测准确度之间的trade-off曲线。我们的UMEC算法持续优于以Group Lasso为baseline方法的实验结果。

 

图4: 在Criteo AI Labs Ad Kaggle dataset开源数据集上的关于特征长度的压缩任务。

 

我们的UMEC算法持续优于以混合维度嵌套层(MDEL) [4] 为baseline的实验结果。

 

我们同时也在另一个体量更大的开源数据集上验证了我们的方法的优越性:

 

 

 

图5: 在Criteo AI Labs Ad Terabyte dataset 开源数据集上的特征选择和模型压缩联合任务(上图),以及特征选择任务(下图)。

 

我们的UMEC算法持续优于相应baseline的实验结果。

 

| 总结

 

本文提出了UMEC压缩框架,将推荐模型的输入特征模块和预测神经网络模块的联合压缩问题整合为一个带有资源约束的优化问题进行求解。 多项系统性的实验结果证明了我们的压缩算法的优越性。 更多关于算法的细节分析以及实验的结果分析,请参考我们的论文。

 

| 参考文献

 

[1] Maxim Naumov et al. Deep learning recommendation model for personalization and recommendation systems. CoRR, abs/1906.00091, 2019. URL https://arxiv.org/abs/1906.00091.

 

[2] Atsushi Nitanda. Stochastic proximal gradient descent with acceleration techniques. In Advances in Neural Information Processing Systems, pp. 1574–1582, 2014.

 

[3] Yoshua Bengio, Nicholas Léonard, and Aaron Courville. Estimating or propagating gradients through stochastic neurons for conditional computation. arXiv preprint arXiv:1308.3432, 2013.

 

[4] Antonio Ginart, Maxim Naumov, Dheevatsa Mudigere, Jiyan Yang, and James Zou. Mixed dimension embeddings with application to memory-efficient recommendation systems. arXiv preprint arXiv:1909.11810, 2019.

 

[5] Haichuan Yang, Yuhao Zhu, Ji Liu. ECC: Platform-Independent Energy-Constrained Deep Neural Network Compression via a Bilinear Regression Model. CVPR, 2019.

 

Illustrastion by Natasha Remarchuk from Icons8

Be First to Comment

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注