[论文解读] Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems

Antonio Ginart, Maxim Naumov+3 more|arXiv (Cornell University)|Sep 25, 2019

Recommender Systems and Techniques参考文献 77被引用 32

一句话总结

本文提出了一种混合维度（MD）嵌入方法用于推荐系统，其中嵌入向量的维度根据特征的流行度进行动态调整，而非固定为统一维度。通过为高频查询的项目分配更多参数，而为低频项目分配更少参数，MD嵌入在Criteo数据集上实现了高达16倍的内存压缩，且准确率无损失，同时在GPU上训练速度提升2倍以上。

ABSTRACT

Embedding representations power machine intelligence in many applications, including recommendation systems, but they are space intensive -- potentially occupying hundreds of gigabytes in large-scale settings. To help manage this outsized memory consumption, we explore mixed dimension embeddings, an embedding layer architecture in which a particular embedding vector's dimension scales with its query frequency. Through theoretical analysis and systematic experiments, we demonstrate that using mixed dimensions can drastically reduce the memory usage, while maintaining and even improving the ML performance. Empirically, we show that the proposed mixed dimension layers improve accuracy by 0.1% using half as many parameters or maintain it using 16X fewer parameters for click-through rate prediction task on the Criteo Kaggle dataset.

研究动机与目标

• 解决大规模推荐系统中嵌入层的高内存消耗问题，其存储量可超过模型总存储的99.9%。
• 克服统一嵌入维度带来的低效问题，即在低频特征上浪费参数，而在高频特征上则参数不足。
• 设计一种参数高效的嵌入架构，根据特征查询频率动态调整嵌入维度。
• 通过理论与实证验证，证明MD嵌入在数据分布高度偏斜的情况下，能同时提升内存效率与模型泛化能力。
• 相比现有非统一嵌入方法，实现更快的训练速度与更低的调参复杂度。

提出的方法

• 提出一种混合维度（MD）嵌入层，其中每个嵌入向量的维度基于其查询频率设定，而非固定统一维度。
• 将嵌入维度选择建模为一个凸优化问题，以在总参数预算下最小化重构失真。
• 采用低秩近似框架，将嵌入矩阵表示为多个块的和，每个块对应一对类别特征。
• 通过拉格朗日松弛法推导出最优维度分配规则，其中每个特征对的最优维度由其流行度、谱衰减特性及参数预算共同决定。
• 对于幂律谱衰减情形，推导出闭式解：$ d^{*}_{ij} = \lambda \zeta_{ij} \Pi_{ij}^{1/(2\beta)} $，其中 $ \lambda $ 为拉格朗日乘子，$ \zeta_{ij} $ 依赖于特征大小与谱衰减特性。
• 将连续的最优维度向下取整为整数，以确保可行性，且近似误差有理论保证（通常仅损失1%–10%的谱质量）。

Figure 1 : Matrix Architecture for UD and MD Embedding Layers.

实验结果

研究问题

• RQ1 能否根据特征流行度动态分配嵌入维度，在不牺牲模型性能的前提下减少内存使用？
• RQ2 在内存受限条件下，MD嵌入是否在重构失真与泛化能力方面优于统一维度嵌入？
• RQ3 在给定特征流行度与整体参数预算的前提下，每个特征的最优嵌入维度是什么？
• RQ4 与现有非统一嵌入架构相比，该方法在训练速度与调参复杂度方面表现如何？
• RQ5 在何种条件下（如谱衰减特性、流行度偏斜程度），MD嵌入能对统一嵌入实现可证明的性能提升？

主要发现

• 在Criteo Kaggle数据集上，MD嵌入在仅使用一半参数量的情况下，准确率比统一嵌入高出0.1%。
• MD嵌入在保持与统一嵌入相同准确率的同时，将参数量减少高达16倍，展现出极强的参数效率。
• 由于减少了内存带宽压力并提升了计算效率，使用MD嵌入进行训练在GPU上速度提升2倍以上。
• 理论分析证明，在数据受限且流行度分布足够偏斜的情况下，MD嵌入的重构失真低于统一嵌入。
• 对于幂律谱衰减，最优嵌入维度按 $ \Pi_{ij}^{1/(2\beta)} $ 的方式缩放，其中 $ \Pi_{ij} $ 为特征对的流行度。
• 将连续最优维度向下取整为整数所带来的近似误差较小（通常仅损失1%–10%的谱质量），确保了实际可行性。

Figure 2 : CTR prediction results for MD embeddings on Criteo dataset using DLRM. Implementation is available as part of an open-source project on GitHub: facebookresearch/dlrm . Fig. 2a (left): Learning curves for selected emb. arch. Fig. 2b (center): Loss vs. # param. for varying $\alpha$ . Fig 2c