Code for the ICML 2023 paper "Dataset Distillation with Convexified Implicit Gradients"
Contact: Noel Loo
We propose a new dataset distillation algorithm using reparameterization and convexification of implicit gradients (RCIG), that substantially improves the state-of-the-art. To this end, we first formulate dataset distillation as a bi-level optimization problem. Then, we show how implicit gradients can be effectively used to compute meta-gradient updates. We further equip the algorithm with a convexified approximation that corresponds to learning on top of a frozen finite-width neural tangent kernel. Finally, we improve bias in implicit gradients by parameterizing the neural network to enable analytical computation of final-layer parameters given the body parameters. RCIG establishes the new state-of-the-art on a diverse series of dataset distillation tasks. Notably, with one image per class, on resized ImageNet, RCIG sees on average a 108% improvement over the previous state-of-the-art distillation algorithm. Similarly, we observed a 66% gain over SOTA on Tiny-ImageNet and 37% on CIFAR-100.
To distill on CIFAR-10 with 10 images/cls:
python3 distill_dataset.py --dataset_name cifar10 --n_images 10 --output_dir ./output_dir/ --max_steps 10000 --config_path ./configs_final/depth_3.txt --random_seed 0
To eval on CIFAR-10 with 10 images/cls (using pre-distilled datasets provided):
python3 eval.py --dataset_name cifar10 --checkpoint_path ./distilled_images_final/0/cifar10_10/checkpoint_10000 --config_path ./configs_final/depth_3.txt --random_seed 0
We include pre-distilled datasets for CIFAR-10 with 10 img/cls, CIFAR-10 1/cls, CIFAR-100 1/cls and Tiny-Imagenet 1/cls.
We recommend using the configs located in ./configs_final/ when distilling datasets as these were what we used to create the results in the paper. These contain all the hyperparameters used for distillation. For a description of these hyperparameters, see ./configs_final/example_config.txt
Specifically the configs we use are:
- For MNIST, and Fashion-MNIST, we use the config
depth_3_no_flip.txt - For CIFAR-10, CIFAR-100, and CUB-200, we use config
depth_3.txt - For CIFAR-100 with 50 ipc we use config
cifar100_50.txt - For Tiny-ImageNet 1 ipc, we use config
depth_4_200.txt - For Tiny-ImageNet 10 ipc and resized Imagenet, we use config
depth_4_big.txt - For ImageNette and ImageWoof, we use config
depth_5.txt
For the number of training iterations (i.e. the --max_steps argument, see section S.4.2 in the appendix)
Please cite our paper as
@misc{https://doi.org/10.48550/arxiv.2302.06755,
doi = {10.48550/ARXIV.2302.06755},
url = {https://arxiv.org/abs/2302.06755},
author = {Loo, Noel and Hasani, Ramin and Lechner, Mathias and Rus, Daniela},
title = {Dataset Distillation with Convexified Implicit Gradients},
publisher = {arXiv},
year = {2023},
}```