[Paper Review] Knowledge Distillation with Feature Maps for Image Classification

Wei‐Chun Chen, Chia-Che Chang+2 more|arXiv (Cornell University)|Dec 3, 2018

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes KDFM, a knowledge distillation method that improves student model performance by transferring feature maps from a teacher network using a shared classifier and generative adversarial networks (GANs). Experiments show KDFM enables a 4-layer CNN to match DenseNet-40's accuracy with 2x faster inference, and MobileNet to achieve 77.20% accuracy on CIFAR-100 while being 6x faster and half the size of DenseNet-100.

ABSTRACT

The model reduction problem that eases the computation costs and latency of complex deep learning architectures has received an increasing number of investigations owing to its importance in model deployment. One promising method is knowledge distillation (KD), which creates a fast-to-execute student model to mimic a large teacher network. In this paper, we propose a method, called KDFM (Knowledge Distillation with Feature Maps), which improves the effectiveness of KD by learning the feature maps from the teacher network. Two major techniques used in KDFM are shared classifier and generative adversarial network. Experimental results show that KDFM can use a four layers CNN to mimic DenseNet-40 and use MobileNet to mimic DenseNet-100. Both student networks have less than 1\% accuracy loss comparing to their teacher models for CIFAR-100 datasets. The student networks are 2-6 times faster than their teacher models for inference, and the model size of MobileNet is less than half of DenseNet-100's.

Motivation & Objective

• To address the challenge of deploying large, computationally expensive deep learning models on resource-constrained devices such as mobile and IoT systems.
• To improve knowledge distillation by leveraging intermediate feature maps, not just logits, for more effective knowledge transfer.
• To develop a method that enables simple, shallow, and wide student networks to match the performance of complex, deep teacher networks like DenseNet.
• To enhance distillation effectiveness through adversarial training that aligns student feature maps with those of the teacher.
• To demonstrate that the proposed method generalizes across different architectures and datasets, including CIFAR-100 and ImageNet.

Proposed method

• KDFM uses a shared classifier between the teacher and student networks to align feature representations during training.
• It applies a generative adversarial network (GAN) framework where the student acts as a generator and the teacher's feature maps are the real data.
• The discriminator in the GAN is trained to distinguish between real feature maps (from the teacher) and generated feature maps (from the student), forcing the student to produce more realistic feature maps.
• The loss function combines knowledge distillation loss on logits and adversarial loss on feature maps to improve generalization.
• The method is applied to student networks that are shallow and wide (e.g., 4-8 convolutional layers), which are more amenable to hardware acceleration than deep, thin networks.
• The framework is evaluated on multiple teacher models (DenseNet-40, DenseNet-100, CondenseNet-86, ResNet-152) and student architectures (simple CNNs, MobileNet).

Experimental results

Research questions

• RQ1 Can distilling intermediate feature maps, rather than just logits, significantly improve the performance of student models in knowledge distillation?
• RQ2 How effective is using a GAN-based framework to align feature maps between teacher and student networks in image classification?
• RQ3 Can a simple, shallow, and wide student network achieve performance comparable to a deep, complex teacher network like DenseNet-100?
• RQ4 Does the combination of shared classifier and adversarial training on feature maps lead to better accuracy and faster inference than standard KD methods?
• RQ5 Can KDFM be generalized across different datasets (CIFAR-100, ImageNet) and model architectures?

Key findings

• A 4-layer CNN student model achieved 74.23% accuracy on CIFAR-100 when mimicking DenseNet-40, with less than 1% accuracy drop and 2x faster inference time.
• MobileNet, when trained with KDFM to mimic DenseNet-100, achieved 77.20% accuracy—only 0.74% below the teacher—while being 6x faster and using less than half the model size.
• On the ImageNet dataset, MobileNet v2 trained with KDFM achieved 71.82% top-1 accuracy, outperforming standard KD (70.16%) and the baseline (68.01%).
• The ablation study showed that removing the adversarial loss component reduced accuracy to 71.32%, confirming its importance in improving feature map alignment.
• Even simple CNNs with 20.2M or 28.1M parameters achieved 74.36% and 75.25% accuracy respectively when trained with KDFM, outperforming the baseline without KD.
• The method demonstrated generalization across architectures: CondenseNet-86 student models trained with KDFM improved from 74.13% to 75.01% accuracy on CIFAR-100.