Flow Matching for Posterior Inference with Simulator Feedback

TL;DR

This work pretrain the flow network and include feedback from the simulator exclusively for finetuning, therefore requiring only a small amount of additional parameters and compute, and demonstrates that including feedback from the simulator improves the accuracy, making it competitive with traditional techniques while being up to $67$x faster for inference.

Abstract

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 a simulator. Control signals can include gradients and a problem-specific cost function if the simulator is differentiable, or they can be fully learned from the simulator output. In our proposed method, we pretrain the flow network and include feedback from the simulator exclusively for finetuning, therefore requiring only a small amount of additional parameters and compute. We motivate our design choices on several benchmark problems for simulation-based inference and evaluate flow matching with simulator feedback against classical MCMC methods for modeling strong gravitational lens systems, a challenging inverse problem in astronomy. We demonstrate that including feedback from the simulator improves the accuracy by $53\%$, making it competitive with traditional techniques while being up to $67$x faster for inference.

Figures (17)

• Figure 1: An overview of our proposed framework. We consider a pretrained flow network $v_\phi$ and use the predicted flow for the trajectory point $\boldsymbol{\theta}_t$ at time $t$ to estimate $\hat{\boldsymbol{\theta}}_1$. On the right, we show a gradient-based control signal with a differentiable simulator and cost function $C$ for improving $\hat{\boldsymbol{\theta}}_1$. An additional network learns to combine the predicted flow with feedback via the control signal to give a new controlled flow. By combining learning-based updates with suitable controls, we avoid local optima and obtain high-accuracy samples with low inference times.
• Figure 2: Gradient-based control signal
• Figure 3: Learning-based control signal
• Figure 5: Evaluation of SBI tasks using different variants of flow matching training. Lower C2ST scores are better.
• Figure 6: Evaluation of simulator feedback for LV.