PaperlayerView exact revision on arXiv

Source-linked AI summary

Dataset Distillation with Infinitely Wide Convolutional Networks

Timothy Nguyen, Roman Novak, Lechao Xiao, Jaehoon Lee

arXiv:2107.13034v3cs.LG

TL;DR

Large datasets enable useful representations but are expensive to acquire and train on, creating a need for compact, informative datasets. The paper applies KIP and LS through distributed meta-learning for infinitely wide convolutional networks and obtains state-of-the-art distillation results across several image datasets. It also analyzes how the distilled data differ from natural images and identifies scaling computational burden as an important boundary.

  • Problem

    Large datasets support useful representations but can be prohibitively costly to acquire and train on, motivating smaller informative datasets.

  • Method

    The paper applies KIP and LS to infinitely wide convolutional networks using a distributed meta-learning framework tailored to expensive neural-kernel computation.

  • Results

    The method achieves state-of-the-art dataset-distillation results across MNIST, Fashion-MNIST, SVHN, CIFAR-10, and CIFAR-100, including over 25% and 37% absolute accuracy gains on CIFAR-10 and SVHN with 10 images.

  • Takeaways & Limitations

    Distilled datasets can be more effective than natural datasets two orders of magnitude larger, while their unusual visual and spectral properties motivate further study of feature learning.

  • Takeaways & Limitations

    The approach has substantial computational burden, and scaling it to larger datasets remains a proposed direction requiring efficient kernel-approximation methods.

Abstract

from arXiv · show

The effectiveness of machine learning algorithms arises from being able to extract useful features from large amounts of data. As model and dataset sizes increase, dataset distillation methods that compress large datasets into significantly smaller yet highly performant ones will become valuable in terms of training efficiency and useful feature extraction. To that end, we apply a novel distributed kernel based meta-learning framework to achieve state-of-the-art results for dataset distillation using infinitely wide convolutional neural networks. For instance, using only 10 datapoints (0.02% of original dataset), we obtain over 65% test accuracy on CIFAR-10 image classification task, a dramatic improvement over the previous best test accuracy of 40%. Our state-of-the-art results extend across many other settings for MNIST, Fashion-MNIST, CIFAR-10, CIFAR-100, and SVHN. Furthermore, we perform some preliminary analyses of our distilled datasets to shed light on how they differ from naturally occurring data.

1 Introduction

Large datasets support useful learned representations but are costly to acquire and train on, motivating distilled datasets that retain task-relevant information with far fewer examples. This paper combines distributed kernel-based meta-learning with infinitely wide convolutional networks, reports state-of-the-art results across image datasets, and analyzes the resulting distilled data.

  • Large datasets enable useful task-adapted representations but can be prohibitively costly to acquire and use for neural-network training.
  • Dataset distillation synthesizes datasets intended to be more informative than natural datasets of equal size, although their learned images differ from natural-image distributions.
  • The paper applies KIP and LS to infinitely wide convolutional networks through a distributed meta-learning framework using hundreds of accelerators per training.
  • The paper analyzes distilled images and labels, finding surprising properties involving interpretability, dimensionality, and spectral structure.
  • Across MNIST, Fashion-MNIST, SVHN, CIFAR-10, and CIFAR-100, the method achieves state-of-the-art dataset-distillation results for kernel ridge regression and neural-network training.
  • Using only 10 images, the method gains over 25% absolute accuracy on CIFAR-10 and 37% on SVHN relative to prior art.

2 Setup

