[VLDB 2024] 大规模可学习向量存储压缩的实验分析

关键词

Embedding memory compression 

Deep learning recommendation model

Retrieval-augmented large language model

导 读

大规模可学习向量存储压缩

可学习嵌入向量是机器学习中最重要的应用之一，广泛应用于各种数据库相关领域，特别是推荐任务和检索任务。然而，推荐任务中稀疏数据的高维性和检索相关任务中的巨量语料库导致嵌入表的内存占用过大，对模型的训练和部署提出了巨大的挑战。最近很多研究提出了各种压缩嵌入的方法，但代价是模型质量略有下降或引入其他开销。虽然相关工作很多，但现有的实验仅涵盖部分方法，并且侧重于有限的评价指标，所以目前各类方法在各个评价指标上的相对效果仍不明确。在这篇论文中，我们对嵌入压缩进行了全面的比较分析和实验评估。我们提出了一种新的分类法，根据这些压缩方法的范式进行分类，并进一步开发了一个集成了14种代表性方法的模块化基准测试框架。在统一的测试环境下，我们公平地评估了每种方法，展示它们在不同内存预算下的效果，并分析了各个方法的实用场景。我们的研究还揭示了当前方法的局限性，并为未来研究提出了潜在的方向。

论文链接：

https://arxiv.org/abs/2311.15578

代码地址：

https://github.com/HugoZHL/Hetu/tree/embedmem/tools/EmbeddingMemoryCompression

01

问 题 背 景

近年来，嵌入技术受到了数据库社区的极大关注，它们被广泛应用于多个数据库相关的研究领域，例如基数估计、查询优化、语言理解、实体解析，文档检索，图学习和广告推荐等。这些应用通常依赖大量的嵌入向量来学习语义表示并提取有意义的模式和相似性。然而，这些可学习向量数量庞大，在实际部署场景中带来了相当大的存储和检索挑战。例如，Meta提出的大型深度学习推荐模型，配备了数十亿个嵌入向量，需要占用 96 TB 的内存来提供服务[1]。克服大型嵌入模型可扩展性问题的一个方向是使用多个分布式实例，这会带来显著的通信开销。另一种方法是压缩嵌入向量以减少其存储占用，且尽量不影响推荐模型的准确性或实用性。在过去的几年里，人们提出了各种压缩方法，包括哈希、量化、降维等。然而，这些技术在不同场景下的性能和有效性仍未明确。当存储受限时，数据科学家难以从现有的压缩技术中选择合适的方法。 

在这篇论文中，我们主要关注推荐和检索任务中的嵌入压缩方法。这两个任务分别包含高维稀疏数据和巨大的文本数据量，从而具有最严重的嵌入向量存储压力。深度学习推荐模型中，实际应用中稀疏数据的维度不断增大，导致嵌入数量不断增多。其中嵌入流行度遵循幂律分布，并且训练过程需要频繁访问和更新多条向量，因此传统的键值数据库存储压缩方法并不适用。检索模型现有的工作更多地侧重于设计高效的嵌入搜索算法，但检索增强大规模语言模型的出现[2]给嵌入向量存储带来了新的挑战。由于大规模语言模型已经消耗了大量的内存，嵌入表无法放入GPU或其他加速器中，导致搜索延迟很高。

02

方 法 分 类

我们首先将所有嵌入压缩方法分为特征间压缩和特征内压缩，前者使特征共享参数，后者对每个特征的向量本身进行压缩。然后，我们根据方法的性质和技术进一步细分。

为了解决特征数量爆炸性增长造成的内存瓶颈，一种直接的方法是仅保留少量嵌入以供不同的特征共享。特征间压缩需要维护从特征到嵌入的新映射。早期的哈希方法[3]利用哈希函数将特征映射到多热向量中，然后从哈希嵌入表中查找子嵌入来构造最终嵌入。根据训练时编码函数是否固定，这些方法可以进一步分为静态编码[4-8]和动态编码[9-10]两大类。相关方法使用的技术、复杂度、以及其他细节如下表所示。

受特征重要性不同的启发，特征内压缩为每个特征分配单独压缩的嵌入。根据压缩范式，特征内压缩可以进一步分为量化[11-12]、降维[13-15]和剪枝[16-17]。量化是深度学习模型中常见的压缩方法，它使用位数较少的数据类型。降维和剪枝为不同重要性的特征提供不同维度和稀疏度的嵌入向量。相关方法使用的技术、复杂度、以及其他细节如下表所示。

03

评 估 框 架

