Hierarchical Koopman Diffusion: Fast Generation with Interpretable Diffusion Trajectory
Hanru Bai, Weiyang Ding, Difan Zou
TL;DR
This work tackles the slow sampling inherent in diffusion models while preserving interpretability of the generation process. It introduces Hierarchical Koopman Diffusion (HKD), a framework that lifts diffusion dynamics into a multi-scale Koopman latent space where the evolution is linear, enabling one-step sampling and explicit access to intermediate states. HKD uses a U-Net-like encoder to project, per-scale linear Koopman operators to evolve the latent features, and a shared decoder to reconstruct images, with a trajectory-consistency loss that supervises intermediate states. Spectral analysis of the learned Koopman operators provides interpretable insights into generation dynamics and enables frequency-aware editing. Empirically, HKD achieves competitive one-step generation performance on CIFAR-10 and FFHQ, while offering a principled mechanism for interpretability and controllable synthesis via spectral components, all with improved training stability and efficiency.
Abstract
Diffusion models have achieved impressive success in high-fidelity image generation but suffer from slow sampling due to their inherently iterative denoising process. While recent one-step methods accelerate inference by learning direct noise-to-image mappings, they sacrifice the interpretability and fine-grained control intrinsic to diffusion dynamics, key advantages that enable applications like editable generation. To resolve this dichotomy, we introduce \textbf{Hierarchical Koopman Diffusion}, a novel framework that achieves both one-step sampling and interpretable generative trajectories. Grounded in Koopman operator theory, our method lifts the nonlinear diffusion dynamics into a latent space where evolution is governed by globally linear operators, enabling closed-form trajectory solutions. This formulation not only eliminates iterative sampling but also provides full access to intermediate states, allowing manual intervention during generation. To model the multi-scale nature of images, we design a hierarchical architecture that disentangles generative dynamics across spatial resolutions via scale-specific Koopman subspaces, capturing coarse-to-fine details systematically. We empirically show that the Hierarchical Koopman Diffusion not only achieves competitive one-step generation performance but also provides a principled mechanism for interpreting and manipulating the generative process through spectral analysis. Our framework bridges the gap between fast sampling and interpretability in diffusion models, paving the way for explainable image synthesis in generative modeling.