From Reconstruction to Compression: Reinforcement Learning in Diffusion-Based Distillation

Zhouyang Jiang, Chen Yuzheng, Yijun Liu, Luo Chengguo, Yucheng Hu, Jiawei Du, Yew-Soon Ong, Joey Tianyi Zhou

Abstract:

Dataset distillation synthesizes compact datasets that retain
the training utility of much larger ones. While diffusion models are natural candidates for this task due to their generative
capabilities, there are few methods that adopt them in dataset
distillation compared to the matching-based approaches and
label-relaxation approaches. A key reason is the fundamental
mismatch between diffusion objectives and distillation goals:
diffusion models are trained to reconstruct high-fidelity data,
whereas distillation requires compressed, task-relevant representations. We address this gap by proposing a reinforcement
learning (RL)-guided framework that steers diffusion models from reconstruction toward compression. By formulating
sampling as a decision process, we optimize the generative
trajectory using rewards derived from student model performance. This enables the generation of synthetic samples that
maximize learning utility under strict compression budgets.
Unlike prior static modifications of the diffusion process, our
method dynamically adapts generation based on downstream
outcomes. Experiments on standard benchmarks show that
our RL-guided diffusion approach consistently improves both
performance and efficiency, advancing the frontier of generative dataset distillation.