[论文解读] KroneckerBERT: Learning Kronecker Decomposition for Pre-trained Language Models via Knowledge Distillation

Marzieh S. Tahaei, Ella Charlaix+3 more|arXiv (Cornell University)|Sep 13, 2021

Topic Modeling参考文献 35被引用 10

一句话总结

KroneckerBERT 提出了一种新颖的框架，通过在嵌入层和 Transformer 块中的所有线性变换上应用 Kronecker 分解，并结合中间层知识蒸馏，实现了对基于 BERT 模型的极端压缩。在 19 倍压缩比（BERT_BASE 大小的 5%）下，KroneckerBERT 在 GLUE 和 SQuAD 上实现了最先进性能，尽管压缩率极高，仍优于先前方法。

ABSTRACT

The development of over-parameterized pre-trained language models has made a significant contribution toward the success of natural language processing. While over-parameterization of these models is the key to their generalization power, it makes them unsuitable for deployment on low-capacity devices. We push the limits of state-of-the-art Transformer-based pre-trained language model compression using Kronecker decomposition. We use this decomposition for compression of the embedding layer, all linear mappings in the multi-head attention, and the feed-forward network modules in the Transformer layer. We perform intermediate-layer knowledge distillation using the uncompressed model as the teacher to improve the performance of the compressed model. We present our KroneckerBERT, a compressed version of the BERT_BASE model obtained using this framework. We evaluate the performance of KroneckerBERT on well-known NLP benchmarks and show that for a high compression factor of 19 (5% of the size of the BERT_BASE model), our KroneckerBERT outperforms state-of-the-art compression methods on the GLUE. Our experiments indicate that the proposed model has promising out-of-distribution robustness and is superior to the state-of-the-art compression methods on SQuAD.

研究动机与目标

• 解决在低容量边缘设备上部署大型、过参数化的预训练语言模型的挑战。
• 开发一种压缩框架，实现在不显著降低性能的前提下实现极端压缩（压缩因子 >10）。
• 将 Kronecker 分解的应用从标准矩阵分解扩展至 Transformer 块中的嵌入层和所有线性层。
• 通过使用教师模型的中间层监督的两阶段知识蒸馏策略，提升压缩模型的性能。
• 在标准 NLP 基准上评估模型大小、推理速度和准确率之间的权衡。

提出的方法

• Kronecker 分解被应用于嵌入层、多头注意力机制和前馈网络中的权重矩阵，将其表示为更小矩阵的 Kronecker 积，从而显著减少参数量。
• 该方法将每个权重矩阵 W 分解为 W = A ⊗ B，其中 A 和 B 为更小的矩阵，从而实现显著的模型大小缩减。
• 在预训练和微调阶段均应用知识蒸馏，使用原始的 BERT_BASE 作为教师模型，指导学生模型的中间层表示。
• 该框架被应用于将 BERT_BASE 压缩为 KroneckerBERT，并通过消融实验验证了知识蒸馏的影响。
• 在智能手机和高性能 CPU 上评估推理速度，以衡量压缩带来的实际延迟改进。
• 使用 TorchScript 序列化压缩后的模型，以便在边缘设备上进行部署和延迟测量。

实验结果

研究问题

• RQ1 Kronecker 分解能否有效应用于基于 Transformer 的模型的嵌入层和所有线性层，以实现极端压缩？
• RQ2 中间层知识蒸馏如何缓解 Kronecker 压缩模型中的性能退化？
• RQ3 在高倍压缩因子（例如 19 倍）下，KroneckerBERT 在标准 NLP 基准上的表现如何？
• RQ4 Kronecker 分解带来的 FLOPs 减少是否能实际转化为边缘设备上的推理加速？
• RQ5 在分布外数据上，KroneckerBERT 与最先进压缩模型相比，在准确率和鲁棒性方面表现如何？

主要发现

• KroneckerBERT 实现了 19 倍的压缩因子（BERT_BASE 大小的 5%），在该压缩水平下显著优于现有方法。
• 在 GLUE 基准上，KroneckerBERT 在 19 倍压缩下优于最先进压缩 BERT 模型，展现出更优的压缩-性能权衡。
• 在 SQuAD 上，KroneckerBERT 表现优于最先进压缩方法，表明其具备强大的鲁棒性和泛化能力。
• 消融实验确认，预训练和微调阶段均应用知识蒸馏至关重要，两阶段 KD 能够带来最高性能。
• 在边缘设备（智能手机）上，KroneckerBERT 在 19 倍压缩下实现了最高 2.4 倍的推理加速，证实了其实际推理优势。
• 在高性能 CPU 上未观察到加速，这是由于 Kronecker 矩阵乘法的串行特性，凸显了硬件感知优化的重要性。