Atomic Feature Mimicking

Unified Low-rank Compression Framework for Large-scale Recommendation Systems

Authors: Hao Yu, Minghao Fu, Jiandong Ding, Yusheng Zhou, Jianxin Wu

[arXiv] [bibtex]

Introduction: This repository provides an implementation for the KDD2024 ADS track paper: "Unified Low-rank Compression Framework for Large-scale Recommendation Systems" based on DeepCTR-Torch. We propose a unified and efficient low-rank decomposition framework to compress the embedding tables and MLP layers of CTR prediction models.

Note: Due to copyright restrictions, we can only provide evaluation code, including eval_criteo.py and eval_avazu.py, and the core compression code demo (afm_emb.py and afm_mlp.py). We cannot provide the detailed engineering code.

Compress MLP Layers:

First given a pre-trained CTR prediction model, calculate the MLP output of each layer of the model and calculate $U$, $V$ and $AVG$ (step2()). Then merge weight (step3()) and save checkpoint.

Compress Embedding Tables:

First given a pre-trained model, count the output of each sparse embedding feature of the model and calculate $U$, $V$ and $AVG$ (step2()). Then merge the embedding table $D$ and compresion weight $U$ (step3()) and retain the checkpoint. Note that fusion will fail when the interaction layer exists between the compressed weight $U^\top$ and the first FC layer weight $W$.

Please view our paper for more informations.

Citation

@article{yu2024Unified,
  title={Unified Low-rank Compression Framework for Large-scale Recommendation Systems},
  journal={arXiv preprint arXiv:2405.18146},
  author={Hao Yu,Minghao Fu,Jiandong Ding,Yusheng Zhou,Jianxin Wu},
  year={2024}
}

Contact

If you have any questions about our work, feel free to contact us through email (Hao Yu: yuh@lamda.nju.edu.cn) or Github issues.

yuhao318/Atomic_Feature_Mimicking

Atomic Feature Mimicking

Citation

Contact