LightPAFF: Efficient Distillation for Language Models

• LightPAFF is a lightweight pre-training and fine-tuning framework that employs a two-stage knowledge distillation process to compress large Transformer models into efficient student models with minimal performance loss.
• It integrates distillation during both unsupervised pre-training and supervised fine-tuning by combining standard likelihood losses with KL divergence to match teacher outputs.
• Empirical evaluations show up to 7× inference speed improvements and minor performance drops, making it ideal for real-time applications in resource-constrained environments.

LightPAFF is a lightweight pre-training and fine-tuning framework for neural LLMs that employs a two-stage knowledge distillation process to compress large pre-trained Transformers into smaller, deployable student models while maintaining high predictive accuracy. Designed to address the prohibitive memory and inference speed requirements of state-of-the-art LLMs such as BERT, GPT-2, and MASS, LightPAFF injects structural flexibility into standard model pipelines by supporting distillation in both unsupervised pre-training and supervised fine-tuning stages. This enables effective deployment of LLMs in resource-constrained, real-time online settings with minimal loss in performance relative to their large teacher counterparts (Song et al., 2020).

1. Framework Structure and Principles

LightPAFF transfers knowledge from a large, pre-trained “teacher” Transformer to a smaller “student” model by explicitly distilling information at two critical stages: pre-training (unsupervised learning of general language representations) and fine-tuning (supervised task adaptation). The student model retains the architectural design of the teacher but with reduced depth, hidden size, and number of attention heads, resulting in approximately 1/5–1/6 the number of parameters.

Standard pre-training and fine-tuning protocols such as masked language modeling (BERT), causal language modeling (GPT-2), and masked sequence-to-sequence modeling (MASS) are extended: in each, the student matches probabilistic token or label distributions produced by the teacher for both the masked tokens during pre-training and ground-truth task outputs during fine-tuning. By performing distillation during pre-training, LightPAFF ensures the student inherits broad linguistic competence rather than merely task-specific behavior.

2. Two-Stage Knowledge Distillation Methodology

Both pre-training and fine-tuning in LightPAFF optimize a weighted combination of the standard likelihood loss on ground-truth data and a Kullback-Leibler (KL) divergence regularization term, guiding the student to approximate the teacher’s output distributions. The general loss for each stage is:

$$L(\theta) = \sum_{(x,y)} (1 - \lambda) \cdot \mathrm{MLE}(x, y; \theta) + \lambda \cdot \mathrm{KL}[P(\cdot|x; \theta_T) \| P(\cdot|x; \theta)]$$

where $\lambda \in [0,1]$ controls the influence of teacher predictions.

Specific loss instantiations for common paradigms:

• BERT (masked language modeling):

$$L_{MLM}(\theta) = -\sum_{x \in D} \sum_{m \in M} \sum_{k = 1}^V \left[ (1 - \lambda)\,\mathbf{1}\{x_m = k\} + \lambda Q(x_m = k \mid x^{\setminus M}; \theta_T) \right] \log P(x_m = k \mid x^{\setminus M}; \theta)$$

• GPT-2 (causal language modeling):

$$L_{CLM}(\theta) = -\sum_{x \in D} \sum_{m=1}^{|x|} \sum_{k = 1}^V \left[ (1 - \lambda)\,\mathbf{1}\{x_m = k\} + \lambda Q(x_m = k \mid x_{<m}; \theta_T) \right] \log P(x_m = k \mid x_{<m}; \theta)$$

• MASS (masked seq2seq):

$$L_{MSSM}(\theta) = -\sum_{x \in D} \sum_{m=s}^t \sum_{k=1}^V \left[ (1-\lambda)\,\mathbf{1}\{x_m=k\} + \lambda\,Q(x_m=k\mid x^{s:t}_{<m};\theta_T) \right] \log P(x_m=k\mid x^{s:t}_{<m}, x^{\setminus s:t}; \theta)$$

The fine-tuning stage applies the same general formulation, with adaptation to the specific downstream task (classification, language modeling, or sequence generation).

3. Model Architectures and Training Configuration

Teacher and student architectures for LightPAFF reflect reductions in model depth (L), hidden size (H), and attention heads (A), resulting in substantially smaller parameter counts. Examples include:

