Flow Annealed Kalman Inversion for Gradient-Free Inference in Bayesian Inverse Problems

Richard D.P. Grumitt,Minas Karamanis,Uroš Seljak
DOI: https://doi.org/10.48550/arXiv.2309.11490
2023-09-21
Abstract:For many scientific inverse problems we are required to evaluate an expensive forward model. Moreover, the model is often given in such a form that it is unrealistic to access its gradients. In such a scenario, standard Markov Chain Monte Carlo algorithms quickly become impractical, requiring a large number of serial model evaluations to converge on the target distribution. In this paper we introduce Flow Annealed Kalman Inversion (FAKI). This is a generalization of Ensemble Kalman Inversion (EKI), where we embed the Kalman filter updates in a temperature annealing scheme, and use normalizing flows (NF) to map the intermediate measures corresponding to each temperature level to the standard Gaussian. In doing so, we relax the Gaussian ansatz for the intermediate measures used in standard EKI, allowing us to achieve higher fidelity approximations to non-Gaussian targets. We demonstrate the performance of FAKI on two numerical benchmarks, showing dramatic improvements over standard EKI in terms of accuracy whilst accelerating its already rapid convergence properties (typically in $\mathcal{O}(10)$ steps).
Computation,Machine Learning
What problem does this paper attempt to address?
The problem that this paper attempts to solve is how to conduct effective Bayesian inference in scientific inverse problems when the forward model is expensive and its gradient is inaccessible. Specifically, the standard Markov Chain Monte Carlo (MCMC) algorithms become impractical in this situation because they require a large number of serial model evaluations to converge to the target distribution. While the existing Ensemble Kalman Inversion (EKI) method performs well on some problems, it has limitations when dealing with non - Gaussian target geometries. To overcome these limitations, the authors propose the Flow Annealed Kalman Inversion (FAKI) method. FAKI relaxes the Gaussian assumption in the standard EKI by embedding Kalman filter updates in a temperature annealing scheme and using Normalizing Flows (NF) to map intermediate measures to the standard Gaussian space. This enables FAKI to approximate non - Gaussian target distributions more accurately and accelerate the convergence speed. ### Main Contributions 1. **Introducing FAKI**: A new gradient - free Bayesian inference algorithm suitable for expensive forward models. 2. **Relaxing the Gaussian assumption**: By using normalizing flows, FAKI can better capture non - Gaussian target geometries. 3. **Accelerating convergence**: Compared with the standard EKI, FAKI significantly reduces the number of required iterations while maintaining or improving accuracy. ### Application Examples The paper demonstrates the performance of FAKI through two numerical benchmark tests: 1. **Two - dimensional Rosenbrock distribution**: This is a highly non - Gaussian target distribution, and the standard EKI has difficulty accurately capturing its geometric features, while FAKI can effectively solve this problem. 2. **Stochastic Lorenz system**: A complex system with a 94 - dimensional parameter space. FAKI performs well in dealing with nonlinearity and non - Gaussianity and can estimate the mean and standard deviation more accurately. ### Conclusion By introducing normalizing flows, FAKI successfully overcomes the limitations of the standard EKI in dealing with non - Gaussian target geometries and improves the accuracy and efficiency of Bayesian inference. Future work can further explore the combination with other unbiased MCMC or importance sampling methods to deal with more complex forward models.