Hessian-Informed Flow Matching
This addresses the limitation of isotropic flow-based models in capturing anisotropic covariance structures for applications like molecular dynamics and robotic control, representing an incremental advancement.
The paper tackles the problem of modeling complex systems with anisotropic equilibrium distributions by introducing Hessian-Informed Flow Matching (HI-FM), which integrates the Hessian of an energy function into flow matching to capture local curvature, resulting in improved test sample likelihood on MNIST and Lennard-Jones datasets.
Modeling complex systems that evolve toward equilibrium distributions is important in various physical applications, including molecular dynamics and robotic control. These systems often follow the stochastic gradient descent of an underlying energy function, converging to stationary distributions around energy minima. The local covariance of these distributions is shaped by the energy landscape's curvature, often resulting in anisotropic characteristics. While flow-based generative models have gained traction in generating samples from equilibrium distributions in such applications, they predominately employ isotropic conditional probability paths, limiting their ability to capture such covariance structures. In this paper, we introduce Hessian-Informed Flow Matching (HI-FM), a novel approach that integrates the Hessian of an energy function into conditional flows within the flow matching framework. This integration allows HI-FM to account for local curvature and anisotropic covariance structures. Our approach leverages the linearization theorem from dynamical systems and incorporates additional considerations such as time transformations and equivariance. Empirical evaluations on the MNIST and Lennard-Jones particles datasets demonstrate that HI-FM improves the likelihood of test samples.