中文

将预训练代码模型压缩至 3 MB

软件工程 2022-09-07 v2

摘要

尽管大型预训练代码模型在各种代码处理任务中取得了显著进展,但在软件开发人员的日常工作流中广泛且流畅地采用这些强大模型仍存在障碍:这些大模型消耗数百兆字节内存,且在个人设备上运行缓慢,导致模型部署困难并严重降低用户体验。这促使我们提出 Compressor,一种能够将预训练代码模型压缩为极小模型且性能损失可忽略的新方法。我们提出的方法将微小模型的设计表述为简化预训练模型架构:搜索一个显著更小且遵循与原预训练模型相似架构设计的模型。Compressor 提出一种基于遗传算法(GA)的策略来指导简化过程。先前研究发现,计算成本更高的模型往往更强大。受此启发,GA 算法被设计为最大化模型的 Giga 浮点运算次数(GFLOPs,模型计算成本的指标),以满足目标模型大小的约束。随后,我们使用知识蒸馏技术训练小模型:将无标签数据输入大模型,并将其输出作为标签来训练小模型。我们在两个最先进的预训练模型(即 CodeBERT 和 GraphCodeBERT)上,针对两个重要任务(即漏洞预测和克隆检测)评估了 Compressor。我们使用该方法将预训练模型压缩至 3 MB 的大小,比原始大小小 160×\times。结果表明,压缩后的 CodeBERT 和 GraphCodeBERT 在推理时分别比原始模型快 4.31×\times 和 4.15×\times。更重要的是,……

关键词

引用

@article{arxiv.2208.07120,
  title  = {Compressing Pre-trained Models of Code into 3 MB},
  author = {Jieke Shi and Zhou Yang and Bowen Xu and Hong Jin Kang and David Lo},
  journal= {arXiv preprint arXiv:2208.07120},
  year   = {2022}
}

备注

Accepted by the Research Papers Track of 37th IEEE/ACM International Conference on Automated Software Engineering (ASE '22)