Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

One-Step Offline Distillation of Diffusion-based Models via Koopman Modeling

Nimrod Berman, Ilan Naiman, Moshe Eliasof, Hedi Zisling, Omri Azencot

TL;DR

The paper tackles the high computational cost of diffusion-model sampling by introducing Koopman Distillation Model (KDM), an offline distillation framework that learns a finite-dimensional Koopman operator in a latent observable space to map noisy inputs to clean outputs in a single step. By training encoders, a linear Koopman operator, and a decoder with reconstruction, latent consistency, and prediction losses (plus adversarial and conditioning components), KDM achieves fast, semantically faithful generation while preserving the teacher's dynamics. The authors provide theoretical guarantees for finite-dimensional Koopman representations and semantic proximity, and demonstrate competitive, and often superior, one-step generation results across CIFAR-10, FFHQ, and AFHQv2 in the offline distillation setting, with favorable scalability and efficiency. This work offers a practical, architecture-agnostic path to efficient diffusion distillation and lays the groundwork for extending Koopman-based dynamics to broader generative modeling tasks.

Abstract

Diffusion-based generative models have demonstrated exceptional performance, yet their iterative sampling procedures remain computationally expensive. A prominent strategy to mitigate this cost is distillation, with offline distillation offering particular advantages in terms of efficiency, modularity, and flexibility. In this work, we identify two key observations that motivate a principled distillation framework: (1) while diffusion models have been viewed through the lens of dynamical systems theory, powerful and underexplored tools can be further leveraged; and (2) diffusion models inherently impose structured, semantically coherent trajectories in latent space. Building on these observations, we introduce the Koopman Distillation Model (KDM), a novel offline distillation approach grounded in Koopman theory - a classical framework for representing nonlinear dynamics linearly in a transformed space. KDM encodes noisy inputs into an embedded space where a learned linear operator propagates them forward, followed by a decoder that reconstructs clean samples. This enables single-step generation while preserving semantic fidelity. We provide theoretical justification for our approach: (1) under mild assumptions, the learned diffusion dynamics admit a finite-dimensional Koopman representation; and (2) proximity in the Koopman latent space correlates with semantic similarity in the generated outputs, allowing for effective trajectory alignment. KDM achieves highly competitive performance across standard offline distillation benchmarks.

One-Step Offline Distillation of Diffusion-based Models via Koopman Modeling

TL;DR

The paper tackles the high computational cost of diffusion-model sampling by introducing Koopman Distillation Model (KDM), an offline distillation framework that learns a finite-dimensional Koopman operator in a latent observable space to map noisy inputs to clean outputs in a single step. By training encoders, a linear Koopman operator, and a decoder with reconstruction, latent consistency, and prediction losses (plus adversarial and conditioning components), KDM achieves fast, semantically faithful generation while preserving the teacher's dynamics. The authors provide theoretical guarantees for finite-dimensional Koopman representations and semantic proximity, and demonstrate competitive, and often superior, one-step generation results across CIFAR-10, FFHQ, and AFHQv2 in the offline distillation setting, with favorable scalability and efficiency. This work offers a practical, architecture-agnostic path to efficient diffusion distillation and lays the groundwork for extending Koopman-based dynamics to broader generative modeling tasks.

Abstract

Diffusion-based generative models have demonstrated exceptional performance, yet their iterative sampling procedures remain computationally expensive. A prominent strategy to mitigate this cost is distillation, with offline distillation offering particular advantages in terms of efficiency, modularity, and flexibility. In this work, we identify two key observations that motivate a principled distillation framework: (1) while diffusion models have been viewed through the lens of dynamical systems theory, powerful and underexplored tools can be further leveraged; and (2) diffusion models inherently impose structured, semantically coherent trajectories in latent space. Building on these observations, we introduce the Koopman Distillation Model (KDM), a novel offline distillation approach grounded in Koopman theory - a classical framework for representing nonlinear dynamics linearly in a transformed space. KDM encodes noisy inputs into an embedded space where a learned linear operator propagates them forward, followed by a decoder that reconstructs clean samples. This enables single-step generation while preserving semantic fidelity. We provide theoretical justification for our approach: (1) under mild assumptions, the learned diffusion dynamics admit a finite-dimensional Koopman representation; and (2) proximity in the Koopman latent space correlates with semantic similarity in the generated outputs, allowing for effective trajectory alignment. KDM achieves highly competitive performance across standard offline distillation benchmarks.
Paper Structure (57 sections, 4 theorems, 34 equations, 14 figures, 9 tables, 2 algorithms)

This paper contains 57 sections, 4 theorems, 34 equations, 14 figures, 9 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Generative Modeling.
  4. Accelerating Diffusion Sampling.
  5. Koopman-Based Modeling.
  6. Mathematical Background
  7. Deterministic Diffusion Sampling.
  8. Koopman Operator Theory.
  9. Koopman Distillation of Diffusion Models
  10. Dynamical System Perspective on Diffusion Models.
  11. Koopman-based Distillation.
  12. Koopman's Loss Functions.
  13. Adversarial Loss.
  14. Control Modeling for Conditional Dynamics.
  15. Theoretical Properties of KDM
  16. ...and 42 more sections

Key Result

Theorem 5.1

Let $\Phi : \mathbb{R}^n \to \mathbb{R}^n$ be an analytic map, and let $x_T \sim \mathcal{N}(0, I_n)$. Then, for any $\epsilon > 0$, there exists a linear operator $C \in \mathbb{R}^{d \times d}$ such that where $\xi : \mathbb{R}^n \rightarrow \mathbb{R}^d$ denotes the observable lifting map. Moreover, the required dimension $d$ grows at most polynomially with $1/\epsilon$, and the mapping $\Phi$

Figures (14)

  • Figure 1: A Gaussian sample (bottom left) evolves into a clean image (bottom right) via nonlinear dynamics $\Phi_T$. Leveraging the Koopman framework, learned encoders $E_\varphi$ and $E_\phi$ transform noise and image into an embedding space where the evolution becomes linear under $C_\mu$ (top).
  • Figure 2: Untrained noise distributions are semantically unstructured (left). Training EDM on the data reveals the emergence of coherent clusters in noise space, i.e., at time $T$ (right).
  • Figure 3: Visualization of vicinity comparison across different noise levels $\sigma$. Left: EDM karras2022elucidating. Right: our method. Each image corresponds to a small perturbation around the original noise sample.
  • Figure 4: Sampled images using 79 NFEs with EDM karras2022elucidating (left), and using 1 NFE with KDM (right).
  • Figure 5: A) Model size vs FID. B) Data Size vs FID. C) Outlier analysis
  • ...and 9 more figures

Theorems & Definitions (6)

  • Theorem 5.1: Finite Koopman Operator for Analytic Dynamics
  • Theorem 5.2: Semantic Proximity via Koopman-Invariant Coordinates
  • Theorem : Finite Koopman Approximation for Analytic Dynamics
  • proof
  • Theorem : Semantic Proximity via Koopman-Invariant Coordinates
  • proof