Model	Teacher Config	Student Config	Teacher Params	Student Params
BERT	L=12, H=768, A=12	L=3, H=512, A=8	110M	25M (EN), 20M (ZH)
GPT-2	L=24, H=1024, A=16	L=4, H=768, A=12	345M	67M
MASS	Enc/Dec L=6/H=1024/A=16	Enc L=6/Dec L=4/H=512/A=8	~307M	67M (Zh–En), 42M (En–De, En–Fr)

Training hyperparameters:

• Pre-training $\lambda$: BERT = 1.0, MASS = 0.7, GPT-2 = 0.4.
• Fine-tuning $\lambda$: tuned per task (approx. 0.4–0.8).
• Distributed infrastructure scales per model and stage, e.g., BERT PT on 8 × V100 GPUs, MASS PT on 8 GPUs with large batch sizes.

Inference latency (batch=1) is reduced by 4.5x–7x on both GPU and CPU compared to teacher models.

4. Empirical Evaluation and Results

LightPAFF is evaluated on representative benchmarks for three model paradigms:

• BERT: SST-2 (sentiment analysis), QQP (duplicate question detection), Polyphone disambiguation (Chinese)
• GPT-2: Language modeling on WikiText-2, PTB, WikiText-103
• MASS: Machine translation WMT17 Zh–En, WMT16 En–De, WMT14 En–Fr

Key quantitative outcomes include:

Task/Metric	Teacher	Student	Δ	Speedup
SST-2 Acc. (BERT)	93.5%	92.9%	-0.6	6–7× GPU/CPU
QQP F1/Acc. (BERT)	71.2/89.2	70.6/88.6	-0.6/-0.6
Polyphone Acc. (BERT)	95.9%	95.4%	-0.5
WT-2 Perp. (GPT-2)	15.5	18.8	+3.3	5.5–6.9× GPU/CPU
Wiki103 Perp. (GPT-2)	13.0	16.4	+3.4
Zh–En BLEU (MASS)	25.2	24.9	-0.3	4.5–5.2× GPU/CPU
En–De BLEU (MASS)	33.1	32.2	-0.9
En–Fr BLEU (MASS)	26.7	25.7	-1.0

Performance drops relative to teacher models are limited to ~0.5–1% absolute in accuracy or BLEU.

Ablation analysis reveals degradation if either pre-training or fine-tuning distillation is omitted, confirming the necessity of the two-stage process. Robustness experiments with random parameter perturbation demonstrate students trained via LightPAFF settle in wider minima and show improved resilience, suggesting better generalization properties.

5. Distillation Dynamics and Task-Specific Considerations

The optimal choice of $\lambda$ is empirically correlated with the reliability of the teacher’s predictions—a larger $\lambda$ is beneficial when the pre-training task is easier (e.g., BERT’s 15% token masking; ~73% token prediction accuracy on held-out data). For harder tasks, where the teacher is less certain (e.g., GPT-2 with only 50% of context visible yields ~42% held-out prediction accuracy), smaller $\lambda$ is preferred to avoid over-regularization toward uncertain soft labels.

In practice, $\lambda$ values for pre-training are set to BERT = 1.0, MASS = 0.7, GPT-2 = 0.4, with fine-tuning $\lambda$ values adapted per task between 0.4–0.8.

6. Practical Implications and Deployment

LightPAFF achieves a ≈5× reduction in model parameters—and 5–7× improvement in inference speed—while incurring only minor accuracy or BLEU loss, making it suitable for deployment in scenarios with constrained memory or low-latency requirements. Empirical trade-offs are explicit:

Model (Polyphone Disamb.)	#Params	Acc (%)
Teacher	110M	95.9
Student (30M)	30M	95.5
Student (20M)	20M	95.4
Student (8M)	8M	93.9

LightPAFF is model-agnostic, supporting various pre-training strategies and downstream tasks, and can be further integrated with pruning or quantization for more aggressive compression.

7. Significance and Outlook

By incorporating knowledge distillation into both pre-training and fine-tuning, LightPAFF extends the effectiveness of distillation beyond task adaptation to encompass general language understanding, addressing limitations of prior works restricted to the supervised stage. It demonstrates that minimal student models can inherit both general linguistic and task-specific competencies from large-scale pre-trained teachers, with minor loss in fidelity and substantial gains in efficiency. This suggests a practical pathway for deploying high-performance NLP models in online and resource-limited environments (Song et al., 2020).