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

Why Offline Distillation?

Offline distillation presents significant benefits compared to online methods. By training on pre-calculated noise-image pairs, it bypasses the necessity for complex, multi-stage training procedures. This approach also minimizes the computational demands on the teacher model, as it doesn't require extensive neural function evaluations during the student's learning process. Furthermore, offline distillation provides a crucial opportunity to filter out any potentially privacy-violating synthetic data generated by the teacher *before* the distillation process even begins. Finally, it removes the resource-intensive requirement of keeping the large teacher model in memory throughout the student's optimization. These advantages culminate in a training paradigm that is more privacy-preserving, highly scalable, inherently asynchronous, and enables broader reuse of the distilled data across various student model architectures and different application scenarios .

It is well established that denoising generative models—such as diffusion models and flow matching models—can be described by stochastic differential equations (SDEs) or ordinary differential equations (ODEs). During training, these models define a forward noising process via a known dynamical system and learn its corresponding reverse process to gradually denoise the data. Once training is complete, sampling new data points becomes possible by simulating the learned reverse dynamics.

From this viewpoint, these models effectively learn a dynamical system that transports data from one distribution to another. Motivated by this perspective, we initiate a novel investigation into the geometric structure of the learned dynamics. Specifically, we ask: Does the initial noisy state uniquely determine the final data point? Our findings indicate that it does. As shown above, before training, noisy samples are scattered arbitrarily. After training, a coherent structure emerges—highlighting the structure of the learned flow.

Koopman's theory proposes that a nonlinear dynamical system can be represented as a linear system when expressed in an embedded space. To move data into and out of this space, we can learn/use linear/nonlinear transformation functions—typically referred to as an encoder (“E”) and a decoder (“D”). Once in this transformed space, the system’s evolution becomes linear and can be modeled using simple operations like matrix multiplication.

Previous research has shown that in certain types of dynamical systems, the Koopman operator preserves the structural properties of the original nonlinear dynamics. In our setup, we theoretically prove that this structure-preserving behavior also holds under mild conditions. The consistent latent structure we observe, along with insights from prior work and our own theoretical findings, motivates the use of Koopman theory in our approach.

But one might ask: is such a Koopman operator guaranteed to exist?

The answer is yes. Although our setup requires a finite Koopman matrix, whereas the theory suggests  that the matrix is infinite, in our work, we provide theoretical proof that a finite approximation of the Koopman operator does indeed exist. This result offers a solid foundation for applying Koopman-based modeling in our setting and justifies its practical and theoretical use.

Can you tell who is the teacher and who is the student ? Well, the left example is the teacher and the right is the student.  Furthermore, we can see that the KDM has learned the dynamical structure of the original model pretty well since all corresponding images started from the same noise (initial state) and ending almost in the same ending state.