Feature-Based & Relation-Based Distillation

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. FitNets: Intermediate Layer Matching
5. Attention Transfer
6. PKT: Probabilistic Knowledge Transfer
7. CRD: Contrastive Representation Distillation
8. RKD: Relational Knowledge Distillation
9. When to Use Feature vs Logit Distillation
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Design feature matching losses with appropriate projections.
2. Derive attention transfer and its connection to feature importance.
3. Explain CRD's contrastive objective and why it maximizes mutual information.
4. Implement relational distillation (distance and angle losses).
5. Choose between feature-based and logit-based distillation for different scenarios.

Notation

• $F_T^l, F_S^l$ — teacher/student features at layer $l$
• $\phi$ — projection/adaptation function
• $A^l$ — attention map at layer $l$

Core Intuition

Logit distillation only transfers the teacher's FINAL output. But much of the teacher's knowledge is in its INTERMEDIATE representations — how it processes input through layers, what features it attends to, how it structures its hidden space. Feature distillation transfers this richer knowledge by matching intermediate layers, attention patterns, or inter-sample relationships.

FitNets: Intermediate Layer Matching

Romero et al. (2015): Match student's intermediate layer to teacher's:

where $\phi: \mathbb{R}^{d_S} \to \mathbb{R}^{d_T}$ is a learned linear projection (since $d_S \neq d_T$ in general).

Layer selection: Match the "hint layer" (teacher's intermediate) to the "guided layer" (student's intermediate):

• Typically: teacher's layer at 2/3 depth → student's layer at 1/2 depth.
• Deeper teacher layers contain more abstract features.

Two-stage training:

1. Stage 1: Train only the student layers up to the guided layer (match intermediate).
2. Stage 2: Train the full student end-to-end (logit distillation + classification loss).

Challenge: Which teacher layer should guide which student layer? Misalignment hurts performance.

Attention Transfer

Zagoruyko & Komodakis (2017): Transfer what the teacher ATTENDS to, not the full feature map.

Attention map: Summarize a feature map $F \in \mathbb{R}^{C \times H \times W}$ by channel-wise aggregation:

With $p=2$: squared attention (emphasizes strong activations).

Loss:

Advantage: Resolution-independent (teacher and student can have different spatial sizes after normalization). Focuses on WHERE the model attends, not the full representation.

PKT: Probabilistic Knowledge Transfer

Passalis & Tefas (2018): Model features as probability distributions in representation space.

Teacher feature distribution: $p_T(\mathbf{f}) = \frac{1}{N}\sum_i K(\mathbf{f}, \mathbf{f}_T^i)$ (kernel density estimate).

Student feature distribution: $p_S(\mathbf{f}) = \frac{1}{N}\sum_i K(\mathbf{f}, \mathbf{f}_S^i)$.

Loss: KL divergence between pairwise similarity distributions:

Intuition: The student's feature space should have the same geometry as the teacher's — similar samples should be close, dissimilar ones far.

CRD: Contrastive Representation Distillation

Tian et al. (2020): Maximize mutual information between teacher and student representations via contrastive learning.

Objective: For a sample $x$, its teacher feature $f_T(x)$ and student feature $f_S(x)$ should be more similar than random pairings:

where $h$ is a learned similarity function and $x_k^-$ are negative samples.

Why contrastive beats MSE:

• MSE requires exact matching (sensitive to irrelevant details).
• Contrastive only requires RELATIVE similarity (preserves structure, ignores noise).
• Invariant to linear transformations of the feature space.

Result: State-of-the-art single-method distillation on CIFAR-100 and ImageNet.

RKD: Relational Knowledge Distillation

Park et al. (2019): Don't match individual features — match RELATIONSHIPS between samples.

Distance loss: Preserve pairwise distances:

where $\hat{d}$ is the normalized pairwise distance: $\hat{d}(i,j) = d(i,j) / \mu_d$.

Angle loss: Preserve angular relationships (triplets):

where $\angle(i,j,k)$ is the angle at vertex $j$ in the triangle formed by features of samples $i, j, k$.

Advantage: Completely agnostic to dimensionality mismatch. No projection layer needed. Works even when teacher and student have wildly different architectures.

When to Use Feature vs Logit Distillation

Use logit distillation when:

• Teacher and student have similar architectures.
• The task is classification with many classes (rich logit signal).
• Simplicity is desired (no layer matching decisions).

Use feature distillation when:

• Large architecture gap (ResNet teacher → MobileNet student).
• Dense prediction tasks (detection, segmentation) — intermediate features matter.
• The student needs to learn specific spatial patterns.

Use relational distillation when:

• Extreme architecture differences (CNN teacher → Transformer student).
• Retrieval/embedding tasks (preserving metric structure).
• No clear layer correspondence between teacher and student.

Best practice: Combine: $\mathcal{L} = \alpha\mathcal{L}_{\text{KD}} + \beta\mathcal{L}_{\text{feature}} + \gamma\mathcal{L}_{\text{task}}$.

Common Pitfalls

Pitfall 1. Matching ALL intermediate layers simultaneously. Too many constraints overconstrain the student. Match 1-3 carefully chosen layers.

Pitfall 2. Forgetting the projection layer when dimensions differ. Direct MSE between $d_S=256$ and $d_T=2048$ features is meaningless without a learned adapter.

Pitfall 3. Using MSE for feature matching when CRD would be better. MSE penalizes absolute values; if the teacher uses a different scale or has irrelevant noise, CRD's relative matching is more robust.

Summary

• FitNets: Match intermediate layers with projection; two-stage training.
• Attention transfer: Match spatial attention maps (WHERE the model looks).
• CRD: Contrastive learning to maximize teacher-student mutual information.
• RKD: Match pairwise distances and angles (architecture-agnostic).
• Choice depends on: architecture similarity, task type, computational budget.
• Best results: combine logit + feature + relational losses.

Exercises

Exercise 1. For a teacher with layers $[512, 1024, 2048]$ and student with $[128, 256, 512]$: design the projection layers for FitNet matching.

Exercise 2. Compute the attention map (equation 2, $p=2$) for a $64 \times 8 \times 8$ feature map with known values.

Exercise 3. For CRD with $K=16$ negatives and $\tau=0.1$: compute the loss for a positive pair with similarity 0.9 and negatives with similarities $[0.1, 0.2, \ldots, 0.3]$.

Exercise 4. Prove that RKD angle loss is invariant to uniform scaling of features.

Exercise 5. Design a combined distillation loss for distilling ViT-L (teacher) into EfficientNet-B0 (student) for ImageNet. Specify which layers to match and the loss weights.