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

2025-05-19

AI Summary

Diffusion models suffer from low sampling efficiency and high computational overhead. To address this, we propose the first offline distillation framework grounded in Koopman operator theory, which models the multi-step nonlinear denoising process as a single-step linear evolution in latent space—enabling semantically faithful one-step generation. We theoretically prove that diffusion trajectories admit a finite-dimensional Koopman representation and that latent-space distances quantify semantic similarity. Methodologically, we design an encoder–Koopman linear propagator–decoder architecture, jointly optimizing latent trajectory alignment and Koopman operator learning. Evaluated on standard offline distillation benchmarks, our approach achieves state-of-the-art performance: one-step generation improves FID by up to 40% over prior methods. All code and configuration files are publicly released.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs

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. Empirically, KDM achieves state-of-the-art performance across standard offline distillation benchmarks, improving FID scores by up to 40% in a single generation step. All implementation details and code for the experimental setups are provided in our GitHub - https://github.com/azencot-group/KDM, or in our project page - https://sites.google.com/view/koopman-distillation-model.

Problem

• Reduce computational cost of diffusion models via distillation
• Leverage Koopman theory for linear dynamics representation
• Achieve single-step generation with semantic fidelity

Innovation

• Koopman Distillation Model for single-step generation
• Linear operator propagates noisy inputs efficiently
• Preserves semantic fidelity in latent space