Disentangled Embedding Optimization for Diffusion-Based Dataset Distillation

arXiv:2611.16224v1 Announce Type: new Abstract: Dataset distillation aims to compress large-scale datasets into compact synthetic subsets that enable efficient model training with comparable performance to the original data. Recent advances show that diffusion models are highly effective for this task due to their strong generative capability. Existing diffusion-based approaches typically perform class-wise clustering and use the obtained cluster prototypes as diffusion inputs to generate distilled samples. However, the uniformly distributed original samples in a category inevitably lead to entanglement near cluster boundaries, pulling the cluster prototypes closer and reducing the inter-cluster separation. Consequently, the diffusion model receives overly similar latent inputs, resulting in synthetic data with limited diversity. To address this issue, we propose Disentangled Embedding Optimization (DEO), a simple yet effective framework that replaces static clustering with trainable embeddings. These embeddings are optimized to match the Gaussian statistics (i.e., mean and covariance) of the original data distribution, while a diversity regularizer is designed to encourage separation among these embeddings to improve the intra-class variability. Extensive experiments demonstrate that our proposed DEO can generate more diversified and representative distilled data than existing cluster-based diffusion methods, achieving state-of-the-art performance and superior generalization across multiple datasets.