Toward Theoretical Insights into Diffusion Trajectory Distillation via Operator Merging
This work addresses a theoretical gap for researchers in diffusion models, providing insights to optimize distillation strategies, though it is incremental as it builds on existing trajectory distillation methods.
The paper tackles the lack of theoretical insights into diffusion trajectory distillation methods, which accelerate sampling in diffusion models but face trade-offs between strategies and generative quality. It reinterprets distillation as an operator merging problem, proposes a dynamic programming algorithm for optimal merging, and demonstrates a sharp phase transition in strategies based on data covariance, enhancing theoretical understanding and offering practical improvements.
Diffusion trajectory distillation methods aim to accelerate sampling in diffusion models, which produce high-quality outputs but suffer from slow sampling speeds. These methods train a student model to approximate the multi-step denoising process of a pretrained teacher model in a single step, enabling one-shot generation. However, theoretical insights into the trade-off between different distillation strategies and generative quality remain limited, complicating their optimization and selection. In this work, we take a first step toward addressing this gap. Specifically, we reinterpret trajectory distillation as an operator merging problem in the linear regime, where each step of the teacher model is represented as a linear operator acting on noisy data. These operators admit a clear geometric interpretation as projections and rescalings corresponding to the noise schedule. During merging, signal shrinkage occurs as a convex combination of operators, arising from both discretization and limited optimization time of the student model. We propose a dynamic programming algorithm to compute the optimal merging strategy that maximally preserves signal fidelity. Additionally, we demonstrate the existence of a sharp phase transition in the optimal strategy, governed by data covariance structures. Our findings enhance the theoretical understanding of diffusion trajectory distillation and offer practical insights for improving distillation strategies.