ERNIE-Tiny Progressive Compression

• The paper demonstrates that a four-stage progressive distillation approach effectively compresses large PLMs into compact student models with over 98% accuracy retention.
• It employs a systematic transition through general distillation, general-enhanced, task-adaptive, and task-specific stages to smoothly bridge pretraining and finetuning regimes.
• The method achieves a 9.4× speedup and significant parameter reduction, outperforming comparable compact models on both English GLUE and Chinese NLP benchmarks.

ERNIE-Tiny Progressive Compression refers to a four-stage progressive distillation framework for compressing large pretrained LLMs (PLMs) such as BERT and ERNIE into significantly smaller, faster student models while maintaining high task performance. The core insight is to mitigate the commonly observed pretrain-finetune distillation discrepancy by gradually morphing the teacher model, training data, and learning objective from general, unsupervised pretraining conditions to task-specific, supervised finetuning conditions. At each of the four stages, only one component is changed, enabling smoother student optimization and minimizing performance degradation. Empirical results demonstrate that ERNIE-Tiny achieves state-of-the-art performance among compact models on English (GLUE) and Chinese NLP tasks, with over 9× inference speedup and substantial parameter reduction (Su et al., 2021).

1. Motivation and Problem Setting

PLMs such as BERT_base (12 layers, 768-d hidden state, 109M parameters) achieve state-of-the-art performance across numerous NLP tasks. However, their large parameter counts and high inference latencies preclude deployment on resource-constrained environments such as CPU servers and edge devices. Knowledge distillation (KD) is a standard approach for compressing PLMs; it involves training a small student to mimic a larger teacher, typically using latent (feature-based) and/or soft-label (probabilistic output) matching.

Conventional KD techniques for PLMs either perform distillation only during pretraining or only during finetuning, or else transition abruptly from a general pretraining regime to a specific finetuning regime. This abrupt change across teacher model, data distribution, and objective introduces an optimization gap and can hinder student adaptation. ERNIE-Tiny addresses this problem by defining a curriculum-inspired progressive distillation sequence, in which each transition modifies only one component, thus enabling smoother optimization and improved generalization, especially in low-resource or out-of-domain settings (Su et al., 2021).

2. Progressive Distillation Framework

The ERNIE-Tiny framework consists of four stages—General Distillation (GD), General-Enhanced Distillation (GED), Task-Adaptive Distillation (TAD), and Task-Specific Distillation (TSD)—and systematically transitions through variations in teacher, data, and loss.

Notation

Symbol	Meaning
$T_g$, $T_f$	Pretrained, finetuned teacher models
$S$	Student model
$D_g$, $D_t$	General (unlabeled), task (labeled) data
$H^T_l(x)$, $H^S_l(x)$	Teacher/student hidden states at layer $l$
$A^T_{l,a}(x)$, $A^S_{l,a}(x)$	Attention probabilities
$N_l$, $M_l$	Projection operators for alignment
$z^T(x)$, $z^S(x)$	Logits
$y$	Ground-truth label
$\alpha, \beta$	Loss weighting hyperparameters

The Four Distillation Stages

