PaperlayerView exact revision on arXiv

Source-linked AI summary

Dataset Distillation using Neural Feature Regression

Yongchao Zhou, Ehsan Nezhadarya, Jimmy Ba

arXiv:2206.00719v2cs.LGcs.CV

TL;DR

Dataset distillation seeks compact synthetic datasets, but differentiating through inner training is expensive and can amplify overfitting. FRePo uses neural feature regression with a model pool to compute efficient meta-gradients and improve generalization. It reports state-of-the-art performance with substantially lower training cost and memory, alongside benefits in downstream applications.

  • Problem

    Dataset distillation requires meta-gradients through inner optimization, while distilled data can overfit to learning algorithms, optimization stages, models, or architectures.

  • Method

    FRePo computes meta-gradients through kernel ridge regression on neural features and maintains a diverse pool of models to mitigate overfitting.

  • Results

    FRePo achieves state-of-the-art results with a 100x reduction in training time and a 10x reduction in GPU memory requirement across various datasets.

  • Takeaways & Limitations

    FRePo distilled data generalizes across architectures and improves downstream applications including continual learning and membership inference defense.

  • Takeaways & Limitations

    The reported simple model pool uses one architecture and no data augmentation, varying initialization, optimization stage, and distilled-data time step.

Abstract

from arXiv · show

Dataset distillation aims to learn a small synthetic dataset that preserves most of the information from the original dataset. Dataset distillation can be formulated as a bi-level meta-learning problem where the outer loop optimizes the meta-dataset and the inner loop trains a model on the distilled data. Meta-gradient computation is one of the key challenges in this formulation, as differentiating through the inner loop learning procedure introduces significant computation and memory costs. In this paper, we address these challenges using neural Feature Regression with Pooling (FRePo), achieving the state-of-the-art performance with an order of magnitude less memory requirement and two orders of magnitude faster training than previous methods. The proposed algorithm is analogous to truncated backpropagation through time with a pool of models to alleviate various types of overfitting in dataset distillation. FRePo significantly outperforms the previous methods on CIFAR100, Tiny ImageNet, and ImageNet-1K. Furthermore, we show that high-quality distilled data can greatly improve various downstream applications, such as continual learning and membership inference defense. Please check out our webpage at https://sites.google.com/view/frepo.

1 Introduction

Dataset distillation learns compact synthetic data through bi-level optimization, but meta-gradient computation and overfitting make training difficult. FRePo addresses these challenges with efficient feature-based gradients and a diverse model pool, achieving strong efficiency and benchmark performance.

  • Motivation: Dataset distillation learns a small synthetic dataset that preserves information from a large dataset while reducing storage and training costs.The distilled data can also support continual learning, neural architecture search, and privacy-preserving tasks.
  • Problem formulation: The bi-level formulation optimizes distilled data in the outer loop while training a model on that data in the inner loop.The outer objective evaluates models trained on the distilled dataset against real data.
  • Challenges: Unrolled meta-gradient computation introduces huge compute and memory requirements, training instability, and truncation bias.The full training trajectory must be stored when differentiating through iterative inner optimization.
  • FRePo: FRePo computes meta-gradients through a kernel and fixed feature extractor, while a diverse model pool addresses overfitting across algorithms, optimization stages, and architectures.The approach is analogous to one-step TBPTT but avoids backpropagating through the inner optimization.
  • Results: 100x reduction in training time and 10x reduction in GPU memory requirement accompany state-of-the-art results across benchmark datasets.The distilled data also transfers across architectures and supports continual learning and membership inference defense.
  • Results: 7.5% top1 accuracy on ImageNet-1K using one image per class exceeds the 1.1% accuracy obtained from a random subset of real images.Previous methods struggle in this setting because of large memory and compute requirements.

2 Method

