Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

DistillKac: Few-Step Image Generation via Damped Wave Equations

Weiqiao Han, Chenlin Meng, Christopher D. Manning, Stefano Ermon

TL;DR

DistillKac introduces a finite-speed alternative to diffusion-based image generation by leveraging the damped wave (telegrapher) equation and its stochastic Kac representation, ensuring probability mass moves with a speed cap $c$ and globally bounded kinetic energy. The method extends classifier-free guidance to velocity space and presents endpoint-only distillation to produce few-step samplers with a formal endpoint-to-trajectory stability guarantee, addressing stiffness and instability seen in diffusion models. Empirical results on CIFAR-10, CelebA-64, and LSUN Bedroom-256 show that DistillKac achieves competitive image quality at substantially fewer function evaluations, with multi-stage distillation further enhancing efficiency. Overall, the work offers a theoretically grounded, stable framework for fast image generation via hyperbolic PDE-based flows and structured distillation, and suggests avenues for genuinely multi-dimensional Kac processes and broader applications.

Abstract

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.

DistillKac: Few-Step Image Generation via Damped Wave Equations

TL;DR

DistillKac introduces a finite-speed alternative to diffusion-based image generation by leveraging the damped wave (telegrapher) equation and its stochastic Kac representation, ensuring probability mass moves with a speed cap and globally bounded kinetic energy. The method extends classifier-free guidance to velocity space and presents endpoint-only distillation to produce few-step samplers with a formal endpoint-to-trajectory stability guarantee, addressing stiffness and instability seen in diffusion models. Empirical results on CIFAR-10, CelebA-64, and LSUN Bedroom-256 show that DistillKac achieves competitive image quality at substantially fewer function evaluations, with multi-stage distillation further enhancing efficiency. Overall, the work offers a theoretically grounded, stable framework for fast image generation via hyperbolic PDE-based flows and structured distillation, and suggests avenues for genuinely multi-dimensional Kac processes and broader applications.

Abstract

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.

Paper Structure

This paper contains 25 sections, 7 theorems, 29 equations, 12 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Theoretical Foundations
  3. Partial Differential Equations
  4. Kac Process and Random Flights
  5. Image Generation via Component-wise 1-D Telegrapher Equation
  6. Notations
  7. Classifier Free Guidance in Velocity Space
  8. Distillation
  9. Experiments
  10. Setup
  11. Main Results
  12. Related Work
  13. Discussion and Future Work
  14. Extra Proofs
  15. Proof of Theorem \ref{['thm:cfg']}
  16. ...and 10 more sections

Key Result

Theorem 0

orsingher2007random For $d = 1, 2$, and $4$, the distribution of the state $X(t)$ can be obtained analytically, while for any other positive integer $d$ there is no explicit formula.

Figures (12)

  • Figure 1: Uncurated CIFAR-10 (left) and CelebA-64 (right) samples generated by Guided Kac Flow (midpoint integrator, 100 steps, guidance $w = 3$ for CIFAR-10 and $w = 0$ for CelebA-64).
  • Figure 2: Left: CIFAR-10 FID vs. guidance strength $w$ for the 100-step Guided Kac Flow. Second-order integrators (midpoint and AB-2) outperform Euler. Right: CIFAR-10 FID vs. integration steps. DistillKac substantially reduces FID at 20, 4, 2, and 1 steps relative to Guided Kac Flow.
  • Figure 3: Time (seconds) to sample 50k images with one H100 GPU at 1, 2, 4, and 20 steps (NFE).
  • Figure 4: Uncurated LSUN Bedroom-256 samples generated by Guided Kac Flow (midpoint integrator, 100 steps, guidance $w = 0$).
  • Figure 5: Uncurated LSUN Bedroom-256 samples generated by DistillKac (Euler integrator, 20 NFE).
  • ...and 7 more figures

Theorems & Definitions (8)

  • Theorem 0
  • Theorem 0
  • Theorem 0
  • Lemma 0
  • Lemma 0
  • Theorem 0
  • proof
  • Corollary 0