Dataset Distillation using Neural Feature Regression

What came before

Responds directly to kip-2021's exact-NTK closed-form inner solver, which the paper states requires "thousands of GPU hours and sophisticated implementation of the distributed kernel computation framework" for modern architectures, and to dd-2018-style unrolled BPTT, which needs the whole training trajectory in memory. Also responds to surrogate-matching alternatives (dc-2021, dsa-2021 gradient matching; CAFE/DM distribution matching; mtt-2022 trajectory matching), arguing each surrogate can introduce its own bias and typically overfits to a single learning algorithm, training stage, or architecture.

The problem

Meta-gradient computation for dataset distillation is expensive (unrolled BPTT) or requires an intractable exact kernel (KIP-style NTK for deep nets); separately, distilled data tends to overfit to the specific algorithm, training stage, or architecture used during synthesis, hurting cross-architecture transfer.

How it works

Splits the matching network into a feature extractor $f(\cdot,\theta)$ and a linear head, and instead of learning the head by gradient descent, solves for it via kernel ridge regression on the extractor's own features (a "conjugate kernel," not the NTK): $\mathcal{L}=\tfrac12\|Y_t-K^\theta_{X_tX_s}(K^\theta_{X_sX_s}+\lambda I)^{-1}Y_s\|_2^2$. Backpropagating this loss through the fixed feature extractor gives the meta-gradient for the synthetic images $X_s$ (and, since the loss is smooth in $Y_s$, learned labels), with no need to unroll the extractor's own training. A "model pool" of $m$ feature extractors at different initializations and training stages is maintained; each step samples one pool member, uses it for the KRR meta-gradient, takes one online training step on it with the current synthetic data, and reinitializes it once it has been updated more than $K$ steps -- diversifying which extractor is matched against, unlike a single periodically-reset model. Stored artifact: only the synthetic images and (optionally) learned labels; the pool is discarded after training.

Evidence

CIFAR-100 x 1 IPC x FRePo x learned labels x 28.7% (Table 1) vs. MTT 24.3%, KIP 15.7%, DSA 16.8%, DM 12.2%; Tiny-ImageNet x 1 IPC: 15.4% (FRePo) vs. MTT 8.8%, DSA 6.6%; CUB-200 x 1 IPC: 12.4% (FRePo) vs. MTT 2.2% -- the paper's headline claim that FRePo is exceptionally strong on complex, many-class label spaces at 1 IPC. Cross-architecture transfer (Table 2, CIFAR-10 x 10 IPC, trained on Conv-BN): FRePo beats DSA/DM/MTT/KIP on every unseen evaluation architecture (Conv-NN 65.5% vs. MTT's 41.6%; AlexNet 61.9% vs. KIP's 57.2%; VGG-BN 59.4% vs. MTT's 46.6%), and the paper notes instance normalization -- used by DSA/DM/MTT -- specifically hurts their cross-architecture transfer, while FRePo's images "look natural" and carry less architectural bias. Resized ImageNet-1K (64x64): 1-2 IPC (1k/2k images) reaches 7.5%/9.7% top-1 vs. 1.1%/1.4% for an equally sized random real subset (no comparison to MTT is given at this scale in-paper). Cost: on CIFAR-100 at 1 IPC, FRePo matches MTT's accuracy (~24%) in 38 seconds vs. MTT's 3,805 seconds (~100x faster) and reaches 92% of its own final accuracy in 385 seconds, with roughly 10x lower peak GPU memory (Fig. 3). Downstream: 5-/10-step class-incremental continual learning on CIFAR-100 (41.6%/38.0% final accuracy vs. DM's 33.9%/34.0%) and membership- inference defense (attack AUC near 0.5, i.e., near random guessing, on MNIST/FashionMNIST at 500 images) both improve over DSA/DM at matched budget. Ablation: label learning is unnecessary on simple, few-class datasets (MNIST) but becomes crucial on complex many-class datasets (CIFAR-100, Tiny-ImageNet), isolating label richness as the mechanism behind the many-class gains.

Limitations

Requires a wider network architecture than prior methods for the KRR component to behave well (a narrow feature dimension causes a significant accuracy drop, per Appendix C.6), so comparisons are not always at matched architecture width. No data augmentation is used during synthesis (only at evaluation), unlike dsa-2021/mtt-2022. Diversifying the model pool with multiple architectures (rather than just different seeds/stages of one architecture) was tried but found unstable, since meta-gradients from different architectures can disagree sharply -- so the reported cross-architecture generalization comes from a single-architecture pool, not explicit multi-architecture training. Splitting per-class distillation across sequential continual-learning steps was observed to underperform distilling the same total budget in one pass, suggesting the joint, whole-dataset view is important. ImageNet-1K results are at only 1-2 IPC and use hard evaluation without any teacher-relabeling or soft-label scheme (unlike tesla-2023's contemporaneous ImageNet-1K soft-label approach), leaving unclear how far the closed-form-KRR proxy scales past very low IPC on datasets with 1,000 classes.