FRePo replaces costly differentiation through unrolled inner optimization with kernel ridge regression over neural features and uses model pooling to reduce overfitting. Its procedure alternates distilled-data updates with online model training.

  • Bi-level optimization: Dataset distillation minimizes expected validation loss on real data after models are trained on a small distilled dataset.The expected objective accounts for stochasticity from model initialization and the inner learning algorithm.
  • Meta-gradient computation: Unrolled optimization computes meta-gradients by differentiating through inner training, requiring substantial computation and memory.The entire training trajectory must be retained in memory.
  • Training procedure: FRePo updates distilled images and labels using the meta-gradient, then trains a sampled pool model for one step on the updated distilled data.The algorithm repeatedly samples models and real-data target batches until convergence.
  • Meta-gradient computation: FRePo trains a linear classifier to convergence with a fixed feature extractor, expressing predictions through kernel ridge regression with a conjugate kernel.This gives the inner optimization a closed-form approximation suitable for efficient meta-gradient computation.
  • Meta-gradient computation: FRePo backpropagates through the conjugate kernel and fixed feature extractor, decoupling meta-gradient computation from the online optimizer.The distilled data therefore remains agnostic to the specific learning algorithm used for online model updates.
  • Model pool: The model pool contains diverse parameters from different initializations, optimization stages, and distilled-data time steps.Diversity targets overfitting to a particular learning algorithm, optimization stage, or architecture.
  • Pool diversity: The implementation uses one architecture without data augmentation, varying only initialization, optimization stage, and distilled-data time step within the model pool.The paper notes that larger pool diversity, augmentation, or different architectures could increase regularization strength, but does not use them in its simple configuration.

3 Related Work

Prior dataset-distillation methods compute meta-gradients through unrolled optimization, surrogate objectives, or closed-form kernel approximations. These alternatives trade off computational cost, bias, or fidelity to the original objective.

  • Unrolled optimization: Unrolled optimization differentiates through the entire inner training sequence but incurs large computation and memory costs.Long unrolls can also produce exploding or vanishing gradients and poorly conditioned loss landscapes, while short unrolls introduce truncation bias.
  • Surrogate objectives: Surrogate objectives match gradients, feature distributions, or expert training trajectories instead of directly differentiating through inner optimization.These methods can introduce bias because the surrogate may not accurately reflect the true objective.
  • Closed-form approximation: Kernel-based closed-form approximations such as KIP use the Neural Tangent Kernel to approximate inner optimization and compute meta-gradients through the kernel.Computing analytical NTK for modern neural networks is described as extremely expensive.

4 Dataset Distillation

FRePo achieves strong dataset-distillation performance across standard, cross-architecture, high-resolution, and complex-label benchmarks while substantially reducing training cost. Its distilled images are natural, and the method generalizes across architectures and challenging datasets.

  • Standard benchmarks: FRePo improves one-image-per-class accuracy on CIFAR100, Tiny ImageNet, and CUB-200 from 24.3%, 8.8%, and 2.2% to 28.7%, 15.4%, and 12.4%, respectively.These gains are reported despite potentially suboptimal hyperparameters.
  • Qualitative results: FRePo’s distilled images look natural, and higher image quality is strongly correlated with higher test accuracy.The reported qualitative results suggest that condensed data can reflect common dataset patterns rather than appearing unlike real data.
  • Training efficiency: 38 seconds: FRePo reaches 23.4% CIFAR100 accuracy, close to MTT’s 24.0% in 3805 seconds, making it roughly two orders of magnitude faster.FRePo reaches 92% of its final test accuracy in 385 seconds.
  • Training efficiency: FRePo requires much less GPU memory and uses faster gradient steps, enabling more optimization steps or larger, more complex evaluation models.The efficiency comparison reports one-tenth of the GPU memory requirement relative to previous methods.
  • Cross-architecture generalization: On unseen CIFAR10 architectures, FRePo outperforms previous methods and transfers well to models without normalization or with different normalization layers.The comparison includes AlexNet, VGG, and ResNet; other methods show degradation when normalization changes or is removed.
  • High-resolution and complex-label datasets: FRePo outperforms MTT on 128x128 ImageNette and ImageWoof settings and scales to 64x64 ImageNet-1K with complex labels.Using one image per class on ImageNet-1K yields 7.5% top-1 accuracy versus 1.1% from an equally sized random real subset.

5 Application

The paper applies distilled data to continual learning and membership-inference defense. FRePo improves class-incremental accuracy, while distilled-data training keeps attack AUC near random guessing with a smaller test-accuracy cost than equally sized real subsets.

  • Continual learning: Distilled data is proposed as a compact replay-buffer component because continual-learning performance depends on replay-buffer quality in the GDumb setup.The experiments use class-incremental CIFAR100 learning with increasing buffer sizes.
  • Continual learning: 41.6% versus 33.9%: FRePo achieves higher final all-class accuracy than DM in 5-step class-incremental learning.In 10-step learning, the corresponding final accuracies are 38.0% and 34.0%.
  • Continual learning: Performance drops as continual-learning steps increase, suggesting that independently distilling separate class groups can create redundant information or miss distinguishable features.Distilling 2000 images continually reaches 41.6%, similar to 41.3% from distilling 1000 images from the whole dataset.
  • Membership-inference defense: Models trained on distilled data have attack AUCs close to random guessing, whereas models trained on real data remain vulnerable to membership-inference attacks.The evaluation uses five black-box attackers, including threshold, logistic-regression, multilayer-perceptron, random-forest, and k-nearest-neighbor attacks.
  • Membership-inference defense: Distilled-data training causes a small test-accuracy drop versus full-data training but achieves much better test performance than an equally sized real subset.Figure 5 presents the trade-off between test accuracy and attack effectiveness as training steps increase.

