AutoDim，如何节省70%的存储空间同时还能大幅提效？

作者：一元，四品炼丹师

Memory-efficient Embedding for Recommendations (WWW21)！

简 介

在工业界使用深度学习算法做大规模的搜索推荐问题，不可避免的都会碰到下面的几个问题：

1. 过大的embedding给机器带来过大内存负担的问题，例如一些成熟的平台，都有上千万的用户，还有上几十亿的商品，这些embedding带来的负担是非常巨大的；
2. 我们有成百上千的特征字段，每个特征的基数大小不一，而我们经常会将其一起编码为一样的维度，这可能是不合理的，如何根据特征场的重要性和可预测性，为不同的特征场分配不同的Embedding维数是非常重要的。

本文我们重点讨论第二个问题，因为大量的特征Filed以及Embedding维数与特征分布和神经网络结构之间的微妙关系，在实际的推荐系统中手工分Embedding的维数是非常困难的，本文提出了一个基于AutoML的框架AutoDim，该框架能够以数据驱动的方式自动选择不同特征Filed的维数 。

框 架

01 AutoDim

本文提出的框架如上所示，一共分为两个阶段：

1. 维度搜索阶段

它的目的是为每个特征Field寻找最佳的Embedding维数。

更具体地说，我们首先通过embedding lookup 步将具有不同维度的候选嵌入分配给特定的分类特征；然后，我们通过一个变换步骤来统一这些候选嵌入的维数，这是因为第一个MLP层的输入维数是固定的；

接下来，我们通过计算所有变换后的候选嵌入的加权和，得到该分类特征的形式化嵌入，并将其输入到MLP组件中。在训练集上学习DLRS参数，包括嵌入和MLP层，而在验证集上优化候选嵌入的体系结构权重，从而防止框架选择过拟合训练集的嵌入维度[24，29]。

在实践中，在DLRS参数和体系结构权重的替代训练之前，我们首先为所有候选embeddings分配等效的体系结构权重。固定这些体系结构权重并预训练DLRS，包括所有候选embedding。

2. 参数再训练阶段

根据在维度搜索阶段学习到的框架的权重，为每个特征字段选择embedding维度，并在训练数据集上以端到端的方式重新训练DLRS参数的参数（即MLPs和选定的embedding）。注意：

• 再训练阶段是必要的，因为在维数搜索阶段，模型的性能也受到次优embedding维数的影响，而在实际的推荐系统中，次优embedding维数是不理想的；
• 由于大多数现有的深度推荐模型通过嵌入向量上的交互操作（如内积）捕获两个特征字段之间的交互，新的embedding仍然统一到同一维度。这些交互操作限制了embedding向量需要拥有相同的维数。

请注意，数值特征将通过bucketing转换为类别特征。

02 维度搜索

不同的特征域具有不同的基数和对最终预测的不同贡献。所以我们应该为不同的特征域提供不同的embedding维数。然而，由于特征域的数量庞大，以及embedding维数与特征分布和神经网络结构之间的复杂关系，传统的降维方法很难人工选择emebdding维数。

• 一个直观方法是为每个特征域分配多个不同维数的embedding空间，然后DLRS自动为每个特征域选择最优的嵌入维数。

1. Embedding Lookup Trick

对一个特征，我们赋予不同的emebdding维度，最终的emebdding维度就有，但是这么做会因为需要的emebdding空间巨大而无法忍受，训练的代价也是巨大的。

2. 权重共享embedding框架

对一个特征，我们仅赋予维度的emebdding，这么做的好处和坏处是:

• 大大节省了存储空间；
• 增加了训练时间；

03 大量维度的融合

由于现有DLRSs中第一MLP层的输入维数通常是固定的，因此它们很难处理各种候选维数。此处我们使用两种方式来处理

1. 线性变换

通过线性变化，我们将所有不同维度的向量映射到同一维度，然后我们再做BatchNormalization，并最终相加。

2. Zero填充

我们先做BatchNormalization， 然后通过使用0填充，得到相同的维度。

04 维度选择

此处引入Gumbel-softmax操作来模拟从一个类别分布中进行不可微的采样。假设我们有权重，是不同维度的类的Giallo，然后hard选择就是：

此处我们使用softmax来近似：

其中是参数，用来控制gumbel-softmax操作的平滑性。

最终我们将所有的emebdding进行concat得到并输入到层的MLP模型。最终使用cross entropy作为我们的Loss。

05 参数重新训练

由于维度搜索阶段的次优Embedding维数也会影响模型的训练，因此需要一个只训练最优维数模型的再训练阶段来消除这些次优影响。

1. 离散维数的推导

在在此训练的阶段，gumbel-softmax操作就不需要了，

2. 模型再训练

此处我们注意：

• 现有的大多数深度推荐算法通过交互操作（如内积和Hadamard积）捕获特征字段之间的交互。这些交互操作要求所有字段的嵌入向量具有相同的维数。

所以我们仍然需要转化到相同的维度，此处类似的，我们使用：

• 线性变换：来自一个特征Field的所有嵌入共享相同的权重矩阵和偏差向量；
• 零填充：不引入额外的可训练参数

实 验

• AutoDim可显著提高推荐性能，并节省70%∼ 80% embedding参数。

小 结

本文提出的AutoDim，通过数据驱动的方式自动为不同的特征字段分配不同的嵌入维数。在现实世界的推荐系统中，由于特征场的数量庞大，以及嵌入维数、特征分布和神经网络结构之间的高度复杂关系，很难在可能的情况下将不同的维数手动分配给不同的特征场。因此，我们提出了一个基于AutoML的框架来自动选择不同的嵌入维度。

具体地，本文首先提出了一种端到端的可微模型，该模型以软连续的形式同时计算不同特征域在不同维度上的权重，并提出了一种基于AutoML的优化算法；然后根据最大权值，导出离散嵌入结构，并对DLRS参数进行重新训练。基于广泛使用的基准数据集，我们通过大量实验对AutoDim框架进行了评估。结果表明，我们的框架可以保持或实现稍好的性能，同时embedding空间也小了很多。

参考文献

1. Memory-efficient Embedding for Recommendations