Scaling Up Dataset Distillation to ImageNet-1K with Constant Memory

What came before

Responds directly to mtt-2022, whose memory use scales with the number of unrolled student steps and which the paper shows was never actually run at ImageNet-1K scale; also compares against frepo-2022 and kip-2021 (kernel-closed-form), which reach ImageNet-1K but only up to 2 IPC, and notes frepo-2022's own finding that soft labels are what let a kernel method work with many classes.

The problem

Trajectory matching backpropagates through an unrolled sequence of student SGD steps, and naive implementations store activations for every step, so memory grows with the number of unrolled steps and the synthetic batch -- this is why no trajectory-matching method had scaled past small low-resolution subsets. Separately, and independently of memory, the paper finds MTT-style matching converges poorly on datasets with many classes (ImageNet-1K's 1,000), because one-hot hard labels give the matching loss no way to share information across classes.

The idea

(1) Reformulate the trajectory-matching gradient algebraically so it can be computed exactly, in constant memory, by processing the unrolled steps' gradient contributions sequentially rather than keeping the whole computation graph live -- an exact reformulation, not an approximation. (2) Assign each synthetic image a soft label for free, by reusing the same expert ("teacher") checkpoint already sampled for that iteration's trajectory-matching target to produce the label, instead of learning labels as extra optimized parameters.

How it works

Same trajectory-matching objective as mtt-2022 (student trained T steps on synthetic images from expert start point $\theta^*_t$, loss = normalized distance to expert target $\theta^*_{t+M}$), but the gradient of this loss with respect to each synthetic image is derived in closed form as a sum over the T steps (Eqs. 1-6) and accumulated one step at a time, so peak memory no longer scales with T. Soft Label Assignment (SLA) is train-free and hyperparameter-free: at each iteration, after sampling the target-step teacher checkpoint $\theta^*_{t+M}$, the synthetic images are passed through it to produce soft labels, which are then used (fixed for that iteration) in the matching loss; labels are not themselves optimized as free parameters. What is stored: per-class synthetic images (pixels) plus one soft-label vector per image (from the target-step teacher, not re-derived at every downstream training epoch); expert trajectories remain a synthesis-time-only cache.

Evidence

Table 1 (soft labels used on ImageNet-1K, ConvNet, * = soft label): CIFAR-10 IPC1/10/50 = 48.5+/-0.8 / 66.4+/-0.8 / 72.6+/-0.7% vs. full-data 86.0+/-0.1%; CIFAR-100 IPC1/10/50 = 24.8+/-0.4 / 41.7+/-0.3 / 47.9+/-0.3% vs. full-data 56.7+/-0.2%; ImageNet-1K IPC1/2/10/50 = 7.7+/-0.2* / 10.5+/-0.3* / 17.8+/-1.3* / 27.9+/-1.2% vs. full-data 33.8+/-0.3% -- a 5.9-point gap at IPC50 while using 4.2% of ImageNet-1K's data points. Crucially, in this same table MTT has no ImageNet-1K entries at all (every IPC is "-"): the strongest prior ImageNet-1K result is frepo-2022's kernel-closed-form method, which only reaches IPC2 (9.7+/-0.2%); TESLA's IPC50 27.9% is the reported "18.2% absolute gain over prior SOTA" (27.9 - 9.7), not a gain over MTT, which was never scaled to ImageNet-1K at all. Table 3 ablation (ImageNet-1K, hard vs. soft labels, IPC1/2/10/50): hard-label MTT-with-constant-memory reaches only 0.7/1.1/4.4/18.1%, vs. TESLA's soft-label 7.7/10.5/17.8/27.9% -- soft labels alone add 7-13.4 absolute points across IPCs, isolating label richness (not the memory fix) as necessary for many-class scaling. Cross-architecture (IPC10, ConvNet-trained data evaluated unseen, Table 2): ImageNet-1K ConvNet 17.8%, ResNet18 7.7%, ViT 11.0%, all exceeding Random's 3.6/1.4/3.2%; CIFAR-10 ConvNet 66.4%, ResNet18 48.9%, ViT 34.8% (vs. MTT 65.3/46.1/34.6%). Cost (Table 4): on CIFAR-100 MTT uses 17.1+/-0.1GB vs. TESLA's 3.6+/-0.1GB (4.75x reduction) at 11.5 vs. 14.5 sec/iteration (1.25x slower); on ImageNet-1K, MTT-scale memory would be 79.9+/-0.1GB vs. TESLA's 13.9+/-0.1GB (5.75x reduction) at 45.0 vs. 46.0 sec/iteration (1.02x slower) -- the constant-memory reformulation costs almost no extra time.

Limitations

Soft labels are a fixed per-image assignment produced once during distillation from a target-step teacher checkpoint, not a teacher-relabeling scheme recomputed every downstream training epoch (the decoupled-synthesis family's later, stronger approach); the paper itself found using a fully-trained (rather than target-step) teacher for label assignment performs much worse, so the choice of which checkpoint generates labels matters and is not obviously optimal. Introducing soft labels changes what is being compared: ImageNet-1K numbers are not directly comparable to any hard-label trajectory-matching result at the same IPC, since part of the accuracy gain is attributable to the label channel rather than the images (the paper's own hard-vs-soft ablation shows this gap is large, 7-13.4 points). Still requires the same cached-expert-trajectory up-front training cost as mtt-2022; the exact-gradient reformulation trades memory for a modest (2-25%) time increase per iteration rather than reducing wall-clock cost.