Ergodic Generative Flows

Abstract: Generative Flow Networks (GFNs) were initially introduced on directed acyclic graphs to sample from an unnormalized distribution density. Recent works have extended the theoretical framework for generative methods allowing more flexibility and enhancing application range. However, many challenges remain in training GFNs in continuous settings and for imitation learning (IL), including intractability of flow-matching loss, limited tests of non-acyclic training, and the need for a separate reward model in imitation learning. The present work proposes a family of generative flows called Ergodic Generative Flows (EGFs) which are used to address the aforementioned issues. First, we leverage ergodicity to build simple generative flows with finitely many globally defined transformations (diffeomorphisms) with universality guarantees and tractable flow-matching loss (FM loss). Second, we introduce a new loss involving cross-entropy coupled to weak flow-matching control, coined KL-weakFM loss. It is designed for IL training without a separate reward model. We evaluate IL-EGFs on toy 2D tasks and real-world datasets from NASA on the sphere, using the KL-weakFM loss. Additionally, we conduct toy 2D reinforcement learning experiments with a target reward, using the FM loss.

Ergodic Generative Flows

In this paper, the authors introduce a novel family of generative flow models, termed Ergodic Generative Flows (EGFs), designed to address existing challenges in Generative Flow Networks (GFNs) such as training complexity in continuous settings and the intractability of flow-matching losses. EGFs employ ergodicity to construct generative models with a finite set of globally defined transformations (diffeomorphisms) ensuring universality and tractable flow-matching loss. Additionally, the paper proposes a new loss function, the KL-weakFM loss, enabling imitation learning (IL) without necessitating a separate reward model.

Theoretical Contributions

The paper extends the theoretical framework of generative flows by providing a quantitative version of the sampling theorem for non-acyclic flows. This is achieved through the introduction of a virtual initial and terminal flow concept that provides control over the total variation during sampling. Further, the universality of EGFs is proved using ergodicity properties. Specifically, if the ergodic policy within an EGF parameterized family possesses summable $L^2$-mixing coefficients, then the family is universal on tori and spheres. This universality theorem implies that EGFs can achieve flow-matching for any non-zero distributions, which is critical for both theoretical and practical implementations across various state spaces.

Strong Numerical Results

The numerical experiments conducted support the efficacy of EGFs. The experiments demonstrate significant improvements in training stability and expressivity in benchmark tests, such as on toy 2D tasks and real-world datasets from NASA. Specifically, EGFs show superior performance compared to Moser Flows and other traditional models in generating distributions accurately using significantly fewer parameters. On the NASA volcano dataset, EGFs achieve better negative log-likelihood scores compared to existing models, illustrating their practical efficiency and robustness.

Implications and Future Work

The implications of EGFs extend to advancing generative modeling capabilities in machine learning applications. By ensuring model expressivity with limited parameters and improved training dynamics, EGFs hold potential for deployment in resource-constrained environments, such as real-time and mobile applications. Future research may focus on exploring EGFs in higher-dimensional state spaces, ensuring scalability and performance across more complex tasks. Additionally, integration with techniques from Neural Ordinary Differential Equations (NeuralODEs) could enhance control over mixing properties, potentially improving training stabilization and convergence speed.

In conclusion, the paper presents EGFs as a substantial development in generative flow modeling, offering robust theoretical guarantees and compelling empirical performance while addressing significant limitations encountered in existing generative models. The introduction of EGFs particularly anticipates further advancements in the development of compact, efficient, and scalable generative models.