Dataset Distillation with Probabilistic Latent Features

TL;DR

This paper introduces a probabilistic latent feature approach for dataset distillation, enabling the creation of compact synthetic datasets that effectively train models, with improved diversity and spatial structure capture.

Contribution

It proposes a novel stochastic method modeling joint latent feature distribution with a low-rank normal distribution, enhancing diversity and spatial structure in synthetic data.

Findings

Achieves state-of-the-art performance on ImageNet, CIFAR-10, and MedMNIST.

Maintains low computational complexity with a lightweight network.

Demonstrates strong cross-architecture generalization.

Abstract

As deep learning models grow in complexity and the volume of training data increases, reducing storage and computational costs becomes increasingly important. Dataset distillation addresses this challenge by synthesizing a compact set of synthetic data that can effectively replace the original dataset in downstream classification tasks. While existing methods typically rely on mapping data from pixel space to the latent space of a generative model, we propose a novel stochastic approach that models the joint distribution of latent features. This allows our method to better capture spatial structures and produce diverse synthetic samples, which benefits model training. Specifically, we introduce a low-rank multivariate normal distribution parameterized by a lightweight network. This design maintains low computational complexity and is compatible with various matching networks used in dataset distillation. After distillation, synthetic images are generated by feeding the learned latent features into a pretrained generator. These synthetic images are then used to train classification models, and performance is evaluated on real test set. We validate our method on several benchmarks, including ImageNet subsets, CIFAR-10, and the MedMNIST histopathological dataset. Our approach achieves state-of-the-art cross architecture performance across a range of backbone architectures, demonstrating its generality and effectiveness.