1. General Distillation (GD)
   • Teacher: Pretrained PLM ($T_g$)
   • Data: General, unlabeled corpus ($D_g$, e.g., Wikipedia + BookCorpus)
   • Loss: Latent distillation. The student matches selected teacher hidden states and attention distributions over general data:

   $$\mathcal{L}_{GD} = \mathbb{E}_{x\sim D_g}\left[ \sum_{l=1}^{L_S} \| H^{T_g}_{k(l)}(x) - H^S_l(x)N_l \|_2^2 + \sum_{l=1}^{L_S} \sum_{a=1}^{h'} \| A^{T_g}_{k(l),a}(x) - M_l A^S_{l,a}(x) \|_2^2 \right]$$

   where $k(l) = l \cdot (L_T / L_S)$.

2. General-Enhanced Distillation (GED)
   • Teacher: Finetuned teacher ($T_f$), trained on $D_t$
   • Data: Still $D_g$
   • Loss: Latent distillation

   $$\mathcal{L}_{GED} = \mathbb{E}_{x\sim D_g}[\mathcal{L}_{latent}^{T_f}(x)]$$

• Significance: Only the teacher changes; allows injection of task-specific knowledge into student representations over large, general corpora, enhancing robustness.

1. Task-Adaptive Distillation (TAD)
   • Teacher: $T_f$
   • Data: Switch to labeled task-specific data ($D_t$)
   • Loss: Latent distillation

   $$\mathcal{L}_{TAD} = \mathbb{E}_{x\sim D_t}[\mathcal{L}_{latent}^{T_f}(x)]$$

• Significance: Only the data changes; prepares the student for supervised objectives without abrupt introduction of classification loss.

1. Task-Specific Distillation (TSD)
   • Teacher: $T_f$
   • Data: $D_t$
   • Loss: Combined latent distillation, soft-label KD, and hard-label cross-entropy:

   $$\mathcal{L}_{TSD} = \beta\,\mathcal{L}_{latent}^{T_f} + \alpha\,\mathcal{L}_{soft} + (1-\alpha)\,\mathcal{L}_{hard}$$

   where

   $$\mathcal{L}_{soft} = -\sum_{i=1}^C p^T_i\log p^S_i,\quad p^T_i = \mathrm{softmax}(z^{T_f}_i/\tau),\ p^S_i = \mathrm{softmax}(z^S_i/\tau)$$

   $\mathcal{L}_{hard} = -\sum_{i=1}^C y_i\log p^S_i$

   Typically $\alpha=1$, $\beta=1$.

Component Transitions Table

Stage Transition	What Changes
GD → GED	Teacher ($T_g\to T_f$)
GED → TAD	Data ($D_g\to D_t$)
TAD → TSD	Loss (latent→latent+soft+hard)

3. Training Procedure and Hyperparameters

A pseudocode encapsulation of the procedure:

for step in 1...N1: x = sample(D_g) L_lat = latent_loss(x; teacher=T_g, student=S) S.params -= η1 * ∇L_lat for step in 1...N2: x = sample(D_g) L_lat = latent_loss(x; teacher=T_f, student=S) S.params -= η2 * ∇L_lat for epoch in 1...E3: for (x, y) in batches(D_t): L_lat = latent_loss(x; teacher=T_f, student=S) S.params -= η3 * ∇L_lat for epoch in 1...E4: for (x, y) in batches(D_t): L_lat = latent_loss(x; teacher=T_f, student=S) L_soft = soft_loss(x; teacher=T_f, student=S) L_hard = hard_loss(y, student_logits(x)) L_TSD = β * L_lat + α * L_soft + (1−α) * L_hard S.params -= η4 * ∇L_TSD

Model specification:

• Teacher $T_g$: BERT_base (12L, 768-d, FFN 3072, 109M params)
• Student $S$: 4 layers, 312-d hidden, FFN 1200, 14.5M params

Data:

• $D_g$: Wikipedia + BookCorpus
• $D_t$ (English): GLUE suite (MNLI, QQP, QNLI, SST-2, CoLA, RTE, MRPC, STS-B)
• $D_t$ (Chinese): XNLI, LCQMC, ChnSentiCorp, NLPCC-DBQA, MSRA-NER

Hyperparameters (for GLUE tasks; analogous for Chinese):

• GD & GED: batch 1280, lr $2\times 10^{-4}$, 500K steps, Adam, warmup 500, dropout 0.0
• TAD: batch 256/128, lr $5\times 10^{-5}$, 5–50 epochs, warmup 10%, dropout 0.0
• TSD: batch 256/128, lr $1\times 10^{-5}$/ $3\times 10^{-5}$, 3 epochs, warmup 10%, 0.0 dropout

4. Empirical Results

English GLUE (test set, official GLUE server):

Method	Params	Speedup	Avg. Score
BERT<sub>base</sub>	109M	1×	79.6
DistilBERT	52M	3×	71.9
BERT-PKD	52M	3×	72.6
BERT-EMD	14.5M	9.4×	74.7
TinyBERT	14.5M	9.4×	77.0
MobileBERT	15.1M	8.6×	77.0
ERNIE-Tiny	14.5M	9.4×	78.0

ERNIE-Tiny retains over 98% of BERT_base’s performance, surpassing prior 4-layer distillation SOTAs (TinyBERT, MobileBERT) by approximately +1% GLUE score (Su et al., 2021).

Chinese NLP Tasks (dev-set averages over 5 runs):

Method	Params	Speedup	Avg. Accuracy
ERNIE2.0<sub>base</sub>	109M	1×	88.4
BERT<sub>base</sub>	109M	1×	86.3
TinyBERT (repro’d)	14.5M	9.4×	85.3
ERNIE-Tiny	14.5M	9.4×	86.7

ERNIE-Tiny slightly exceeds the performance of the full 12-layer BERT_base on Chinese suite tasks (by ≈ +0.4%), while being approximately 7.5× smaller and 9.4× faster.

Ablation Analysis

• Removing both GD & GED reduces average GLUE from 75.7→55.0 (for MRPC/CoLA/MNLI).
• Skipping GED (keeping teacher as $T_g$ through TAD) decreases performance by ≈ 3.3 points.
• Skipping TAD (immediate supervised loss in TSD) results in ≈ 4.1-point drop.

5. Mechanistic Analysis and Insights

Progressive component transition in ERNIE-Tiny contributes to its empirical effectiveness through several factors:

1. Smoother optimization: Gradual introduction of teacher specialization, data specificity, and classification loss ensures the student is not destabilized by simultaneous abrupt shifts, leading to improved optimization trajectories and final accuracy.
2. Task-knowledge augmentation through GED: Finetuned teacher guidance (on general data) exposes the student to task-specific features on a large corpus, reducing overfitting (when $D_t$ is small) and improving out-of-domain generalization (demonstrated by boosted performance on SNLI/RTE when distilled using MNLI).
3. Adaptive bridging via TAD: Preceding classification loss with latent KD on labeled data enables the student to internalize teacher representations specific to the task, reducing loss shocks and enhancing adaptation.

6. Limitations and Future Directions

Key limitations are the increased training cost associated with four sequential distillation stages (∼2 days on 4× V100 GPUs for English GLUE), and reliance on access to large unlabeled corpora in GED (same as BERT pretraining data).

Potential research avenues include:

• Scaling GED with larger corpora (e.g., C4).
• Designing dynamic curricula where the number of GED/TAD iterations is task-adaptive.
• Jointly searching for optimal student architectures (e.g., via NAS) within the progressive distillation paradigm.

ERNIE-Tiny’s progressive distillation sequence demonstrates that carefully staged transitions in teacher, data, and loss are instrumental for attaining nearly teacher-level accuracy in highly compressed student architectures, setting new state-of-the-art results for 4-layer transformer-based PLMs on both English and Chinese language benchmarks (Su et al., 2021).