[论文解读] KDLSQ-BERT: A Quantized Bert Combining Knowledge Distillation with Learned Step Size Quantization

Jing Jin, Liang Cai+3 more|arXiv (Cornell University)|Jan 15, 2021

Topic Modeling参考文献 40被引用 13

一句话总结

KDLSQ-BERT 提出了一种新颖的 BERT 量化方法，结合知识蒸馏（KD）与可学习步长量化（LSQ），实现高精度、超低比特推理。通过在 LSQ 训练过程中将知识从全精度教师 BERT 转移到量化学生模型，该方法在 2-bit 精度下实现 14.9 倍模型压缩，且精度损失极小，在 GLUE 和 SQuAD 基准测试中表现优于当前最先进量化方法。

ABSTRACT

Recently, transformer-based language models such as BERT have shown tremendous performance improvement for a range of natural language processing tasks. However, these language models usually are computation expensive and memory intensive during inference. As a result, it is difficult to deploy them on resource-restricted devices. To improve the inference performance, as well as reduce the model size while maintaining the model accuracy, we propose a novel quantization method named KDLSQ-BERT that combines knowledge distillation (KD) with learned step size quantization (LSQ) for language model quantization. The main idea of our method is that the KD technique is leveraged to transfer the knowledge from a "teacher" model to a "student" model when exploiting LSQ to quantize that "student" model during the quantization training process. Extensive experiment results on GLUE benchmark and SQuAD demonstrate that our proposed KDLSQ-BERT not only performs effectively when doing different bit (e.g. 2-bit $\sim$ 8-bit) quantization, but also outperforms the existing BERT quantization methods, and even achieves comparable performance as the full-precision base-line model while obtaining 14.9x compression ratio. Our code will be public available.

研究动机与目标

• 解决在资源受限设备上部署计算密集型大型 BERT 模型的挑战。
• 克服 BERT 在超低比特（如 2-bit）量化时通常出现的显著精度下降问题。
• 通过结合知识蒸馏与可学习步长量化（LSQ），提升量化性能。
• 在不改变模型架构的前提下，实现高效且硬件友好的量化。
• 在保持与全精度 BERT 竞争性性能的同时，实现高模型压缩比。

提出的方法

• 在训练过程中利用知识蒸馏，将全精度教师 BERT 的知识转移到量化学生 BERT。
• 对学生模型的权重和激活应用可学习步长量化（LSQ），实现具有可学习缩放因子的可微量化。
• 提出一种新颖的 LSQ 缩放因子初始化方法，以加速训练并改善收敛性。
• 在反向传播中同时结合蒸馏损失与真实标签损失，以稳定训练并提升性能。
• 通过端到端训练，结合蒸馏与监督信号，实现在低比特宽度下有效的知识迁移。
• 通过设计支持任意比特宽度量化（如 2-bit、4-bit），避免了三值或混合精度方法的局限性。

实验结果

研究问题

• RQ1 当与可学习步长量化结合时，知识蒸馏是否能显著提升量化 BERT 模型的性能？
• RQ2 与现有量化方法相比，KDLSQ-BERT 在超低比特宽度（如 2-bit）下保持高精度的能力如何？
• RQ3 结合蒸馏损失与真实标签损失是否能带来比单独使用任一损失更好的收敛性与精度？
• RQ4 在不牺牲性能的前提下，模型压缩能达到何种程度，特别是在参数量和推理速度方面？
• RQ5 该方法能否在包括 GLUE 和 SQuAD 基准在内的多样化 NLP 任务中实现良好泛化？

主要发现

• KDLSQ-BERT 在 2-bit 量化下实现 14.9 倍模型压缩（从 418 MB 压缩至 28 MB），同时保持 83.987 的平均 GLUE 得分。
• 在 4-bit 量化下，KDLSQ-BERT 实现 84.414 的平均 GLUE 得分（54 MB，7.7 倍压缩），优于全精度 TinyBERT（258 MB，1.6 倍压缩）。
• 在 2-2-4 量化设置下（权重、激活、嵌入分别 2-bit、2-bit、4-bit），KDLSQ-BERT 在 GLUE 上实现 85.013 的 MNLI 得分与 89.989 的 RTE 得分，接近全精度 BERT 性能。
• 在 SQuAD 1.1 上，KDLSQ-BERT 使用 2-2-4 量化实现 88.996 的 EM 得分，仅比全精度 BERT（88.696）低 3.699 分，展现出强大的零样本泛化能力。
• 在 SQuAD 2.0 上，KDLSQ-BERT 使用 4-4-8 量化实现 84.086 的平均 F1 得分，优于全精度 TinyBERT 1.361 分。
• 蒸馏损失与真实标签损失的结合显著提升了精度，尤其在低比特宽度下，证实了双重监督的有效性。