[Paper Review] 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 Techniques4 citations

TL;DR

MiniRBT is a small, efficient Chinese pre-trained language model that uses a narrow and deep architecture with two-stage knowledge distillation and dynamic whole word masking to achieve 94% of RoBERTa's performance while reducing model size to 10% and delivering a 6.8x inference speedup. It is optimized for practical deployment in downstream NLP tasks with strong efficiency and effectiveness.

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.

Motivation & Objective

• To address the lack of efficient, small-scale pre-trained models for Chinese NLP.
• To reduce inference latency and model size while maintaining high performance on downstream tasks.
• To improve distillation efficiency by introducing a two-stage distillation process with a teacher-assistant intermediate model.
• To investigate whether narrow and deep architectures outperform wide and shallow ones in small models of similar parameter count.
• To enable practical deployment of Chinese NLP models in resource-constrained environments.

Proposed method

• A narrow and deep student model architecture is used to improve parameter efficiency and inference speed.
• Dynamic whole word masking (WWM) is applied during pre-training to preserve semantic coherence by masking entire Chinese words rather than individual characters.
• Two-stage knowledge distillation is employed: first distilling from the teacher model (Chinese RoBERTa-wwm) to a teacher-assistant model, then from the teacher-assistant to the student model.
• The student model is trained using distillation loss to mimic the attention maps and logits of the teacher models.
• Pre-training is conducted on a large-scale Chinese corpus using a batch size of 4096 and 100K training steps with a temperature of 8.
• Fine-tuning is performed on downstream tasks with learning rates of 5e-5 or 1e-4 over 2–10 epochs, with results averaged across multiple random seeds.

Experimental results

Research questions

• RQ1 Does a narrow and deep model structure outperform a wide and shallow one when both have similar parameter counts in small Chinese pre-trained models?
• RQ2 Can two-stage knowledge distillation improve performance compared to one-stage distillation in small model compression?
• RQ3 How effective is dynamic whole word masking in enhancing the semantic modeling capability of small Chinese pre-trained models?
• RQ4 To what extent can a 10% parameter model match the performance of RoBERTa in downstream NLP tasks?
• RQ5 Can a distilled small model achieve significant speedups without sacrificing accuracy in real-world deployment scenarios?

Key findings

• MiniRBT achieves 94% average relative performance compared to RoBERTa, with a 6.8x speedup on a single NVIDIA M40 GPU.
• MiniRBT with 10% of RoBERTa’s parameters outperforms RBT3, which has 3–4 times more parameters, on machine reading comprehension tasks.
• The two-stage distillation method improves performance by 0.4–0.5 F1 points on average compared to one-stage distillation.
• MiniRBT achieves 98% of RoBERTa’s accuracy on text classification tasks, with an average relative performance of 95.3% across multiple benchmarks.
• The narrow and deep architecture outperforms wide and shallow counterparts of similar size, confirming its effectiveness for small model design.
• MiniRBT surpasses larger models like RBT4-H312 in performance despite having fewer parameters, demonstrating the advantage of architectural design and distillation.