Ergodic Generative Flows

TL;DR

This work introduces Ergodic Generative Flows (EGFs), a family of generative flows built from a small set of globally defined diffeomorphisms whose ergodicity yields universality in continuous RL/IL settings. EGFs provide tractable forward/inflow densities and enable a KL-weakFM loss that supports imitation learning without a separate reward model, addressing key FM-related limitations. Theoretical results include a Master Universality Theorem and universality on tori and spheres, plus a quantitative sampling bound linking initial/terminal mismatches to sampling error. Empirically, EGFs achieve stable training and strong performance on toy 2D tasks and NASA Earth datasets, while using far smaller models than competitive baselines. The framework offers a modular, scalable approach for efficient, reward-free IL and RL in non-acyclic, continuous spaces, with clear directions for future theory and extensions.

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.

Figures (10)

• Figure 1: Checkerboard RL task.
• Figure 2: Comparison of imitation learning on standard toy distributions using tiny models. Background filter is applied to EGF, a similar filter on Moser flow would yield worse results.
• Figure 3: EGF generated point with reward field $\ftermhat$ (left) and whole dataset point with KDE field (right) for the volcano task.
• Figure 4: NLL estimation for a range of values of $k$ for a given model on the volcano dataset.
• Figure 5: Comparison of the EGF's samples using the same model with and without a density filter.

Theorems & Definitions (19)

• Theorem 2.1: bengio2021flowbrunswic2024theory
• Definition 3.1: Ergodic Generative Flows
• Definition 3.2
• Definition 3.3
• Theorem 3.4
• Theorem 3.5
• Theorem 3.6
• Definition 3.7: Virtual initial and terminal flow
• Theorem 3.8: Quantitative Sampling of Generative Flows
• Corollary 3.9