[论文解读] Extreme Compression for Pre-trained Transformers Made Simple and Efficient

Xiaoxia Wu, Zhewei Yao+3 more|arXiv (Cornell University)|Jun 4, 2022

Advanced Image and Video Retrieval Techniques被引用 13

一句话总结

本文提出XTC，一种通过超低比特量化（1-bit/2-bit）和模型蒸馏实现预训练Transformer模型极端压缩的简单而高效的方法。通过系统分析超参数和训练策略，作者发现先前的最先进方法存在训练不足的问题，并证明XTC在GLUE任务上实现了最先进性能，模型尺寸缩小50倍——甚至超越6层TinyBERT，且无需多阶段知识蒸馏或复杂优化。

ABSTRACT

Extreme compression, particularly ultra-low bit precision (binary/ternary) quantization, has been proposed to fit large NLP models on resource-constraint devices. However, to preserve the accuracy for such aggressive compression schemes, cutting-edge methods usually introduce complicated compression pipelines, e.g., multi-stage expensive knowledge distillation with extensive hyperparameter tuning. Also, they oftentimes focus less on smaller transformer models that have already been heavily compressed via knowledge distillation and lack a systematic study to show the effectiveness of their methods. In this paper, we perform a very comprehensive systematic study to measure the impact of many key hyperparameters and training strategies from previous works. As a result, we find out that previous baselines for ultra-low bit precision quantization are significantly under-trained. Based on our study, we propose a simple yet effective compression pipeline for extreme compression, named XTC. XTC demonstrates that (1) we can skip the pre-training knowledge distillation to obtain a 5-layer BERT while achieving better performance than previous state-of-the-art methods, e.g., the 6-layer TinyBERT; (2) extreme quantization plus layer reduction is able to reduce the model size by 50x, resulting in new state-of-the-art results on GLUE tasks.

研究动机与目标

• 探究为何先前BERT的极端量化方法需要复杂的多阶段流水线及大量超参数调优。
• 确定昂贵的知识蒸馏和权重拆分是否对二值化（1-bit）和三值化（2-bit）量化确实必要。
• 评估极端量化在已小型化、蒸馏后的模型（如TinyBERT和MiniLM）上的有效性。
• 开发一种简化、高效的压缩流水线，在最小化训练开销的同时保持高精度。
• 为超低比特量化提供系统化的超参数与训练策略基准。

提出的方法

• 提出XTC，一种简化压缩流水线，跳过预训练蒸馏，直接对小型、蒸馏后的模型应用量化感知训练（QAT）。
• 采用单阶段训练流程，使用直通估计器（STE）实现1-bit或2-bit权重量化，并进行梯度反向传播。
• 引入层剪枝（如6层SkipBERT）以进一步压缩模型尺寸，同时保持性能。
• 使用标准微调策略，结合混合精度训练与学习率调度，避免多阶段蒸馏和数据增强。
• 在部分消融实验中应用LoRa（低秩适应）实现参数高效的微调。
• 通过超过1,000种模型变体的超参数消融实验，识别出最优配置。

实验结果

研究问题

• RQ1 在1-bit量化BERT模型中，是否真的需要多阶段知识蒸馏与数据增强才能实现高精度？
• RQ2 更简单的单阶段训练流水线能否实现与复杂多阶段基线相当或更优的性能？
• RQ3 极端量化与已压缩的蒸馏模型（如TinyBERT或MiniLM）之间有何交互作用？
• RQ4 哪些超参数与训练策略对超低比特量化模型的性能影响最大？
• RQ5 我们能否在无需预训练蒸馏或权重拆分的情况下实现极端压缩的最先进结果？

主要发现

• 先前在1-bit和2-bit量化方面的最先进方法存在严重训练不足，导致性能未达最优。
• XTC在1-bit量化6层SkipBERT上实现81.43的平均GLUE分数，优于无预训练蒸馏的6层TinyBERT（80.56）。
• 通过极端量化与层剪枝，模型尺寸缩小50倍（从16MB降至0.3MB），在GLUE基准上达到最先进结果。
• 对于1-bit量化，XTC仅用8.1MB模型尺寸即实现81.43的平均GLUE分数，超越采用复杂多阶段流水线的先前方法。
• 该方法在不同随机种子下表现稳定，平均GLUE分数的标准差仅为0.09–0.11，表明训练过程高度稳定。
• XTC消除了对权重拆分和多阶段蒸馏的需求，降低了训练复杂度，同时提升或保持了精度。