SIGMOD会议是数据库领域最具影响力的顶级国际学术会议之一，与VLDB和ICDE并称为数据库领域的三大顶级会议。PKU-DAIR实验室发表于SIGMOD 2024的研究成果《CAFE: Towards Compact, Adaptive, and Fast Embedding for Large-scale Recommendation Models》荣获了SIGMOD24 Honorable Mention for Best Artifact，该奖项每年仅授予至多三篇文章，旨在表彰那些在可复现性、灵活性和可移植性方面表现卓越的研究工作。

导读

近年来，深度学习推荐模型（DLRM）中嵌入表的内存需求不断增长，给模型训练和部署带来了巨大的挑战。现有的嵌入压缩解决方案无法同时满足三个关键设计要求：高内存效率、低延迟和动态数据分布的适应性。本文提出了CAFE，一种紧凑、自适应和低延迟的嵌入压缩框架，可以同时满足上述要求。CAFE的设计理念是动态地为重要的特征（称为热特征）分配更多的内存资源，为不重要的特征分配更少的内存。在CAFE中，我们提出了一种快速且轻量级的草图数据结构，名为HotSketch，用于捕获特征重要性并实时识别热特征。对于每个热特征，我们为其分配唯一的嵌入；对于非热门特征，我们使用哈希嵌入技术允许多个特征共享一个嵌入。在该设计理念下，我们进一步提出了多级哈希嵌入框架来优化非热门特征的嵌入表。我们从理论上分析了HotSketch的准确性，并分析了模型的收敛性。实验表明，CAFE显着优于现有的嵌入压缩方法，在10000倍的压缩比下，在Criteo Kaggle数据集和CriteoTB数据集上的测试AUC分别提高了3.92%和3.68%。

一、问题背景

近年来，嵌入技术广泛应用于数据库领域，其中深度学习推荐模型（DLRM）是嵌入技术最重要的应用之一。典型的DLRM将分类特征向量化为可学习的嵌入，然后将嵌入与其他数值特征一起喂给下游神经网络。最近，嵌入表的内存需求随着DLRM分类特征数量的指数级增长而增长，给各种应用带来了巨大的存储挑战。在本文中，我们重点关注压缩超大规模DLRM的嵌入表。DLRM有离线训练和在线训练两种范式。离线训练时，预先收集训练数据，并在整个训练过程后部署模型以供使用。在线训练中，训练数据是实时生成的，模型同时更新参数并服务请求。本文主要针对在线训练的场景，因为在线训练的难度更大，且在线训练的压缩方法可以直接应用于离线训练。在线训练中的嵌入压缩有三个重要的设计要求：

• 高内存效率。为了在内存限制内保持模型质量，一种方案是使用分布式实例，但它们会带来大量的通信开销。此外，嵌入表的训练和部署通常发生在存储容量较小的边缘或终端设备上，使得内存问题更加严重。另一方面，由于模型质量直接影响利润，即使AUC（ROC 曲线下面积）0.001的微小变化也会带来相当大的对利润的影响。因此，我们需要保持模型质量的内存高效压缩方法。
• 低延迟。延迟是服务质量的关键指标，因此嵌入的压缩方法必须足够快，以避免引入显着的延迟。
• 对数据分布变化的适应性。在线训练中，数据分布并不像离线训练那样固定，一般会随着时间动态变化。我们需要适合在线训练的自适应压缩方法。

现有的行压缩方法可以分为基于哈希和自适应方法两类。基于哈希的方法利用简单的哈希函数将特征映射到具有冲突的嵌入中[1]，尽管它们简单易用，但预先确定的哈希函数会扭曲特征的语义信息，导致模型精度大幅下降；集成特征频率信息[2]可以提高离线训练中基于哈希方法的模型质量，但不能应用于在线训练。自适应方法在训练过程中识别重要特征以适应在线训练[3]，但现有方法的压缩率受到重要性分数存储的限制，导致内存效率较低，且额外的检查、采样操作会增加整体延迟。总之，现有方法无法同时满足所有的三个关键要求。

二、CAFE设计方案

1. 总体流程