6 Conclusion

FRePo addresses meta-gradient computation and overfitting in dataset distillation, achieving strong benchmark performance with substantially lower training cost. The distilled data also supports continual learning and membership inference defense.

  • FRePo achieves state-of-the-art performance with 100x less training time and 10x less GPU memory.
  • The method produces distilled data that looks real and generalizes across a wide range of architectures.
  • FRePo-generated distilled data is demonstrated for continual learning and membership inference defense.
  • The paper frames synthetic data generation as a downstream classification direction with potential interpretability benefits.

A.1 Implementation Details

The implementation uses a convolutional model, regularized preprocessing, and a model-pool algorithm for efficient dataset distillation. Experiments vary label learning, update ordering, initialization, and implementation hyperparameters while documenting architecture-related effects.

  • Data Preprocessing: Regularized ZCA preprocessing is added for RGB datasets, with checkerboard ZCA used for high-resolution ImageNette and ImageWoof.The checkerboard transformation causes artifacts in the corresponding distilled images.
  • Models: FRePo uses a convolutional network with 3, 4, or 5 blocks for 32x32, 64x64, or 128x128 inputs, respectively.Blocks combine 3x3 convolution, normalization, ReLU, and 2x2 average pooling.
  • FRePo Algorithm: Algorithm 2 samples models from a pool, updates distilled images and labels using a meta-training loss, and periodically reinitializes models after K steps.
  • Architecture Effects: Model architecture can affect distilled data, motivating an architecture ablation study.
  • Initialization and Labels: Real-image initialization is used for convergence speed, while label learning is treated as a Boolean experimental hyperparameter.Gaussian-noise image initialization has little effect on final performance but changes convergence speed.
  • Training and Evaluation: The default implementation uses v2, which computes the meta-gradient before the online model update; the two update orders show no empirical difference.The method also uses distilled images and their flipped versions to mitigate mirroring effects for RGB datasets.

A.2 Experimental Setups

The experiments visualize distilled data, measure training efficiency and downstream behavior, and compare FRePo with prior methods across resolutions and datasets. Evaluations use reproduced baselines, controlled preprocessing, and specified accuracy, privacy, and continual-learning protocols.

  • Visualizations: Figure 1 visualizes selected FRePo images for CIFAR100, Tiny ImageNet, ImageNette, and ImageWoof under specified image-per-class and preprocessing settings.
  • Efficiency: Efficiency curves measure average wall-clock time per step and evaluate checkpoints at different training times using default models and hyperparameters.
  • Visualizations: Figures 4a–4d compare FRePo, MTT, and DSA distilled images and show FRePo distilled labels for CIFAR100 at 1 Img/Cls.The FRePo labels correspond to sunflower, girl, and bear, and are unnormalized logits.
  • Continual Learning: Continual-learning plots report multi-class accuracies across classes observed by a given time point over five class-split runs.The primary variance source is the class split.
  • Membership Inference: Privacy experiments track test accuracy and attack AUC as training steps increase, without data augmentation to expose distilled-data effects on membership inference.
  • Comparisons: Baseline comparisons reproduce prior methods across preprocessing and architectures, while FRePo results report the best value with or without label learning.

B Additional Results

Additional results evaluate FRePo on resized ImageNet-1K and membership inference defense. The results show useful accuracy from very small distilled sets and privacy preservation with good performance, while identifying an evaluation-loss limitation for complex data.

  • Resized ImageNet-1K: 7.5% and 9.7% Top1 accuracy are achieved on resized ImageNet-1K using 1k and 2k distilled examples, versus 1.1% and 1.4% for equally sized real subsets.The full dataset reaches 19.8% with MSE evaluation, while cross-entropy reaches 32.2% for the same model.
  • Membership Inference Defense: Training on distilled data can preserve privacy while maintaining good performance on membership inference experiments.FashionMNIST shows a similar trend to the reported MNIST results.

C Ablation Study

