[论文解读] TernaryBERT: Distillation-aware Ultra-low Bit BERT

Wei Zhang, Lu Hou+5 more|arXiv (Cornell University)|Sep 27, 2020

Topic Modeling参考文献 37被引用 10

一句话总结

TernaryBERT 提出了一种感知蒸馏的超低比特量化方法，通过混合粒度对 BERT 权重进行三值化处理，将权重映射为 {-1, 0, +1}，其中词嵌入部分采用头粒度，Transformer 层采用层粒度，结合基于近似的三值化与损失感知的三值化策略。该方法在 GLUE 和 SQuAD 基准测试中，性能接近全精度 BERT，同时将模型大小缩小了 14.9 倍，优于现有的量化与压缩方法。

ABSTRACT

Transformer-based pre-training models like BERT have achieved remarkable performance in many natural language processing tasks.However, these models are both computation and memory expensive, hindering their deployment to resource-constrained devices. In this work, we propose TernaryBERT, which ternarizes the weights in a fine-tuned BERT model. Specifically, we use both approximation-based and loss-aware ternarization methods and empirically investigate the ternarization granularity of different parts of BERT. Moreover, to reduce the accuracy degradation caused by the lower capacity of low bits, we leverage the knowledge distillation technique in the training process. Experiments on the GLUE benchmark and SQuAD show that our proposed TernaryBERT outperforms the other BERT quantization methods, and even achieves comparable performance as the full-precision model while being 14.9x smaller.

研究动机与目标

• 为解决全精度 BERT 模型在资源受限设备上的高计算与内存开销问题。
• 通过结合三值化与知识蒸馏，减轻超低比特量化（如 1–2 比特）带来的精度下降问题。
• 探究 BERT 不同组件（如词嵌入与 Transformer 层）中三值化的最优粒度。
• 开发一种在三值权重约束下仍能保持高性能的方法，即使模型容量因三值化而减小。

提出的方法

• 采用基于近似与损失感知的量化策略，将 BERT 权重三值化为 {-1, 0, +1}。
• 对不同组件采用不同的三值化粒度：Transformer 权重使用层粒度，词嵌入使用头粒度。
• 从全精度教师模型中引入知识蒸馏，指导三值化学生模型的训练。
• 在中间表示与注意力图上也应用蒸馏损失，而不仅限于 logits，以更好地保留知识。
• 采用混合训练目标，结合交叉熵损失与蒸馏损失，以提升泛化能力。
• 通过损失感知量化（LAQ）方法将该方法扩展至 3 比特量化，用于消融实验与对比分析。

实验结果

研究问题

• RQ1 BERT 权重的三值化是否能在显著减小模型大小的同时，实现接近全精度 BERT 的性能？
• RQ2 三值化粒度的选择（层粒度 vs. 头粒度）如何影响不同 BERT 组件的模型精度？
• RQ3 知识蒸馏在多大程度上缓解了超低比特量化 BERT 模型中的精度下降问题？
• RQ4 不同蒸馏损失（logits、中间特征、注意力图）在三值化 BERT 中对性能提升的贡献如何？
• RQ5 TernaryBERT 在精度与模型大小方面，与当前最先进的量化与压缩方法相比表现如何？

主要发现

• TernaryBERT 在仅 28 MB 模型大小下实现了 83.5% 的 MNLI-m 准确率，相比全精度 BERT 缩小了 14.9 倍。
• 该模型在 MNLI 和 SQuAD 基准测试中均优于混合精度 Q-BERT 与 GOBO（2 比特权重），准确率更高。
• 知识蒸馏显著提升了性能：若从 Transformer 层与 logits 中均移除蒸馏，MNLI-m 准确率将下降超过 5%。
• 在中间表示与注意力图上应用蒸馏，其性能提升效果优于仅依赖 logits 的蒸馏方式。
• 数据增强与微调后的 BERT 初始化策略可进一步提升性能，尤其在 CoLA 与 RTE 上，两者结合时 CoLA 准确率最高提升达 9.7%。
• TernaryBERT 在 SQuAD v1.1 上的性能（F1 为 87.4/87.4）与全精度 BERT 相当，但模型大小仅为后者的 14.9 倍，展现出极强的压缩效率。