Vector-Quantized Soft Label Compression for Dataset Distillation

TL;DR

A vector-quantized autoencoder (VQAE) for compressing soft labels is introduced for compressing soft labels, achieving substantial compression while preserving the effectiveness of the distilled data.

Abstract

Dataset distillation is an emerging technique for reducing the computational and storage costs of training machine learning models by synthesizing a small, informative subset of data that captures the essential characteristics of a much larger dataset. Recent methods pair synthetic samples and their augmentations with soft labels from a teacher model, enabling student models to generalize effectively despite the small size of the distilled dataset. While soft labels are critical for effective distillation, the storage and communication overhead they incur, especially when accounting for augmentations, is often overlooked. In practice, each distilled sample is associated with multiple soft labels, making them the dominant contributor to storage costs, particularly in large-class settings such as ImageNet-1K. In this paper, we present a rigorous analysis of bit requirements across dataset distillation frameworks, quantifying the storage demands of both distilled samples and their soft labels. To address the overhead, we introduce a vector-quantized autoencoder (VQAE) for compressing soft labels, achieving substantial compression while preserving the effectiveness of the distilled data. We validate our method on both vision and language distillation benchmarks. On ImageNet-1K, our proposed VQAE achieves 30--40x additional compression over RDED, LPLD, SRE2L, and CDA baselines while retaining over $90\%$ of their original performance.

Figures (3)

• Figure 1: (a) We begin with a standard dataset distillation pipeline, where a small number of distilled inputs per class (IPC) are passed through a teacher model to produce soft labels across multiple augmentations. These soft labels, although highly informative, impose a significant memory and communication burden, especially for large numbers of classes and training epochs. (b) To mitigate this overhead, we introduce a vector-quantized autoencoder (VQ-AE) that encodes soft labels into compact latent representations segmented and quantized using a learned dictionary. Only the code indices, dictionary, and decoder need to be stored and transmitted. At distillation time, soft labels are reconstructed from these compact codes and used to train student models via KL divergence. This approach dramatically reduces storage while preserving the fidelity of the teacher’s knowledge.
• Figure 2: Performance ratio of training with compressed over uncompressed soft labels at different compression levels, averaged across image distillation techniques presented in Table \ref{['tab: main_table']} at each IPC level. Our method consistently outperforms LPLD across all compression factors.
• Figure 3: Accuracy vs. storage across codebook size $k$ and block size $d_c$.