Flow Matching for Posterior Inference with Simulator Feedback

Benjamin Holzschuh,Nils Thuerey

2024-10-30

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.

Machine Learning

What problem does this paper attempt to address?

### What problem does this paper attempt to solve?

This paper aims to solve inverse problems in physical sciences, especially by introducing simulator feedback to improve flow - based generative models. Specifically, the author proposes a method to fine - tune a pre - trained flow model with control signals to improve the accuracy and efficiency of parameter inference.

#### Main problem description:

1. **Challenges of inverse problems**:
- In physical sciences, obtaining the posterior distribution of given observational data is an important scientific problem, involving practical applications from particle physics, gravitational - wave inference to weather forecasting.
- Traditional likelihood - based methods are computationally expensive when dealing with high - dimensional data and costly likelihood evaluations.

2. **Limitations of existing methods**:
- Flow models (such as normalizing flows) perform well in generative modeling and density estimation, but lack a direct feedback mechanism during the training process, making it difficult to generate highly accurate samples solely by learning.
- Simulator feedback can provide additional information, but how to effectively integrate this information into the flow model remains a challenge.

3. **Proposed solution**:
- The author proposes a strategy to improve the pre - trained flow model by introducing simulator feedback. This method includes:
- Using a pre - trained flow network and fine - tuning it with control signals.
- The control signals can be based on gradients and cost functions (if the simulator is differentiable), or directly learned from the simulator output.
- By combining learning updates and appropriate control signals, avoid local optimal solutions and obtain high - precision samples while maintaining a low inference time.

4. **Application scenarios**:
- The author chooses strong gravitational lensing systems as experimental objects, which is an inverse problem in astronomy and requires an accurate posterior distribution for accurate observational modeling.
- Experimental results show that using flow matching and simulator feedback can obtain a posterior distribution competitive with the MCMC method, but is 67 times faster in inference speed.

5. **Summary of contributions**:
- Proposes a general strategy to improve pre - trained flow models through simulator feedback.
- Evaluates different types of flow - matching variants and shows a significant performance improvement with simulator feedback.
- Also verifies the effectiveness of control signals in low - dimensional problems, especially more obvious in high - dimensional problems and observations with less noise.

Through these improvements, the author hopes to significantly improve the accuracy and applicability of flow - based generative models in complex inverse problems while maintaining high efficiency.