我们针对可学习向量压缩方法设计并实现了一个模块化的评估框架，如下图所示。该框架包含四个部分：编码函数、嵌入层、调度器、辅助模块。编码函数负责将特征映射到独热向量或者多热向量；嵌入层存储了实际的可学习向量的参数，支持各种可能的形式；调度器负责管理整个训练流程，支持训练阶段的切换、以及用合适的数据训练合适的部分参数；辅助模块存储其他需要的参数或者频率信息等。现有的压缩方法都可以由上述四个模块组成的流程表示。在此基础上，我们选择了14个具有代表性的压缩方法进行实现，并用于后续实验。

04

实 验

对于深度学习推荐模型，我们选择了14个具有代表性的压缩方法，包括静态编码方法CompoEmb[4]、TT-Rec[5]、DHE[6]、ROBE[7]、Dedup[8]，动态编码方法MGQE[9]、AdaptEmb[10]，量化方法INT8/16[11]、ALPT[12]，降维方法MDE[13]、AutoDim[14]、OptEmbed[15]，剪枝方法DeepLight[16]、AutoSrh[17]。我们在两个公开数据集Avazu[18]、Criteo[19]和一个更大的工业数据集Company上做实验，使用的模型是DLRM[20]、WDL[21]、DCN[22]。我们对推断时的存储设置了4个压缩率，分别为50%，10%，1%，0.1%的原存储大小。实验结果如下：

基于实验结果，我们有如下发现：

（1）在压缩能力上，基于哈希的方法（包括Dedup）和剪枝方法能达到我们设置的所有压缩率，而其他方法只能支持一定范围内或者特定的压缩率。

（2）在模型质量上，静态编码方法、量化方法和剪枝方法有最好的模型质量；降维方法也能在特定场景下达到很好的模型质量。没有任何方法能在所有场景下达到最优。

（3）在训练所需存储上，静态编码方法（除了Dedup）和MDE方法使用了最少的存储，因为它们没有引入其他额外的结构；而其他方法都或多或少引入了额外的结构，部分方法需要在训练中保持原本的嵌入表。

（4）在训练时间上，简单的哈希方法（包括Dedup）和简单量化方法INT8/16能最快达到收敛，因为它们没有改变原本的训练范式；而其他方法都一定程度上改变了原本的训练范式，有的还引入了复杂的计算或者多个训练阶段。

（5）在推理延迟上，Dedup、ROBE、OptEmbed、量化方法、剪枝方法有着最低的推理延迟，因为它们可以简单的将所有特征存储于一个嵌入表中，而其他压缩方法都需要维护多个嵌入表、或者引入额外的计算。

（6）我们的方法分类是根据方法的压缩范式来确定的，同类方法的实现方式基本相同。同类方法在不同指标上的表现有一定的相似性，但也存在一定差异，因为它们使用的具体技术不同。我们的分析从压缩范式、使用的技术两个角度出发，能够得到更加综合的结论。

（7）在工业数据集上，我们能够得到和公开数据集类似的结论，原因在于公开数据集也是从真实场景中得到的，且我们的结论有一定的泛化性。

另外两个模型的实验结果如图4所示，三个模型的表现差别不大，因为我们的研究对象是即插即用、只改变嵌入层的压缩方法，与后面的模型结构是正交的。

我们还尝试使用不同的向量维度，如图5，发现在压缩时考虑了特征重要性的方法能够在更高维度（相同存储）下达到更好的效果，因为它们能为重要的特征分配更多的存储；而对于不考虑特征重要性的方法，在相同存储下增加维度基本没什么效果。我们在图6中可视化了其中一个方法在不同维度时各频率的特征实际的参数量，发现频率和重要性有一定正相关，且重要的特征在维度较大时参数量更大。

对于检索增强大规模语言模型，我们选择了8个方法（或者相应的变体），包括TT（tensor-train分解，类似TT-Rec），Dedup，PQ（MGQE的变体，不考虑频率），MagPQ（MGQE的变体，使用向量的范数代替频率），INT8/16，SVD（降维的变体），MagSVD（MDE的变体，使用向量的范数代替频率），Pruning（剪枝绝对值较小的值）。我们用RAG[2]进行实验，它使用DPR[23]作为检索器、BART[24]作为生成模型，在开放域问答数据集Natural Questions[25]上做实验。对于检索相关的任务，可学习向量是神经网络的输出，只有在推断的时候才会生成所有的向量，因此我们考察模型的推断能力。与推荐任务相同，我们对存储设置了4个压缩率。实验结果如表6所示：

在实验中，我们有如下发现：

（1）在压缩能力上，TT、Dedup、Pruning能够达到我们设置的所有压缩率，而其他方法只能支持一定范围内或者特定的压缩率。

（2）在模型质量上，INT8/16、(Mag)SVD、Pruning在存储预算较大时有较好的模型质量，(Mag)PQ在存储预算较小时有较好的模型质量。其中INT8/16和PQ因为具有良好的性能，被广泛用于实际应用中[26]。

