DiM: Distilling Dataset into Generative Model

What came before

Bilevel/surrogate DiI ("distill into images") methods — DC, DSA, DM, CAFE, MTT, FRePo, IDC, HaBa — all store the dataset as a fixed set of images, distilled with a matching objective computed through a 3-layer ConvNet. IT-GAN had already shown that optimizing a frozen GAN's latents with a matching loss beats pixel optimization, but its stored artifact is still one latent per image, so it must be re-optimized whenever the image budget (IPC) changes.

The problem

DiI methods have two coupled failures the paper attributes to storing images rather than a generative process: synthetic images distilled with a small ConvNet transfer poorly to larger unseen architectures (ResNet, VGG, DenseNet), and any change to the IPC/distillation-ratio budget requires re-running the whole distillation (8.3 GPU-hours for MTT, 18.2 for IDC to go from IPC=1 to 10 on CIFAR-10).

The idea

Store the dataset in the weights of a conditional generator instead of in images: train the generator so that, for a batch of real and generator-sampled images with matched labels, a randomly chosen classifier from a pool gives it the same logits as it gives the real images. Because the artifact is a generator, any number of samples at any budget can be drawn from it after one training run.

How it works

A conditional GAN $\mathcal{G}(\cdot;\beta)$ is first trained for $N{=}120$ epochs with the vanilla GAN loss $L_g$ alone (Adam, lr 1e-4) so it produces valid images; then a Logits Matching loss $L_m=\mathrm{MSE}(m(\mathcal{S}_{\text{DiM}}), m(I_r))$ is added, where $m$ is a model freshly sampled each epoch from a "models pool" of randomly-initialized architectures (ConvNet-3, ResNet-10, ResNet-18) trained on the target dataset, and $\mathcal{S}_{\text{DiM}}=\mathcal{G}([\mathcal{Z}\oplus\mathcal{Y}];\beta)$ is a batch of generated images conditioned on random noise $\mathcal{Z}$ and one-hot labels $\mathcal{Y}$. The generator is updated by $\beta \leftarrow \arg\min_\beta (L_g+\lambda L_m)$ with $\lambda{=}0.01$; nothing besides the models pool and the frozen GAN's own training data is used — no external pretrained teacher, no diffusion or LDM. Because logits are class-wise soft targets from real images, they double as soft supervision, but a fresh model is trained from these generated images at evaluation with hard one-hot labels (the paper never keeps or reuses the pool's logits as teacher labels at deployment). At deployment, images are sampled from the trained generator on the fly for any chosen Image-Number-Per-Class (INPC), the paper's budget unit since IPC undercounts a generator's capacity.

Evidence

MNIST/FashionMNIST/SVHN/CIFAR-10, INPC=1/10/50, ConvNet-3-trained pool, hard labels, Table 2: DiM beats Conditional GAN and every listed DiI/coreset baseline at every INPC, e.g. CIFAR-10 INPC=1 DiM 51.3% vs. MTT 46.3% vs. FRePo 46.8% vs. Conditional GAN 46.4% (real-data upper bound 84.8%); CIFAR-10 INPC=10 DiM 66.2% vs. FRePo 65.5%; SVHN INPC=1 DiM 80.9% vs. best DiI baseline ~58.5% (MTT), a ~10-22 point gain the abstract highlights. Cross-architecture (Table 7, CIFAR-10 INPC=10, distilled with ConvNet-3): DiM is the first method where ResNet-18 (69.2%) beats ConvNet-3 (66.2%) on its own synthetic data, vs. MTT dropping from 64.3% (ConvNet-3) to 46.4% (ResNet-18). Redeployment cost (Fig. 1d): re-deploying IPC 1->10 costs DiM ~0.1 GPU-hours vs. MTT 8.3 and IDC 18.2 (13x-160x cheaper), because DiM need only sample more images, not re-optimize. Key ablation (Table 4): replacing logits matching with feature matching or gradient matching inside the same DiM pipeline drops CIFAR-10 INPC=50 ResNet-18 accuracy from 75.0% to 73.8% (feature) or 71.7% (gradient), isolating logits matching (not the generator alone) as the mechanism; Table 6 further shows accuracy rises monotonically with the diversity of the models pool (single arch 73.1-74.2% avg vs. all three archs 75.0%/73.9% avg).

Limitations

Tested only at MNIST/FashionMNIST/SVHN/CIFAR-10 resolution; the conditional GAN is trained from scratch on the exact target dataset for each experiment, so (like IT-GAN) none of the reported gain isolates a generator not trained on the target distribution; evaluation is hard-label only, with the models-pool logits used purely as a matching signal during synthesis, not carried forward as soft teacher labels at deployment; the paper's own stated limitation is that DiM must generate images on the fly at deployment, adding a small but nonzero per-batch cost (Table 9: 2ms generation vs. 8-80ms training per batch, 2%-20% overhead); it does not test ImageNet-scale data (explicitly named as future work), and the "13x-160x" redeployment saving excludes the upfront GAN-pretraining cost (120 GAN epochs) from the comparison.