Dataset Distillation by Matching Training Trajectories

What came before

Responds to dd-2018's expensive, unstable full backprop-through-training and to dc-2021/dsa-2021's single-step gradient matching, which the paper argues drifts from the true training trajectory because it only matches one local gradient direction rather than where training actually ends up many steps later; also compares against kip-2021's infinite-width kernel approach.

The problem

Single-step gradient matching is a short-horizon proxy: matching one step's gradient direction does not guarantee the resulting long trajectory of parameters resembles real training, so accumulated drift over many steps caps accuracy. Full bi-level BPTT (dd-2018) would fix this but is too memory-hungry and unstable to unroll over realistic numbers of steps.

The idea

Precompute many "expert" trajectories from ordinary training runs on real data, then optimize the synthetic images so that a few steps of training on them move a network's parameters to (approximately) the same place an expert's parameters reach many steps later, starting from the same point on the expert trajectory -- matching long-range parameter movement instead of single-step gradients or the full trajectory.

How it works

Optimizes a per-class synthetic image set (pixels) plus a learned synthetic-step learning rate. In the loop: expert trajectories $\{\theta^*_t\}$ are trained once on real data and cached to disk (not part of the stored distilled artifact). Each distillation step samples a start point $\theta^*_t$ from a random cached expert and epoch, initializes a student at $\hat\theta_t=\theta^*_t$, takes $N$ differentiable-augmentation SGD steps on the synthetic images to reach $\hat\theta_{t+N}$, and minimizes the normalized squared distance $\mathcal{L}=\|\hat\theta_{t+N}-\theta^*_{t+M}\|_2^2 / \|\theta^*_t-\theta^*_{t+M}\|_2^2$ to the expert's parameters $M>N$ real-training epochs later (Eq. 2), backpropagating through the $N$ student steps into the images. A per-class batching scheme (Eq. 3) samples a fresh mini-batch of synthetic images at every inner step to control memory when the synthetic set is large. Labels are fixed/hard.

Evidence

Table 1 (ConvNet-128, ZCA whitening, hard labels): CIFAR-10 IPC1 46.3+/-0.8%, IPC10 65.3+/-0.7%, IPC50 71.6+/-0.2% vs. full-data 84.8+/-0.1%; CIFAR-100 IPC1 24.3+/-0.3%, IPC10 40.1+/-0.4%, IPC50 47.7+/-0.2% vs. full-data 56.2+/-0.3%; Tiny-ImageNet (64x64) IPC1 8.8+/-0.3%, IPC10 23.2+/-0.2%, IPC50 28.0+/-0.3% vs. full-data 37.6+/-0.4%, beating DC/DSA/DM/CAFE/CAFE+DSA at every matched IPC. Table 2 vs. kip-2021 on the same 128-width ConvNet: CIFAR-10 IPC1 46.3% vs. KIP 38.3% (49.9% on a 1024-width net), IPC10 65.3% vs. 57.6%, IPC50 71.5% vs. 65.8%; CIFAR-100 IPC1 24.3% vs. 18.2%, IPC10 39.4% vs. 32.8%. Table 4, 128x128 ImageNet subsets, IPC1/10 vs. full data: ImageNette 47.7/63.0% (full 87.4%), ImageWoof 28.6/35.8% (full 67.0%), ImageFruit 26.6/40.3% (full 63.9%), ImageMeow 30.7/40.4% (full 66.7%), ImageSquawk 39.4/52.3% (full 87.5%), ImageYellow 45.2/60.0% (full 84.4%) -- the first trajectory-matching result above 64x64 resolution. Cross-architecture, CIFAR-10 IPC10, ConvNet-trained data evaluated on unseen nets: ResNet 46.4%, VGG 50.3%, AlexNet 34.2% (vs. DSA 42.8/43.2/35.9% and KIP 36.8/42.1/24.4%). Cost: expert trajectories took ~3 sec/epoch on CIFAR (8 GPU-hours total for 200 CIFAR experts, ~60MB/expert) and ~11 sec/epoch on the ImageNet subsets (15 GPU-hours for 100 experts, ~120MB/expert); the largest single distillation run used up to 144GB VRAM over 6xRTX6000 GPUs. Ablation (Table 5, CIFAR-100 IPC1): removing ZCA whitening, real-image initialization, or the learned synthetic step size each degrades accuracy below the full method's 24.3%, isolating the combination as necessary; the paper separately shows varying the matched expert-epoch offset $M$ and student steps $N$ both materially change accuracy, and that cached (precomputed) experts match live-training experts in accuracy at far lower per-run cost.

Limitations

States explicitly that pre-computing and storing expert trajectories is a real up-front cost, both in compute (8-15 GPU-hours per dataset's expert pool) and disk (tens to hundreds of MB per expert), amortized across but not eliminated by reuse. Backpropagating through the $N$ synthetic-training steps is memory-intensive and scales with $N$ and with the number of images distilled jointly, requiring a per-class or per-batch chunking workaround (Eq. 3) and, in the largest configurations, up to 144GB of GPU memory (a bottleneck tesla-2023 directly targets next). Only hard labels and only ConvNet-family experts are used; no ImageNet-1K result (reached at 128x128 crops of 10-class subsets only); no soft-label or teacher-relabeling scheme.