Efficient Dataset Distillation Using Random Feature Approximation

Noel Loo; Ramin Hasani; Alexander Amini; Daniela Rus

Source-linked AI summary

Efficient Dataset Distillation Using Random Feature Approximation

Noel Loo, Ramin Hasani, Alexander Amini, Daniela Rus

arXiv:2210.12067v1Oct 21, 2022cs.LGcs.AIcs.NEstat.ML

TL;DR

Dataset distillation needs a faster way to compress large datasets into useful synthetic coresets because KIP’s exact kernel computation scales quadratically with support-set size. RFAD approximates the NNGP kernel with random features, reducing complexity to linear in support-set size; it achieves a 100-fold speedup over KIP while retaining accuracy and supports interpretability and privacy applications.

Problem
KIP’s exact NTK kernel computation scales as O(|S|^2), making accurate dataset distillation computationally expensive as support-set size grows.
Method
RFAD uses random neural networks to approximate the NNGP kernel, reducing dataset-distillation complexity from O(|S|^2) to O(|S|).
Results
RFAD provides a 100-fold speedup over KIP while retaining accuracy across dataset-distillation evaluations.
Takeaways & Limitations
RFAD enables efficient dataset distillation for model interpretability and privacy preservation.
Takeaways & Limitations
Distilled datasets transfer poorly across finite networks with and without instance normalization, reflecting differing feature representations.

Abstract

from arXiv · show

Dataset distillation compresses large datasets into smaller synthetic coresets which retain performance with the aim of reducing the storage and computational burden of processing the entire dataset. Today's best-performing algorithm, \textit{Kernel Inducing Points} (KIP), which makes use of the correspondence between infinite-width neural networks and kernel-ridge regression, is prohibitively slow due to the exact computation of the neural tangent kernel matrix, scaling $O(|S|^2)$, with $|S|$ being the coreset size. To improve this, we propose a novel algorithm that uses a random feature approximation (RFA) of the Neural Network Gaussian Process (NNGP) kernel, which reduces the kernel matrix computation to $O(|S|)$. Our algorithm provides at least a 100-fold speedup over KIP and can run on a single GPU. Our new method, termed an RFA Distillation (RFAD), performs competitively with KIP and other dataset condensation algorithms in accuracy over a range of large-scale datasets, both in kernel regression and finite-width network training. We demonstrate the effectiveness of our approach on tasks involving model interpretability and privacy preservation.

1 Introduction

Dataset distillation compresses datasets into optimized synthetic coresets, but KIP’s exact NTK computation limits applicability. RFAD addresses this bottleneck with a faster approximation while retaining competitive accuracy and supporting interpretability and privacy applications.

1 Introduction: Synthetic coresets summarize large datasets for downstream tasks and permit continuous optimization beyond selecting existing images and labels.Applications include faster training, reduced catastrophic forgetting, and improved interpretability.
1 Introduction: Over 100-fold speedup over KIP is reported, with RFAD exceeding KIP’s CIFAR-10 performance.Figure 1 presents the speed and CIFAR-10 comparison.
1 Introduction: KIP achieves state-of-the-art dataset-distillation accuracy but its exact NTK computation is prohibitively expensive and limits applicability.KIP computes infinite-width network outputs through NTK ridge regression.
1 Introduction: RFAD reduces KIP’s support-set complexity from O(|S|^2) to O(|S|) while retaining similar accuracy and substantially improving performance.The method uses a random-feature approximation of the NNGP kernel.
1 Introduction: RFAD is presented as fast, accurate, and scalable, with applications in model interpretability and privacy preservation.The paper’s contributions explicitly include both application areas.

2 Background and Related Work

Prior work spans representative coresets, synthetic dataset distillation, infinite-width neural-network kernels, and random-feature approximations. These lines of work motivate replacing expensive exact kernel calculations with finite feature representations.

