Task-Specific Distillation: Detection, Segmentation & NLP

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Object Detection Distillation
5. Semantic Segmentation Distillation
6. NLP Distillation: DistilBERT & TinyBERT
7. Generative Model Distillation
8. Multi-Task Distillation
9. Common Pitfalls
10. Summary
11. Exercises

Learning Objectives

1. Design distillation losses for detection (foreground/background imbalance).
2. Apply pixel-level distillation for dense prediction.
3. Describe DistilBERT and TinyBERT architectures and training.
4. Explain progressive distillation for generative models.
5. Handle multi-task scenarios (conflicting teacher signals).

Notation

• FPN — feature pyramid network
• $\mathcal{R}$ — region proposals
• $\omega_{ij}$ — pixel importance weight

Core Intuition

Classification distillation (match softmax outputs) doesn't directly transfer to other tasks. Detection needs spatial localization; segmentation needs pixel-level decisions; NLP needs token-level AND sentence-level understanding. Each task requires custom distillation losses that transfer the right type of knowledge at the right granularity.

Object Detection Distillation

Challenges unique to detection:

1. Multi-scale features (FPN levels).
2. Foreground/background imbalance (1:100 ratio).
3. Both classification AND regression outputs.
4. Region proposal quality matters.

Feature Pyramid Distillation:

Match teacher's FPN features at each pyramid level (P2-P5).

Foreground-aware distillation: Weight by foreground probability:

Focus distillation signal on object regions (where it matters most).

Classification + regression:

Semantic Segmentation Distillation

Pixel-level KD: Apply distillation at every spatial position:

Structured distillation: Beyond pixel-level, transfer spatial structure:

• Pair-wise: Encourage same relative distances between pixel features.
• Holistic: Match global scene statistics (Gram matrices).

Boundary-aware: Increase weight near class boundaries (where errors matter most):

Channel-wise distillation: Some channels encode object presence, others encode position. Transfer both types.

NLP Distillation: DistilBERT & TinyBERT

DistilBERT (Sanh et al., 2019):

• Architecture: 6 layers (half of BERT-base's 12).
• Distillation: Token-level soft label + masked language model loss.
• Result: 97% of BERT quality at 60% size, 60% faster.
• Training: Initialize from alternating BERT layers (skip every other).

TinyBERT (Jiao et al., 2020): More aggressive compression with multi-layer distillation:

Three types of knowledge transferred:

1. Attention matrices: $\mathcal{L}_{\text{attn}} = \text{MSE}(A_S^l, A_T^{g(l)})$ (attention distribution matching).
2. Hidden states: $\mathcal{L}_{\text{hidden}} = \text{MSE}(H_S^l W, H_T^{g(l)})$ (with projection $W$).
3. Predictions: $\mathcal{L}_{\text{pred}} = \text{KL}(p_T \| p_S)$ (standard logit KD).

Layer mapping $g(l)$: Maps student layer $l$ to teacher layer $g(l)$. For 4-layer student from 12-layer teacher: $g = [3, 6, 9, 12]$.

Generative Model Distillation

Progressive distillation for diffusion (Salimans & Ho, 2022):

1. Teacher: 1000-step diffusion model.
2. Student: Trained to match 2 teacher steps in 1 student step.
3. Repeat: Student becomes teacher for next round.
4. Result: 4-step student matches 1000-step teacher quality.

Loss: Student predicts the SAME denoised output as two teacher steps combined:

GAN distillation: Compress generator via standard KD + adversarial loss:

Multi-Task Distillation

Challenge: Teacher excels at all tasks; student has limited capacity and must balance them.

Multi-teacher distillation: Different teacher per task:

Task routing: Not all tasks benefit equally from distillation. Some tasks are better learned from data directly; others from the teacher.

Gradnorm for task balancing: Dynamically adjust $\alpha_k$ based on task gradient magnitudes.

Common Pitfalls

Pitfall 1. Using classification-style KD for detection without handling foreground/background imbalance. Background regions dominate and the student learns mostly "not an object."

Pitfall 2. Matching ALL attention heads in TinyBERT. Many teacher heads are redundant; matching all of them overconstrain the student. Select the most informative heads.

Pitfall 3. Progressive distillation with too large step reduction per round. Going from 2 steps to 1 step in one round often fails; prefer 4→2→1 progression.

Summary

• Detection: FPN feature matching + foreground-aware weighting.
• Segmentation: Pixel-level KD + boundary-aware + structural losses.
• NLP: Multi-layer (attention + hidden + prediction) distillation.
• Generative: Progressive step-halving distillation.
• Multi-task: Balance task weights dynamically; use task-specific teachers.
• Each task needs custom losses matching the nature of its predictions.

Exercises

Exercise 1. Design a distillation loss for YOLO (single-stage detector): specify which outputs to match and how to handle scale.

Exercise 2. For TinyBERT with 4 layers from BERT-12: propose an alternative layer mapping to $[3,6,9,12]$ and argue why it might work better.

Exercise 3. Compute the compression factor of DistilBERT (6 layers, 768 hidden) vs BERT-base (12 layers, 768 hidden) in parameters, FLOPs, and latency.

Exercise 4. Design a progressive distillation schedule for a 50-step flow model targeting 2-step generation.

Exercise 5. For multi-task distillation (classification + detection + segmentation): how should you allocate student capacity across tasks?