[Paper Review] Why Skip If You Can Combine: A Simple Knowledge Distillation Technique for Intermediate Layers

Yimeng Wu, Peyman Passban+2 more|arXiv (Cornell University)|Oct 6, 2020

Natural Language Processing Techniques19 references4 citations

TL;DR

This paper proposes a novel knowledge distillation technique, Combinatorial Knowledge Distillation (CKD), that combines intermediate-layer features from a deep teacher network into a compact student network without skipping any layers, addressing the 'skip problem' in existing methods. CKD achieves state-of-the-art performance with 50% fewer parameters than 12-layer teachers, matching their translation quality on low-resource NMT tasks.

ABSTRACT

With the growth of computing power neural machine translation (NMT) models also grow accordingly and become better. However, they also become harder to deploy on edge devices due to memory constraints. To cope with this problem, a common practice is to distill knowledge from a large and accurately-trained teacher network (T) into a compact student network (S). Although knowledge distillation (KD) is useful in most cases, our study shows that existing KD techniques might not be suitable enough for deep NMT engines, so we propose a novel alternative. In our model, besides matching T and S predictions we have a combinatorial mechanism to inject layer-level supervision from T to S. In this paper, we target low-resource settings and evaluate our translation engines for Portuguese--English, Turkish--English, and English--German directions. Students trained using our technique have 50% fewer parameters and can still deliver comparable results to those of 12-layer teachers.

Motivation & Objective

• To address the 'skip problem' in knowledge distillation where intermediate layers of a deeper teacher model cannot be matched to a shallower student due to architectural mismatch.
• To improve knowledge distillation in neural machine translation by leveraging not only final predictions but also intermediate-layer representations.
• To develop a method that enables full utilization of all teacher layers without selective pruning or skipping, enhancing student model performance.
• To demonstrate the effectiveness of combining intermediate-layer features via concatenation in a novel combinatorial mechanism for distillation.
• To validate the approach on low-resource NMT settings, showing that compact students can match the performance of large teachers.

Proposed method

• Proposes a combinatorial mechanism that concatenates intermediate-layer features from multiple encoder layers of the teacher network into a single representation for distillation.
• Introduces a new loss function that combines the standard cross-entropy loss (hard loss), the standard knowledge distillation loss (soft loss), and a new intermediate-layer distillation loss.
• The intermediate-layer loss matches the concatenated output of multiple teacher encoder layers with the corresponding student encoder layer outputs.
• Uses a weighted combination of losses: total loss = β×hard_loss + η×soft_loss + λ×intermediate_loss, with λ=0.7 and η=0.1, optimized via hyperparameter search.
• Applies distillation only on the encoder side, keeping decoder layers unconnected to the teacher to avoid performance degradation.
• Employs a multi-head attention-based architecture with residual connections, and trains students with 4 encoder layers against 12-layer teachers.

Experimental results

Research questions

• RQ1 Can a knowledge distillation method that combines multiple intermediate-layer features from a deep teacher model outperform existing methods that skip layers?
• RQ2 Does a combinatorial fusion of intermediate representations improve student model performance compared to layer-wise matching or prediction-only distillation?
• RQ3 Can a compact student model with 50% fewer parameters match the translation quality of a 12-layer teacher model in low-resource NMT settings?
• RQ4 How does the proposed method perform in large-scale NMT datasets compared to existing techniques like PKD?
• RQ5 Is the performance gain from combining intermediate layers statistically significant and robust across different language directions?

Key findings

• CKD achieves a BLEU score of 43.78 on the Portuguese-to-English translation task, outperforming PKD (43.28) and the original RKD baseline.
• On English-to-German translation, CKD achieves 24.14 BLEU, surpassing PKD (23.38) and No-KD (24.31), demonstrating consistent gains across tasks.
• For English-to-French translation, CKD achieves 36.10 BLEU, exceeding both No-KD (35.45) and PKD (34.97), confirming effectiveness on large-scale datasets.
• The method enables students with 50% fewer parameters than 12-layer teachers to match their performance, validating model efficiency.
• Experiments show that matching self-attention components in PKD improves performance slightly, but CKD still outperforms even these enhanced variants.
• The high weight (λ=0.7) assigned to the intermediate-layer loss indicates its critical role in training high-quality compact models.