一、MDE [2020]

《Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems》

1. 标准的 embedding 做法是：将 objects 以固定的统一纬度（uniform dimension: UD） $d$$d$ 嵌入到 ${\mathbb{R}}^d$$\mathbb R^d$ 中。

   • 当 embedding 维度 $d$$d$ 太小时，下游任务的预测性能会受到影响。

   • 当 embedding 维度 $d$$d$ 太大、并且需要表示的 objects 数量很大时，内存消耗就成为一个问题。例如，在推荐模型中， embedding layer 可以占到存储模型所需内存的 99.9% 以上，而在 large-scale setting 中，它可能会消耗数百 GB 甚至数百TB 。

     不仅内存消耗是一个瓶颈，模型太大也容易过拟合。

   因此，寻找新颖 embedding representation 是一个重要的挑战，其中该 embedding representation 使用更少的参数，同时保留下游模型的预测性能。

   在现实世界的applications 中，object 的频率往往是严重倾斜的。例如：

   • 对于 full MovieLens 数据集，top 10% 的用户收到的 query 数量和剩余 90% 的用户一样多、top 1% 的 item 收到的 query 数量和剩余 99% 的 item 一样多。

   • 此外，在 Criteo Kaggle 数据集上， top 0:0003% 的 indices 收到的 query 数量与剩余 3200 万个 indices 一样多。

   为了在推荐中利用 heterogeneous object popularity ，论文 《Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems》提出了 mixed dimension(MD) embedding layer ，其中一个 specific-object 的 embedding 维度随着该object 的 popularity 而变化，而不是保持全局统一。论文的案例研究和理论分析表明：MD embedding 的效果很好，因为它们不会在 rare embedding 上浪费参数，同时也不会欠拟合 popular embedding 。此外，在测试期间，MD embedding 通过有效地分配参数从而最小化 popularity-weighted loss。

   论文贡献：

   • 论文提出了一种用于推荐系统的 mixed dimension embedding layer ，并提供了一种新颖的数学方法来确定具有可变 popularity 的特征的尺寸。这种方法训练速度快、易于调优、并且在实验中表现良好。

   • 通过矩阵补全（matrix completion）和 factorization model ，论文证明了在有足够的 popularity 倾斜的情况下， mixed dimension embedding 在内存有限的情况下会产生较低的失真，在数据有限的情况下会有更好的泛化效果。

   • 对于内存有限的领域，论文推导出最佳特征维度。这个维度只取决于特征的 popularity、参数预算、以及pairwise 交互的奇异值谱。

1.1 背景

1. 与典型的协同过滤（collaborative filtering: CF）相比，CTR 预测任务包括额外的上下文。这些上下文特征通过索引集合（categorical）和浮点数（continuous）来表达。这些特征可以代表关于点击事件或个性化事件的上下文的任意细节。第 $i$$i$ 个 categorical 特征可以用一个索引 $x_i\in \{1,2,\cdots ,n_i\}$$x_i\in \{1,2,\cdots,n_i\}$ 来表达，$i=1,\cdots ,\kappa$$i=1,\cdots,\kappa$ 。 除了 categorical 特征，我们还有 $s$$s$ 个标量特征，这些标量一起产生一个稠密的特征向量 ${\overrightarrow{\mathbf{x}}}^{\prime }\in {\mathbb{R}}^s$$\mathbf{\vec x}^\prime\in \mathbb R^s$ 。因此，给定 $({\overrightarrow{\mathbf{x}}}_1,{\overrightarrow{\mathbf{x}}}_2,\cdots ,{\overrightarrow{\mathbf{x}}}_{\kappa },{\overrightarrow{\mathbf{x}}}^{\prime })\in {\mathbb{R}}^{n_1+n_2+\cdots +n_{\kappa }+s}$$\left(\mathbf{\vec x}_1,\mathbf{\vec x}_2,\cdots,\mathbf{\vec x}_{\kappa},\mathbf{\vec x}^\prime\right) \in \mathbb R^{n_1+n_2+\cdots+n_\kappa+s}$ ，我们想预测点击事件$y\in \{0,1\}$$y\in \{0,1\}$ ，其中 ${\overrightarrow{\mathbf{x}}}_i$$\mathbf{\vec x}_i$ 为 $x_i$$x_i$ 的 one-hot 向量。

   我们使用 SOTA 的 deep learning recommendation model: DLRM 作为一个现成的深度模型。各种深度 CTR 预测模型都是由内存密集型的 embedding layer 驱动的。 embedding layer 的大小和预测性能之间的权衡似乎是一个不可避免的权衡。对于一个给定的模型 $f_{\theta }$$f_\theta$ （$\theta$$\theta$ 为参数），标准的做法是用 embedding layer $\mathbf{E}$$\mathbf E$ 来表示 categorical 特征。通常，每个 categorical 特征都有自己的独立的 embedding matrix：$E[(x_1,\cdots ,x_{\kappa })]=({\mathbf{E}}^{(1)}{\overrightarrow{\mathbf{x}}}_1,\cdots ,{\mathbf{E}}^{(\kappa )}{\overrightarrow{\mathbf{x}}}_{\kappa })$$E[(x_1,\cdots,x_\kappa)] = \left(\mathbf E^{(1)}\mathbf{\vec x}_1,\cdots,\mathbf E^{(\kappa)}\mathbf{\vec x}_\kappa\right)$ ，其中 $E[\cdot ]$$E[\cdot]$ 为 embedding layer，${\mathbf{E}}^{(i)}$$\mathbf E^{(i)}$ 为第 $i$$i$ 个 categorical 的 embedding matrix 。

