Improving Flow Matching for Posterior Inference with Physics-based Controls

Flow-based generative modeling is a powerful tool for solving inverse problems in physical sciences that can be used for sampling and likelihood evaluation with much lower inference times than traditional methods. We propose to refine flows with additional control signals based on an underlying physics model. In our experiments, this control signal is represented by gradients with respect to a differentiable cost function. We train a neural network to aggregate a pretrained flow and physics-based control signal to yield a hybrid update. We evaluate the refinements against classical MCMC methods for modeling strong gravitational lens systems, a challenging inverse problem in astronomy. We demonstrate that including physics-based controls improves the accuracy by $57$%, making them competitive with MCMC methods while being 12x to 83x faster for inference.