2 Background and Related Work: Coresets are smaller data subsets designed to preserve model performance relative to training on the full dataset.Standard selection methods commonly use importance sampling.
2 Background and Related Work: Dataset distillation generates synthetic datasets whose training performance aims to match training on the original data.Related methods use gradient matching, differentiable augmentation, and distribution matching.
2 Background and Related Work: Infinite-width neural networks connect Gaussian processes and kernel methods, while gradient-descent training corresponds to kernel ridge regression with the NTK.The NNGP forms the leading term of the NTK in the cited formulation.
2 Background and Related Work: Random-feature methods approximate kernel feature maps with finite-dimensional vectors, limiting kernel rank and enabling faster matrix inversion.This provides a general basis for scaling kernel methods to larger datasets.

3 Algorithm Setup and Design

RFAD modifies KIP by replacing the exact NTK with an NNGP random-feature approximation and by using Platt scaling. The resulting pipeline constructs kernel matrices more efficiently while optimizing synthetic support sets against training data.

3 Algorithm Setup and Design: KIP computes infinite-width network outputs as f(X_T) = K_TS(K_SS + λI)^−1y_S and optimizes support images and labels through a training loss.Its kernel matrices use the NTK and require exact computation.
3 Algorithm Setup and Design: The NNGP approximation is motivated because NNGP is the leading NTK term, while the remaining fully connected contributions decay for ReLU activations.The authors verify performance under NTK and finite networks trained by gradient descent.
3 Algorithm Setup and Design: RFAD samples finite features by evaluating randomly initialized neural networks on training batches and support sets.The resulting features form empirical kernel matrices used for kernel-ridge predictions.
3 Algorithm Setup and Design: Random features reduce kernel construction to O(|T| + |S|) for K_TS and O(|S|) for K_SS, versus O(|T||S|) and O(|S|^2) in KIP.The runtime becomes linear in coreset size because feature construction dominates matrix inversion.
3 Algorithm Setup and Design: Platt scaling replaces MSE with cross-entropy using a jointly learned positive temperature parameter, addressing loss behavior in constrained distillation.The method uses the same kernel-ridge output before temperature scaling and softmax conversion.
3 Algorithm Setup and Design: RFAD combines NNGP replacement, random-feature approximation, and Platt scaling into one dataset-distillation algorithm.These three changes are implemented in Algorithm 1.

4 Experiments with RFAD

RFAD is evaluated across dataset-distillation benchmarks, efficiency, kernel transfer, finite-network transfer, and ablations. It retains competitive accuracy while substantially reducing training time relative to KIP.

Benchmarks: RFAD outperforms other methods in almost every fixed-label dataset setting and is at least two orders of magnitude faster than KIP.It exceeds KIP by up to 6.1% on CIFAR-10 with 10 images per class; learned-label performance is slightly lower than KIP.
Time Savings during training: RFAD training time is linear in coreset size and the number of random models, whereas KIP’s exact kernel computation is quadratic in coreset size.Both methods converge in 3000–15000 iterations, but RFAD takes 1–14 hours versus several hundred GPU hours for KIP.
NTK Kernel and Finite Network Transfer: NNGP-trained RFAD coresets generally transfer to NTK with little or no performance loss, with the largest gap reaching 8% on SVHN with 10 images per class.On some datasets, including CIFAR-100 with 10 images per class, NTK evaluation outperforms NNGP evaluation.
NTK Kernel and Finite Network Transfer: RFAD coresets achieve the best finite-network transfer performance in almost all benchmarks against KIP and often against DC/DSA.The authors attribute this gain primarily to centering and label scaling during training.
Empirical NNGP Performance at Inference: Performance depends mainly on total random-feature count rather than individual network width, with 70% accuracy achieved using one convolutional channel.The standard training configuration used 8 models with 256 convolutional channels each.
Objective Function Sensitivity: Cross-entropy with Platt scaling substantially improves performance, including an 8% gain on Fashion-MNIST with one image per class.The ablation varied the loss and the number of models used during training on Fashion-MNIST and CIFAR-10.

