Probability Measure Gradient Flows for Diffusion-Based Amortized Optimization

Nonconvex optimization problems are ubiquitous but can be computationally expensive to solve repeatedly. Recent work on amortized optimization accelerates repeated optimization by learning a mapping from problem parameters to their solutions. However, existing methods typically learn a deterministic solution mapping, limiting their ability to represent and explore multiple promising regions of a nonconvex solution landscape. In this work, we introduce DiFlowOpt, a self-supervised amortized optimization method that models solutions as a distribution parameterized by a diffusion model. We lift amortized optimization to the space of probability measures and study its optimization dynamics under Wasserstein and Fisher–Rao gradient flows. We show that these two geometries induce complementary gradient estimators for training diffusion-based amortized optimizers, leveraging first- and zeroth-order objective information, respectively. Motivated by this connection, DiFlowOpt updates the diffusion model along a Wasserstein–Fisher–Rao flow that combines both sources of information. To accommodate constraints, we further develop a primal-dual formulation in which instance-specific dual variables are learned by a neural network. Across three standard high-dimensional nonconvex optimization problems and two real-world applications, DiFlowOpt achieves better solution quality than classical solvers and existing amortized optimization methods, with substantially lower solve time than classical solvers.