The setup uses infinitely wide convolutional models and KIP or LS to optimize small support sets against larger target datasets. A distributed client-server workflow addresses the expensive kernel-gradient computations required by this optimization.

  • Background on infinitely wide convolutional networks: Infinitely wide convolutional networks converge to Gaussian Processes and provide convolutional inductive biases such as locality and translation invariance.
  • Specific models considered: The central model is a simple four-layer convolutional network with average-pooling layers, alongside convolutional and fully connected variants.
  • Background on algorithms: KIP minimizes kernel ridge-regression loss over a support set while sampling target batches and updating support images and optionally labels with gradients.
  • Background on algorithms: LS solves for the least-norm support-label minimizer, using the whole target dataset to determine the labels.
  • Background on algorithms: The target dataset is substantially larger than the support set, so learned support images and solved labels serve as distilled target representations.
  • Client-Server Distributed Workflow: Distributed KIP separates kernel-matrix computation from support-set gradient computation, partitioning image pairs and derivative blocks across workers.
  • Client-Server Distributed Workflow: The main computational novelty uses backward-mode differentiation with JAX's jax.vjp because naive forward-mode differentiation is infeasible for large image inputs and expensive kernel elements.

3 Experimental Results

Experiments show strong dataset distillation across five datasets and support sizes, with robustness from kernel sampling and effective transfer from kernel-optimized data to finite-width networks.

  • Kernel Distillation Results: Strong results span MNIST, Fashion-MNIST, SVHN, CIFAR-10, and CIFAR-100 across support sets of 1, 10, or 50 images per class.Regularized ZCA preprocessing is crucial for the RGB dataset results, while augmentation benefits are generally larger for sufficiently large support sets.
  • Kernel Distillation Results: 37% absolute test-accuracy gain is the largest margin over prior art, with especially strong performance in the small-support regime.The reported results outperform all prior baselines across dataset settings, and 1-image-per-class KRR results outperform over 100 times as many natural images.
  • Kernel Transfer: Sampling KIP remains effective when evaluated with ConvNet, Conv-Vec3, or Conv-Vec8, whereas single-kernel training suffers a significant drop under kernel mismatch.The sampled-kernel approach is designed to improve robustness across the family of kernels used during training and testing.
  • Neural Network Transfer: Finite-width transfer produces mostly moderate or small performance losses despite KIP and LS datasets being optimized for kernel ridge-regression.The transfer setting uses KIP datasets formed with the ConvNet kernel and evaluates them in corresponding finite-width ConvNet networks.
  • Neural Network Transfer: Finite-width results outperform DC/DSA for all RGB datasets with 1 image per class and for CIFAR-10 in all settings.Transfer performance generally improves with KIP training, although learned-label settings can deteriorate after longer training.
  • Neural Network Transfer: KIP images outperform natural images consistently across hyperparameters and checkpoints, suggesting potential use in accelerated hyperparameter search.The comparison follows performance over training for a single run and a range of hyperparameters.

4 Understanding KIP Images and Labels

The paper investigates why KIP improves generalization by analyzing learned images, labels, dimensionality, and spectral contributions. Across these analyses, KIP generally produces more complex representations that combine features from many target images, while jointly learning labels improves KRR but can reduce image interpretability.

  • Transfer robustness: Transfer performance improves generally with KIP training, while normalization layers harm MSE performance but can benefit cross-entropy performance.The transfer studies vary network depth, normalization, loss, augmentation, label learning, and network width.
  • Hyperparameter robustness: KIP images consistently outperform natural images across many hyperparameters and checkpoints, suggesting usefulness for accelerated hyperparameter search.The comparison covers both MSE and cross-entropy training settings.
  • Spectral Analysis: KIP images leverage the spectrum more evenly than natural images, whose generalization performance is concentrated in a small number of top eigendirections.The spectral comparison evaluates contiguous eigenvalue bands using kernel ridge-regression test accuracy.
  • Complexity analysis: KIP generally increases image complexity by producing richer textures and contours, higher intrinsic dimensionality, and more even use of the spectral components.These observations combine visual, dimensional, and spectral analyses of learned support images.
  • Visual Analysis: Learned KIP images can be visually heterogeneous, ranging from nested or merged object depictions to classes that are visually indistinct.Examples include apples represented by multiple nested apples, bottles transforming from two objects into one, and indistinct beaver images.
  • Labels: Joint image-label optimization makes labels more informative and images less interpretable, improving KRR results but appearing less efficient for neural-network transfer.The comparison uses fixed-label versus jointly optimized support images and labels.
  • Dimensional Analysis: KIP image intrinsic dimension increases during training across datasets and configurations, contrary to the expectation that distillation should reduce dimensionality.The analysis uses the Two-NN estimator, which requires relatively few assumptions and two nearest-neighbor distances per datapoint.

