哈工大 NLP 在线哈尔滨工业大学nlp

转载

mob6454cc70863a 2024-01-16 09:14:50

文章标签 哈工大 NLP 在线压缩率性能提升基线 文章分类 NLP 人工智能

每天给你送来NLP技术干货！

哈工大讯飞联合实验室（HFL）2022年终技术发布第三弹，推出一种结合了知识蒸馏和内部注意力裁剪的预训练模型压缩方案GRAIN（Gradient-based Intra-attention Pruning）。该方案针对预训练模型的注意力机制，设计了细粒度裁剪单元，并结合模型结构正规化机制和知识蒸馏技术，实现了预训练模型的在下游任务上的高效压缩。在多个NLP任务上的实验结果表明，在同等压缩率下，GRAIN相比其他业界领先的蒸馏裁剪方法获得更优的任务效果。

论文标题：Gradient-based Intra-attention Pruning on Pre-trained Language Models

论文作者：Ziqing Yang, Yiming Cui, Xin Yao, Shijin Wang

预训练语言模型（PLM）的压缩与加速一直是NLP研究社区关注的热点问题之一。在众多压缩方法中，模型裁剪和知识蒸馏是两种主流的途径。在针对预训练模型的结构化裁剪方法中，人们通常通过移除前馈网络（FFN）层的隐层维度和多头注意力（MHA）层的注意力头实现减少模型参数量的效果。在本文中，我们提出一种结合了知识蒸馏和内部注意力裁剪的压缩方案GRAIN。通过考察注意力头的内部计算机制，本文发现可以在比注意力头的层级更细的粒度上进行结构化裁剪，即内部注意力裁剪（Intra-attention Pruning），从而显著扩大模型结构的搜索空间，使得裁剪过程能找到更优的模型结构。为了解决内部注意力裁剪带来的结构碎片化问题，本文提出了结构正规化以使模型结构趋于规整并提升推理速度。此外，GRAIN结合知识蒸馏可以获得更好的压缩效果。在文本分类、阅读理解、命名实体识别等任务上的实验结果表明，在多种压缩率下，GRAIN相比可比基线方法带来明显的性能提升。在仅保留3% transformer参数的极端压缩率下，通过GRAIN得到的模型在下游任务上仍能达到具有竞争力的效果。

主要贡献：

本文提出了一种压缩方案GRAIN，一个结合了知识蒸馏和内部注意力裁剪的预训练模型压缩方法。
在多种任务上、多个模型大小下的实验表明，GRAIN能够取得相对可比基线方法的显著性能提升。

重点技术

蒸馏裁剪一体化方案GRAIN主要涵盖了以下重点技术。

内部注意力裁剪

传统的针对PLM的模型裁剪方法中，裁剪单元为FFN隐层维度和MHA层的的注意力头。GRAIN采用了比注意力头更细粒度的裁剪单元。注意到每个注意力头可以是不同的大小，因此可以对每个注意力头的输出维度进行裁剪；进一步，在自注意力机制中，query和key的参数矩阵与value和output参数矩阵大小也可以不同，因此裁剪后模型的结构会更丰富。

哈工大 NLP 在线哈尔滨工业大学nlp_基线

左：注意力头裁剪；右：内部注意力裁剪

为了减少模型裁剪后产生零碎的子结构以提升推理速度，本文设计了一种结构正规化（Structure Regularization）机制，鼓励算法优先裁剪掉零碎的单元，从而在同样参数量的条件下生成结构更规整的模型。

哈工大 NLP 在线哈尔滨工业大学nlp_性能提升_02

左：无结构正规化；右：使用结构正规化。虚线代表空注意力头

知识蒸馏与梯度分离

通过与知识蒸馏技术的结合，模型裁剪的效果能得到进一步提升。然而，知识蒸馏的训练过程可能会对模型裁剪产生干扰。具体来说，知识蒸馏的中间隐层损失并不适合作为计算神经元重要性指标进而帮助裁剪模型。因此本文提出梯度分离方法（Gradient Separation），使用来自知识蒸馏中不同损失的梯度，分别优化模型参数和模型结构，从而减少知识蒸馏对裁剪的负面干扰。

哈工大 NLP 在线哈尔滨工业大学nlp_性能提升_03

知识蒸馏中的来自不同损失的梯度流

训练流程

GRAIN的裁剪算法基于梯度裁剪实现。总体的流程如下图所示，在一次完整的训练过程中同时实现模型的裁剪和蒸馏。

哈工大 NLP 在线哈尔滨工业大学nlp_基线_04

GRAIN训练流程

实验效果

为了验证裁剪方案GRAIN的效果，本文在文本分类（QNLI、MNLI）、阅读理解（SQuAD）、命名实体识别（CoNLL-2003）等多个NLU标准任务上进行了实验，并与其他业界领先的模型压缩方法进行比较。结果表明，在多种压缩率（仅保留5%-20%的参数量）下，GRAIN相比可比基线方法带来明显的性能提升。

哈工大 NLP 在线哈尔滨工业大学nlp_压缩率_05