[Paper Review] Graph-based Knowledge Distillation by Multi-head Attention Network

Seunghyun Lee, Byung Cheol Song|arXiv (Cornell University)|Jul 4, 2019

Advanced Neural Network Applications31 references39 citations

TL;DR

The paper introduces MHGD, a graph-based knowledge distillation framework that uses multi-head attention to distill dataset-embedding knowledge from a teacher network to a student network, improving performance on CIFAR100 and TinyImageNet.

ABSTRACT

Knowledge distillation (KD) is a technique to derive optimal performance from a small student network (SN) by distilling knowledge of a large teacher network (TN) and transferring the distilled knowledge to the small SN. Since a role of convolutional neural network (CNN) in KD is to embed a dataset so as to perform a given task well, it is very important to acquire knowledge that considers intra-data relations. Conventional KD methods have concentrated on distilling knowledge in data units. To our knowledge, any KD methods for distilling information in dataset units have not yet been proposed. Therefore, this paper proposes a novel method that enables distillation of dataset-based knowledge from the TN using an attention network. The knowledge of the embedding procedure of the TN is distilled to graph by multi-head attention (MHA), and multi-task learning is performed to give relational inductive bias to the SN. The MHA can provide clear information about the source dataset, which can greatly improves the performance of the SN. Experimental results show that the proposed method is 7.05% higher than the SN alone for CIFAR100, which is 2.46% higher than the state-of-the-art.

Motivation & Objective

• Motivate the need to distill knowledge about intra-data relations, not just feature vectors, to improve dataset embedding in CNNs.
• Propose a graph-based distillation method using multi-head attention to capture dataset embedding procedures.
• Enable the student network to inherit relational inductive bias via a transfer task in a multi-task learning setup.

Proposed method

• Compress feature maps from two sensing points into feature vectors via KD-SVD.
• Compute relations between front-end and back-end feature vectors using a multi-head attention network (MHAN).
• Distill embedding knowledge by training multiple attention heads to produce graph-based relations.
• Transfer the distilled graph-based knowledge to the student via multi-task learning with a transfer loss.
• Use a smoothed attention map and KL-divergence between teacher and student graphs to guide learning.

Experimental results

Research questions

• RQ1 Can graph-based embedding knowledge captured by multi-head attention improve KD beyond traditional feature-vector based methods?
• RQ2 How does the number of attention heads affect the quality of distilled knowledge and SN performance?
• RQ3 Does multi-task learning with graph-based knowledge transfer provide advantages over initialization-based or single-task KD methods?

Key findings

• MHGD improves SN performance by up to about 7% on CIFAR100 and about 4% on TinyImageNet compared to the SN alone.
• MHGD outperforms KD-SVD and state-of-the-art methods in several setups (e.g., VGG and WResNet backbones).
• Increasing the number of attention heads generally improves performance up to a point, after which gains may saturate or degrade due to over-complexity.
• Multi-task learning with graph-based knowledge transfer maintains performance improvements during training, unlike some initialization-based KD methods.
• The method yields architecture-agnostic gains, enhancing SN performance across VGG, MobileNet, and ResNet backbones.