Rethinking Dataset Quantization: Efficient Coreset Selection via Semantically-Aware Data Augmentation

Code release for the TMLR 2026 paper "Rethinking Dataset Quantization: Efficient Coreset Selection via Semantically-Aware Data Augmentation."

Overview

This codebase implements DQ_v2, a coreset-selection pipeline that improves the original Dataset Quantization (DQ, Zhou et al., NeurIPS 2023) by replacing its MAE pixel-reconstruction step with a Semantically-aware Data Augmentation (SDA) step. A randomly initialized CNN produces a spatial foreground mask for each training image; the mask is used to crop a semantically focused SDA view, and the original images plus their SDA views form a 1.5x candidate pool. The pool is partitioned into bins via GraphCut submodular maximization and sampled at the target ratio. Across CIFAR-10, ImageNette, ImageNet, CUB-200, Food-101, TinyImageNet, and a COVID-19 chest X-ray benchmark, DQ_v2 matches or exceeds DQ while removing the dependence on a pre-trained MAE encoder.

Repository structure

release/
├── README.md                    This file
├── REPRODUCE.md                 Detailed figure / table to script mapping
├── requirements.txt
├── sda/                         SDA pipeline (mask generation, SDA view rendering, mixing)
│   ├── tools/                   Per-dataset mask generators (CIFAR, ImageNet, CUB, Food, COVID)
│   ├── moco/                    Dynamic transform datasets used during mask generation
│   ├── save_image_sda_parallel.py
│   └── mix_original_and_sda{,_parallel}.py
├── scripts_imagenet/            ImageNet-scale optimized variants (batched mask cache, threaded SDA)
├── benchmarks/                  End-to-end drivers, one per dataset
│   ├── benchmark_cifar10.py
│   ├── benchmark_imagenette_vit.py
│   ├── benchmark_cub200.py
│   ├── benchmark_food101.py
│   ├── benchmark_food101_server.py
│   ├── benchmark_efficiency.py
│   ├── covid19/
│   │   ├── benchmark_covid19.py
│   │   ├── covid_resnet18_baseline.py
│   │   └── validate_covid19_random_subset.py
│   └── imagenet_c/
│       └── imagenet_c_complete_pipeline.py
├── ablations/                   Ablation drivers and analysis utilities
│   └── run_ablations.py
├── experiments/                 Rebuttal-era multi-seed / multi-GPU drivers and training helpers
│   ├── run_cub_fair224_*.py
│   ├── run_food_train_only_multigpu.py
│   ├── run_point3_multidataset.py
│   ├── run_imagenet_ratio20_vitb*.py
│   ├── run_tinyim_point3_randinit.py
│   ├── train_subset_cls*.py
│   └── fast_train_food101.py
└── third_party/
    └── dataset_quantization/    Fork of upstream DQ (NeurIPS 2023) with our patches
        ├── quantize_sample.py
        ├── quantize_bin.py
        ├── quantize_pixel.py
        ├── mae_models.py
        ├── validate_{cifar,imagenette,imagenette_seed_vit,COVID-19}.py
        ├── dq/, util/, pytorch_image_models/
        ├── pretrained/          (empty; place MAE checkpoint here)
        └── README.md            Upstream DQ README, preserved for attribution

Quick start

The COVID-19 chest X-ray benchmark is the smallest end-to-end pipeline in this release and the recommended smoke test. It exercises the full chain (mask generation -> SDA rendering -> mixed candidate pool -> GraphCut binning -> ratio sampling -> ResNet-18 training).

1. Install dependencies (Python 3.9, CUDA 11.7 tested):

   pip install -r requirements.txt

2. Prepare data. The COVID-19 dataset is the public COVIDx CXR-3 split; see REPRODUCE.md for the exact source and directory layout used in the paper.

3. (Optional, only for the DQ baseline branch) Download the MAE checkpoint mae_visualize_vit_large_ganloss.pth from the official MAE release and place it under third_party/dataset_quantization/pretrained/. DQ_v2 itself does not require this checkpoint.

4. Run the end-to-end COVID-19 benchmark:

   python benchmarks/covid19/benchmark_covid19.py \
       --data /path/to/covidx_cxr3 \
       --mae-checkpoint third_party/dataset_quantization/pretrained/mae_visualize_vit_large_ganloss.pth \
       --ratio 0.3

Each driver exposes a benchmark_dq() branch (original DQ with MAE) and a benchmark_dq_v2() branch (our method); flags are documented in the individual scripts. Larger-scale datasets follow the same pattern but should be launched with the scripts_imagenet/ variants when the candidate pool no longer fits in memory.

Reproducing the paper

Detailed per-figure and per-table instructions, including all cross-check warnings, live in REPRODUCE.md. The mapping below is a high-level index.

Script	Paper reference
benchmarks/benchmark_cifar10.py	Figure 1 (CIFAR-10 panel)
benchmarks/benchmark_imagenette_vit.py	Section 6.2, ImageNette ViT cross-architecture
benchmarks/benchmark_cub200.py	Section 6.2, CUB-200
benchmarks/benchmark_food101.py	Section 6.2, Food-101
benchmarks/covid19/benchmark_covid19.py	Table 2, COVID-19 X-ray
benchmarks/benchmark_efficiency.py	Table 1, runtime / memory
benchmarks/imagenet_c/imagenet_c_complete_pipeline.py	ImageNet-C mean Corruption Error
ablations/run_ablations.py	Table 3, Figures 5-7
experiments/run_*.py	Multi-seed runs reported in the rebuttal

Citation

@article{liu2026rethinking,
  title={Rethinking Dataset Quantization: Efficient Coreset Selection via Semantically-Aware Data Augmentation},
  author={Liu, Yangze and Liu, Hong},
  journal={Transactions on Machine Learning Research},
  year={2026},
  url={https://openreview.net/forum?id=Mb2nn1yx66}
}

Attribution

third_party/dataset_quantization/ is a fork of the official Dataset Quantization codebase released with Zhou et al., Dataset Quantization, NeurIPS 2023, with patches applied for compatibility with the DQ_v2 pipeline. The original upstream README is preserved at third_party/dataset_quantization/README.md. pytorch_image_models/ inside that fork is itself a vendored copy of timm and retains its own license.

License

Code in this repository (outside third_party/) is released under the MIT License. The vendored upstream DQ fork retains its original license; see third_party/dataset_quantization/ for details.