Gradient amplification for gradient matching based dataset distillation.

Abstract

Dataset distillation (DD) aims to construct a smaller dataset compared to the original cumbersome one. Models trained on both datasets are expected to achieve almost the same accuracy on the test set. Previous work using a gradient matching (GM) framework achieved suboptimal performance because it only matched the gradient information for correct labels, neglecting to account for the model's surprise regarding incorrect answers. In this paper, we aim to produce more informative gradient information during the matching process and present a novel framework for DD by leveraging label cycle shifting. Specifically, it involves using pre-trained neural networks to process mismatched image-label pairs, resulting in the generation of diverse and substantial gradients during the backpropagation of the cross-entropy loss. Furthermore, GM with larger gradients tends to converge more rapidly compared to conventional GM approaches, prompting us to propose an early exit mechanism. To enhance the performance further, we employ an ensemble approach by applying an exponential moving average to the distilled dataset and introduce distribution matching to the total matching function. We demonstrate that the model implicitly considers gradient experiences from past rounds, and we have delved into mechanisms where gradient matching and distribution matching mutually enhance each other. Our design can outperform most previous DD methods with fewer training iterations. Experiments on the benchmark datasets (CIFAR10, CIFAR100, TinyImageNet, and a subset of ImageNet) present the effectiveness of our method.