Dataset Distillation with Convexified Implicit Gradients

What came before

Responds directly to frepo-2022 (whose loss it explicitly generalizes) and to kip-2021/rfad-2022's exact/approximate NTK approach, and to dd-2018-style BPTT and gradient/trajectory-matching surrogates (dc-2021, mtt-2022), framing all of them as either expensive-to-unroll or introducing their own matching bias.

The problem

Bi-level dataset distillation needs the meta-gradient of an outer loss with respect to the inner learner's parameters; implicit gradients (via the implicit function theorem) offer a way to get this without unrolling, but the paper shows "vanilla" implicit gradients (VIG) perform poorly -- sometimes worse than random real images -- because deep networks are non-convex and the implicit-function theorem's unique-minimizer assumption fails, and diagnoses this as the specific reason implicit gradients had not previously worked for distillation.

The idea

Make implicit gradients work by making the inner problem actually convex: linearize the network (first-order Taylor expansion around initialization, i.e. train on the frozen finite-width NTK/ conjugate kernel), which convexifies the inner objective; then further reduce bias by reparameterizing so the final layer's optimal weights are computed analytically from the body parameters at every step, rather than left to imperfect implicit-gradient optimization.

How it works

Splits network parameters into body $\theta_B$ and final layer $\theta_F$; under linearized (NTK) dynamics, $\theta_F^*$ has a closed form given $\theta_B$ (Eq. 3-4, computed via KRR on the network's own hidden features, offset by the body's contribution), so only $\theta_B$ is treated as the inner variable. The meta-gradient combines a direct term ($\partial\mathcal{L}_o/\partial\psi$, non-zero because $\psi$ affects $\theta_F^*$ -- this direct term reduces to the FRePo loss when using MSE) and an implicit term computed by approximating $v=H_S^{-1}g_T$ via SGD on a Neumann-series-equivalent auxiliary loss (Hessian-vector products via the Pearlmutter trick). A pool of $m=30$ partially trained models (max $K=100$ steps before reinit, matching FRePo's scheme) supplies diversity in $\theta_B$ across steps. A Platt-scaling (learnable-temperature cross-entropy) outer loss replaces MSE. For support sets exceeding 1,000 images, a bias-free subsampling trick runs the full forward pass but backpropagates through only a random subset of nodes, since gradient contributions from each support image are interchangeable. Ablation architecture: VIG (naive implicit gradients) -> CIG (+ linearization/ convexification) -> RCIG (+ reparameterized analytic final layer), isolating each component's contribution (Table 1).

Evidence

Ablation (Table 1, CIFAR-100 1 IPC): random subset 6.2%, VIG 17.8%, CIG 19.7%, RCIG 39.6% -- showing convexification alone helps modestly but the reparameterized analytic final layer drives most of the gain. Headline (Table 2): CIFAR-100 x 1 IPC x 39.3+/-0.4% vs. prior SOTA (FRePo) 28.7% (37% relative gain, matches abstract); Tiny-ImageNet x 1 IPC x 25.6+/-0.3% vs. FRePo 15.4% (66% relative gain); resized ImageNet-1K x 1 IPC x 15.6+/-0.2% vs. FRePo 7.5% (108% relative gain, i.e. roughly doubling SOTA, matching the abstract's headline number). RCIG wins 13/16 standard-benchmark settings (Table 2) but underperforms FRePo on CUB-200 (12.1% vs. 12.4% at 1 IPC), which the paper attributes to overfitting on CUB-200's small training set (5,994 images across 200 classes) since RCIG uses no data augmentation. Cross-architecture transfer (Table 3, CIFAR-10 x 10 IPC): RCIG trained with BatchNorm reaches 54.4-62.1% on unseen ResNet/VGG/AlexNet architectures, comparable to or better than FRePo (52.4-61.6%) and clearly ahead of DSA/DM/MTT/KIP; trained with no normalization (the paper's default), RCIG transfers less well to BN/VGG architectures (46.2-51.3%) than the BN-trained variant, isolating training-time normalization choice (not the RCIG mechanism itself) as what drives cross-architecture robustness. ImageNet subsets (Table 4): ImageNette/ImageWoof 1 IPC 53.0%/33.9% vs. FRePo 48.1%/29.7%, though gains shrink at 10 IPC (roughly matching FRePo). Cost: reported on RTX A6000 GPUs (48GB) and RTX 4090; memory is higher than other methods because linearization roughly doubles memory, and for large support sets (e.g. CIFAR-100 at 50 IPC = 5,000 images) the subsampling trick is required to fit in memory (e.g. n_subset=2000 uses ~20Gb-scale memory per the appendix's Figure 5 sweep, with only a 1.4-point accuracy cost between the smallest and largest subsample sizes tested).

Limitations

States explicitly that gains shrink as the support-set size grows (e.g. only a 14-point gap over SOTA at Tiny-ImageNet 10 IPC vs. a much larger gap at 1 IPC), hypothesized to be because the inner objective is harder to optimize at larger set sizes and would need a lower learning rate and more inner steps; this is the same "matching methods help most at very small IPC" pattern later diagnosed more generally by datm-2024. Requires full-batch (non-stochastic) gradients through the entire support set for both the inner objective and the Hessian-vector product, and linearization roughly doubles memory, so very large support sets are memory-constrained even with the subsampling workaround (which still needs a full forward pass on the whole set). Overfits on small, many-class datasets (CUB-200) because it uses no data augmentation during distillation, unlike dsa-2021/mtt-2022. Its cross-architecture robustness depends on using BatchNorm at training time, not on the RCIG mechanism per se. No result beyond resized 64x64 ImageNet-1K at 1-2 IPC, and no soft-label/teacher-relabeling scheme.