Wasserstein-Based Feature Map Knowledge Transfer to Improve the Performance of Small Deep Neural Networks

Abstract

The Convolution Neural Network (CNN) is a class of Deep Neural Networks, specifically used for computer vision applications. The CNN architecture consists of several hidden layers. This hidden layer information, which is referred to as a feature map, contains rich spatial and semantic information about the input data. A few techniques like Knowledge Distillation, Deep Mutual Learning, and Adversarial Deep Mutual Learning have been introduced to transfer the knowledge of a large DNN to a smaller DNN for improving its prediction accuracy. However, the existing methods do not effectively use feature maps as a source of information that can be exploited to improve the performance of the small DNN. We propose an adversarial learning-based approach that consists of a simple generator and discriminator network for transferring feature map information from a large pre-trained DNN to a smaller DNN using the Wasserstein metric. Our approach helps the smaller DNNs generate feature maps similar to a large pre-trained DNN, thereby improving the accuracy and generalization ability. Our experiments show that a variety of small DNN networks benefit from the proposed approach and achieve compelling results on the CIFAR-100 and CIFAR-10 benchmark datasets. Moreover, experimental results on the MobileNet architecture illustrate that the proposed approach is particularly effective for relatively small DNN networks.