2. 相关工作：最近的工作提出了类似、但实质上不同的 non-uniform embedding 架构的技术，尤其是针对自然语言处理领域（《Groupreduce: Block-wise low-rank approximation for neural language model shrinking》、《Adaptive input representations for neural language modeling》）。这些方法都不适合用于 CTR 预测，因为它们忽略了 CTR 中固有的 feature-level 结构。

   还有一些方法是为 RecSys embedding layer 提出神经架构搜索（neural architecture search : NAS）（《Neural input search for large scale recommendation models》），其中使用强化学习算法来构建 embedding layer 。与计算昂贵的 NAS 相比，我们表明 non-uniform embedding layer 的架构搜索可以简化为调优一个超参数，而不需要 NAS 。由于我们的理论框架，模型搜索的这种简化是可能的。此外，与以前所有的 non-uniform embedding 的工作相比，我们从理论上分析了我们的方法。此外，过去的工作并没有从经验上验证他们的方法所推测的工作机制。

3. 从 popularity 倾斜的角度来看， embedding normalization 实际上隐含了对 rare embeddings 的惩罚，而不是对 popular embeddings 的惩罚。这是因为更大一部分的 training updates 仅包含 rare embeddings 的 norm penalty signal (而很少包含 rare 出现的事件）。

4. 如下图所示，图 (a) 表示 Criteo Kaggle 数据集中所有特征的 access 的直方图；图 (b), (c) 分别为 MovieLens 数据集的 user feature, item feature 的 access 的直方图。

   这些图都是典型的长尾分布。

   

1.2 模型

1. MD embedding layer 架构 $\overline{\mathbf{E}}$$\bar{\mathbf E}$ 由 $\kappa$$\kappa$ 个 block 组成，每个 block 对应一个 categorical field 。这些 block 定义为 $2\kappa$$2\kappa$ 个矩阵：$\overline{\mathbf{E}}=({\overline{\mathbf{E}}}^{(1)},\cdots ,{\overline{\mathbf{E}}}^{(\kappa )},{\overline{\mathbf{P}}}^{(1)},\cdots ,{\overline{\mathbf{P}}}^{(\kappa )})$$\bar{\mathbf E} = \left(\bar{\mathbf E}^{(1)},\cdots,\bar{\mathbf E}^{(\kappa)},\bar{\mathbf P}^{(1)},\cdots,\bar{\mathbf P}^{(\kappa)}\right)$，其中：

   • ${\overline{\mathbf{E}}}^{(i)}\in {\mathbb{R}}^{n_i\times d_i},{\overline{\mathbf{P}}}^{(i)}\in {\mathbb{R}}^{d_i\times \overline{d}}$$\bar{\mathbf E}^{(i)}\in \mathbb R^{n_i\times d_i}, \bar{\mathbf P}^{(i)}\in \mathbb R^{d_i\times \bar d}$ 为第 $i$$i$ 个 block 的矩阵；$d_i$$d_i$ 为第 $i$$i$ 个 categorical 特征的 embeding size ；$n_i$$n_i$ 为第 $i$$i$ 个 categorical 特征的词表大小；$\overline{d}$$\bar d$ 为所有 categorical 特征共享的 base dimension ，且满足 $\overline{d}\ge d_i$$\bar d \ge d_i$ ；$i=1,\cdots ,\kappa$$i=1,\cdots,\kappa$ 。

   • ${\overline{\mathbf{E}}}^{(i)}$$\bar{\mathbf E}^{(i)}$ 作为 MD embedding，然后 ${\overline{\mathbf{P}}}^{(i)}$$\bar{\mathbf P}^{(i)}$ 作为投影矩阵从投影到维度 $\overline{d}$$\bar d$ 的公共空间。

2. 一个关键问题是如何确定 embedding size $\overrightarrow{\mathbf{d}}=(d_1,d_2,\cdots ,d_{\kappa })$$\mathbf{\vec d} = (d_1,d_2,\cdots,d_\kappa)$ 。

   Power-law Sizing Scheme：我们定义 block-level 概率 $\overrightarrow{\mathbf{p}}\in {\mathbb{R}}^{\kappa }$$\mathbf{\vec p}\in \mathbb R^\kappa$ ，其中 $p_i$$p_i$ 为第 $i$$i$ 个 block 中的 embedding 的平均查询概率。当 block 与特征一一对应（我们这里的情况），那么 $p_i=\frac{1}{n_i}$$p_i = \frac{1}{n_i}$ 。

   假设 categorical feature 都是单值（而不是多值的），那么对于词表大小为 $n_i$$n_i$ 的 categorical feature，每个取值出现的平均概率为 $1/n_i$$1/n_i$ 。

   更一般的情况下，${\overrightarrow{\mathbf{p}}}_i=\sum _j\prod _{i,j}$$\mathbf{\vec p}_i = \sum_j \prod_{i,j}$ ，其中 $\prod$$\prod$ 为 block-wise 伯努利采样矩阵。那么 Popularity-Based Dimension Sizing 为：

   $$\overrightarrow{\mathbf{d}}\leftarrow \overline{d}\times \frac{{\overrightarrow{\mathbf{p}}}^{\alpha }}{{\Vert \overrightarrow{\mathbf{p}}\Vert }_{\mathrm{\infty }}^{-\alpha }}$$

   其中：无穷范数是元素绝对值中的最大值；$0\le \alpha \le 1$$0\le \alpha\le 1$ 为温度系数，当 $\alpha =0$$\alpha = 0$ 时 embedding size 都等于 $\overline{d}$$\bar d$ ，当 $\alpha =1$$\alpha = 1$ 时 embedding size 与它们的 popularity 成正比。

   理论分析见原始论文。

   注意：这里仅考虑 field-level 的 popularity，而没有考虑 value-level 的 popularity 。例如，“学历” 这个 field 中，低学历的 value 占比更大，需要更高的维度。

   论文中的这种维度分配方式，使得词表越大的 field，其维度越小；词表越小的 field，其维度越大。例如，“性别” 只有两个取值，该 field 被分配的 embedding 维度最大。

   论文中的这种维度分配是启发式的规则，并不是从数据中学到的。

1.3 实验

1. baseline 方法：统一的 $d=32$$d=32$ 的 DLRM 。

2. 实验结果：

   • $\alpha =0.3$$\alpha = 0.3$ 的 MD embedding 产生的 learning curve 与 $d=32$$d=32$ 的 UD embedding 相当（下图 (a)），但是参数规模减少了 16 倍。

   • MD embedding layer 改善了每种 parameter budget 下的 memory-performance （下图 (b)）。

   • 最佳温度 $\alpha$$\alpha$ 取决于 parameter budget ，对于较小的预算，较高的温度会导致更低的 loss 。

   • $\alpha =0.2$$\alpha = 0.2$ 的 embedding 获得的性能以 0.1% 的准确性优势略微优于 UD embedding，但是参数规模约为 UD embedding 的一半。

   • MD embedding 的训练速度比 UD embedding 快两倍（下图 (c)）。