[Paper Review] Momentum Adversarial Distillation: Handling Large Distribution Shifts in Data-Free Knowledge Distillation

Kien Do, Hung Lê+6 more|arXiv (Cornell University)|Sep 21, 2022

COVID-19 diagnosis using AI18 citations

TL;DR

This paper proposes Momentum Adversarial Distillation (MAD), a method that mitigates large distribution shifts in data-free knowledge distillation by maintaining an exponential moving average (EMA) copy of the generator to stabilize knowledge transfer. By training the student on synthetic data from both the current and EMA generator, MAD reduces catastrophic forgetting and achieves state-of-the-art performance on benchmarks including ImageNet and Places365.

ABSTRACT

Data-free Knowledge Distillation (DFKD) has attracted attention recently thanks to its appealing capability of transferring knowledge from a teacher network to a student network without using training data. The main idea is to use a generator to synthesize data for training the student. As the generator gets updated, the distribution of synthetic data will change. Such distribution shift could be large if the generator and the student are trained adversarially, causing the student to forget the knowledge it acquired at previous steps. To alleviate this problem, we propose a simple yet effective method called Momentum Adversarial Distillation (MAD) which maintains an exponential moving average (EMA) copy of the generator and uses synthetic samples from both the generator and the EMA generator to train the student. Since the EMA generator can be considered as an ensemble of the generator's old versions and often undergoes a smaller change in updates compared to the generator, training on its synthetic samples can help the student recall the past knowledge and prevent the student from adapting too quickly to new updates of the generator. Our experiments on six benchmark datasets including big datasets like ImageNet and Places365 demonstrate the superior performance of MAD over competing methods for handling the large distribution shift problem. Our method also compares favorably to existing DFKD methods and even achieves state-of-the-art results in some cases.

Motivation & Objective

• To address the problem of large distribution shifts in adversarial data-free knowledge distillation (DFKD), where rapid generator updates cause student model forgetting.
• To improve knowledge transfer stability and performance in DFKD without access to original training data, especially on large-scale datasets like ImageNet.
• To reduce catastrophic forgetting in the student model during joint training with a generator by leveraging historical generator states.
• To enable effective DFKD on large datasets using a class-conditional generator with a novel loss that suppresses spurious samples.

Proposed method

• Maintain an exponential moving average (EMA) copy of the generator to preserve historical knowledge and reduce distribution shift.
• Train the student network on synthetic data generated by both the main generator and the EMA generator, using a combined loss function.
• Use a class-conditional generator that takes a noise vector and a class embedding as input, improving sample quality on large-scale datasets.
• Introduce a new loss function for the generator that suppresses spurious solutions, enhancing the quality of synthetic data.
• Control the relative contribution of EMA-generated and main generator-generated samples via learnable coefficients in the student loss.
• Apply momentum-based EMA updates with a decay hyperparameter α to balance stability and adaptability of the EMA generator.

Experimental results

Research questions

• RQ1 Can maintaining an EMA copy of the generator reduce catastrophic forgetting in data-free knowledge distillation under large distribution shifts?
• RQ2 How does the use of synthetic data from both the current and EMA generator affect student performance compared to using only the current generator?
• RQ3 Can a class-conditional generator with a novel loss function enable effective DFKD on large-scale datasets like ImageNet?
• RQ4 What is the optimal balance between current and EMA-generated data in the student training objective?
• RQ5 Does MAD outperform existing SOTA methods in data-free knowledge distillation, especially on challenging benchmarks?

Key findings

• MAD achieves state-of-the-art performance on several benchmarks, including ImageNet and Places365, outperforming existing data-free knowledge distillation methods.
• On CIFAR100, MAD achieves 64.05% test accuracy with λ₀ = λ₁ = 1.0, the optimal coefficient setting, and performance degrades when either coefficient is reduced to zero.
• The momentum hyperparameter α has a significant impact: performance peaks at α = 0.99 and drops sharply at α = 1.0, indicating that too slow generator updates harm learning.
• Increasing the number of student update steps improves performance, confirming that more training steps lead to better knowledge alignment with the teacher.
• The class-conditional generator with the new loss function enables stable training on large datasets, where unconditional generators often fail due to spurious samples.
• MAD reduces catastrophic forgetting by allowing the student to recall past knowledge via the EMA generator, which is less prone to rapid distribution shifts than the main generator.