Path Gradients after Flow Matching

TL;DR

The paper addresses improving Boltzmann Generators by combining Flow Matching (FM) pre-training with Path Gradient (PG) fine-tuning, leveraging gradient information from the target energy without extra sampling. It develops a memory-efficient method (augmented adjoint) to apply PG to forward KL training and demonstrates that a short PG fine-tuning phase can substantially boost importance-sampling efficiency (ESS) and log-likelihood on Lennard-Jones systems and alanine dipeptide, without substantially altering the learned flow dynamics. The approach also shows promise for transferability, improving performance on unseen dipeptides. Overall, the work offers a practical path to higher-quality samples under the same computational budget, with careful attention to model structure preservation and potential applicability to diffusion models and distillation tasks.

Abstract

Boltzmann Generators have emerged as a promising machine learning tool for generating samples from equilibrium distributions of molecular systems using Normalizing Flows and importance weighting. Recently, Flow Matching has helped speed up Continuous Normalizing Flows (CNFs), scale them to more complex molecular systems, and minimize the length of the flow integration trajectories. We investigate the benefits of using path gradients to fine-tune CNFs initially trained by Flow Matching, in the setting where a target energy is known. Our experiments show that this hybrid approach yields up to a threefold increase in sampling efficiency for molecular systems, all while using the same model, a similar computational budget and without the need for additional sampling. Furthermore, by measuring the length of the flow trajectories during fine-tuning, we show that path gradients largely preserve the learned structure of the flow.

Figures (8)

• Figure 1: Training a CNF on a simple 2D Gaussian Mixture Model. Comparison between pure training with Flow Matching, pre-training with Flow Matching and fine-tuning with Path Gradients and pure training with Path Gradients. We see that given the same wall-time hybrid training performs best in terms of forward KL divergence. The bottom row shows the target and the final model after training.
• Figure 2: FM loss objective \ref{['eq:CFM']} during training on 2D GMM averaged on three runs. Training with Path gradients leaves the MSE largely unchanged.
• Figure 3: Reverse ESS, NLL and trajectory length for Flows trained with standard FM, Optimal Transport and Equivariant Transport during fine-tuning with Path Gradients on LJ13. We can observe that fine-tuning largely leaves the trajectory length unchanged, while substantially improving performance. Mean $\pm$ sterr over three runs.
• Figure 4: Alanine dipeptide results for the TBG model and the classical force field with and without Path Gradient finetuning. Left: Ramachandran plots for the dihedral angel distribution of a reference MD simulation and non reweighted samples from the different TBG models. Right: Corresponding energy distributions of generated samples.
• Figure 5: Loss in space after training as done in \ref{['fig:2d-GMM']}