ICLR2021 | 推荐系统中可学习的嵌入维度

| 作者：YEN

| 单位：东北大学

| 研究方向：推荐系统、计算广告

本文分享一篇发表在ICLR’21的推荐系统方向的文章：推荐系统中可学习的嵌入维度。

论文链接：https://arxiv.org/abs/2101.07577

论文核心内容：基于Embedding的表示学习方法是推荐模型中常用的方式，传统的方法为每个特征分配相同大小的嵌入维度（embedding size），这篇文章通过可学习的剪枝操作为每个特征分配不同的嵌入维度。

简介

基于嵌入的表示学习（embedding-based representation learning）方法广泛应用于推荐模型中，它将原始的高维稀疏特征映射为低维的稠密向量。然后将学习到的向量输入预测模型，如FM 的内积、 AutoInt的自注意网络，以获得预测结果。然而，传统嵌入方式为所有特征分配一个相同的嵌入维度（Embedding size），这种方式有两个问题。

• 首先，由于有大量的原始特性，不可避免地会导致一个巨大的嵌入表(Embedding table)，从而导致较高的存储成本。（特征嵌入表占据了推荐模型中最大比例的存储成本，一般在嵌入表的参数量占据整个推荐模型的以上。）
• 其次，相同的特征嵌入维度可能很难处理不同特征之间的异质性。例如一些特性极其稀疏，因此当嵌入维度对所有特征都是一致的时，很可能会导致那些不需要太大表示容量的特征出现过拟合问题，导致推荐模型往往是次优的。

现有的一些试图解决这个问题的工作要么会导致推荐性能的显著下降，要么会遭受无法承受的训练时间成本的限制。

在这篇文章中，作者提出了一种新的方法，称为PEP（Plug-in Embedding Pruning的缩写），以减少嵌入表的大小，同时避免精度的下降和优化成本的上升。PEP通过阈值修剪嵌入表，其中修剪阈值（s）可以自适应地从数据中进行学习。因此，可以通过剪枝每个特征的冗余参数来自动获得一个混合维度嵌入方案。PEP是一个可以插入各种推荐模型的通用框架。大量的实验表明，它可以有效地减少嵌入参数，提高基模型的性能。具体来说，它在减少的参数的同时还保证了强大的推荐性能。至于计算成本，与基模型相比，PEP只增加了的时间成本。

问题定义

基于特征的推荐系统（Feature-based recommender system）广泛用于个性化信息服务平台中。一般来说，深度学习推荐模型以各种原始特征（包括用户和物品的属性、交互上下文等）作为输入，进而预测用户喜欢物品的概率。具体地说，推荐模型将用户和物品的属性的组合作为它的输入向量，表示为 ，其中 可以定义为如下的所有字段的连接：

其中 表示特征域的数量， 是第个域的特征表示（通常表示为one-hot向量）。然后对于， 基于嵌入的推荐模型通过以下公式生成相应的嵌入向量：

其中是第i个特征域的嵌入矩阵，表示这个特征域共包含多少特征，表示嵌入维度。该模型对所 有特征域的嵌入矩阵 可以表述如下：

模型预测时通过 和模型的其它参数 进行计算，如下所示：

表示预测的概率，表示预测模型，如FM、AutoInt、DeepFM等。

在模型训练中，为了学习模型参数，优化器将训练损失最小化如下：

其中，表示输入到模型中的数据，表示输入特征，表示真实标签，是损失函数。CTR预估问题中，LogLoss是最常用的损失函数。

其中，是训练样本的总数，这里省略了正则化项。

PEP模型：通过剪枝可获得可学习的嵌入尺寸

如上所述，内存高效的嵌入学习（memory-efficient embedding learning）的一个可行解决方案是自动为不同的特征嵌入 自动分配不同嵌入大小的维度 。

由于其离散性和极大的优化空间，直接学习 是不可行的。为了解决这个问题，作者提出了一个新的想法，在 上强制执行列稀疏，它等价地缩小了嵌入的维度。

如图1所示，嵌入 中的第一个值被剪裁并设置为零，从而导致一个 的嵌入大小。此外，还有一些不重要的特征嵌入，如 ，通过设置所有值为零可以进行丢弃，即。因此，这样的方法可以显著地减少嵌入参数。另外，稀疏矩阵存储技术有助于我们显著节省内存使用量。

因此，作者以这种方式将嵌入矩阵 的嵌入大小选择问题重新转换为学习列稀疏矩阵问题。为了实现这一点， 作者对 的稀疏约束如下：

其中表示范数，即非零元数量。是参数预算，即对嵌入表参数总数的约束。

然而，由于范数约束的非凸性，上面的方程直接优化是 NP 困难的。为了解决这个问题，一些研究对范数进行了凸松弛，即范数。然而，这样的求解方法依旧有如下挑战：

• 首先，将最优值投影到 球面上的过程需要太多的计算成本，特别是当推荐模型有数百万个参数时。
• 其次，参数预算 需要专家手动设置。

考虑到特征对任务具有不同的重要性，这种操作显然是次优的。为了解决这两个挑战，受软阈值重参数化（Soft Threshold Reparameterization）的启发，作者通过梯度下降来更新的可学习的阈值进而对 进行剪枝操作。

的重参数化可以表述如下：

这里 表示重参数化的矩阵，作为一个剪枝阈值，其中sigmoid函数是一个简单而有效的解决方案。可训练参数 的初始值（称为）需要进行设置，以确保阈值在最开始时接近于零。符号函数将正输入转换为 ，负输入转换为 ，零输入将保持不变。

由于 应用于 的每个元素， 因此推荐模型的优化问题可以重新定义如下：

然后通过标准的反向传播将可训练剪枝参数与推荐模型的参数联合优化。得益于自动微分框架tensorflow, pytorch等，可以通过端到端的方式自动训练上述参数。

当然，上面的剪枝阈值是全局的，也可进行细粒度的剪枝：

• Global: 剪枝阈值为。
• Dimension Wise: 对每个维度进行剪枝，剪枝阈值为。
• Feature Wise: 对每个特征进行剪枝，剪枝阈值为。
• Feature-Dimension Wise: 对每个特征的每个维度进行剪枝，剪枝阈值为。

PEP实现

PEP与基模型（FM，DeepFM, AutoInt...）结合时，只需要把传统的嵌入表换成.

实验结果

PEP推荐精度

PEP模块显著地减少了参数的数量，特别是对于较大的数据集而言。

• 在Criteo和Avazu数据集中，与最好的基线相比，作者的PEP-0可以减少(从 到 )。
• PEP的方法始终优于基于统一嵌入维度的模型，并在大多数情况下比其他方法获得更好的精度。例如，对于 Criteo 数据集上的 FM 模型，PEP 比 UE 的相对性能改进为 ， 比 DartsEmb 的相对性能改进为 。

PEP训练效率

与基模型相比，PEP需要引入额外的时间来为不同的特征找到合适的尺寸。作者对比了三个不同模型的训练时间。我们可以观察到，PEP 的额外计算成本是 到 ，这相比于基模型是可以接受的。而DartsEmb模型需要近一倍的计算时间才能在其双层优化过程中搜索一个良好的嵌入大小。