Hessian-Informed Flow Matching
Christopher Iliffe Sprague, Mert Can Kurucu, Riccardo Tedoldi, Alessandro Tibo, Jon Paul Janet, Hossein Azizpour
原文摘要
Data from physical systems often concentrate near stable states corresponding to minima of an energy function. Around these states, the curvature of the energy landscape shapes the local covariance structure of the data, which is often highly anisotropic: small perturbations around a data point raise its energy more sharply in some directions than in others. However, standard flow matching commonly connects noise to target data through isotropic conditional paths that do not explicitly follow this physical structure. We introduce Hessian-Informed Flow Matching (HI-FM), which uses the energy Hessian at each data sample to encode this local curvature in the conditional probability paths used for training. By locally linearizing energy-driven dynamics, we derive anisotropic paths that move samples toward their targets at different rates in different directions, without changing the standard flow matching objective. We evaluate HI-FM through molecular generation on Hessian-QM9 and GEOM-Drugs, using SemlaFlow as a shared backbone. With Hessians calculated offline only for the training data, across both datasets, HI-FM improves the physical and geometric quality of generated molecules while maintaining comparable chemical validity. Under a matched optimizer-update budget, HI-FM trained with only $10\%$ of Hessian-QM9 achieves lower mean energy, strain, and bond-length $W_1$ than the full-data SemlaFlow baseline. The results support local energy curvature as a useful inductive bias for flow matching.
cs.LG · eess.SY