Publication Type
Conference Proceeding Article
Version
publishedVersion
Publication Date
10-2022
Abstract
Although large pre-trained models of code have delivered significant advancements in various code processing tasks, there is an impediment to the wide and fluent adoption of these powerful models in software developers’ daily workflow: these large models consume hundreds of megabytes of memory and run slowly on personal devices, which causes problems in model deployment and greatly degrades the user experience. It motivates us to propose Compressor, a novel approach that can compress the pre-trained models of code into extremely small models with negligible performance sacrifice. Our proposed method formulates the design of tiny models as simplifying the pre-trained model architecture: searching for a significantly smaller model that follows an architectural design similar to the original pre-trained model. Compressor proposes a genetic algorithm (GA)-based strategy to guide the simplification process. Prior studies found that a model with higher computational cost tends to be more powerful. Inspired by this insight, the GA algorithm is designed to maximize a model’s Giga floating-point operations (GFLOPs), an indicator of the model computational cost, to satisfy the constraint of the target model size. Then, we use the knowledge distillation technique to train the small model: unlabelled data is fed into the large model and the outputs are used as labels to train the small model. We evaluate Compressor with two state-of-the-art pre-trained models, i.e., CodeBERT and GraphCodeBERT, on two important tasks, i.e., vulnerability prediction and clone detection. We use our method to compress pre-trained models to a size (3 MB), which is 160× smaller than the original size. The results show that compressed CodeBERT and GraphCodeBERT are 4.31× and 4.15× faster than the original model at inference, respectively. More importantly, they maintain 96.15% and 97.74% of the original performance on the vulnerability prediction task. They even maintain higher ratios (99.20% and 97.52%) of the original performance on the clone detection task
Keywords
Model compression, Genetic algorithm, Pre-trained models
Discipline
Software Engineering
Research Areas
Software and Cyber-Physical Systems
Publication
ASE '22: Proceedings of the 37th IEEE/ACM International Conference on Automated Software Engineering, Rochester, MI, October 10-14
First Page
1
Last Page
12
ISBN
9781450394758
Identifier
10.1145/3551349.3556964
Publisher
ACM
City or Country
New York
Citation
SHI, Jieke; YANG, Zhou; XU, Bowen; KANG, Hong Jin; and LO, David.
Compressing pre-trained models of code into 3 MB. (2022). ASE '22: Proceedings of the 37th IEEE/ACM International Conference on Automated Software Engineering, Rochester, MI, October 10-14. 1-12.
Available at: https://ink.library.smu.edu.sg/sis_research/7725
Copyright Owner and License
Authors
Creative Commons License
This work is licensed under a Creative Commons Attribution 4.0 International License.
Additional URL
https://doi.org/10.1145/3551349.3556964