Source code for AAAI 2022 paper: From Dense to Sparse: Contrastive Pruning for Better Pre-trained Language Model Compression.
Coming soon!
Most model pruning approaches only consider task-specific knowledge towards downstream tasks, but ignore the essential task-agnostic knowledge. Therefore, we propose ContrAstive Pruning (CAP), a general pruning framework under the pre-training and fine-tuning paradigm, which aims at maintaining both task-specific and task-agnostic knowledge during pruning. CAP is designed as a general framework, compatible with both structured and unstructured pruning. Unified in contrastive learning, CAP encourage the pruned model to learn from the pre-trained model, the snapshots (intermediate models during pruning), and the fine-tuned model, through three contrastive modules, PrC, SnC, and FiC, respectively. You can refer to our paper for more details.
For unstructured pruning, you have to download the GLUE and SQuAD v1.1 data and put them into the data folder. We also provide the data here.
Before training, you have to first train a teacher model and put it into the teacher folder, which is used in our FiC module and knowledge distillation. Using the scripts provided by Huggingface, you can easily train a model for GLUE and SQuAD v1.1.
We provide codes for both structured pruning and unstructured pruning based on our proposed Contrastive Pruning framework.
If you use this work or code, please kindly cite the following paper:
@inproceedings{xu-etal-2021-contrastivepruning,
title = "From Dense to Sparse: Contrastive Pruning for Better Pre-trained Language Model Compression",
author = "Runxin Xu and
Fuli Luo and Chengyu Wang and
Baobao Chang and Jun Huang and
Songfang Huang and Fei Huang",
booktitle = "Thirty-Sixth {AAAI} Conference on Artificial Intelligence (AAAI)",
year = "2022"
}