我们提出的嵌入框架CAFE能够同时满足高内存效率、低延迟、适应数据分布动态变化三个要求。CAFE的核心思想是动态区分重要特征（称为热特征）和不重要特征（称为非热特征），并为热特征分配更多的资源。具体来说，根据我们的理论推导，我们使用梯度的L2范数定义特征的重要性得分。我们观察到，在大多数训练数据中，特征重要性遵循高度倾斜的分布，其中大多数特征的重要性得分较小，只有一小部分热门特征非常重要。如图1所示，Criteo Kaggle数据集和CriteoTB数据集中的特征重要性分布分别近似参数1.05和1.1的Zipf分布。因此，CAFE为热点特征的嵌入分配更多的内存，为非热点特征的嵌入分配更少的内存，从而在嵌入表的内存使用量不变的情况下显着提高模型质量。 图2是CAFE框架的示意图，在CAFE 中，我们提出了一种新颖的HotSketch算法，用于实时记录top-k热特征的重要性。在每次训练迭代中，我们首先在HotSketch里查询输入样本中的每个特征，重要性分数超过阈值的视为热门特征；然后我们分别查找热门特征和非热门特征的嵌入，我们给热门特征分配唯一的嵌入，而让多个非热门特征共享一个嵌入；随着训练数据分布的变化，热门特征和非热门特征之间可以互相迁移；在神经网络反向阶段，我们使用特征的梯度范数来更新特征重要性。

图1. 特征重要性分布

图2. CAFE示意图

2. HotSketch算法

HotSketch算法旨在训练中识别高重要性分数的热门特征，这本质上是在流数据中查找top-k频繁项（特征）的问题。目前，Space-Saving[4]是解决该问题最受认可的算法，它在排序的双向链表中维护频繁项，并使用哈希表来索引。然而，该哈希表不仅使内存使用量增加了一倍，而且指针操作引起的大量内存访问导致时间效率低下。我们提出的HotSketch删除了哈希表，同时仍然保持o(1)时间复杂度，并很好地继承了Space-Saving的理论结果。

数据结构：如图3所示，HotSketch由多个桶组成，我们使用哈希函数将特征映射到桶中。 每个桶包含多个槽，其中每个槽都存储一个特征及其重要性分数。

插入：对于每个传入特征，我们首先计算哈希函数来定位桶，并遇到三种可能的情况：(1) 该特征记录在桶中，我们将新分数添加到其重要性得分上；(2)桶中没有记录该特征且存在空槽，那么我们将该特征和分数插入到空槽中；(3)桶已满且未记录该特征，那么我们找到得分最小的特征替换为插入的特征并添加分数。

讨论：HotSketch具有以下优点：（1）HotSketch插入速度快，具有o(1) 时间复杂度，并且只有一次内存访问；(2) HotSketch内存效率高，它不存储指针，并且短暂冷启动后不会出现空槽；(3)HotSketch对硬件友好，可以通过多线程和SIMD进行加速，从而实现数据并行性。

图3. HotSketch示意图

3. 迁移策略

在DLRM的在线训练过程中，特征重要性的分布随着数据分布的变化而波动，这意味着在整个训练过程中热点特征并不是恒定的。HotSketch已经记录了特征重要性，自然可以通过嵌入表之间的嵌入迁移来支持动态热特征。我们设置了一个阈值来区分热点特征，几乎每个首次达到HotSketch的特征都被视为非热门特征。当非热门特征的重要性得分超过阈值时，它会转变为热门特征，并且其嵌入将从共享表迁移出来，确保特征的表示在整个训练过程中保持平滑；相反，HotSketch中特征的重要性会随时间进行衰减，当热门特征的重要性得分低于阈值时，它就变成非热门特征，考虑到新迁移的非热特征不再重要，因此我们简单地忽略其原始的独占嵌入，而使用共享嵌入。通过在HotSketch中设置合适的阈值，我们能够保持适度的迁移频率从而使模型适应分布的变化，而不会对收敛和延迟产生负面影响。

4. 多级哈希嵌入

非热特征的数量远超热门特征，鉴于这些非热特征的重要性得分也符合高度倾斜的Zipf分布，因此我们可以根据其重要性进一步分隔并分配不同的哈希嵌入表。如图4所示，我们将非热门特征划分为不同重要性级别，并为它们分配来自多个表的不同数量的嵌入。简单起见，我们关注2级哈希嵌入，其中非热特征分为中等特征和冷特征。我们扩展了HotSketch的功能，以估计中等特征的重要性分数。由于中等特征比冷特征更重要，我们赋予它们两个不同哈希嵌入表中的两个嵌入向量，而冷特征仅查找单个嵌入向量。

