Dataset Distillation with Infinitely Wide Convolutional Networks

In plain words

Dataset distillation asks how small a training set can get and still train a good model. Instead of picking a subset of real photos, this paper learns a handful of synthetic images, and optionally their labels, that carry most of the useful signal. The method uses a network with infinitely many channels, whose behaviour is described by a kernel. With that kernel, predicting a large target set from the learned images has a closed-form answer, so the images can be improved by ordinary gradient steps against that target set. The authors make this practical for convolutional networks by splitting the expensive kernel computations across many machines. The learned images also train ordinary finite networks, with some loss in accuracy, and the authors examine how they differ from natural photos.

TL;DR

Learns a few synthetic images per class, plus optional labels, by meta-learning through kernel ridge regression (KRR) with the kernel of a 4-layer convolutional network of infinite width §2 app. A. Author-reported: CIFAR-10 with one image per class reaches 64.7% KRR test accuracy without augmentation tab. 1. Transfer to finite-width ConvNets loses accuracy tab. 2. The distributed trainer is the expensive part §2.1 app. B; the released repository does not contain it (see Paper vs code).

Problem

Large datasets make training costly §1. Coresets and low-dimensional projections keep a natural subset or a coarse summary, and trade accuracy for size §1. The earlier KIP and Label Solve work used fully connected kernels; this paper extends it to convolutional kernels at large scale §1. Convolutional and pooling layers make kernel computation expensive: each added convolution costs O(d) and each global pooling layer O(d²), where d is the number of spatial positions app. B. The method assumes that data optimised for KRR under the infinite-width limit also helps finite networks trained by gradient descent §2; this is tested empirically §3.3.

Method

Models. The main model, ConvNet, is a 4-layer convolutional network with average pooling and no normalisation layers, so that its infinite-width limit exists. It is a modification of the ConvNet used in the DC/DSA baselines, which has three layers and instance normalisation §2 app. A. Variants are ConvVec (no pooling before the readout) and FC (no convolutions) §2.

Meta-objective. Given a kernel K, the KRR loss on a support set (X_s, y_s), evaluated on a target set (X_t, y_t), is eq. 1:

L(X_s, y_s) = \tfrac{1}{2}\left\| y_t - K_{X_tX_s}\,(K_{X_sX_s}+\lambda I)^{-1} y_s \right\|_2^2 eq. 1

Here λ > 0 is the regulariser. In the experiments it is set to λ = λ0 · tr̄(K_{X_sX_s}), with λ0 = 10⁻⁶ and tr̄(C) = tr(C)/len(C) app. A.

KIP (Kernel Inducing Points). Minimise eq. 1 over the support images, and optionally the support labels, with gradient steps. Each step draws a fresh target batch §2. Optional variants augment the target batch or sample a kernel from a family, which gives the robust "sampling KIP" variant §2 §3.2.

Label Solve (LS). Keep the support images fixed and solve for the labels that minimise eq. 1 in the least-norm sense, using the whole target set eq. 2:

y_s^{*} = \left( K_{X_tX_s}\,(K_{X_sX_s}+\lambda I)^{-1} \right)^{+} y_t eq. 2

where ^+ is the pseudo-inverse §2. The support set must be smaller than the target set §2.

Initialisation. Both methods start from class-balanced subsets of the training data, with the same preprocessing; random noise also works and gives similar final results app. A.

Preprocessing. Standard per-channel normalisation is the default. Regularised ZCA whitening is used for the RGB datasets in the main results app. A tab. 1. The regulariser is scaled by tr̄ as above app. A.

Distributed gradient computation §2.1. The support-support and target-support kernel blocks are computed in B×B tiles by workers and summed on a server. For the gradient, the server computes ∂L/∂K and cuts it into the same tiles. Each worker applies a vector-Jacobian product to its tile to get the input gradient. The server sums the tile contributions to update X_s.

# KIP step on the support images  [§2] [eq. 1] [app. A]
# Input:  support images (S, H, W, C); labels (S, K); target set
# Output: updated support images (and labels, if label training is on)

