Feynman-Kac-Flow: Inference Steering of Conditional Flow Matching to an Energy-Tilted Posterior

Mark, Konstantin, Galustian, Leonard, Kovar, Maximilian P. -P., Heid, Esther

arXiv.org Artificial Intelligence 

Institute of Materials Chemistry, TU Wien, A-1060 Vienna, Austria Conditional Flow Matching(CFM) represents a fast and high-quality approach to generative modelling, but in many applications it is of interest to steer the generated samples towards precise requirements. While steering approaches like gradient-based guidance, sequential Monte Carlo steering or Feynman-Kac steering are well established for diffusion models, they have not been extended to flow matching approaches yet. In this work, we formulate this requirement as tilting the output with an energy potential. We derive, for the first time, Feynman-Kac steering for CFM. We evaluate our approach on a set of synthetic tasks, including the generation of tilted distributions in a high-dimensional space, which is a particularly challenging case for steering approaches. We then demonstrate the impact of Feynman-Kac steered CFM on the previously unsolved challenge of generated transition states of chemical reactions with the correct chirality, where the reactants or products can have a different handedness, leading to geometric constraints of the viable reaction pathways connecting reactants and products. Code to reproduce this study is avaiable open-source at https://github.com/heid-lab/fkflow. I. INTRODUCTION Since its introduction by Lipman et al. [1], Conditional Flow Matching (CFM) has seen several interesting applications, ranging from image [1], audio [2] and video [3] generation to decision-making [4], time series modelling [5], protein modelling [6, 7] or molecular structure design [8], amongst others. CFM transforms samples from a source distribution (such as random noise) to samples following a given target distribution (such as images or molecular structures) by modelling probability paths via vector fields. It largely improves on diffusion-based methods both in quality and speed, establishing CFM as a popular generative method [1].