5 Related Work

Prior dataset-distillation methods progressively expanded from synthetic data and label distillation to gradient matching and augmentation-based condensation. This paper builds on kernel methods motivated by the correspondence between infinitely wide neural networks and kernel methods, including finite-width corrections.

  • Prior dataset distillation: Dataset distillation originated with synthetic dataset construction and was extended through label distillation, gradient matching, and augmentation strategies.The cited progression includes Wang et al., Sucholutsky and Schonlau, Bohdal et al., Zhao et al., and Zhao and Bilen.
  • Prior dataset distillation: Zhao and Bilen’s augmentation-based method is limited because it applies only a single augmentation per training iteration.The passage identifies augmentation expressiveness as a limitation of that prior approach.
  • Spectral analysis: Spectral analysis evaluates how contiguous 20% eigenvalue bands affect kernel ridge-regression test accuracy after projection onto corresponding eigenspaces.The analysis uses CIFAR-10 with 500 images.
  • Kernel correspondence: The paper’s kernel approach relies on the correspondence between infinitely wide neural networks and kernel methods, extended by analyses of finite-width corrections.These correspondences underlie transferability from kernel ridge-regression results to neural networks.

6 Conclusion

The paper reports state-of-the-art dataset distillation across convolutional architectures and multiple image-classification datasets, while identifying interpretability, label optimization, and scaling as important directions.

  • Efficient kernel-approximation methods are proposed as a way to reduce computational burden and support scaling to larger datasets.The authors specifically identify understanding compression as dataset size grows as an open question.
  • Optimizing labels is described as a powerful tool for dataset distillation, although the learned labels differ substantially from standard interpretable labels.The authors suggest that optimal label design warrants further study.
  • Learned datasets may provide insights into interpretability and sample-efficient representations through their novel visual features.The paper contrasts spurious-looking pollen features in distilled bee images with essential bicycle contours.

A Experimental Details

The experiments evaluate KIP and LS across several datasets using varied preprocessing, augmentation, label-training, optimization, architecture, and transfer configurations.

  • Augmentations: RGB datasets use a learned AutoAugment policy followed by horizontal flips, while grayscale datasets use rotations and image shifts.Crops and cutout were found to harm performance in the reported experiments.
  • Preprocessing: Regularized ZCA preprocessing standardizes flattened features and applies a regularized whitening transformation, with dataset-specific λ values for RGB data.The paper uses λ = 100 for SVHN and 0.1 for CIFAR-10 and CIFAR-100 unless otherwise stated.
  • Label training: KIP label training either optimizes support labels or keeps them fixed at centered one-hot initialization.The choice is treated as a boolean hyperparameter in KIP training experiments.
  • Optimization and transfer: KIP training uses NTK parameterization, target batches of 5K, up to 50K iterations, and Adam with learning rates of 0.04 or 0.01 depending on dataset.Transfer experiments evaluate multiple KIP checkpoints and tune learning rate and L2 regularization over a small grid.
  • Architecture: The ConvNet architecture removes instance normalization and adds an initial convolution-ReLU layer, yielding four convolutional layers.This modification supports the infinite-width setting and remains comparable to the referenced architecture.

B Computational Costs

