Path Gradients after Flow Matching

TL;DR

This paper explores a hybrid method combining Flow Matching and path gradients to improve sampling efficiency in molecular systems, maintaining learned structures without extra sampling or significant computational costs.

Contribution

It introduces a novel hybrid approach that fine-tunes CNFs with path gradients after Flow Matching, significantly boosting sampling efficiency while preserving flow structure.

Findings

Up to threefold increase in sampling efficiency.

Path gradients preserve learned flow structure.

Method requires no additional sampling or significant computation.

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.