Compressing Pre-trained Models of Code into 3 MB

Jieke Shi1, Zhou Yang1, Bowen Xu1, Hong Jin Kang1, David Lo1

1Singapore Management University, Singapore

引用

Shi J, Yang Z, Xu B, et al. Compressing pre-trained models of code into 3 mb[C]//Proceedings of the 37th IEEE/ACM International Conference on Automated Software Engineering. 2022: 1-12.

论文：https://dl.acm.org/doi/abs/10.1145/3551349.3556964

仓库：https://github.com/soarsmu/Compressor.git

摘要

本文提出了Compressor，一种能够将预训练代码模型压缩到极小模型的新方法，且性能损失可忽略不计。此方法将小型模型的设计形式化为简化预训练模型架构：寻找一个显著更小的模型，其架构设计与原始预训练模型类似。Compressor提出了一种基于遗传算法（GA）的策略来指导简化过程。GA算法旨在最大化模型的Giga浮点运算（GFLOPs），这是模型计算成本的指标，以满足目标模型大小的约束。然后，使用知识蒸馏技术训练小模型：将未标记数据输入大模型，并使用输出作为标签来训练小模型。

1 引言

预训练的代码模型在各种程序理解和生成任务中取得了巨大成功。然而，这些模型体积庞大，参数量达到数十亿，推理速度较低，因此不适合消费者级设备或需要严格低延迟的应用。以CodeBERT为例，这是一种处理多种编程语言的最先进预训练模型，拥有1.25亿个参数，模型大小为476 MB。在消费者级笔记本电脑上进行推理时，响应延迟高达1.5秒。根据Visual Studio团队在为大量用户部署集成开发环境（IDE）时的经验，Svyatkovskiy等人建议，IDE组件或编辑器插件的合理上限大小为50 MB，而3 MB的模型更为理想，因为它可以在极为受限的环境中部署（例如，用于教学学生编程的低端硬件）。同时，Aye等人认为推理延迟应尽可能减少到0.1秒。最近，研究人员也在努力减少代码模型的响应时间。预训练代码模型旨在为编码任务服务，有望成为现代IDE中的一个组成部分，以辅助开发人员进行软件开发和维护。然而，巨大的存储和运行内存消耗，加上高推理延迟，使得这些大模型难以集成到现代IDE中。因此，在部署这些预训练模型之前，必须减少它们的大小。

迄今为止，一些方法被提出用于压缩自然语言处理或其他任务中的预训练模型。这些现有研究主要分为三类：模型剪枝、模型量化和知识蒸馏。模型剪枝通过将部分模型参数替换为零或删除架构组件（如网络层或注意力头）来工作。模型量化技术将模型参数从32位浮点数转换为低位定点数。不幸的是，模型剪枝无法将CodeBERT和GraphCodeBERT压缩到50 MB以下——即使删除预训练模型中的所有网络层，模型仍会留下约150 MB的嵌入表。虽然模型量化可以通过使用低位数存储参数显著减少模型大小，但压缩后的模型在推理速度上并不会更快，也不会减少CPU内存的消耗。相反，运行时需要专门的硬件或专门的底层处理库。我们不能期望IDE用户的设备除了消费者级CPU外还拥有专门的硬件。特别是压缩后的模型预计将在CPU上带来显著的效率提升。

与模型剪枝和模型量化技术不同，知识蒸馏有望生成足够小且在CPU上具有效率提升的压缩模型。为了将知识从大模型传递到小模型，知识蒸馏训练一个较小的学生模型来模仿较大教师模型的行为。然而，在知识蒸馏后保留知识是具有挑战性的。正如Cho等人指出的那样，大预训练模型之间的模型容量差距使得知识蒸馏难以取得良好效果。Vapnik等人提出了模型容量的概念，用来描述机器学习模型或一组函数的复杂性和表达能力。低模型容量的小型学生模型不能有效吸收大模型的知识，这使得小型模型的性能远不及大型模型。与预训练模型相比，小型模型具有更薄更浅的网络架构，因此模型容量相应减少。因此，为小型学生模型选择具有足够容量的适当架构非常重要。之前的研究主要集中在将预训练模型的知识转移到具有预定义架构的小型模型上，而对于如何特别设计小型学生模型的架构，文献中的探索较为有限。

