Hierarchical Koopman Diffusion: Fast Generation with Interpretable Diffusion Trajectory

AI summary

Overview

• Research area: Generative machine learning, specifically fast sampling for diffusion models and the interpretability of their denoising dynamics, using Koopman operator theory from dynamical systems.
• Technical level: Advanced. The paper assumes familiarity with continuous-time diffusion models, probability flow ODEs, U-Net architectures, and Koopman operator theory.
• Scope: The paper proposes Hierarchical Koopman Diffusion (HKD), a framework that lifts the diffusion trajectory into scale-specific latent spaces governed by linear operators, enabling one-step generation while keeping intermediate generative states accessible for analysis and editing.

What This Paper Is About

Diffusion models generate high-fidelity images but require an expensive iterative denoising process, while existing one-step accelerations (distillation, consistency models) collapse generation into a direct noise-to-image mapping and thereby lose the intermediate trajectory that makes diffusion interpretable and controllable. This paper asks whether a single framework can deliver one-step sampling and still expose an interpretable, manipulable generative trajectory. The authors answer by mapping the deterministic diffusion trajectory into Koopman spaces where evolution is linear, so that closed-form trajectory solutions exist and every intermediate state remains analytically accessible.

Key Contributions

1. An interpretable one-step generation paradigm (HKD) built on hierarchical Koopman dynamics. It integrates explicit intermediate generative states into the process, enabling control along the diffusion trajectory, and provides spectral tools from dynamical systems theory for analyzing the underlying generative mechanism.
2. A theoretical expressiveness argument showing that the Koopman explicit formulation is provably more expressive than directly learning a black-box noise-to-image mapping with standard neural networks, when compared at the same total number of activation functions.
3. Empirical validation on CIFAR-10 and FFHQ, demonstrating competitive one-step generation, plus a spectral analysis linking Koopman spectral components to semantic image attributes.
4. A frequency-aware image editing experiment, intervening at an intermediate stage of the diffusion trajectory to demonstrate interpretability and controllable generation.

Main Findings

• One-step generation on CIFAR-10: HKD reaches an FID of 3.30 at 1 NFE (FID-50k). For context in the same table, iCT-deep reports 2.51, iCT 2.83, 2-Rectified flow++ 3.38, CD (LPIPS) 3.55, ECM 3.60, DMD 3.77, 2-Rectified flow 4.85, 3-Rectified flow 5.21, 1-Rectified flow 6.18, PD 8.34, CT (LPIPS) 8.70, and KD 9.36, all at 1 NFE.
• Multi-step baselines for comparison: DDPM (1000 NFE, 3.17), Score SDE (2000 NFE, 2.38), DDIM (100 NFE, 4.16; 10 NFE, 13.36), EDM (35 NFE, 1.97; 15 NFE, 5.62).
• One-step generation on FFHQ: HKD achieves an FID of 5.70 at 1 NFE, compared with ECM at 5.99 (1 NFE), EDM at 9.85 (15 NFE) and 2.47 (79 NFE), and DDIM at 18.30 (10 NFE). The text refers to this as FFHQ 64x64.
• Training efficiency and stability: HKD achieved comparable performance within 2-3 days on 8×V100 GPUs, whereas the paper states state-of-the-art consistency training typically requires nearly a week of training on 8 GPUs and remains unstable in practice. The paper also notes that retraining with consistency distillation results in a degraded FID of 10.53, as reported in the literature it cites.
• Spectral-semantic correspondence: Sorting eigenvalue pairs by real part and masking to smallest, intermediate, or largest ranges shows a semantic hierarchy: low-range modes capture global structure, mid-range modes recover overall shape and pose, and high-range modes refine local details.
• Frequency-aware editing: Injecting high-frequency features from a reference image into the lower-left half of a generated image at mixing ratios of 10%, 20%, 50%, 80%, and 90% at the midpoint of the Koopman trajectory gradually reveals more facial details from the reference. Frequency-agnostic editing, which mixes all frequency bands at a 90% ratio, disrupts global structures.
• Image recuperation: Inpainting and coloring experiments on CIFAR-10 were performed following Algorithm 4 of the cited consistency-model work, iteratively mixing a reference image along the generative trajectory by adding and removing noise at each time step t as t decreases from T to 0.
• Ablation on CIFAR-10: Removing Koopman dynamics (keeping hierarchical design) gives FID 5.72; adding Koopman evolution to skips and bottleneck gives 5.57; using Koopman and the trajectory consistency loss without the hierarchical design gives 4.78; the full configuration gives 3.30.
• Theoretical bound: Theorem 3.1 states err_HKD ≤ err_one-step + O(κ) at equal total activation counts, with κ = O(N^{-1/2}) + o(m^{-r/3}), where N is dataset size and m is the dimension of the Koopman space (a fuller expression for κ involving σ, δ, ρ_inf, ρ_sup appears in the theorem).

