[Paper Review] A mean-field games laboratory for generative modeling

B. Zhang, Markos A. Katsoulakis|arXiv (Cornell University)|Apr 26, 2023

Sports Analytics and Performance4 citations

TL;DR

This paper introduces a mean-field games (MFG) framework as a unifying mathematical foundation for generative modeling, demonstrating that continuous normalizing flows, score-based generative models (SGMs), and Wasserstein gradient flows all arise as solutions to MFG optimality conditions—coupled forward-backward PDEs. The key contribution is the derivation of Hamilton-Jacobi-Bellman (HJB) regularizers from MFG optimality, which enhance training stability and performance, as validated by an HJB-regularized SGM outperforming standard SGMs.

ABSTRACT

We demonstrate the versatility of mean-field games (MFGs) as a mathematical framework for explaining, enhancing, and designing generative models. In generative flows, a Lagrangian formulation is used where each particle (generated sample) aims to minimize a loss function over its simulated path. The loss, however, is dependent on the paths of other particles, which leads to a competition among the population of particles. The asymptotic behavior of this competition yields a mean-field game. We establish connections between MFGs and major classes of generative flows and diffusions including continuous-time normalizing flows, score-based generative models (SGM), and Wasserstein gradient flows. Furthermore, we study the mathematical properties of each generative model by studying their associated MFG's optimality condition, which is a set of coupled forward-backward nonlinear partial differential equations. The mathematical structure described by the MFG optimality conditions identifies the inductive biases of generative flows. We investigate the well-posedness and structure of normalizing flows, unravel the mathematical structure of SGMs, and derive a MFG formulation of Wasserstein gradient flows. From an algorithmic perspective, the optimality conditions yields Hamilton-Jacobi-Bellman (HJB) regularizers for enhanced training of generative models. In particular, we propose and demonstrate an HJB-regularized SGM with improved performance over standard SGMs. We present this framework as an MFG laboratory which serves as a platform for revealing new avenues of experimentation and invention of generative models.

Motivation & Objective

• To unify diverse generative modeling approaches—continuous normalizing flows, score-based models, and Wasserstein gradient flows—under a single mathematical framework.
• To reveal the inductive biases of generative models by analyzing their associated mean-field game (MFG) optimality conditions.
• To develop HJB regularizers from MFG structure to enhance training stability and performance of generative models.
• To establish a systematic, modular laboratory for inventing new generative models by manipulating MFG cost functions and dynamics.

Proposed method

• Formulate generative flows as mean-field games where each particle minimizes a loss dependent on the empirical distribution of the population, leading to coupled forward-backward PDEs as optimality conditions.
• Derive the MFG optimality conditions (a system of forward-backward nonlinear PDEs) for continuous normalizing flows, SGMs, and Wasserstein gradient flows.
• Construct Hamilton-Jacobi-Bellman (HJB) regularizers from the MFG optimality conditions to stabilize training in generative models.
• Apply HJB regularizers to score-based generative models, demonstrating improved performance over standard SGMs in numerical experiments.
• Use the MFG framework to systematically generate new models by modifying cost functions and dynamics, including relaxations of Wasserstein gradient flows.
• Leverage physics-informed neural networks (PINNs) to solve the forward-backward PDE systems arising from MFGs for score-matching.

Figure 1 : Flow chart describing how mean-field games are related to and provide new insights for flow-based generative modeling. Blue boxes denote the current understanding of flow-based generative models. Green boxes are the new perspective provided by mean-field games.

Experimental results

Research questions

• RQ1 How can mean-field games (MFGs) serve as a unifying framework for continuous-time generative models such as normalizing flows and score-based generative models?
• RQ2 What mathematical structure underlies the inductive biases of generative flows, and how can it be revealed through MFG optimality conditions?
• RQ3 Can the optimality conditions of MFGs be used to derive effective regularizers that improve training stability and performance in generative modeling?
• RQ4 How do relaxations of Wasserstein gradient flows, formulated as MFGs, interpolate between geodesics and gradient flows, and what is their potential in generative modeling?
• RQ5 What new classes of generative models can be systematically invented by manipulating MFG cost functions and dynamics?

Key findings

• The paper establishes that continuous normalizing flows, score-based generative models (SGMs), and Wasserstein gradient flows all emerge as solutions to mean-field game (MFG) optimality conditions, which are coupled forward-backward nonlinear PDEs.
• The MFG framework reveals the inductive biases of generative models through the structure of their associated optimality conditions, providing a deeper mathematical understanding of their behavior.
• Hamilton-Jacobi-Bellman (HJB) regularizers derived from MFG optimality conditions are shown to enhance training stability and performance, with an HJB-regularized SGM outperforming standard SGMs in numerical experiments.
• The MFG formulation of Wasserstein gradient flows is derived, and relaxations of these flows are shown to interpolate between Wasserstein geodesics and gradient flows, suggesting new computational and theoretical avenues.
• The framework enables systematic invention of new generative models by modifying cost functions and dynamics, such as introducing interaction terms or using nonlocal diffusions and jump processes.
• Physics-informed neural networks (PINNs) are successfully applied to solve the forward-backward PDE systems of MFGs, enabling score-matching via MFG-based optimization.