Hessian-Informed Flow Matching

TL;DR

This paper introduces 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 and leverages the linearization theorem from dynamical systems and incorporates additional considerations such as time transformations and equivariance.

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.

Figures (4)

• Figure 1: A schematic of HI-FM, where conditional flows are defined via a linear approximation of the dynamics of equilibrium (via the Hessian). These approximations incorporate the underlying anisotropy of the system, which may lead to generalization to unseen parts of the energy landscape (green). These flows are radially unbounded, flowing faster than plateauing potentials (e.g. Lennard-Jones).
• Figure 2: A depiction of the Eigen subspaces of the Hessian of a group-invariant energy function. This energy function is invariant to rotations in $\mathrm{SO}(2)$.
• Figure 3: MNIST: The negative log-likelihood of the test set during training (a) and the number of function evaluations to compute it (b).
• Figure 4: LJ13: The negative log-likelihood of the test set during training (a) and the number of function evaluations to compute it (b).