5 RFAD Application I: Interpretability

RFAD enables influence-based explanations by measuring how coreset elements affect predictions and connecting source-training examples to those coreset influences. It also supports semantically meaningful queries and highly corrupted but accurate coresets.

Influence-based explanations: Influence scores identify coreset elements whose removal most changes a test prediction, while influence embeddings connect those effects to original training examples.The method uses prediction differences and cosine similarity for relatively fast queries.
Influence-based explanations: On CIFAR-10, influence-based queries for both correct and incorrect predictions are visually similar to the test examples.These queries can help explain incorrect predictions and identify harmful items or data-collection gaps.
Corrupted coresets: RFAD can produce corrupted coresets whose pixels contain little human-understandable source information while retaining high test performance.Corruption is implemented by initializing images with random noise and preserving a random fraction of pixels at initialization.

6 RFAD Application II: Privacy

RFAD supports privacy-oriented distillation by retaining task performance while replacing much of the coreset content with random noise. The reported results show this effect on CIFAR-10 and CelebA.

CIFAR-10: At 90% corruption, CIFAR-10 reaches 40.6% accuracy with one sample per class, above the 16.1% natural baseline.The evaluation used corruption ratios up to 0.9 and 1, 10, or 50 images per class.
CelebA: CelebA reaches 81% accuracy with only two images while 95% of their pixels are random noise.The task is male/female classification using one image per class.

7 Conclusions

RFAD achieves a 100-fold speedup over KIP while retaining accuracy, and remains effective under substantial corruption. Its limitations include normalization-specific transfer and possible overfitting with low compression.

7 Conclusions: RFAD performance degrades very gradually with noise, retaining high performance at 90% corruption on CIFAR-10 and CelebA.These privacy-related experiments used varying support-set sizes and corruption ratios.
7 Conclusions: 100-fold speedup over KIP is achieved while retaining accuracy.The speedup comes from replacing exact NTK computation with an approximate NNGP kernel, reducing complexity from O(|S|^2) to O(|S|).
7 Conclusions: Datasets distilled with or without instancenorm transfer poorly across networks that differ in instancenorm usage.The authors attribute this to differing features between normalized and unnormalized networks.
7 Conclusions: Platt loss may overfit simple datasets or large coresets, making it potentially detrimental when the compression ratio is low.The authors observed near-100% training accuracy in such cases and hypothesize that adding images would not improve performance.

Checklist

The paper’s checklist reports compliance for contributions, limitations, reproducibility, training details, compute reporting, and asset citations, while marking inapplicable ethics items as N/A.

Checklist: The paper reports that its main claims and limitations are accurately described.The checklist marks both contribution accuracy and discussion of limitations as satisfied.
Checklist: Reproducibility materials, training details, error bars, and compute resources are reported as included.The checklist marks each of these experimental-reporting items as satisfied.
Checklist: Existing assets are cited, while licensing and new-asset disclosure are marked not applicable.The checklist marks citation of existing-asset creators as satisfied and the other two items as N/A.
Checklist: Consent, personally identifiable information, offensive content, and human-subject procedures are marked not applicable.The checklist marks the relevant data-governance and participant-risk items as N/A.
Checklist: RFAD’s time complexity is O(|S|) when the coreset is no larger than the training batch, versus KIP’s O(|S|^2) in the corresponding regime.For larger batches, RFAD scales as O(λ1|S| + λ2|B|), while KIP uses O(|B||S|) and can be effectively worse than quadratic scaling.
Checklist: The experiments use standard preprocessing, convolutional architectures, Gaussian initialization, and no normalization layers unless otherwise stated.The setup covers black/white and color preprocessing, three-layer ConvNets, and Gaussian weight and bias initialization.
Checklist: Removing the final fully connected layer for NNGP feature computation did not affect performance.RFAD instead uses representations after the final ReLU layer, which also reduces memory and final-layer variance.

