Unfolding Generative Flows with Koopman Operators: Trajectory-Preserving Linearization
Abstract
Continuous Normalizing Flows (CNFs) enable elegant generative modeling but remain bottlenecked by their iterative nature requiring costly sampling and lacking interpretability of the intermediate states. Recent approaches accelerate sampling by straightening trajectories or distilling endpoints, yet they treat the original generative process as a black box, discarding the teacher’s intermediate dynamics. We propose a fundamentally different perspective: globally linearizing flow dynamics via Koopman theory to achieve trajectory-preserving linearization. By lifting Conditional Flow Matching (CFM) into a higher-dimensional Koopman space, we represent its evolution with a single linear operator. Crucially, unlike boundary-only distillation, our method enforces infinitesimal consistency with the teacher's vector field along the full generative path. We derive a practical, simulation-free training objective that ensures this global alignment and yields two key benefits. First, sampling becomes one-step and parallelizable. Second, because the linearization is faithful to the dynamics, the Koopman operator provides unique insights on the generation. We demonstrate that this structure enables novel applications unavailable in prior approaches, including discovery of semantically coherent editing directions, inversion with a teacher-aligned linear operator and class-conditional spectral signatures. Empirically, our approach achieves competitive sample quality, while enabling spectral analysis and control of the entire trajectories of generative flows.
Lay Summary
Image generators create pictures by gradually transforming random noise into a coherent image. This works remarkably well, but the process is slow and acts like a "black box": we cannot easily see what is happening inside it, and we have little control over the result. Existing methods that speed up generation discard this internal information. We adapted a century-old mathematical formalism called Koopman theory, which turns a complicated, twisting process into a simple straight-line one by viewing it from a new mathematical perspective. Just as Fourier analysis breaks a sound into the individual pure tones that compose it, our approach breaks the image-generation process into independent, simple components. Crucially, we ensure that our simplified version remains faithful to the original at every step, not only at the final picture. The benefit is twofold. Generation becomes a single fast step rather than many. More importantly, the process becomes transparent: we can isolate these components to edit images, for example, adding sunglasses or changing hair color, and examine how the generative model acquires its behavior during training. This offers a new way to understand and steer generative models.