[Paper | BibTex | 🤗Dataset | 📂Logs]
Official Implementation for "Unifying Dataset Pruning and Distillation for Efficient Large-scale Compression".
Lingao Xiao, Songhua Liu, Yang He*, Xinchao Wang
Abstract: Dataset pruning (DP) and dataset distillation (DD) fundamentally differ in their outputs: DP selects original image subsets, while DD generates synthetic images. Recently, DD’s increasing reliance on original images suggests a convergence of the two directions. To investigate this convergence trend, we propose a unified dataset compression (DC) benchmark. This benchmark reveals an interesting trade-off for soft-label-DD: while soft labels provide valuable information, they can make the distillation process less essential, as distilled images may not always outperform random subsets. In addition, the benchmark reveals that in current stages, dataset pruning outperforms dataset distillation at small dataset sizes. Given these observations, we explore hard-labelDC as a complementary approach that emphasizes image quality while offering substantial storage efficiency. Our PCA (Prune, Combine, and Augment) is the first framework that does not rely on soft labels but instead focuses on image quality. (1) “P” means selecting easy samples based on dataset pruning metrics, (2) “C” indicates combining these samples effectively, and (3) “A” is to apply constrained image augmentation during training. Extensive experiments validate that PCA significantly outperforms existing DD and DP methods without soft labels.
- release large-scale benchmark
- release SOTA datasets
- release PCA framework
- release PCA datasets
*Note: for soft label benchmark, we use fast evaluation code without relabeling.
Datasets (🤗Hugging Face)
SOTA datasets used in our experiments are available at 🤗Hugging Face. We have preprocessed all images into fixed 224x224 resolutioins and creates the datasets for a fair storage comparison.
| Type | DD |
DP |
|---|---|---|
| Explain | Dataset Distillation | Dataset Pruning |
| Method | Type | Venue | Dataset Key | Avaiable IPCs |
|---|---|---|---|---|
| random | - | - | he-yang/2025-rethinkdc-imagenet-random-ipc-[IPC] |
[1,10,20,50,100,200] |
| SRe2L | DD |
NeurIPS'23 | he-yang/2025-rethinkdc-imagenet-sre2l-ipc-[IPC] |
[10,50,100] |
| CDA | DD |
TMLR'24 | he-yang/2025-rethinkdc-imagenet-cda-ipc-[IPC] |
[10,50,100] |
| G-VBSM | DD |
CVPR'24 | he-yang/2025-rethinkdc-imagenet-gvbsm-ipc-[IPC] |
[10,50,100] |
| LPLD | DD |
NeurIPS'24 | he-yang/2025-rethinkdc-imagenet-lpld-ipc-[IPC] |
[10,50,100] |
| RDED | DD |
CVPR'24 | he-yang/2025-rethinkdc-imagenet-rded-ipc-[IPC] |
[10,50,100] |
| DWA | DD |
NeurIPS'24 | he-yang/2025-rethinkdc-imagenet-dwa-ipc-[IPC] |
[10,50,100] |
| Forgetting | DP |
ICLR'19 | he-yang/2025-rethinkdc-imagenet-forgetting-ipc-[IPC] |
[10,50,100] |
| EL2N | DP |
NeurIPS'21 | he-yang/2025-rethinkdc-imagenet-el2n-ipc-[IPC] |
[10,50,100] |
| AUM | DP |
NeurIPS'20 | he-yang/2025-rethinkdc-imagenet-aum-ipc-[IPC] |
[10,50,100] |
| CCS | DP |
ICLR'23 | he-yang/2025-rethinkdc-imagenet-ccs-ipc-[IPC] |
[10,50,100] |
1. Install from pip (tested on python=3.12)
pip install rethinkdc2. Or install from source
Step 1: Clone Repo,
git clone https://github.com/ArmandXiao/Unifying-Dataset-Pruning-and-Distillation.git
cd Rethinking-Dataset-CompressionStep 2: Create Environment,
conda env create -f environment.yml
conda activate rethinkdcStep 3: Install Benchmark,
make build
make install1. Prepare ImageNet Validation Folder:
# download and prepare ImageNet Val (skip if you have)
wget -qO- https://github.com/ArmandXiao/Unifying-Dataset-Pruning-and-Distillation/script/download_val.sh | bash
# set environment (IMPORTANT!)
export IMAGENET_VAL_DIR="Your ImageNet Val Path"2. Hyper-parameter for "rethinkdc"
rethinkdc --help📘 Manual
Rethinking Large-scale Dataset Compression
usage: rethinkdc [-h] [--soft | --hard | --yaml YAML] [--batch-size BATCH_SIZE] [--gradient-accumulation-steps GRADIENT_ACCUMULATION_STEPS] [-j WORKERS]
[--val-dir VAL_DIR] [--output-dir OUTPUT_DIR] [--hf-cache-dir HF_CACHE_DIR] [--mode MODE] [--cos] [--adamw-lr ADAMW_LR]
[--adamw-weight-decay ADAMW_WEIGHT_DECAY] [--sgd-setting] [--hard-label] [--start-epoch START_EPOCH] [--epochs EPOCHS] [--model MODEL]
[--teacher-model TEACHER_MODEL] [-T TEMPERATURE] [--mix-type MIX_TYPE] [--mixup MIXUP] [--cutmix CUTMIX] [--ipc IPC] [--wandb-project WANDB_PROJECT]
[--wandb-api-key WANDB_API_KEY]
Example Usage:
rethinkdc he-yang/2025-rethinkdc-imagenet-random-ipc-10 --soft --ipc 10 --output-dir ./random_ipc10_soft
options
-h, --help ┃ show this help message and exit ┃ str ┃ ==SUPPRESS==
Configuration Options
--soft ┃ Use standard_soft_config.yaml (Example: rethinkdc PATH ┃ str ┃ False
┃ --soft) ┃ ┃
--hard ┃ Use standard_hard_config.yaml (Example: rethinkdc PATH ┃ str ┃ False
┃ --hard) ┃ ┃
--yaml ┃ Custom config file (Exmpale: rethinkdc ┃ str ┃
┃ YOUR_PATH_TO_CONFIG.yaml) ┃ ┃
Data Options
train_dir ┃ path to training dataset or huggingface dataset key ┃ str ┃
--batch-size ┃ batch size ┃ int ┃ 1024
--gradient-accumul… ┃ gradient accumulation steps for small gpu memory ┃ int ┃ 1
-j, --workers ┃ number of data loading workers ┃ int ┃ 16
--val-dir ┃ path to validation dataset ┃ str ┃ /path/to/imagenet/val
--output-dir ┃ path to output dir ┃ str ┃ ./save/1024
--hf-cache-dir ┃ cache dir for huggingface dataset ┃ str ┃ ./hf_cache
--mode ┃ mode for training ┃ str ┃ fkd_save
Training Options
--cos ┃ cosine lr scheduler ┃ str ┃ False
--adamw-lr ┃ adamw learning rate ┃ float ┃ 0.001
--adamw-weight-dec… ┃ adamw weight decay ┃ float ┃ 0.01
--sgd-setting ┃ using sgd evaluation settting (lr=0.1, scheduler=cos) ┃ str ┃ False
--hard-label ┃ use hard label ┃ str ┃ False
--start-epoch ┃ start epoch ┃ int ┃ 0
--epochs ┃ total epoch ┃ int ┃ 300
Model Options
--model ┃ student model name ┃ str ┃ resnet18
--teacher-model ┃ teacher model name ┃ str ┃
-T, --temperature ┃ temperature for distillation loss ┃ float ┃ 3.0
Mixup/CutMix Options
--mix-type ┃ choices in {mixup, cutmix, None} ┃ str ┃
--mixup ┃ mixup alpha, mixup enabled if > 0. (default: 0.8) ┃ float ┃ 0.8
--cutmix ┃ cutmix alpha, cutmix enabled if > 0. (default: 1.0) ┃ float ┃ 1.0
--ipc ┃ number of images per class ┃ int ┃ 50
Wandb Options
--wandb-project ┃ wandb project name ┃ str ┃ Temperature
--wandb-api-key ┃ wandb api key ┃ str ┃
For more information, please visit the project repository: https://github.com/ArmandXiao/Unifying-Dataset-Pruning-and-Distillation
3. Example Usage (more examples can be found in folder script):
rethinkdc [YOUR_PATH_TO_DATASET] [*ARGS]
# example (test random subset)
rethinkdc he-yang/2025-rethinkdc-imagenet-random-ipc-10 --soft --ipc 10 --output-dir ./random_ipc10_softMain Table Result (📂Google Drive)
Logs for main tables are results provided in google drive for reference.
| Table | Explanation |
|---|---|
| Table 3 | Random baselines in soft label setting with standard evaluation |
| Table 4 & Table 18 | SOTA methods in soft label setting with std |
| Table 5 & Table 19 | SOTA methods in hard label setting with std |
| Table 6 | SOTA Pruning Rules |
| Table 7 | Ablation Study of PCA |
| Table 8 | Cross-architecture Performance of PCA |
| Table 12 & Table 22 | Regularization-based Data Augmentation |
| Table 20 | Pure Noise as Input |
| Table 24 | PCA using Different Pruning Methods |
- https://github.com/VILA-Lab/SRe2L
- https://github.com/he-y/soft-label-pruning-for-dataset-distillation
- https://github.com/haizhongzheng/Coverage-centric-coreset-selection
- https://github.com/NUS-HPC-AI-Lab/DD-Ranking
@inproceedings{xiaounifying,
title={Unifying Dataset Pruning and Distillation for Efficient Large-scale Compression},
author={Xiao, Lingao and Liu, Songhua and He, Yang and Wang, Xinchao},
booktitle={Forty-third International Conference on Machine Learning}
}