获得具有足够容量以保留从大模型蒸馏知识的适当模型架构是具有挑战性的。首先，找到合适的架构本质上是一个具有大量搜索空间的组合问题。存在大量网络结构可供选择（例如，CNN、LSTM），并且这些结构有各自相关的可调整超参数（例如，网络层数）。其次，通过训练和测试评估每个可能的候选模型在计算上是不可行的，这使得难以指导搜索。因此，针对这个难题，需要设计一个易于计算且有效的预测指标。

针对上述挑战，我们提出了一种新技术Compressor，能够有效地找到一个适当的小型学生模型，该模型可以有效吸收大模型的知识。Compressor首先利用预训练模型的架构设计来缩小搜索空间。特别地，它尝试找到一个与教师模型具有相似架构但规模更小的学生模型，我们称之为模型简化问题。Gao等人通过计算成本来衡量模型容量，计算成本被定义为Giga浮点运算（GFLOPs）的单位。我们采用这种度量方法，即候选模型的计算成本，作为指导搜索的预测指标。我们设计了一个定制的基于遗传算法（GA）的算法，并使用找到的解决方案的计算成本作为适应度函数。在使用基于GA的算法进行搜索后，Compressor应用知识蒸馏来训练找到的小模型。它将大量未标记输入数据输入到预训练模型中，并获取每个输入的预测结果。然后，这些知识（即输入和预测对）被用来训练小模型。

总的来说，本文的贡献如下：

2 技术介绍

本节解释了我们提出的工具Compressor的设计，包括搜索问题的形式化、基于GA（遗传算法）引导的模型简化算法，以及如何进行知识蒸馏以训练小模型。

图1：模型压缩方法的流程图

2.1 概述

知识蒸馏是一种将大模型压缩为小模型的技术。首先，我们查询大模型并获得其输出，这些输出被视为大模型的“知识”。然后，使用蒸馏得到的知识来训练一个较小的模型，使后者能够保持（大部分）大模型的行为但具有更小的尺寸。然而，小模型的架构会显著影响知识蒸馏后的模型性能。寻找小模型的最优架构的过程本质上是一个搜索问题：在给定模型大小约束的情况下，我们希望搜索一个在知识蒸馏后具有最佳性能的模型。

为了解决这个搜索问题，需要应对两个主要挑战。第一个挑战是庞大的搜索空间，其中有无数可能的组合。一个模型堆叠了许多层，每层包含不同数量的参数。即使是架构的任何元素的小变化也可能导致新的神经网络，即使在相同的数据集上训练，性能也可能大不相同。模型开发者通常需要投入大量的工程努力来找到合适的小模型架构，这既耗时又耗费计算资源。第二个挑战是这个搜索问题的目标，即小模型在蒸馏后的性能，计算成本非常高。在搜索过程中训练和评估每个找到的模型是不切实际和不可行的。因此，需要一个易于计算且有效的预测指标来适应这个困难的搜索问题。

在这项工作中，我们提出了Compressor，这是一种通过引导模型简化来压缩预训练代码模型的新方法。Compressor旨在解决上述两个挑战。为了克服第一个挑战，Compressor试图遵循预训练模型（即教师模型）的类似架构设计来缩小搜索空间。特别是，它找到一个与教师模型具有相似架构但宽度和深度较小的学生模型，我们称之为模型简化问题。候选架构由一组与架构相关的超参数决定，例如层数、网络维度等。为了克服第二个挑战，我们定制了一种基于遗传算法（GA）的策略来引导模型简化过程，以便找到能够最大限度地放大搜索到的小模型容量并同时符合给定大小限制的有希望的与架构相关的超参数。最后，我们采用知识蒸馏来获得一个性能良好的小模型。

应用范围：相近年来，在不同的学术场合中出现了十几种预训练代码模型。它们中的大多数属于BERT家族，因为它们采用与BERT或RoBERTa相同的架构，例如CodeBERT、GraphCodeBERT和CuBERT。少数例外也保持了基于Transformer的架构，例如CodeT5、GPT-C和CodeGPT。我们认为，我们的研究可以直接应用于BERT家族的预训练模型，而对于其他基于Transformer的模型仅需要少量修改。Compressor的当前实现主要集中于CodeBERT和GraphCodeBERT的压缩，因其受欢迎程度和优越性能。然而，我们的方法可以推广到其他基于Transformer的模型。我们将压缩其他预训练模型的评估留作未来工作。

