[论文解读] MiniRBT: A Two-stage Distilled Small Chinese Pre-trained Model

Xin Yao, Ziqing Yang+2 more|arXiv (Cornell University)|Apr 3, 2023

Natural Language Processing Techniques被引用 4

一句话总结

MiniRBT 是一个小型、高效的中文预训练语言模型，采用窄而深的网络结构，结合两阶段知识蒸馏与动态全词掩码技术，在模型参数量仅为 RoBERTa 的 10% 的同时，实现了 94% 的 RoBERTa 性能，并带来 6.8 倍的推理速度提升。该模型在下游 NLP 任务中具备出色的效率与效果，专为实际部署而优化。

ABSTRACT

In natural language processing, pre-trained language models have become essential infrastructures. However, these models often suffer from issues such as large size, long inference time, and challenging deployment. Moreover, most mainstream pre-trained models focus on English, and there are insufficient studies on small Chinese pre-trained models. In this paper, we introduce MiniRBT, a small Chinese pre-trained model that aims to advance research in Chinese natural language processing. MiniRBT employs a narrow and deep student model and incorporates whole word masking and two-stage distillation during pre-training to make it well-suited for most downstream tasks. Our experiments on machine reading comprehension and text classification tasks reveal that MiniRBT achieves 94% performance relative to RoBERTa, while providing a 6.8x speedup, demonstrating its effectiveness and efficiency.

研究动机与目标

• 为解决中文 NLP 领域缺乏高效、小型化的预训练模型的问题。
• 在保持下游任务高性能的前提下，降低推理延迟与模型大小。
• 通过引入教师-助教中间模型的两阶段蒸馏过程，提升蒸馏效率。
• 探究在参数量相近的小型模型中，窄而深的结构是否优于宽而浅的结构。
• 实现在资源受限环境下的中文 NLP 模型实用化部署。

提出的方法

• 采用窄而深的学生模型架构，以提升参数效率与推理速度。
• 在预训练阶段应用动态全词掩码（WWM），通过掩码整个中文词语而非单个字符，以保持语义连贯性。
• 采用两阶段知识蒸馏：首先从教师模型（中文 RoBERTa-wwm）蒸馏到教师-助教模型，再从教师-助教模型蒸馏到学生模型。
• 使用蒸馏损失训练学生模型，使其模仿教师模型的注意力图与输出 logits。
• 在大规模中文语料上进行预训练，使用 4096 的批量大小与 100K 次训练步数，温度系数为 8。
• 在下游任务上进行微调，学习率设为 5e-5 或 1e-4，训练 2–10 个周期，结果取多个随机种子的平均值。

实验结果

研究问题

• RQ1 当小型中文预训练模型的参数量相近时，窄而深的模型结构是否优于宽而浅的结构？
• RQ2 与单阶段蒸馏相比，两阶段知识蒸馏是否能提升小型模型压缩的性能？
• RQ3 动态全词掩码在提升小型中文预训练模型语义建模能力方面的有效性如何？
• RQ4 参数量仅为 10% 的模型在下游 NLP 任务中，能在多大程度上匹配 RoBERTa 的性能？
• RQ5 在真实部署场景中，蒸馏后的轻量模型是否能在不损失准确率的前提下实现显著的速度提升？

主要发现

• MiniRBT 在单张 NVIDIA M40 GPU 上实现了 6.8 倍的推理速度提升，平均相对性能达到 RoBERTa 的 94%。
• MiniRBT 参数量仅为 RoBERTa 的 10%，但在机器阅读理解任务上优于参数量多出 3–4 倍的 RBT3 模型。
• 与单阶段蒸馏相比，两阶段蒸馏在平均 F1 分数上提升了 0.4–0.5 分。
• MiniRBT 在文本分类任务上达到 RoBERTa 98% 的准确率，多个基准测试的平均相对性能为 95.3%。
• 窄而深的架构在参数量相近的模型中优于宽而浅的结构，证实其在小型模型设计中的有效性。
• 尽管参数量更少，MiniRBT 在性能上仍超越了更大的模型（如 RBT4-H312），凸显了架构设计与蒸馏技术的优势。