Methodology in Plain English

The authors start from the probability flow ODE, the deterministic trajectory that a diffusion model follows from noise to data. Rather than learning a single black-box jump from noise to image, they train an encoder to lift images at any time t into a set of latent feature maps organized as a U-Net-style hierarchy of spatial resolutions. At each level and each spatial location, the latent features are assumed to evolve under a local linear operator, so the state at one time can be written in closed form as an exponential of that operator times the elapsed time. This is the Koopman idea: nonlinear dynamics become linear in a lifted observable space, and the linear structure can be diagonalized, here into blocks of size 2×2 corresponding to pairs of complex conjugate eigenvalues.

Because the trajectory has a closed form, the model can jump directly from noise at time T to a sample at time ε in a single step. It can also compute any intermediate state, which enables a new "trajectory consistency loss": an encoded intermediate state is evolved to the final time and decoded, and the decoded image is compared with the ground-truth clean image. The total objective combines this trajectory consistency loss with a reconstruction loss supervising the direct one-step noise-to-image mapping. The distance function combines MSE and LPIPS, with the LPIPS weight fixed at 1 and the MSE weight annealed so training shifts from coarse alignment to perceptual refinement. Encoder, decoder, and the Koopman matrices are trained end-to-end, with encoder and decoder initialized from a pretrained EDM U-Net.

For the theory, the authors measure network size by the number of activation functions, introduce a "simplicial error" measuring deviation in simplicial complex estimation of the data and noise spaces, and use these to bound the error of HKD relative to an end-to-end one-step model of equal size.

Why This Matters

• Research impact: The work challenges the assumption that fast one-step sampling must sacrifice trajectory access. It introduces Koopman operator theory into image generation, per the authors' claim to be the first to do so, and offers spectral analysis as a lens for studying how diffusion models build images.
• Interpretability of generative models: Because intermediate states remain computable, the generative process can be probed, masked, and intervened on, which is not possible with purely implicit one-step mappings.
• Real-world applications:
  • Frequency-aware image editing, where reference image high-frequency detail is transferred into a generated image at chosen mixing ratios while preserving the original low-frequency structure.
  • Interactive controllable synthesis, where a user intervenes at a specific point along the trajectory rather than only at the input.
  • Image restoration-style tasks demonstrated on CIFAR-10 for inpainting and coloring by iteratively mixing a reference image along the trajectory.
  • Model diagnostics, using spectral contribution analysis to inspect which modes carry global layout, shape and pose, or fine local texture.
• Industry relevance: One-step inference reduces the compute cost of deploying diffusion models, and the reported 2-3 day training on 8×V100 GPUs, along with the claim of improved training stability relative to consistency models, is relevant to teams that find consistency training slow or unstable. The paper notes that its training paradigm does not rely on adversarial optimization or heavy tuning.

Future Directions

• Integrating advanced training techniques such as adversarial learning to further improve generation quality, since the current paradigm avoids adversarial optimization.
• Exploring high-resolution generation, which the hierarchical design is meant to support but which the paper states remains underexplored.
• Extending spectral interventions beyond frequency-aware manipulation to text-guided editing and attribute-specific control.
• Further theoretical and empirical work on the Koopman approximation itself, given that the error term κ depends on the Koopman space dimension m and dataset size N.

Target Audience

This paper is most useful to researchers and graduate students working on diffusion model acceleration, one-step generative models, or interpretability of generative models, and to practitioners who need fast inference with some ability to intervene in the generation process. Readers without background in continuous-time diffusion formulations, U-Net architectures, and Koopman operator theory will need to consult the background and appendix material before the technical sections and the error-bound theorem are fully accessible.

Authors’ 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.