表1：检索示例BERT家族预训练代码模型的超参数以及GA引导模型简化的搜索空间

2.2 问题搜索空间

此正如Liu等人所指出的，一个与BERT架构相同的学生模型，即来自Transformer的双向编码器架构，仍然可以实现与大BERT模型相当甚至更优的性能。这表明，与教师模型共享相似架构可以提升学生模型在知识蒸馏后的性能。BERT家族模型的大小与一组超参数有关。我们对这些超参数进行了彻底分析，并允许调整以下与模型架构相关的超参数以简化模型：网络层数（L）、网络层的维度（H）、注意力头的数量（A）、前馈层的维度（D）以及词汇表大小（V）。我们将这些超参数的可调值限制在不大于预训练模型中相应参数的范围内，如表1所示。

2.3 GA引导的模型简化

在模型大小约束下找到合适的简化模型本质上是一个组合优化问题，其目标是在给定大小限制下找到架构相关超参数的最佳组合，以最小化小模型大小与给定大小限制之间的差异，并最大化小模型的容量以在随后的蒸馏中实现更好的性能。然而，由于可能的超参数组合数量巨大，识别出一个好的超参数组合来调整预训练模型的结构并非易事。尽管我们已经限定了搜索空间，但仍有11,059,200种可能的超参数组合（小模型的候选架构）需要测试。通过训练和测试评估每个可能的候选模型在计算上是不可行的，因此难以指导搜索过程。因此，我们提出了一种基于遗传算法（GA）的策略来引导模型架构简化过程。算法1展示了GA引导的模型架构简化的概述。它首先初始化种群（第1行），然后执行遗传操作以生成新解（第2到11行）。我们的GA引导策略计算适应度函数并保留适应度值较大的解（第13行）。最后，算法返回具有最高适应度值的解（第15到16行）。

适应度函数。GA使用适应度函数来衡量和比较种群中染色体的质量。较高的适应度值表明染色体更接近问题的目标。如第1节所述，一个小模型预计具有足够的模型容量以吸收预训练模型的知识，从而在极小的模型尺寸下保证与大模型相当的性能。不成比例的模型容量将使小模型无法很好地捕获和蒸馏知识。根据Gao等人的建议，我们通过计算成本（以Giga浮点运算GFLOPs为单位）来衡量模型容量。GFLOPs计算模型在前向传递中需要进行多少乘加操作。较大的GFLOPs意味着模型具有更大的容量。在给定目标模型大小的情况下，Compressor旨在找到一个满足模型大小且具有最大GFLOPs的小模型。

操作和选择。在每次迭代中，我们采用两种标准遗传操作，即交叉和变异，来生成新的染色体。我们以概率r应用交叉（第7行），以1-r的概率应用变异（第9行）。给定两个染色体（c1和c2），交叉操作如下：我们首先随机选择一个切割位置h，并用c2在相应位置的基因替换c1在h位置后的基因。变异操作也是首先随机选择一个切割位置h，然后用一个随机值替换c1在h位置后的基因。在一轮迭代中生成新一代后，我们将它们与当前种群合并并执行选择操作（第13行）。我们始终保持种群大小不变，并丢弃适应度值较低的染色体。

算法1

2.4 使用无标签数据进行蒸馏

我们倡导在知识蒸馏中使用无标签数据，原因如下：(1) Tang等人和Jiao等人表明，不涉及真实标签的知识蒸馏可以取得更好的效果；(2) 标签数据昂贵且难以获得，而无标签数据可以轻松地从代码托管网站的开源项目的原始代码中抓取；(3) 一个现实的情况是，部署的小模型需要使用用户生成的新数据进行训练以增强性能，而这些数据没有标签。我们的研究采用了Hinton等人介绍的知识蒸馏方法，其中学生模型直接从教师模型的输出概率中学习。这种方法属于第2.2节中提到的任务特定蒸馏组。最近的一项研究确认了其在多个任务中从目标大模型向小模型转移知识的有效性。具体来说，我们首先将无标签数据输入到较大的预训练模型中并收集输出概率值。然后，我们按照Hinton等人介绍的目标，训练通过GA引导的模型简化识别出的微型模型架构。

