Learnable Embedding Sizes for Recommender Systems

这篇文章刚被ICLR2021接收，代码开源。主要研究在推荐系统中如何减少Embedding的大小。

基于embedding的表示学习范式通常被应用于深度推荐模型中，传统的模型为每个特征分配相同的embedding size。然而，这种方式面临两个问题：

• 导致embedding table特别大，消耗内存，因为有些特征（比如）年龄可能不需要用户ID一样的Embedding大小
• 导致模型过拟合，有大量的冗余参数

因此，这篇文章提出PEP(Plug-in Embedding Pruning)方法，为embedding动态学习阈值从而进行剪枝。PEP不仅能够减少参数量，还能够提升模型效果。PEP在减少97-99%参数的基础上还能够保持良好的效果，同时只带来20-30%的计算开销。

现有方法

1. Embedding Parameter Sharing： 不同的特征共享部分Embedding参数。比如把Embedding矩阵分解为小片段，或者让某些低频特征共享小部分参数。这类方法需要算法人员精心设计参数共享比例，而且也无法发现冗余参数，往往会降低最终模型效果。
2. Embedding Size Selection： 不同的特征拥有不同的Embedding size。比如基于一些启发式规则，或者基于AutoML和强化学习直接对Embedding size进行搜索。然而，这类方法计算量非常大，而且需要精心设计搜索空间。

本文方法

本文的方法可以通过下面这个图直观表示。左边\(V\)就是原始的Embedding table，我们可以通过一个可学习的非负数阈值\(g(s)\)对\(V\)进行剪枝，即将绝对值小于\(g(s)\)的参数直接mask为0，从而得到新的Embedding table，再进行后续的任务。如果选择合适的阈值能够使得新的Embedding table中0的比例比较大，那么使用稀疏存储技术对Embedding table进行存储时将节省大量空间。

那么怎么实现上面这个想法呢。一般的想法是在Embedding上添加正则化，比如约束其\(L_0\)范数。但\(L_0\)范数是非凸的，难以优化。一种合理的方式是将\(L_0\)范数放松为上界\(L_1\)范数。注意，这里的\(L_1\)范数并不是说在Embedding table上施加\(L_1\)正则化，这会改变最终的Loss。这里是把学习目标转换为带约束的优化问题：

\[min \mathcal{L}, \ s.t. ||\mathbf{V}||_1 \leq k \]

其中，\(\mathcal{L}\)是损失函数，\(k\)是Embedding table中非0元素的数目。

这个带约束优化问题可以通过soft thresholding进行求解，将原来的Embedding table \(\mathbf{V}\)转换为\(\hat{\mathbf{V}}\):

\[\hat{\mathbf{V}} = \mathcal{S}(\mathbf{V}, s) = sign(\mathbf{V})ReLU(|\mathbf{V}| - g(s)) \]

其中，\(g(s)\)就是学习的阈值，\(s\)是参数，\(g\)是将\(s\)映射到\([0, \infty]\)的函数，在论文中用的是sigmoid与常数的乘积。

完整的算法流程如下图所示

整个算法分为两个阶段，1-3行基于上面公式对模型进行优化，第4行根据学习到的阈值生成Embedding的mask，5-8行重新恢复原始的Embedding，与mask结合生成新的Embedding，然后基于这份剪枝的Embedding重新进行普通训练流程。

论文中也对重训练做了实验，使用原始的Embedding进行初始化再重新训练效果是最好的，而且收敛很快。这也可以理解，毕竟这份Embedding的Mask是基于原始的Embedding训练而来的，如果再随机初始化，岂不是与mask不对应了。