Dataset Distillation with Probabilistic Latent Features

TL;DR

This paper addresses the growing costs of large-scale data by introducing Stochastic Latent Feature Distillation (SLFD), a dataset distillation framework that explicitly models spatial uncertainty in latent features via a low-rank multivariate normal distribution. By sampling diverse latent representations and using a pretrained generator, SLFD produces synthetic images whose training dynamics align with those of real data through gradient matching, improving cross-architecture generalization. The method demonstrates state-of-the-art performance on ImageNet subsets, CIFAR-10, and MedMNIST across multiple backbones and resolutions, including strong results in medical imaging domains. Its lightweight stochastic module, compatibility with existing distillation pipelines, and robustness across domains suggest substantial practical impact for efficient data sharing and privacy-preserving training.

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.

Figures (6)

• Figure 1: Overview of the proposed SLFD framework. (a) The distillation backbone is based on GLaD cazenavette2023generalizing, where the generator (shown in gray to indicate frozen weights) is split into two parts: the first 16 layers compute intermediate features $f_n$ from latent codes $w_s$ (black arrows), and the remaining 16 layers generate synthetic images using both $w_s$ and $f_n$ (red arrows). (b) The Stochastic Module models uncertainty in the latent space by estimating a low-rank multivariate normal distribution using three linear layers to compute the mean, variance, and covariance factors. Multiple samples drawn from this distribution are used to represent synthetic probability distributions. (c) The Gradient Matching component computes the matching loss between gradients of real and synthetic data. (d) In the Evaluation phase, classifiers are trained solely on the distilled synthetic images.
• Figure 2: Illustration Eq. \ref{['samplemc']} about how we sample probabilities.
• Figure 3: Example images distilled from ImageNet subsets.
• Figure 4: Synthetic images from MedMNIST yang2023medmnist for different output resolutions.
• Figure 5: Results of ablation study. Distilled images with the diffusion model UViT bao2023all with resolution $(256\times 256)$.