3 实验设置

本节解释了我们评估Compressor的实验设置，包括实现细节、数据集和评估标准。

3.1 实现和配置

在本Compressor使用Python实现，基于HuggingFace Transformers编写了约2000行代码。我们借用了Clark等人的工作来计算模型的FLOPs，并使用PyTorch中的API来计算模型大小。

在实验中，GA算法维护了50个候选模型的种群。交叉率设置为0.6。当迭代次数达到100次时，GA算法终止。所有实验在Ubuntu 18.04服务器上进行，服务器配备了Intel Xeon E5-2698 CPU、504GB RAM和8个Tesla P100 GPU。

3.2 数据集和任务

为了充分评估Compressor，我们选择了两个在一系列下游任务中表现出色的预训练代码模型（CodeBERT和GraphCodeBERT），并使用Compressor来减少它们的模型大小。我们考虑了两个重要的下游任务：漏洞预测和克隆检测。

漏洞预测。该任务旨在预测给定的代码片段是否存在漏洞。将该任务的模型集成到IDE组件中可以极大地帮助开发人员在早期阶段识别潜在的代码缺陷。我们使用由Zhou等人发布的数据集Devign4。Devign包含从两个流行的开源C库（FFmpeg和Qemu）中提取的27,318个函数。这些函数被手动标记为是否包含漏洞。Devign是CodeXGLUE基准测试的一部分，用于研究CodeBERT在预测漏洞方面的有效性。按照CodeXGLUE的设置，我们使用2,732个样本作为验证和测试集。我们将剩余的样本平均分成两部分。一半用于微调CodeBERT和GraphCodeBERT。另一半用于进行知识蒸馏和从头训练压缩模型，其标签被删除。

克隆检测。该任务旨在确定两个给定的函数是否是克隆，即在操作语义上等效。克隆检测可以识别相同功能的冗余实现。找到代码克隆可以帮助开发人员采取措施提高软件质量，例如通过代码重构。BigCloneBench是一个广泛使用的克隆检测基准测试。BigCloneBench从各种开源的Java项目中收集，包含超过6,000,000对克隆的Java方法以及260,000对非克隆对。按照之前工作的设置，我们随机选择90,102个样本用于训练，4,000个用于验证和测试，以保持实验在计算上友好的规模。与漏洞预测数据集类似，训练数据也被平均分为有标签和无标签部分，分别用于获取预训练模型和压缩模型。

为了微调CodeBERT在漏洞预测和克隆检测任务上的表现，我们使用了CodeXGLUE基准测试中采用的相同超参数设置。为了微调GraphCodeBERT，我们遵循GraphCodeBERT论文中的相同超参数设置。我们按照CodeXGLUE论文和Yang等人的做法，使用准确率作为评估模型在下游任务中表现的评估指标。尽管我们使用的有标签训练数据是完整训练集的一半，所有模型在漏洞预测和克隆检测任务上都取得了与CodeXGLUE基准测试和GraphCodeBERT论文中报道的最先进结果相当的结果。数据集的统计信息和微调模型的表现见表2。

表2：数据集统计信息及两个预训练模型的性能

3.3 评估指标和基线

一个好的模型压缩方法应当能够（1）显著减少模型大小，并且（2）对模型性能（例如准确率）的负面影响最小。正如Svyatkovskiy等人所建议的，3 MB大小的模型在各种设备上都易于部署，甚至在一些用于教学生编程的低端硬件上也是如此。我们将预训练模型压缩到3 MB。按照CodeXGLUE基准测试和Yang等人的做法，我们使用准确率作为下游任务的评估指标。我们测量了Compressor获得的小模型的预测准确率，并计算准确率的下降程度。

此外，我们将Tang等人提出的任务特定蒸馏方法BiLSTMsoft作为我们的基线。在我们的实验中，BiLSTMsoft能够将预训练模型压缩到7.5 MB，这一大小接近相关文献中期望的3 MB。我们选择BiLSTMsoft作为基线有两个原因。一方面，他们的工作中压缩的模型也是从预训练模型的输出logits中学习的；进行知识蒸馏不需要其他模型输出（例如，预训练模型中每个编码器层的输出）。这样的设置与我们的设置一致。其次，BiLSTMsoft能够将预训练模型压缩到非常小的尺寸。

