| Abstract: |
| In this work, we present a new generative model-based variational framework that utilizes Normalizing Flows (NFs) to solve Wasserstein Gradient Flows (WGFs), going beyond traditional chain-of-state methods by directly parameterizing the paths using the forward process of NFs as a representation of WGFs. Our approach enhances the capture of the Wasserstein Gradient Structure by incorporating an minimum action cost regularization term during the training of NFs. Additionally, we introduce an innovative adaptive sampling strategy that iteratively generates efficient stochastic collocation points, reflecting the evolving density estimates. This framework naturally incorporates the continuity equation within the network architecture, enabling the efficient estimation of high-dimensional density functions essential for free energy computations and improving the approximation of population dynamics. By integrating seamlessly with stochastic gradient descent techniques prevalent in deep learning, our method demonstrates robust performance on several illustrative problems, showcasing its potential in efficiently modeling the evolution of probability distributions in complex systems. |
|