[Paper Review] Knowledge Distillation by On-the-Fly Native Ensemble

Lan Xu, Xiatian Zhu+1 more|arXiv (Cornell University)|Jun 12, 2018

Advanced Neural Network Applications19 references271 citations

TL;DR

ONE trains a single multi-branch network that creates an on-the-fly teacher from its branches to perform online distillation, boosting generalisation without extra test-time cost.

ABSTRACT

Knowledge distillation is effective to train small and generalisable network models for meeting the low-memory and fast running requirements. Existing offline distillation methods rely on a strong pre-trained teacher, which enables favourable knowledge discovery and transfer but requires a complex two-phase training procedure. Online counterparts address this limitation at the price of lacking a highcapacity teacher. In this work, we present an On-the-fly Native Ensemble (ONE) strategy for one-stage online distillation. Specifically, ONE trains only a single multi-branch network while simultaneously establishing a strong teacher on-the- fly to enhance the learning of target network. Extensive evaluations show that ONE improves the generalisation performance a variety of deep neural networks more significantly than alternative methods on four image classification dataset: CIFAR10, CIFAR100, SVHN, and ImageNet, whilst having the computational efficiency advantages.

Motivation & Objective

• Motivate compact models that generalise well under low-memory and fast-run constraints.
• Eliminate the need for a separate pre-trained teacher by constructing a strong on-the-fly teacher during training.
• Improve target network performance through online distillation within a single training phase.
• Demonstrate efficiency and generalisation gains across standard image classification benchmarks.

Proposed method

• Augment the target network with m auxiliary branches that share low-level layers.
• Construct an on-the-fly ONE teacher by ensembling the branches via a gating component.
• Train each branch with ground-truth cross-entropy loss and a distillation loss from the ONE teacher.
• Compute soft targets with temperature T for both branches and the teacher to guide learning.
• Use a closed-loop knowledge distillation objective combining cross-entropy and KL-divergence with appropriate scaling by T^2.
• During testing, remove auxiliary branches to recover the original single-branch model and retain same test cost; optionally deploy ONE as an ensemble (ONE-E).

Experimental results

Research questions

• RQ1 Can online distillation with an on-the-fly teacher improve generalisation for networks of varying capacity?
• RQ2 Does a multi-branch single-model design provide competitive or superior performance to offline teacher-student distillation and online peer-teaching methods?
• RQ3 What is the effect of gating-based ensemble of branches on teacher quality and downstream student learning?
• RQ4 How does ONE affect training cost, test cost, and scalability on large-scale data (e.g., ImageNet)?

Key findings

• ONE improves generalisation across CIFAR-10/100, SVHN, and ImageNet for several backbones.
• Smaller models gain larger relative improvements from ONE than larger models.
• Training with ONE incurs lower cost than offline teachers or multi-network online methods while achieving competitive or better accuracy.
• The ONE teacher benefits from distillation, and the target model approaches the teacher’s performance over training.
• Ensembling via gating (ONE-E) can yield the best test performance among compared configurations with favorable cost trade-offs.