[Paper Review] Role-Wise Data Augmentation for Knowledge Distillation

Jie Fu, Xue Geng+7 more|arXiv (Cornell University)|Apr 19, 2020

Data Quality and Management43 references18 citations

TL;DR

This paper proposes role-wise data augmentation for knowledge distillation, where distinct data augmentation policies are learned specifically for the teacher and student networks to improve knowledge transfer. By tailoring training data to each model's learning capacity, the method achieves significant accuracy gains—up to 3% over vanilla training and 1.5% over standard KD—especially in low-precision and full-precision settings.

ABSTRACT

Knowledge Distillation (KD) is a common method for transferring the ``knowledge'' learned by one machine learning model (the extit{teacher}) into another model (the extit{student}), where typically, the teacher has a greater capacity (e.g., more parameters or higher bit-widths). To our knowledge, existing methods overlook the fact that although the student absorbs extra knowledge from the teacher, both models share the same input data -- and this data is the only medium by which the teacher's knowledge can be demonstrated. Due to the difference in model capacities, the student may not benefit fully from the same data points on which the teacher is trained. On the other hand, a human teacher may demonstrate a piece of knowledge with individualized examples adapted to a particular student, for instance, in terms of her cultural background and interests. Inspired by this behavior, we design data augmentation agents with distinct roles to facilitate knowledge distillation. Our data augmentation agents generate distinct training data for the teacher and student, respectively. We find empirically that specially tailored data points enable the teacher's knowledge to be demonstrated more effectively to the student. We compare our approach with existing KD methods on training popular neural architectures and demonstrate that role-wise data augmentation improves the effectiveness of KD over strong prior approaches. The code for reproducing our results can be found at https://github.com/bigaidream-projects/role-kd

Motivation & Objective

• Address the limitation in existing knowledge distillation methods that use identical training data for both teacher and student, despite differing model capacities.
• Investigate whether adaptive data augmentation tailored to each model's learning profile can improve knowledge transfer effectiveness.
• Improve knowledge distillation performance in low-bitwidth (quantized) neural networks, where standard KD often leads to significant accuracy degradation.
• Develop a two-stage framework that learns distinct data augmentation schedules for the teacher and student, independent of the distillation objective.
• Demonstrate that custom augmentation policies enhance student performance beyond standard KD, even when the teacher and student have similar architectures.

Proposed method

• Propose a two-stage training framework: first train the teacher with a policy-learned data augmentation schedule, then distill knowledge into the student using a separate, learned augmentation schedule.
• Use a differentiable architecture search (DARTS)-inspired approach to jointly optimize the distillation objective and the data augmentation policy for both teacher and student.
• Learn operation probability and magnitude parameters for data augmentation operations (e.g., rotation, cutout, mixup) over training epochs, forming an epoch-based schedule.
• Introduce a feature-based knowledge distillation variant that captures both intra- and inter-feature map relationships to improve knowledge transfer.
• Decouple the data augmentation policy learning from the distillation objective, enabling the method to be combined with any existing KD approach.
• Apply the method to both low-precision (quantized) and full-precision student networks, validating its generalization across settings.

Experimental results

Research questions

• RQ1 Can learning distinct data augmentation policies for the teacher and student improve knowledge distillation performance compared to using the same data for both?
• RQ2 Does role-wise data augmentation lead to better performance in low-precision model training, where standard KD often fails?
• RQ3 How do the learned augmentation schedules differ between teacher and student models, and what patterns emerge in their evolution over training?
• RQ4 Does the effectiveness of the method depend on the architectural similarity between teacher and student?
• RQ5 Can the proposed method be combined with existing KD techniques without compromising their performance?

Key findings

• The proposed role-wise data augmentation improves knowledge distillation performance by up to 3% over vanilla training and 1.5% over standard KD in full-precision settings.
• On CIFAR-100, the method achieves 76.19% accuracy with WRN-16-2 as student and WRN-28-10 as teacher, outperforming the vanilla baseline (72.68%) and standard KD (74.41%).
• When distilling into a more similar architecture (e.g., PyramidNet-200-240 to WRN-16-2), the method achieves 75.03% accuracy, showing stronger gains when teacher and student are architecturally aligned.
• The learned augmentation schedules for teacher and student differ significantly: the student’s schedule evolves to emphasize operations like rotation and cutout, which are challenging for low-precision models.
• The method generalizes well beyond quantization, improving performance on full-precision networks such as WRN-16-2, PreResNet32, and WRN-16-4 across multiple teacher architectures.
• The two-stage framework is orthogonal to existing KD methods, enabling seamless integration with soft label, feature map, or hybrid distillation strategies.