Imagine Flash: Accelerating Emu Diffusion Models with Backward Distillation

Jonas Kohler,Albert Pumarola,Edgar Schönfeld,Artsiom Sanakoyeu,Roshan Sumbaly,Peter Vajda,Ali Thabet
2024-05-09
Abstract:Diffusion models are a powerful generative framework, but come with expensive inference. Existing acceleration methods often compromise image quality or fail under complex conditioning when operating in an extremely low-step regime. In this work, we propose a novel distillation framework tailored to enable high-fidelity, diverse sample generation using just one to three steps. Our approach comprises three key components: (i) Backward Distillation, which mitigates training-inference discrepancies by calibrating the student on its own backward trajectory; (ii) Shifted Reconstruction Loss that dynamically adapts knowledge transfer based on the current time step; and (iii) Noise Correction, an inference-time technique that enhances sample quality by addressing singularities in noise prediction. Through extensive experiments, we demonstrate that our method outperforms existing competitors in quantitative metrics and human evaluations. Remarkably, it achieves performance comparable to the teacher model using only three denoising steps, enabling efficient high-quality generation.
Computer Vision and Pattern Recognition
What problem does this paper attempt to address?
The paper proposes a solution to the problem of long inference time and low efficiency in the diffusion model, a powerful generative framework. Existing acceleration methods may affect image quality at very low numbers of steps and performance under complex conditions. The paper introduces a new distillation framework, consisting of three key components: reverse distillation to reduce the gap between training and inference distributions, dynamically adjusted knowledge transfer loss to adapt to different time steps, and noise correction to improve sample quality. This approach enables the model to generate high-fidelity and diverse images within 1 to 3 steps, outperforming existing methods and striking a good balance between quality and efficiency.