DistillKac: Few-Step Image Generation via Damped Wave Equations
This work addresses the computational inefficiency in image generation for AI applications, offering a novel approach that is incremental in improving upon existing diffusion models.
The authors tackled the problem of slow image generation in diffusion models by introducing DistillKac, a method based on damped wave equations that enforces finite speed transport to avoid stiffness and unbounded propagation. The result is a fast image generator that produces high-quality samples with very few function evaluations while maintaining numerical stability.
We present DistillKac, a fast image generator that uses the damped wave equation and its stochastic Kac representation to move probability mass at finite speed. In contrast to diffusion models whose reverse time velocities can become stiff and implicitly allow unbounded propagation speed, Kac dynamics enforce finite speed transport and yield globally bounded kinetic energy. Building on this structure, we introduce classifier-free guidance in velocity space that preserves square integrability under mild conditions. We then propose endpoint only distillation that trains a student to match a frozen teacher over long intervals. We prove a stability result that promotes supervision at the endpoints to closeness along the entire path. Experiments demonstrate DistillKac delivers high quality samples with very few function evaluations while retaining the numerical stability benefits of finite speed probability flows.