[Paper Review] On Compressing U-net Using Knowledge Distillation

Karttikeya Mangalam, Mathieu Salzamann|arXiv (Cornell University)|Dec 1, 2018

Energy Efficient Wireless Sensor Networks11 references7 citations

TL;DR

This paper proposes a knowledge distillation approach to compress U-Net architectures by over 1000×, reducing parameters to 0.1% of the original size while maintaining near-identical segmentation accuracy. By integrating batch normalization and class re-weighting into the student network during distillation, the method overcomes training instability and class imbalance, enabling reliable compression where standard distillation fails.

ABSTRACT

We study the use of knowledge distillation to compress the U-net architecture. We show that, while standard distillation is not sufficient to reliably train a compressed U-net, introducing other regularization methods, such as batch normalization and class re-weighting, in knowledge distillation significantly improves the training process. This allows us to compress a U-net by over 1000x, i.e., to 0.1% of its original number of parameters, at a negligible decrease in performance.

Motivation & Objective

• To address the challenge of deploying large U-Net models in resource-constrained environments such as mobile devices.
• To investigate whether knowledge distillation can effectively compress U-Net architectures to minimal sizes.
• To identify and resolve failures of standard distillation in compressing small U-Net variants.
• To improve distillation reliability for small student networks by introducing architectural and loss modifications.

Proposed method

• Introduce batch normalization layers into every convolutional layer of the contracting path in the student U-Net to reduce internal covariate shift and stabilize training.
• Apply class re-weighting to the cross-entropy loss, assigning a foreground weight of 17.8 (based on background-to-foreground pixel ratio) to address class imbalance.
• Use mixed knowledge distillation with a temperature-scaled soft label loss from a 4-Unet teacher network, combining soft labels at T=5 with hard labels.
• Train the student network using both the distillation loss (soft labels) and the original cross-entropy loss (hard labels), then temperature-soften the final output.
• Use a 4-Unet trained from scratch with hard labels as the teacher model, and distill it into a 2-Unet with minimal parameters.
• Scale the soft label loss by T² during training to maintain consistent gradient magnitude across temperatures.

Experimental results

Research questions

• RQ1 Can standard knowledge distillation reliably compress a U-Net to less than 1% of its original parameter count?
• RQ2 Why does standard distillation fail when compressing U-Net to very small sizes, such as 2 initial channels?
• RQ3 How do batch normalization and class re-weighting improve the training of compressed U-Net student networks?
• RQ4 Can a 2-Unet achieve segmentation performance comparable to a 64-Unet when trained with distillation enhanced by regularization techniques?

Key findings

• A 4-Unet with only 122,394 parameters achieves a test loss of 0.0974 and IoU of 0.807, comparable to the original 64-Unet’s 0.804 IoU, demonstrating that parameter reduction without distillation is effective.
• Standard distillation (vanilla or mixed) fails to train a 2-Unet, as shown by test losses of 0.505–0.507 in Table 2, indicating poor generalization and instability.
• With batch normalization and class re-weighting, the 2-Unet achieves an IoU of 0.759 and cross-entropy loss of 0.135 using mixed distillation, outperforming the non-distilled 2-Unet (IoU 0.752, loss 0.134).
• The final 2-Unet model uses only 30,902 parameters, representing 0.1% of the original 64-Unet’s 31 million parameters, while maintaining near-identical performance.
• Training the 2-Unet without distillation but with modifications yields a test loss of 0.307, confirming that distillation remains essential for performance.
• The combination of batch normalization and class re-weighting enables successful distillation at low parameter counts, where standard distillation fails.