The ablations examine FRePo’s meta-gradient computation, model pooling, and robustness to training duration. Compared with TBPTT, FRePo reduces truncation-related overfitting by training part of the network to convergence.

  • C.1 FRePo vs TBPTT: FRePo computes meta-gradients through a kernel and feature extractor rather than backpropagating through inner optimization.This separates meta-gradient computation from the online model update.
  • C.1 FRePo vs TBPTT: TBPTT-generated distilled data loses test accuracy as training continues, indicating overfitting to its k-step training scheme.More truncation steps can alleviate but not eliminate truncation bias without intractable infinite unrolling.
  • C.1 FRePo vs TBPTT: Training a subset of the neural network to convergence allows FRePo to alleviate TBPTT’s truncation bias.FRePo also decouples meta-gradient computation from inner-level optimization, reducing dependence on a specific learning setup.
  • C.2 Model pool and Batchsize: Using more than one online model outperforms using one model, while ten models provide reasonable rather than best performance.The one-model setting corresponds to the training-and-resetting strategy used by previous methods.
  • C.2 Model pool and Batchsize: Both too few and too many online updates hurt performance, whereas the default of 100 updates provides a good starting point.Larger batch sizes may converge faster in training steps but increase time per step.

C.3 Initialization

The initialization ablations show that image and label initialization affect convergence and appearance, while learning labels becomes especially important for datasets with many classes.

  • Image initialization: 27.8% test accuracy results from combining real-image initialization with properly chosen Gaussian noise, versus 27.2% for the default hyperparameter.A noise standard deviation of 0.5 produces the reported 27.8% result.
  • Image initialization: Learning label scales and noise scales can significantly improve convergence speed.For CIFAR100, scaling the mean-centered one-hot label vector by 0.3 is reported as a good choice; 1.0 works well for ten-class datasets.
  • Image initialization: Different image initializations produce distilled images that become quite similar after optimization.This holds even though the images look very different at initialization.
  • C.4 Label Learning: Learning labels is often unnecessary for MNIST, Fashion MNIST, and CIFAR10 but becomes crucial for CIFAR100, Tiny ImageNet, and ImageNet.The paper links learned labels to class-similarity information, easier optimization, and more natural-looking images.
  • C.4 Label Learning: 7.5% test accuracy is achieved on ImageNet when labels are learned, compared with 1.6% when labels are fixed.This comparison is reported for models trained from scratch on distilled data.

C.5 Training Cost Analysis

FRePo’s training cost rises with the amount of distilled data because it uses all distilled examples for gradients and inverts an increasingly large KRR matrix.

  • C.5 Training Cost Analysis: Time per step and GPU memory usage increase as the number of distilled data increases.The paper attributes this to full-dataset gradient computation and KRR matrix inversion with O(N^3) time complexity.
  • C.5 Training Cost Analysis: FRePo may struggle to distill tens of thousands of examples because of kernel-method scaling costs.The paper suggests batch sampling, subset distillation, or class-wise distillation as possible mitigations, while expecting performance trade-offs.
  • C.5 Training Cost Analysis: The reported cost measurements use milliseconds for time per step and gigabytes for peak GPU memory.Tables 8–11 provide numerical values corresponding to the cost figures.

C.6 Model Architectures

Architecture and normalization choices shape distilled-image appearance and transferability. Conv-BN training provides strong cross-architecture behavior, while instance normalization transfers less well, especially to architectures without normalization.

  • Architecture effects: Distilled data can encode an architecture’s inductive bias, producing different images across architectures and normalization layers.The study evaluates Conv, DCConv, AlexNet, VGG, ResNet, and several normalization choices.
  • Qualitative Results: The simplest Conv architecture produces images that look almost like real images, while normalization changes brightness and contrast.No normalization or batch normalization appears to generate the most natural-looking images.
  • Quantitative Results: Training with Conv-BN and evaluating with Conv-NN yields the best performance among the reported architecture choices.Conv-BN also produces the most generalizable and transferable images across other architectures.
  • Quantitative Results: Training with Conv-BN produces distilled data that transfers well to architectures unseen during training.Cross-architecture transfer is evaluated on CIFAR10 with 10 images per class.
  • Instance Normalization: Instance-normalized distilled data transfers less well to other architectures, especially architectures without normalization.Conv-IN performs slightly worse than the default Conv-BN, while DCConv-IN degrades substantially when distilling 1,000 CIFAR100 images.
  • Practical modifications: Dropout, learning-rate schedules, and training-time data augmentation can improve test accuracy but are omitted for simplicity or fair comparison.These modifications are presented as potentially useful to practitioners.
Loading 2206.00719v2…