Code for this paper UMEC: Unified Model and Embedding Compression for Efficient Recommendation Systems

Jiayi Shen, Haotao Wang*, Shupeng Gui*, Jianchao Tan, Zhangyang Wang, and Ji Liu

We propose a unified model and embedding compression (UMEC) framework to hammer an efficient neural network-based recommendation system. Our framework jointly learns input feature selection and neural network compression together, and solve them as an end-to-end resource-constrained optimization problem using ADMM.

We perform the compression process on DLRM, which is a public recommendation model. Our proposed algorithm is mainly implemented inrc_optimizer.py and rc_utils.py. Other files are inherited from the original DLRM code repo, with several lines of modifications, such as joint_train.py, input_selection.py, and finetune.py, in order to plug in our algorithm. To run the code in this repo, you have to first follow the instructions in the original repo to download the dataset, and run the corresponding training part, to finish the data preprocessing process.

To implement joint training and compressing under the resource constraint, please see the script in script/joint_train.sh.

To implement input selection under the resource constraint, please see the script in script/input_selection.sh.

We thank the author of DLRM for providing a recommendation model benchmark.

@inproceedings{
shen2021umec,
title={{\{}UMEC{\}}: Unified model and embedding compression for efficient recommendation systems},
author={Jiayi Shen and Haotao Wang and Shupeng Gui and Jianchao Tan and Zhangyang Wang and Ji Liu},
booktitle={International Conference on Learning Representations},
year={2021},
url={https://openreview.net/forum?id=BM---bH_RSh}
}