Vector-Quantized Soft Label Compression for Dataset Distillation

TL;DR

This paper introduces a vector-quantized autoencoder to compress soft labels in dataset distillation, significantly reducing storage costs while maintaining high performance on vision and language benchmarks.

Contribution

It presents a novel VQAE method for compressing soft labels in dataset distillation, reducing storage overhead by 30-40x without sacrificing effectiveness.

Findings

Achieves 30-40x compression on ImageNet-1K.

Retains over 90% of baseline performance.

Validates effectiveness on vision and language benchmarks.

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.