[Paper Review] Dataset Distillation with Convexified Implicit Gradients

Noel Loo, Ramin Hasani+2 more|arXiv (Cornell University)|Feb 13, 2023

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes RCIG, a new dataset distillation algorithm that leverages reparameterized convexified implicit gradients to achieve state-of-the-art performance. By convexifying the inner optimization via a frozen finite-width neural tangent kernel and reducing bias through analytical parameterization, RCIG achieves 108% higher accuracy than prior SOTA on ImageNet with one image per class.

ABSTRACT

We propose a new dataset distillation algorithm using reparameterization and convexification of implicit gradients (RCIG), that substantially improves the state-of-the-art. To this end, we first formulate dataset distillation as a bi-level optimization problem. Then, we show how implicit gradients can be effectively used to compute meta-gradient updates. We further equip the algorithm with a convexified approximation that corresponds to learning on top of a frozen finite-width neural tangent kernel. Finally, we improve bias in implicit gradients by parameterizing the neural network to enable analytical computation of final-layer parameters given the body parameters. RCIG establishes the new state-of-the-art on a diverse series of dataset distillation tasks. Notably, with one image per class, on resized ImageNet, RCIG sees on average a 108\% improvement over the previous state-of-the-art distillation algorithm. Similarly, we observed a 66\% gain over SOTA on Tiny-ImageNet and 37\% on CIFAR-100.

Motivation & Objective

• Address the challenge of achieving higher performance in dataset distillation, especially with minimal data per class.
• Overcome limitations of existing implicit gradient-based methods, which suffer from non-convexity and high bias in gradient estimation.
• Develop a scalable and accurate distillation algorithm that generalizes across diverse benchmarks including ImageNet, Tiny-ImageNet, and CIFAR-100.
• Improve robustness to membership inference attacks by leveraging distillation as a defense mechanism.

Proposed method

• Formulate dataset distillation as a bi-level optimization problem with model weights as inner parameters and distilled data as outer parameters.
• Apply the implicit function theorem to compute meta-gradients via backpropagation through the inner optimization loop.
• Introduce convexification by approximating the inner optimization using a frozen finite-width neural tangent kernel (NTK), transforming the non-convex problem into a convex one.
• Reparameterize the neural network to enable analytical computation of final-layer parameters from body parameters, reducing bias in implicit gradient estimation.
• Optimize the distilled data using a hybrid approach that combines implicit gradients with linearized training for improved stability and convergence.
• Scale the method efficiently by subsampling backward passes, maintaining low memory usage even on large datasets.

Experimental results

Research questions

• RQ1 Can convexifying the inner optimization problem through a frozen NTK improve the performance of implicit gradient-based dataset distillation?
• RQ2 How does reparameterization of the neural network architecture reduce bias in implicit gradient computation?
• RQ3 To what extent does RCIG outperform existing SOTA methods across diverse benchmarks with varying data complexity and class counts?
• RQ4 Can dataset distillation via RCIG serve as an effective defense against membership inference attacks?
• RQ5 What hyperparameter settings (e.g., inner steps, Hessian inverse steps) yield optimal performance across different datasets?

Key findings

• RCIG achieves a 108% accuracy improvement over the previous SOTA on ImageNet with one image per class, reaching 15.6% top-1 accuracy.
• On Tiny-ImageNet, RCIG achieves a 66% gain over the prior state-of-the-art, demonstrating strong generalization on smaller, more complex datasets.
• On CIFAR-100, RCIG improves performance by 37% compared to the previous SOTA, confirming its effectiveness across multiple benchmarks.
• The method maintains high robustness to membership inference attacks, with models trained on distilled data showing low AUC (e.g., 0.53) across multiple attack types, indicating strong privacy preservation.
• Ablation studies show that using both implicit gradients and Hessian inverse steps (nH−1 > 0) significantly improves accuracy, with optimal performance at ninner = 20 and nH−1 = 20.
• RCIG scales efficiently to large datasets without increasing memory footprint, thanks to subsampled backward passes during training.