图4展示了一个示例。(1)输入特征f1、f2、f3中，f1的得分大于热阈值，f2的得分大于中阈值，f3的得分低于阈值，因此它们分别被分为热特征、中特征和冷特征；(2) 热特征和冷特征查找嵌入向量的方式不变；（3）中特征从两个哈希嵌入表中查找两个嵌入向量，并通过池化过程获得最终的嵌入。为了确保训练过程保持平滑，当一个特征在中类和冷类之间迁移时，它总是从第一个嵌入表中检索相同的嵌入向量。对于池化操作，在实践中，我们发现嵌入的简单求和效果很好，因为特征的嵌入向量总是在相同的方向上更新。多级哈希嵌入的设计基于这样的观察：唯一嵌入是一种没有信息丢失的综合表示，而对于哈希嵌入，涉及的嵌入数量越多，冲突越少，信息量越多。

图4. 多级哈希嵌入示意图

三、实现

我们将CAFE实现为基于PyTorch的插件嵌入层模块，可以直接替换任何基于PyTorch的推荐模型中的原始Embedding模块。

CAFE后端：我们用C++实现HotSketch算法以减少总体延迟，并使用PyTorch实现CAFE的其余部分。我们HotSketch中桶的数量设置为预先确定的热门特征数量，每个桶有4个槽；我们对所有特征字段使用一个HotSketch结构，而不是每个字段一个结构，因为跨字段的热门特征分布不明确，最好直接用重要性分数来处理。

容错性：我们将所有HotSketch的状态注册为CAFE的PyTorch模块中的buffer，使其可以与模型参数一起保存和加载。这种简单的设计不需要额外的修改，使带有CAFE的DLRM能够使用检查点进行训练和推理。

内存管理：考虑到HotSketch不是计算密集型任务，我们将整个HotSketch结构放在CPU上。CAFE的嵌入模块基于PyTorch运算符构建，可以在任何支持PyTorch的加速器（包括 CPU、GPU）上运行。

四、实验结果

我们用三个模型DLRM[5]、WDL[6]、DCN[7]在四个数据集Criteo，CriteoTB，Avazu，KDD12上做了实验。我们将CAFE与Hash Embedding[1]、Q-R Trick[8]、AdaEmbed[3]、MDE[9]进行了比较。

1、模型质量比较

我们用测试AUC和训练loss分别表示离线和在线场景下的模型质量。如图5所示，相比Hash Embedding、Q-R Trick、MDE方法，CAFE能够在同样的压缩率下达到最好的模型质量，且在训练过程中一直保持更低的训练loss。AdaEmbed是一个自适应方法，在低压缩率场景下有很好的效果，而CAFE不仅能够达到更高的压缩率，还能在低压缩率下取得和AdaEmbed相当或者更好的效果。

具体而言，所有方法中只有CAFE和Hash能达到极限的10000x压缩率，而其他方法只能支持较小的压缩率，例如AdaEmbed和MDE只能支持不到100x的压缩率。DLRM模型里，在Criteo数据集上，与Hash和Q-R Trick相比，CAFE的测试AUC平均提高了1.79%和0.55%；在CriteoTB数据集上，改进分别为1.304%和0.427%； 在KDD12数据集上，改进分别为1.86%和3.80%。 考虑训练loss，与Hash和Q-R Trick相比，CAFE在Criteo数据集上降低了2.31%、0.72%，在CriteoTB数据集上降低了1.35%、0.59%，在Avazu数据集上降低了3.34%、0.76%。与AdaEmbed相比，CAFE在Criteo数据集上达到了几乎相同的测试AUC和训练损失，在CriteoTB数据集上提升了0.04%的测试AUC和减少0.12%的训练损失，在KDD12数据集上提升了0.82%的测试 AUC，在Avazu数据集上的训练损失减少了0.83%。

图5. 模型质量对比

我们进一步查看了多级哈希嵌入的效果，如图6所示，其中CAFE-ML表示CAFE结合多级哈希嵌入。在不同的压缩比下，CAFE-ML始终比CAFE表现更好，测试AUC提高了0.08%，训练损失降低了0.25%。CAFE-ML在较小的压缩比下表现尤其出色，在100倍压缩比下几乎没有造成性能下降。这是因为CAFE-ML在小压缩比下为多级哈希嵌入表分配了更多的内存，使得中等特征的表示更加精确。