（3）在压缩时间上，INT8/16的压缩时间最短，其次是Pruning；其他方法都涉及较为复杂的计算，需要较多压缩时间。

（4）在批量解压延迟上，INT8/16的解压时间最短，其次是Pruning、Dedup和(Mag)PQ。

05

现 有 挑 战 和 未 来 研 究 方 向

目前，推荐任务中没有一个方法能够在所有指标上达到最优，未来可能的一个研究方向是对现有方法进行适当的结合，从而达到取长补短的效果。

对于检索任务而言，可学习向量的压缩仍处于早期研究阶段。在存储的挑战下，未来可能出现更多专门针对检索向量的压缩方法，这些方法可以与搜索过程相结合，从而达到更好的存储、时间、效果的权衡。 

在实验中，我们发现在不同数据集上各压缩方法表现不同，目前尚不明确数据集与压缩方法之间的关系。我们认为数据分布对压缩方法有潜在的影响，未来我们期待有研究工作能对推荐、检索数据有更深的理解，这不仅有助于选择合适的压缩方法，也能启发对该领域进一步的研究探索。

详细了解本工作，请访问下方链接地址：

论文链接：

https://arxiv.org/abs/2311.15578

References：

[1] Dheevatsa Mudigere, Yuchen Hao, Jianyu Huang, Zhihao Jia, Andrew Tulloch, Srinivas Sridharan, Xing Liu, Mustafa Ozdal, Jade Nie, Jongsoo Park, Liang Luo, Jie Amy Yang, Leon Gao, Dmytro Ivchenko, Aarti Basant, Yuxi Hu, Jiyan Yang, Ehsan K. Ardestani, Xiaodong Wang, Rakesh Komuravelli, Ching-Hsiang Chu, Serhat Yilmaz, Huayu Li, Jiyuan Qian, Zhuobo Feng, Yinbin Ma, Junjie Yang, Ellie Wen, Hong Li, Lin Yang, Chonglin Sun, Whitney Zhao, Dimitry Melts, Krishna Dhulipala, K. R. Kishore, Tyler Graf, Assaf Eisenman, Kiran Kumar Matam, Adi Gangidi, Guoqiang Jerry Chen, Manoj Krishnan, Avinash Nayak, Krishnakumar Nair, Bharath Muthiah, Mahmoud khorashadi, Pallab Bhattacharya, Petr Lapukhov, Maxim Naumov, Ajit Mathews, Lin Qiao, Mikhail Smelyanskiy, Bill Jia, and Vijay Rao. 2022. Software-hardware co-design for fast and scalable training of deep learning recommendation models. In The 49th

Annual International Symposium on Computer Architecture (ISCA).

[2] Patrick S. H. Lewis, Ethan Perez, Aleksandra Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, and Douwe Kiela. 2020. Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. In Advances in Neural Information Processing Systems 33 (NeurIPS).

[3] Kilian Q. Weinberger, Anirban Dasgupta, John Langford, Alexander J. Smola, and Josh Attenberg. 2009. Feature hashing for large scale multitask learning. In Proceedings of the 26th Annual International Conference on Machine Learning (ICML).

[4] Hao-Jun Michael Shi, Dheevatsa Mudigere, Maxim Naumov, and Jiyan Yang. 2020. Compositional Embeddings Using Complementary Partitions for Memory-Efficient Recommendation Systems. In Proceedings of the 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD)

[5] Wang-Cheng Kang, Derek Zhiyuan Cheng, Tiansheng Yao, Xinyang Yi, Ting Chen, Lichan Hong, and Ed H. Chi. 2021. Learning to Embed Categorical Features without Embedding Tables for Recommendation. In Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD).

[6] Chunxing Yin, Bilge Acun, Carole-JeanWu, and Xing Liu. 2021. TT-Rec: Tensor Train Compression for Deep Learning Recommendation Models. In Proceedings of Machine Learning and Systems (MLSys).

[7] Aditya Desai, Li Chou, and Anshumali Shrivastava. 2022. Random Offset Block Embedding (ROBE) for compressed embedding tables in deep learning recommendation systems. In Proceedings of Machine Learning and Systems (MLSys).

[8] Lixi Zhou, Jiaqing Chen, Amitabh Das, Hong Min, Lei Yu, Ming Zhao, and Jia Zou. 2022. Serving Deep Learning Models with Deduplication from Relational Databases. Proceedings of the VLDB Endowment 15, 10 (2022), 2230–2243.

[9] Wang-Cheng Kang, Derek Zhiyuan Cheng, Ting Chen, Xinyang Yi, Dong Lin, Lichan Hong, and Ed H. Chi. 2020. Learning Multi-granular Quantized Embeddings for Large-Vocab Categorical Features in Recommender Systems. In Companion of The 2020 Web Conference.