Convolutional kernels are substantially more expensive than dot-product kernels because they track pixel correlations, making distributed computation central to KIP training.

  • Kernel complexity: O(d) complexity is added by each convolutional layer, while an additional global pooling layer introduces O(d^2) complexity.Here d denotes the image’s spatial dimension, defined as height times width.
  • GPU constraints: A ConvNet kernel batch size of approximately 10^1 is reported for a Tesla V100 GPU with 16GB RAM.The small batch size limits how much of the kernel matrix can be computed at once.
  • Per-step cost: 460 seconds of work per training step is required for a support set of 100 and target batch size of 5000 under the described V100 computation.Kernel-matrix computation takes about 120 seconds, while gradients take roughly three times longer.
  • Distributed solution: Hundreds of GPUs working in parallel make the distributed training framework computationally feasible.

C.1 KIP Image Analysis

KIP images are jointly learned and correlated, producing subsampling behavior that differs from natural images and motivating visual and dimensional analyses.

  • Subsampling: Natural images degrade more gracefully under subsampling than KIP images in kernel ridge-regression test accuracy.The difference is attributed to correlations among jointly learned KIP images.
  • Visual analysis: Additional visual analyses cover learned images for the remaining datasets in Figure A2.
  • Dimensional analysis: The paper measures intrinsic, linear, and gradient-based dimensions of learned images in supplementary figures.Linear dimension is based on the number of PCA components explaining 90% of variance.

C.2 Natural Image Baselines

Natural-image subsets provide a baseline for evaluating distilled images, but KIP images retain substantially more task-relevant information at very small dataset sizes. Under kernel-ridge regression, one KIP image per class outperforms natural subsets exceeding 100 images per class.

  • Natural-image comparison: One KIP image per class outperforms more than 100 natural images per class under ConvNet-kernel kernel-ridge regression.This comparison evaluates test accuracy after subsampling natural-image subsets of varying sizes.
  • Natural-image comparison: Subsampling KIP images causes a more severe relative test-accuracy drop than subsampling corresponding natural images.Figure A1 compares the relative drops at initialization and also shows hyperparameter choices among KIP-trained images.
  • Natural-image comparison: With 1 image per class, CIFAR-10 reaches 16.3±2.1 accuracy for natural subsets, compared with 76.1 for the full dataset.The table reports mean and standard deviation over twenty random subsets, with the “All” row denoting the entire dataset.

C.3 Ablation Studies

The ablations evaluate preprocessing, label training, and augmentation choices for distilled images. Across the reported sweeps, ZCA regularization and label training provide a clear and consistent benefit for KIP.

  • Ablation findings: ZCA regularization and label training consistently improve KIP performance across the reported hyperparameter sweeps.The sweeps include standard preprocessing and no-label-training settings, with corresponding neural-network transfer results where available.
  • Ablation settings: Examples of learned images are shown for multiple datasets using 10 distilled images, without ZCA, label training, or augmentations.The settings cover learned-image examples beyond the main CIFAR-100 visualization.

C.4 DC/DSA ConvNet

The paper evaluates its modified ConvNet and associated infinite-width model across architecture, preprocessing, training, and representation analyses. Results indicate robustness to a shallower architecture, while learned-dataset dimensionality generally increases during training.

  • C.4 DC/DSA ConvNet: Removing the initial convolution and ReLU changes CIFAR-10 kernel-ridge results typically by less than 1%, with a maximum difference of 1.9%.The shallower ConvNet sometimes outperforms the deeper version, supporting robustness to this architectural change.
  • Representation analysis: The learned dataset’s dimensionality generally grows with training across image, gradient, and infinite-width-gradient measures.Figure A3 reports linear dimension, intrinsic dimension of infinite-width gradients, and linear dimension of gradients.
  • Experimental settings: The supplementary experiments vary ZCA, instance normalization, augmentation, and loss function for neural-network training on CIFAR-10.The corresponding full-data experiment is summarized in Table A2.
  • Representation analysis: Under ZCA preprocessing, dimensionality generally increases during training, but the trend is less robust and initialization can show unusually high dimensionality.Figure A4 reports intrinsic and linear dimensions for datasets and gradients.
Loading 2107.13034v3…