Going Beyond Feature Similarity: Effective Dataset Distillation based on Class-Aware Conditional Mutual Information

TL;DR

This paper introduces a class-aware conditional mutual information approach to improve dataset distillation, resulting in more learnable synthetic datasets that enhance training efficiency and performance.

Contribution

It proposes minimizing class-aware conditional mutual information to better regularize dataset distillation, addressing the difficulty of synthetic datasets for neural network learning.

Findings

Improved training efficiency with smaller synthetic datasets.

Enhanced performance of models trained on distilled datasets.

Versatile regularization method applicable to existing DD techniques.

Abstract

Dataset distillation (DD) aims to minimize the time and memory consumption needed for training deep neural networks on large datasets, by creating a smaller synthetic dataset that has similar performance to that of the full real dataset. However, current dataset distillation methods often result in synthetic datasets that are excessively difficult for networks to learn from, due to the compression of a substantial amount of information from the original data through metrics measuring feature similarity, e,g., distribution matching (DM). In this work, we introduce conditional mutual information (CMI) to assess the class-aware complexity of a dataset and propose a novel method by minimizing CMI. Specifically, we minimize the distillation loss while constraining the class-aware complexity of the synthetic dataset by minimizing its empirical CMI from the feature space of pre-trained networks, simultaneously. Conducting on a thorough set of experiments, we show that our method can serve as a general regularization method to existing DD methods and improve the performance and training efficiency.