Hessian-Informed Flow Matching

TL;DR

Hessian-Informed Flow Matching (HI-FM) enhances flow-based generative models by incorporating the Hessian of an energy function, enabling better modeling of anisotropic covariance structures in complex equilibrium distributions.

Contribution

The paper introduces HI-FM, a novel method that integrates Hessian information into flow matching to better capture local curvature and anisotropic covariance in equilibrium distributions.

Findings

Improves likelihood of test samples on MNIST.

Effectively models anisotropic covariance structures.

Leverages dynamical systems linearization theorem.

Abstract

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.