每天给你送来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
文献地址：https://arxiv.org/abs/2212.07634
简介
预训练语言模型（PLM）的压缩与加速一直是NLP研究社区关注的热点问题之一。在众多压缩方法中，模型裁剪和知识蒸馏是两种主流的途径。在针对预训练模型的结构化裁剪方法中，人们通常通过移除前馈网络（FFN）层的隐层维度和多头注意力（MHA）层的注意力头实现减少模型参数量的效果。在本文中，我们提出一种结合了知识蒸馏和内部注意力裁剪的压缩方案GRAIN。通过考察注意力头的内部计算机制，本文发现可以在比注意力头的层级更细的粒度上进行结构化裁剪，即内部注意力裁剪（Intra-attention Pruning），从而显著扩大模型结构的搜索空间，使得裁剪过程能找到更优的模型结构。为了解决内部注意力裁剪带来的结构碎片化问题，本文提出了结构正规化以使模型结构趋于规整并提升推理速度。此外，GRAIN结合知识蒸馏可以获得更好的压缩效果。在文本分类、阅读理解、命名实体识别等任务上的实验结果表明，在多种压缩率下，GRAIN相比可比基线方法带来明显的性能提升。在仅保留3% transformer参数的极端压缩率下，通过GRAIN得到的模型在下游任务上仍能达到具有竞争力的效果。
主要贡献：
本文提出了一种压缩方案GRAIN，一个结合了知识蒸馏和内部注意力裁剪的预训练模型压缩方法。
在多种任务上、多个模型大小下的实验表明，GRAIN能够取得相对可比基线方法的显著性能提升。
重点技术
蒸馏裁剪一体化方案GRAIN主要涵盖了以下重点技术。
内部注意力裁剪
传统的针对PLM的模型裁剪方法中，裁剪单元为FFN隐层维度和MHA层的的注意力头。GRAIN采用了比注意力头更细粒度的裁剪单元。注意到每个注意力头可以是不同的大小，因此可以对每个注意力头的输出维度进行裁剪；进一步，在自注意力机制中，query和key的参数矩阵与value和output参数矩阵大小也可以不同，因此裁剪后模型的结构会更丰富。
左：注意力头裁剪；右：内部注意力裁剪
为了减少模型裁剪后产生零碎的子结构以提升推理速度，本文设计了一种结构正规化（Structure Regularization）机制，鼓励算法优先裁剪掉零碎的单元，从而在同样参数量的条件下生成结构更规整的模型。
左：无结构正规化；右：使用结构正规化。虚线代表空注意力头
知识蒸馏与梯度分离
通过与知识蒸馏技术的结合，模型裁剪的效果能得到进一步提升。然而，知识蒸馏的训练过程可能会对模型裁剪产生干扰。具体来说，知识蒸馏的中间隐层损失并不适合作为计算神经元重要性指标进而帮助裁剪模型。因此本文提出梯度分离方法（Gradient Separation），使用来自知识蒸馏中不同损失的梯度，分别优化模型参数和模型结构，从而减少知识蒸馏对裁剪的负面干扰。
知识蒸馏中的来自不同损失的梯度流
训练流程
GRAIN的裁剪算法基于梯度裁剪实现。总体的流程如下图所示，在一次完整的训练过程中同时实现模型的裁剪和蒸馏。
GRAIN训练流程
实验效果
为了验证裁剪方案GRAIN的效果，本文在文本分类（QNLI、MNLI）、阅读理解（SQuAD）、命名实体识别（CoNLL-2003）等多个NLU标准任务上进行了实验，并与其他业界领先的模型压缩方法进行比较。结果表明，在多种压缩率（仅保留5%-20%的参数量）下，GRAIN相比可比基线方法带来明显的性能提升。
实验效果（横轴：相对参数量，纵轴：效果）
该项工作是继知识蒸馏工具TextBrewer、模型裁剪工具TextPruner、中文小型预训练模型MiniRBT后，哈工大讯飞联合实验室提出的全新高效模型训练框架，后续将逐步接入以上工具并将代码进行开源以进一步促进相关技术的研究。
更多的技术细节和实验分析，请查阅文末列出的技术报告。
相关文献
Ziqing Yang, Yiming Cui, Xin Yao, Shijin Wang. Gradient-based Intra-attention Pruning on Pre-trained Language Models. arXiv pre-print 2212.07634. https://arxiv.org/abs/2212.07634
其他相关资源
知识蒸馏工具TextBrewer：http://textbrewer.hfl-rc.com
模型裁剪工具TextPruner：http://textpruner.hfl-rc.com
中文小型预训练模型MiniRBT：http://minirbt.hfl-rc.com
📝论文解读投稿，让你的文章被更多不同背景、不同方向的人看到，不被石沉大海，或许还能增加不少引用的呦~ 投稿加下面微信备注“投稿”即可。
最近文章
COLING'22 | SelfMix：针对带噪数据集的半监督学习方法
ACMMM 2022 | 首个针对跨语言跨模态检索的噪声鲁棒研究工作
ACM MM 2022 Oral | PRVR: 新的文本到视频跨模态检索子任务
点击这里进群—>加入NLP交流群
4319
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