BiLSTMsoft使用浅层双向长短期记忆（BiLSTM）网络来蒸馏一个大规模的预训练模型。更具体地说，学生模型具有相当简单的结构，由单层BiLSTM和两层全连接层组成。BiLSTMsoft的训练目标是最小化学生模型的logits与预训练模型的logits之间的均方误差。Tang等人还提出使用数据增强技术来增强学生模型，通过基于词性标签（即每个词的语法属性）替换句子中的词生成大量合成数据。由于源代码没有这样的标签，这种数据增强方法在压缩预训练代码模型上本质上是不可用的。

在我们的实验中，我们使用了他们官方代码库中的代码。Tang等人提供了他们方法的所有超参数设置信息，除了词汇表大小。我们遵循所有提供的超参数设置，并保持词汇表大小与我们搜索到的模型的词汇表大小相同。此外，为了展示压缩模型的效率，我们关注每次查询模型时的推理延迟（即模型返回输入预测所需的时间），这在实际部署中非常重要。我们还考虑了Compressor中模型压缩过程的时间成本。时间成本包括搜索参数以简化模型的时间和知识蒸馏的时间，展示了Compressor在微调后压缩预训练模型所需的时间。

4 实验评估

本节解释了我们评估Compressor的实验设置，包括实现细节、数据集和评估标准。

4.1 实验设置

研究问题。本节展示了我们提出的方法和基线的评估结果。我们的评估旨在回答以下研究问题（RQ）：

RQ1：Compressor在极端压缩预训练模型时能否导致较小的准确率损失？

RQ2：压缩后的模型能获得多大的效率提升？

RQ3：Compressor压缩预训练模型的速度有多快？

RQ1 Compressor的准确率损失

天下没有免费的午餐——压缩预训练模型可能会影响模型的原始效果，即在下游任务中的预测准确率降低。在这个研究问题中，我们旨在展示Compressor可以在准确率损失很小的情况下，将模型压缩到极小的尺寸。

与预训练模型的比较。我们使用Compressor将两个预训练代码模型压缩到每个仅3 MB，并研究压缩后的模型损失了多少预测准确率。表3展示了Compressor获得的压缩模型的预测准确率，以及与原始预训练模型相比性能下降的情况，分别在第5和第8行。我们还在第4行和第7行报告了BiLSTMsoft的结果，以强调Compressor带来的改进。我们首先分析了使用我们的方法压缩的模型与原始预训练模型的准确率相比的准确率损失。结果显示，当模型大小减少到仅3 MB时，压缩模型的预测准确率几乎没有下降，这仅是原始体积的0.6%。Compressor获得的压缩后的CodeBERT在漏洞预测任务中保持了原始准确率的96.15%。同时，Compressor在压缩GraphCodeBERT时达到了更小的准确率下降，压缩后的模型保持了原始准确率的97.74%。在克隆检测任务中，Compressor获得的压缩后的CodeBERT和GraphCodeBERT分别能够达到原始大模型准确率的99.20%和97.52%。这些相对于原始模型的高保真准确率结果表明，Compressor可以在几乎没有准确率损失的情况下显著减少模型大小，暗示了其在压缩预训练代码模型方面的有效性。

表3：预训练模型和压缩模型的准确率比较

与BiLSTMsoft的比较。我们观察到Compressor在两个任务上都大大优于BiLSTMsoft，尤其是在克隆检测任务上。与7.5 MB的BiLSTMsoft相比，BiLSTMsoft的模型大小是Compressor获得模型的两倍多，Compressor压缩的CodeBERT和GraphCodeBERT分别减少了93.72%（即(0.77-12.27)/12.27）和80.86%的准确率损失。在漏洞预测任务上，Compressor获得的模型也大大减少了BiLSTMsoft的准确率下降：压缩后的CodeBERT将BiLSTMsoft的准确率下降减少了39.90%，压缩后的GraphCodeBERT将准确率下降减少了58.63%。这些结果表明Compressor相对于BiLSTMsoft具有显著的优越性。

总而言之，当Compressor将预训练的代码模型压缩到3 MB时，这仅为原始模型大小的0.6%，在两个任务上的准确率下降都是可以忽略的。平均而言，Compressor压缩的模型在漏洞预测和克隆检测任务上仍能分别保持原始大模型96.95%和98.36%的准确率。此外，在减少压缩模型的准确率损失方面，Compressor分别比BiLSTMsoft高出49.27%和87.29%。