图6. 比较多级哈希嵌入

2、吞吐量对比

我们在图7中测试了每种方法的吞吐量，实验是在压缩率为10x的CriteoTB数据集上进行的。由于不同方法的数据加载时间和DNN计算时间是相同的，因此差异仅在于嵌入层的操作。与未压缩的嵌入操作相比，Hash仅需要额外的模运算，因此是训练和推理中最快的方法。Q-R Trick只额外引入了哈希过程和嵌入向量的聚合，而MDE引入了矩阵乘法、但需要更少的内存访问来获取嵌入参数，因此这两个方法也有较高吞吐量。为了适应在线场景，AdaEmbed和CAFE引入了嵌入的动态调整，即重新分配或迁移，从而导致吞吐量降低。AdaEmbed会定期采样数千个数据来确定是否重新分配，带来较大时间开销；相比之下，CAFE在HotSketch中的计算开销可以忽略不计。与AdaEmbed相比，CAFE的训练吞吐量提高了33.97%，推理吞吐量提高了1.20%。

图7. 吞吐量对比

3、HotSketch评测

桶中槽数的影响（图8a-b）：我们记录使用不同槽数的HotSketch的召回率和吞吐量。实验使用Criteo数据集上1000x压缩率下的热门特征数量作为k。在图8a中，召回随着存储变大而增加。根据我们的理论推导（见论文推论3.5），给定参数1.05到1.1的Zipf分布，每个桶的最佳槽数位于11到21。因此，槽数为8和16比4和32表现出更好的召回率。图 8b中所示的序列化插入（写入）和查询（读取）的吞吐量约为1e7，比DLRM本身的吞吐大得多。考虑到我们在实践中可以并行化操作，HotSketch只占训练和推理的一小部分。随着槽数量的增加，吞吐量会下降，因为需要更多的时间在桶内进行比较。权衡召回率和吞吐量，我们在实现中采用每个桶4个槽，此时的模型质量已经足够好。

寻找实时top-k特征（图8c-d）：我们进行实验来评估HotSketch在线训练中寻找两类实时热门特征的性能：最新的top-k特征，以及前一个时间窗口中的top-k特征。这两类top-k特征在在线训练过程中随着数据分布的变化而变化，因此可以有效体现HotSketch适应动态工作负载的能力。实验在 Criteo 上进行，使用6天的在线训练数据，滑动窗口大小为0.5天。图8c-d显示了HotSketch在不同压缩比下在线训练时的实时召回率，始终满足>90%，这意味着它可以很好地跟上不断变化的数据分布。

图8. HotSketch评测

五、总结

在本文中，我们提出了CAFE，一种紧凑、自适应且快速的嵌入压缩方法，它满足三个基本设计要求：高内存效率、低延迟和对动态数据分布的适应性。我们引入了一种轻量数据结构HotSketch记录特征的重要性分数，它产生的时间开销可以忽略不计，并且其内存消耗明显低于原始嵌入表。通过将独占嵌入分配给一小组重要特征，并将共享嵌入分配给其他不太重要的特征，我们在有限的内存约束内实现了卓越的模型质量。为了适应在线训练过程中的动态数据分布，我们采用了基于HotSketch的嵌入迁移策略。我们通过多级哈希嵌入进一步优化CAFE，创建更细粒度的重要性组。实验结果表明，CAFE优于现有方法，在Criteo和CriteoTB数据集上10000x压缩比下，测试AUC分别提高了3.92%、3.68%，训练损失降低了4.61%、3.24%，在离线训练和在线训练中均表现出优异的性能。

实验室简介

北京大学数据与智能实验室（Data And Intelligence Research Lab at Peking Univeristy，PKU-DAIR实验室）由北京大学计算机学院崔斌教授领导，长期从事数据库系统、大数据管理与分析、人工智能等领域的前沿研究，在理论和技术创新以及系统研发上取得多项成果，已在国际顶级学术会议和期刊发表学术论文200余篇，发布多个开源项目。课题组同学曾数十次获得包括CCF优博、ACM中国优博、北大优博、微软学者、苹果奖学金、谷歌奖学金等荣誉。PKU-DAIR实验室持续与工业界展开卓有成效的合作，与腾讯、阿里巴巴、苹果、微软、百度、快手、中兴通讯等多家知名企业开展项目合作和前沿探索，解决实际问题，进行科研成果的转化落地。