[10] DeepRec. 2021. Adaptive Embedding. https://github.com/alibaba/DeepRec/blob/main/docs/docs_en/Adaptive-Embedding.md.

[11] Zhiqiang Xu, Dong Li, Weijie Zhao, Xing Shen, Tianbo Huang, Xiaoyun Li, and Ping Li. 2021. Agile and Accurate CTR Prediction Model Training for Massive-Scale Online Advertising Systems. In International Conference on Management of Data (SIGMOD).

[12] Shiwei Li, Huifeng Guo, Lu Hou, Wei Zhang, Xing Tang, Ruiming Tang, Rui Zhang, and Ruixuan Li. 2023. Adaptive Low-Precision Training for Embeddings in Click-Through Rate Prediction. In Thirty-Seventh AAAI Conference on Artificial Intelligence (AAAI).

[13] Antonio A. Ginart, Maxim Naumov, Dheevatsa Mudigere, Jiyan Yang, and James Zou. 2021. Mixed Dimension Embeddings with Application to Memory-Efficient Recommendation Systems. In IEEE International Symposium on Information Theory (ISIT).

[14] Xiangyu Zhao, Haochen Liu, Hui Liu, Jiliang Tang, Weiwei Guo, Jun Shi, Sida Wang, Huiji Gao, and Bo Long. 2021. AutoDim: Field-aware Embedding Dimension Searchin Recommender Systems. In Proceedings of the Web Conference (WWW).

[15] Fuyuan Lyu, Xing Tang, Hong Zhu, Huifeng Guo, Yingxue Zhang, Ruiming Tang, and Xue Liu. 2022. OptEmbed: Learning Optimal Embedding Table for Click-through Rate Prediction. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management (CIKM).

[16] Wei Deng, Junwei Pan, Tian Zhou, Deguang Kong, Aaron Flores, and Guang Lin. 2021. DeepLight: Deep Lightweight Feature Interactions for Accelerating CTR Predictions in Ad Serving. In The 14th ACM International Conference on Web Search and Data Mining (WSDM).

[17] Shuming Kong,Weiyu Cheng, Yanyan Shen, and Linpeng Huang. 2023. AutoSrh: An Embedding Dimensionality Search Framework for Tabular Data Prediction. IEEE Transactions on Knowledge and Data Engineering 35, 7 (2023), 6673–6686.

[18] Steve Wang and Will Cukierski. 2014. Avazu Click-Through Rate Prediction.  https://kaggle.com/competitions/avazu-ctr-prediction.

[19] Criteo Labs. 2014. Kaggle display advertising challenge dataset. https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset.

[20] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang, Narayanan Sundaraman, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu, Alisson G. Azzolini, Dmytro Dzhulgakov, Andrey Mallevich, Ilia Cherniavskii, Yinghai Lu, Raghuraman Krishnamoorthi, Ansha Yu, Volodymyr Kondratenko, Stephanie Pereira, Xianjie Chen, Wenlin Chen, Vijay Rao, Bill Jia, Liang Xiong, and Misha Smelyanskiy. 2019. Deep Learning Recommendation Model for Personalization and Recommendation Systems. CoRR abs/1906.00091 (2019).

[21] Heng-Tze Cheng, Levent Koc, Jeremiah Harmsen, Tal Shaked, Tushar Chandra, Hrishi Aradhye, Glen Anderson, Greg Corrado, Wei Chai, Mustafa Ispir, Rohan Anil, Zakaria Haque, Lichan Hong, Vihan Jain, Xiaobing Liu, and Hemal Shah. 2016. Wide & Deep Learning for Recommender Systems. In Proceedings of the 1st Workshop on Deep Learning for Recommender Systems (DLRS@RecSys).

[22] Ruoxi Wang, Bin Fu, Gang Fu, and Mingliang Wang. 2017. Deep & Cross Network for Ad Click Predictions. In Proceedings of the ADKDD’17.

[23] Vladimir Karpukhin, Barlas Oguz, Sewon Min, Patrick S. H. Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih. 2020. Dense Passage Retrieval for Open-Domain Question Answering. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP).

[24] Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Veselin Stoyanov, and Luke Zettlemoyer. 2020. BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics (ACL).

[25] Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur P. Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoc Le, and Slav Petrov. 2019. Natural Questions: a Benchmark for Question Answering Research. Transactions of the Association for Computational Linguistics 7 (2019), 452–466.

[26] Jeff Johnson, Matthijs Douze, and Hervé Jégou. 2019. Billion-Scale Similarity Search with GPUs. IEEE Transactions on Big Data 7, 3 (2019), 535–547.

END

欢迎关注本公众号，帮助您更好地了解北京大学数据与智能实验室（PKU-DAIR），第一时间了解PKU-DAIR实验室的最新成果！

实验室简介

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