D.3 Training

The training procedure uses learned transformations, regularization, validation-based early stopping, and controlled experimental protocols to optimize RFAD coresets and evaluate their transfer.

D.3 Training: RFAD training uses Adabelief with learning rate 1e-3, ε = 1e-16, and an adaptive kernel regularizer with λ0 = 5e-3.The regularizer is applied when computing (KSS + λI)^-1.
D.3 Training: A learned transformation matrix speeds convergence without adding coreset variables.It is initialized as the identity and trained with a small learning rate of 5e-5.
D.3 Training: The coreset is initialized with class-balanced real data and zero-mean one-hot labels, with Platt-scale parameter τ learned logarithmically when applicable.The logarithm of τ uses learning rate 1e-2.
D.3 Training: Each gradient step uses 5120 training examples, processed in four batches, with no-gradient feature computation and double-precision matrix inversion.The no-gradient flag saves memory, while double precision reduces sensitivity to rounding errors.
D.3 Training: Training uses validation loss measured every 40 epochs with patience of 1000 iterations for early stopping.The validation set contains 1000 examples and uses 16 fixed random neural networks.
D.3 Training: Runtime experiments average 200 iterations across N = 1, 2, 4, 8 on a single RTX 3090, with a 40% adjustment for V100 comparison.Finite-network evaluations average 12 training runs, while privacy experiments use one run per corruption ratio.

E Time taken additional results

Additional runtime results report original RFAD times and total convergence runtimes, with all distilled datasets converging in under 10 hours on one RTX 3090.

E Time taken additional results: The reported additional runtimes use original iteration times rather than the 40% V100-comparison adjustment.The section also reports total training iterations and runtimes for each distillation result.
E Time taken additional results: Figure 9 compares CIFAR-10 time per training iteration across model counts and support sizes against KIP on an RTX 3090.The figure varies both the number of models and support-set size.
E Time taken additional results: All RFAD distilled datasets converge in under 10 hours on a single RTX 3090.Label-learning configurations usually require longer runtimes.

F Centering and label scaling for finite networks ablations

Label scaling and centering improve finite-network transfer, especially for small support sets, while instance-normalization mismatches impose a large transfer penalty.

F Centering and label scaling for finite networks ablations: Label scaling and centering improve finite-network transfer performance, particularly for smaller support sets.The paper reports no theoretical explanation for label-scaling gains and observes improvement for α > 1.
F Centering and label scaling for finite networks ablations: Transferring RFAD distilled datasets between networks with and without instance normalization incurs a large performance penalty.This applies in both NNGP kernel-ridge regression and finite-network evaluation.
F Centering and label scaling for finite networks ablations: The instance-normalization experiments use an empirical NNGP kernel because the neural-tangents library lacks an exact implementation.

I Empirical NNGP at Inference additional results

Additional experiments show that empirical NNGP inference closely approaches exact-NNGP performance for convolutional architectures with at least 128 channels, including a 70% CIFAR-10 result with one channel.

I Empirical NNGP at Inference additional results: Empirical NNGP inference achieves close to exact-NNGP performance for convolutional architectures with C ≥128 across the RFAD distilled datasets.
I Empirical NNGP at Inference additional results: 70% CIFAR-10 accuracy is achieved with 10 images per class and fixed labels using an empirical NNGP kernel from one-channel random convolutional networks.The experiment uses a distilled coreset and is reported for inference.
I Empirical NNGP at Inference additional results: The qualitative CIFAR-10 figures compare incorrectly and correctly predicted test images with their most relevant coreset and training-set images.
I Empirical NNGP at Inference additional results: The additional inference figures cover MNIST, Fashion-MNIST, SVHN, CIFAR-10, and CIFAR-100 with both fixed and learned labels.