[论文解读] Know What You Don't Need: Single-Shot Meta-Pruning for Attention Heads

Zhengyan Zhang, Fanchao Qi+3 more|arXiv (Cornell University)|Nov 7, 2020

Topic Modeling参考文献 47被引用 5

一句话总结

本文提出单次元剪枝（SMP），一种在微调前剪除预训练 BERT 模型中高达 50% 注意力头的方法，利用元学习准则保留文本表征的分布。SMP 在减少微调和推理成本的同时，实现与微调基线相当或更优的性能，剪枝后的模型在语义相关性任务上表现出更好的泛化能力。

ABSTRACT

Deep pre-trained Transformer models have achieved state-of-the-art results over a variety of natural language processing (NLP) tasks. By learning rich language knowledge with millions of parameters, these models are usually overparameterized and significantly increase the computational overhead in applications. It is intuitive to address this issue by model compression. In this work, we propose a method, called Single-Shot Meta-Pruning, to compress deep pre-trained Transformers before fine-tuning. Specifically, we focus on pruning unnecessary attention heads adaptively for different downstream tasks. To measure the informativeness of attention heads, we train our Single-Shot Meta-Pruner (SMP) with a meta-learning paradigm aiming to maintain the distribution of text representations after pruning. Compared with existing compression methods for pre-trained models, our method can reduce the overhead of both fine-tuning and inference. Experimental results show that our pruner can selectively prune 50% of attention heads with little impact on the performance on downstream tasks and even provide better text representations. The source code will be released in the future.

研究动机与目标

• 减少大型预训练 Transformer 模型在微调和推理中的计算开销。
• 在微调前而非之后，识别并移除预训练模型中不必要的注意力头。
• 开发一种可泛化的单次剪枝方法，在多种下游任务中保持模型性能。
• 通过元学习学习隐式的剪枝规则，使其在不同模型尺寸和任务间具有可迁移性。

提出的方法

• 训练一个元学习器（SMP），通过保持剪枝后文本表征分布的自监督目标，预测注意力头的重要性。
• 在剪枝器训练过程中，利用多样化语料库的元学习范式，模拟各种下游任务。
• 采用单次剪枝策略，在微调前一步骤内移除注意力头，避免迭代优化。
• 通过反映信息量的评分函数衡量头的重要性，高分表示有用的头。
• 使用剪枝器识别并移除低分头，认为其为非信息性头。
• 在下游任务上评估剪枝后的模型，以衡量性能和效率的提升。

实验结果

研究问题

• RQ1 我们能否在微调前有效剪除预训练 Transformer 中的注意力头，以降低计算成本？
• RQ2 元学习剪枝器是否能在不同下游任务和模型架构间实现泛化？
• RQ3 单次剪枝能否实现与迭代微调压缩方法相当或更优的性能？
• RQ4 剪除部分注意力头是否能提升模型在某些 NLP 任务上的泛化能力？
• RQ5 被剪除的注意力头揭示了哪些隐式模式？这些模式是否与人类对注意力机制的解释一致？

主要发现

• SMP 可在 GLUE 基准任务上对 BERT 剪除高达 50% 的注意力头，性能损失极小。
• 剪除 50% 的头可使内存使用减少 30%，从而支持更高批量大小的更高效训练。
• 在某些情况下，剪枝能提升性能，尤其在 SentEval 的语义相关性任务中，SMP 表现优于随机剪枝，并与 HISP-微调相当。
• 即使参数更少，SMP 的剪枝结果也与微调后压缩的基线方法相当或更优。
• SMP 学习到的隐式剪枝规则与人类直觉一致，如注意力矩阵的可视化所示。
• 剪枝器表现出强大的可迁移性：在 BERT_BASE 上训练的模型能良好泛化到 BERT_LARGE，并在新任务（如 QNLI）上表现良好。