View all files | ||||
Official PyTorch implementation of "Fair Dataset Distillation via Cross-Group Barycenter Alignment" (ICML 2026).
COBRA distills a small synthetic training set whose class-conditional representation matches the uniform Wasserstein-style barycenter of the sensitive subgroups, instead of the (biased) population mean used by vanilla distillation. The resulting synthetic data trains downstream classifiers that are simultaneously accurate and fair under Equalized Odds, across four distillation backbones and seven biased benchmarks.
Each train_*.py exposes the same CLI surface, so switching backbones means switching the script name. The fairness behaviour is controlled by a single flag, --mode:
| vanilla | Original distillation loss (population mean) | DC / DM / IDC / CAFE |
| fairdd | FairDD — per-subgroup loss, summed independently | DC / DM / IDC |
| cobra | COBRA — match synthetic mean to subgroup barycenter | DC / DM / IDC / CAFE |
Tested with Python 3.9 / 3.10, PyTorch ≥ 1.13 (CUDA 11.8), and a single NVIDIA A100. CPU runs are supported but slow.
All datasets are loaded automatically from --data_path (default ./data). The first run will download what is downloadable; UTKFace and BFFHQ require a manual one-time download (see data_handler/utkface.py and data_handler/bffhq.py for the expected directory structure).
| CIFAR10_S_90 | CIFAR-10-S | colour skew (0.9) |
| Colored_MNIST_foreground | Colored MNIST (FG) | digit colour |
| Colored_MNIST_background | Colored MNIST (BG) | background colour |
| Colored_FashionMNIST_foreground | Colored F-MNIST | foreground colour |
| Colored_FashionMNIST_background | Colored F-MNIST | background colour |
| UTKface | UTKFace | gender |
| BFFHQ | BFFHQ | age |
The shell scripts under scripts/ accept three positional arguments: <dataset> <ipc> <mode>.
To reproduce the entire main table at ipc=10 (serial, ~1 GPU-day):
Each run writes a checkpoint and a visualisation grid under --save_path (default ./result/):
The checkpoint stores (synthetic_images, synthetic_labels) plus the list of accuracies across evaluation seeds.
Common flags for every train_*.py:
| --mode | vanilla | vanilla / fairdd / cobra (CAFE: vanilla / cobra) |
| --dataset | CIFAR10_S_90 | See dataset table above |
| --ipc | 10 | Images per class |
| --model | ConvNet | ConvNet / AlexNet / VGG11 / ResNet18 |
| --num_exp | 1 | # of independent distillation runs |
| --num_eval | 5 (DC/DM) | # of evaluation seeds per run |
| --Iteration | varies | Outer distillation iterations |
| --lr_img | varies | Synthetic-data learning rate |
| --lr_net | 0.01 | Inner-network learning rate |
| --batch_real | varies | Real-data batch size per class |
| --init | real | real (init from real samples) or noise |
| --data_path | data | Dataset root |
| --save_path | result | Output root |
| --seed | 42 | Top-level RNG seed |
Backbone-specific flags:
Run any script with --help for the full list.
Given a class c and a sensitive attribute taking values g ∈ G, COBRA computes per-subgroup feature means μ_{c,g} = E[ φ(x) | y=c, s=g ] through the embedding network φ, then sets the target for the synthetic batch of class c to the uniform barycenter b_c = (1/|G|) Σ_g μ_{c,g}. The distillation loss aligns the synthetic feature mean to b_c (DM / CAFE) or aligns gradient signals computed against b_c (DC / IDC). For DM, the embedding network is briefly warm-started on the current synthetic set so the target reflects an informative representation. Because the barycenter weights subgroups uniformly regardless of their prior frequency, the resulting synthetic data is balanced in representation space rather than in input space, which empirically yields lower Equalized-Odds gaps without sacrificing accuracy.
This codebase builds on the public implementations of DC/DM, CAFE, and IDC. We thank the original authors for releasing their code.