Model Compression Using Optimal Transport

TL;DR

This paper introduces a novel model compression technique using optimal transport-based loss functions to improve knowledge distillation, resulting in better performance on image classification tasks.

Contribution

It proposes a new optimal transport-based loss for knowledge distillation, enhancing the alignment of student and teacher feature distributions.

Findings

Optimal transport loss performs comparably or better than existing loss functions.

The method improves student network performance on CIFAR-100, SVHN, and ImageNet.

The approach facilitates efficient deployment of deep models in resource-constrained environments.

Abstract

Model compression methods are important to allow for easier deployment of deep learning models in compute, memory and energy-constrained environments such as mobile phones. Knowledge distillation is a class of model compression algorithm where knowledge from a large teacher network is transferred to a smaller student network thereby improving the student's performance. In this paper, we show how optimal transport-based loss functions can be used for training a student network which encourages learning student network parameters that help bring the distribution of student features closer to that of the teacher features. We present image classification results on CIFAR-100, SVHN and ImageNet and show that the proposed optimal transport loss functions perform comparably to or better than other loss functions.