def kip_step(x_support, y_support, target_set, kernel_fn):
    # Draw a fresh class-balanced target batch of 5000  [app. A]
    x_target, y_target = sample_balanced(target_set, 5000)
    # Kernel blocks for the support and target sets  [§2.1]
    k_ss = kernel_fn(x_support, x_support)           # (S, S)
    k_ts = kernel_fn(x_target, x_support)            # (T, S)
    # Ridge term scaled by the mean kernel diagonal  [app. A]
    lam = 1e-6 * trace(k_ss) / S
    # KRR prediction on the target batch  [eq. 1]
    preds = k_ts @ solve(k_ss + lam * I, y_support)
    # Squared loss over the target batch  [eq. 1]
    loss = 0.5 * sum((y_target - preds) ** 2)
    # Adam step on the images: lr 0.04 for CIFAR, 0.01 otherwise  [app. A]
    x_support = adam_update(x_support, grad(loss, x_support))
    # Optional label training, same loss  [app. A]
    if label_training:
        y_support = adam_update(y_support, grad(loss, y_support))
    return x_support, y_support

# Label Solve  [§2] [eq. 2]
# Input:  fixed support images; full target set (X_t, y_t)
# Output: least-norm support labels

def label_solve(x_support, x_target, y_target, kernel_fn):
    # Kernel blocks for the fixed support set  [§2]
    k_ss = kernel_fn(x_support, x_support)           # (S, S)
    k_ts = kernel_fn(x_target, x_support)            # (T, S)
    # Lambda as defined under eq. 1; reused here (see Pitfalls)  [app. A]
    lam = 1e-6 * trace(k_ss) / S
    # Mapping K_ts (K_ss + lam I)^-1  [eq. 2]
    mapping = k_ts @ inverse(k_ss + lam * I)        # (T, S)
    # Least-norm labels via pseudo-inverse  [eq. 2]
    y_support = pinv(mapping) @ y_target            # (S, K)
    return y_support

# Backward pass in tiles  [§2.1]
# Input:  tile of input images x_i, x_j; cotangent dL/dK for that tile
# Output: contribution of this tile to dL/dX_s

# The server sums these tile contributions to form dL/dX_s  [§2.1]

def tile_gradient(x_i, x_j, dL_dK_tile, kernel_fn):
    # Vector-Jacobian product of the tile kernel  [§2.1]
    kernel_of_x = lambda x: kernel_fn(x, x_j)
    _, vjp_fn = vjp(kernel_of_x, x_i)
    return vjp_fn(dL_dK_tile)

Hyperparameters used for the infinite-width distillation, and for finite network transfer, are listed below. Values are as stated in the paper.

Hyperparameters

Results

Author-reported. Setup: ConvNet kernel for distillation, with the dataset settings in each row; baseline numbers are from the cited prior work tab. 1. Table 1 reports KRR accuracy with the ConvNet kernel for KIP and LS, next to neural-network distillation baselines (DC, DSA) and KIP-FC. The KIP entries use label training, and RGB datasets use regularised ZCA tab. 1.

Neural-network transfer. Finite ConvNets (width 1,024, standard parameterisation) trained on the KIP or LS sets, compared with DC/DSA. The DC/DSA column is the better of the two per row. "Perf. change" is the difference from a KRR number for the chosen dataset, which the paper does not list tab. 2.

Additional results, with the paper's own wording condensed:

• Full CIFAR-10 training with the ConvNet, varying ZCA, normalisation and augmentation. The best row is augmentation alone, at 86.6 ± 0.2 (MSE) and 88.9 ± 0.2 (XENT); no ZCA, no normalisation, no augmentation gives 76.9 ± 0.8 (MSE) and 80.4 ± 0.2 (XENT) app. C.2.
• Natural-image KRR baselines: subsets of 1 to 128 images per class, with the full dataset as reference app. C.2.
• Replacing the first conv layer of the ConvNet with the DC/DSA version changes CIFAR-10 KRR results by typically under 1%, at most 1.9%, and sometimes in favour of the shallower net app. C.4.
• Subsampled KIP sets are more correlated than natural sets: they degrade faster under subsampling app. C.1.
• Intrinsic dimension of learned sets grows during training, and ZCA raises dimension §4 fig. 7 app. C.1.
• Spectral analysis: natural images concentrate their performance in the top 20% of eigendirections; KIP images spread it across bands (FC and Conv-Vec8) or skew it to the tail (ConvNet) §4 fig. 8.
• KIP images also improve over natural images across the hyperparameter sweep and checkpoints tested, for CIFAR-10 with 500 images §3.3 fig. 4.