RQ2 压缩模型的效率

为了提升用户体验，IDE或代码编辑器中的插件需要为开发者提供即时的帮助，从而使开发过程更高效。在这个研究问题中，我们调查了Compressor获得的压缩模型的效率。我们分析了模型进行预测时的推理延迟，即模型平均需要多长时间才能返回输入的预测。推理延迟与两个因素有关：模型运行的机器和输入长度（即输入中的标记数量）。为了公平比较，我们使用第4.1节中提到的相同机器来运行模型。我们限制模型仅使用8个CPU核心，以模拟在普通消费者级笔记本电脑上的运行。通常，较长的输入需要更长的计算时间。我们将输入长度设置为与RQ1实验中的相同，即漏洞预测为400，克隆检测为800。对于每个任务，我们从测试集中随机抽取100个样本。我们使用抽取的100个样本查询模型，并计算模型的平均推理延迟。为了减少随机性的影响，我们重复了三次实验。结果如表4所示。

大型预训练模型的推理延迟超过一千毫秒，而我们的压缩模型的延迟仅为几百毫秒。在压缩CodeBERT时，Compressor获得的3 MB模型在漏洞预测任务中的平均推理延迟为347毫秒，在克隆检测任务中的平均推理延迟为625毫秒；后者任务的输入更长。与原始预训练模型相比，压缩模型的推理延迟分别显著减少了76.97%和76.64%。对于压缩后的GraphCodeBERT，压缩后的GraphCodeBERT在漏洞预测任务上的速度比原始预训练模型快64.52%。在克隆检测任务上，压缩后的GraphCodeBERT的推理延迟比原始预训练模型减少了81.77%。这些结果表明，小模型在实际部署中可能带来的效率优势，这是Compressor的另一个关键优势。

表4：原始预训练模型和压缩模型在不同任务上的推理延迟。括号中的百分比是Compressor带来的相对改进

RQ3 Compressor的时间成本

Compressor是一种任务特定的压缩方法，这意味着要压缩的预训练模型应首先在下游任务上进行微调。因此，在模型开发过程中应用Compressor将不可避免地引入额外的时间成本。由于微调一个大型预训练模型已经非常耗时，随后的压缩过程不应增加太多的开销。在这个研究问题中，我们调查了Compressor在压缩预训练模型时花费的时间，并旨在表明相较于微调大型预训练模型，Compressor能够以更少的时间成本获得性能良好的小型模型。我们重复每个实验三次，并计算每个阶段的时间消耗，包括微调和压缩预训练模型。我们在表5中展示了平均时间消耗。需要注意的是，我们仍然将3 MB作为压缩模型的目标大小。表5展示了在不同任务上微调和压缩预训练模型的时间消耗。

我们观察到，在两个任务上，压缩预训练模型所需的时间少于微调模型的时间。括号内的数字是压缩时间相对于微调时间的比率。这里模型简化和知识蒸馏的时间成本作为一个整体计算。在漏洞预测任务中，Compressor仅需微调时间的26.53%和34.25%来压缩CodeBERT和GraphCodeBERT。尽管Compressor在处理显著更大数据集的克隆检测任务时花费了更多时间（如第4.2节所述），但它仍然仅需微调时间的37.90%和36.21%来压缩CodeBERT和GraphCodeBERT。平均而言，Compressor仅增加了微调时间的30.39%和37.06%的额外时间，我们认为这是合理的。需要注意的是，这种开销只是偶尔发生，例如第一次压缩预训练模型用于部署，以及每次按需更新模型以产生良好结果时，这可能是每月或每年进行的。

我们还分析了GA引导的模型简化算法的效率。得益于精心设计的搜索空间和适应度函数，GA引导的模型简化可以快速返回一个合适的模型架构，然后在随后的知识蒸馏中用作学生模型。为了调查GA算法的速度，我们运行了10次GA引导搜索。平均而言，基于GA的算法在搜索模型架构时花费了1.22秒。结果进一步证明了Compressor的效率。

表5：在不同任务上微调和压缩预训练模型的时间成本。括号中的百分比是Compressor所用时间与微调时间的比率

转述：葛一飞