Paper vs code

The repository is google-research/google-research at 0eebfd403ba438bfd14cb70b9b24a7745c997083, under kip/. It holds a README and three Colab notebooks: KIP.ipynb (the earlier ICLR 2021 KIP and LS work), dataset.ipynb (loads released distilled sets from cloud storage) and nn_training.ipynb (finite ConvNet training). Code: kip/README.md README.md. Licence: Apache-2.0 LICENSE.

• Distributed trainer absent. The paper's Conv-KIP trainer, built on Courier and JAX vjp tiles §2.1, is not in the repository. Only the notebooks above are present. No code produces the 10-, 100- or 500-image convolutional sets; they are loaded from cloud storage [code: raw/code/2107.13034/kip/dataset.ipynb#L164].
• Learning rate. The released KIP colab sets Adam to 4e-2 [code: raw/code/2107.13034/kip/KIP.ipynb#L157]. This matches the paper's CIFAR value and does not show the 0.01 used for other datasets app. A.
• Target batch. The released colab uses 5,000 [code: raw/code/2107.13034/kip/KIP.ipynb#L159], which matches app. A.
• Kernel regulariser. Support-kernel regularisation is |reg| · trace(k_ss)/n · I with reg = 1e-6 KIP.ipynb. This matches λ = λ0 · tr̄(K) app. A.
• Label Solve regulariser. The LS routine adds reg · I / n, with no trace factor KIP.ipynb. The paper gives one λ definition for eq. 1 and eq. 2 app. A. The two differ unless the trace factor is absorbed elsewhere.
• ZCA regulariser. Regularisation amount is zca_reg · trace(cov)/dim, which matches tr̄ nn_training.ipynb. The notebook config sets 0.1, the CIFAR value; the SVHN value of 100 is not in the release nn_training.ipynb.
• Empirical NTK for LR. The paper uses min(64, support size) training points app. A. The code uses 50 validation images [code: raw/code/2107.13034/kip/nn_training.ipynb#L126] [code: raw/code/2107.13034/kip/nn_training.ipynb#L597]. The notebook calls nt.predict.max_learning_rate nn_training.ipynb; I did not check that library's formula against the paper's 2/(λ_min+λ_max).
• Learning-rate factor and L2. The notebook fixes the factor c at 1.0 [code: raw/code/2107.13034/kip/nn_training.ipynb#L128] and L2 at 0.001 [code: raw/code/2107.13034/kip/nn_training.ipynb#L129]. The paper grid-searches both app. A. The notebook is an example run, not the sweep.
• Batch and steps. The notebook uses batch 100 [code: raw/code/2107.13034/kip/nn_training.ipynb#L122] and 5,000 steps [code: raw/code/2107.13034/kip/nn_training.ipynb#L124]. The paper uses full batch for support sizes 10 and 100, and 25,000 or 10,000 steps for some settings app. A.
• Architecture and width. A prepended conv layer is on for ConvNet [code: raw/code/2107.13034/kip/nn_training.ipynb#L138], average pooling is used [code: raw/code/2107.13034/kip/nn_training.ipynb#L477], the width is 1,024 [code: raw/code/2107.13034/kip/nn_training.ipynb#L132] and the parameterisation is standard nn_training.ipynb. These match app. A.
• Optimiser. Momentum 0.9 with cosine decay [code: raw/code/2107.13034/kip/nn_training.ipynb#L120-L121] [code: raw/code/2107.13034/kip/nn_training.ipynb#L386-L390] matches app. A.
• Augmentation. The paper uses AutoAugment and horizontal flips app. A. The released notebooks contain no AutoAugment policy. KIP.ipynb has an Augmentor class with horizontal_flip defaulting to False [code: raw/code/2107.13034/kip/KIP.ipynb#L246-L264].

Pitfalls

• Distributed code is yours to write. The paper's main engineering contribution (tiled vjp gradients across hundreds of accelerators §2.1) is not in the repository. Expect to build the tiling, aggregation and checkpointing yourself. The paper gives no code-level detail beyond §2.1.
• Cost. The paper gives about 80 ms per kernel tile, 18 × 18 tiles, and 460 s per step app. B. 120 s for the kernel and three times that for gradients would be 480 s, so the 460 s figure is approximate. At about 460 s per step, 50,000 steps is roughly 266 device-days. This is my arithmetic, not the paper's.
• Abstract vs tables. The abstract claims over 65% CIFAR-10 accuracy from 10 images [abstract]. The highest CIFAR-10 one-image-per-class entry is 64.7 ± 0.2 without augmentation tab. 1 app. C.3. I found no 65%-or-higher entry for that setting.
• Improvement claims. §1 says CIFAR-10 gains exceed 25 points over prior art §1. Against KIP-FC (40.5) the gap is 24.2 points tab. 1; against DSA (28.8) it is 35.9. For SVHN, the 37-point claim matches 64.3 − 27.5 (DSA). Against DC (31.2), the best prior at that setting, it is 33.1 tab. 1.
• Error bars. The ± values in Table 1 are standard deviations over five checkpoints, not over seeds. LS numbers come from one random support set and one seed app. A.
• Selection bias. Table 1 uses checkpoints chosen by lowest training loss app. A. Table 2 takes the best of a checkpoint and hyperparameter grid, chosen on a separate 5K validation set app. A. Table 2 is therefore a best-of-grid number, not a single-run number.
• Implied KRR numbers. Table 2 changes are relative to KRR values that the paper does not list for the chosen KIP sets. For CIFAR-10 at one image per class, the implied KRR is 49.9 + 9.2 = 59.1 tab. 2. Table 1 gives 64.7 or 63.4 for the same setting.
• Shared SVHN settings. SVHN without augmentation uses the CIFAR-10 hyperparameters (ZCA 0.1, LR 0.04), because ZCA 100 looked unstable app. A.
• Architecture gap. Table 1 uses a 4-layer ConvNet, and DC/DSA uses a 3-layer ConvNet with instance normalisation. The paper's ablation shows small differences, under 2% app. C.4. Comparisons with DC/DSA are not like-for-like in architecture.
• Parameterisation mismatch. KRR uses NTK parameterisation and finite-width transfer uses standard parameterisation at width 1,024 app. A. A transfer gap therefore mixes parameterisation with finite-width effects.
• Label training. Learned labels are the best choice for some transfer rows tab. 2. For XENT, the target class is the argmax of the learned labels app. A.
• Memory and kernel sizes. Kernel batch sizes on a 16 GB V100 are about 10⁴ for FC1, 10³ for Conv-Vec1 and 10¹ for ConvNet app. B. ConvNet is the bottleneck, which is why the tiling exists.
• Interpretation is qualitative. The claims about image complexity, intrinsic dimension and spectra are measured on specific settings §4 app. C.1. Measured on a different preprocessing or kernel, they may change.

Lineage

• Builds on: Nguyen et al. 2021, KIP and Label Solve on fully connected kernels, which this paper scales to convolutional kernels §1 §2.
• Compared against: DC §3.1 (Zhao et al. 2020), DSA §3.1 (Zhao et al. 2021), KIP-FC tab. 1 (Nguyen et al. 2021), and LD (Bohdal et al. 2020), which uses AlexNet and distils labels only tab. 1.
• Uses: Neural Tangents for NTK computation §2; the finite-network transfer protocol of Lee et al. 2020 app. A; regularised ZCA from Shankar et al. 2020 app. A.
• Related: infinite-width theory §5; Ansuini et al. 2019 for intrinsic dimension measurements app. A.