Flow Matching for Posterior Inference with Simulator Feedback

## Abstract

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 % percent 53 53\% 53 % , making it competitive with traditional techniques while being up to 67x faster for inference. Code and experiments are available at https://github.com/tum-pbs/sbi-sim .

## 1 Introduction

Acquiring posterior distributions given measurement data is of paramount scientific interest (Cranmer et al., ), with real-world applications ranging from particle physics (Baydin et al., ), over the inference of gravitational waves (Dax et al., ) to predictions of dynamical systems such as weather forecasting (Gneiting & Raftery, ).
For an observation $\bm{x}_{o}$ and model parameters $\bm{\theta}$, the likelihood $p(\bm{x}_{o}|\bm{\theta})$ corresponds to how strongly we believe a model with parameters $\bm{\theta}$ causes $\bm{x}_{o}$ to occur. In Bayesian modeling, we are interested in the posterior $p(\bm{\theta}|\bm{x}_{o})$, which is proportional to the likelihood times the prior $p(\bm{\theta})$ and tells us which parameters most likely explain the observation.

Inferring the posterior based on samples from traditional likelihood-based methods can be expensive for high-dimensional data and when likelihood evaluations are costly. Simulation-based inference (Cranmer et al., , SBI) represents the posterior as a parametric function $q(\bm{\theta}|\bm{x}_{o})$, which is a learnable conditional density estimator that can be trained purely by simulations $\bm{x}\sim p(\bm{x}|\bm{\theta})$ alone.
By investing an upfront cost for training the density estimator, we can sample and compute likelihoods from $q(\bm{\theta}|\bm{x}_{o})$ much faster than other methods, thereby amortizing the training cost over many observations.

Traditionally, normalizing flows (Rezende & Mohamed, ; Dinh et al., ; Papamakarios et al., ) have been a popular class of density estimators used in many areas of science. To compute likelihoods and for sampling, normalizing flows transform a noise distribution to the posterior distribution. With the success of diffusion models Ho et al. (); Liu et al. (); Lipman et al. (), it became clear that the mapping between sampling and posterior distribution can be specified a priori, for example by specifying a corruption process that transforms any data distribution to a normal Gaussian. The resulting continuous-time models outperform discrete architecture in many areas, and training larger models is much more scalable.

Figure: 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 $\bm{\theta}_{t}$ at time $t$ to estimate $\hat{\bm{\theta}}_{1}$. On the right, we show a gradient-based control signal with a differentiable simulator and cost function $C$ for improving $\hat{\bm{\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.
Refer to caption: x1.png

Despite the widespread success of flow-based models for generative modeling and density estimation, there is no direct feedback loop between the model, the observation $\bm{x}_{o}$ and the sample $\bm{\theta}$ during training, which makes it very difficult to produce highly accurate samples based on learning alone.

We propose a simple strategy to reintroduce control signals using simulators into the flow network. We refine an existing pretrained flow-based model with a flexible
control signal by aggregating the learned flow and control signals into a controlled flow. The aggregation network is small compared to the pretrained flow network, and we find that freezing the weights of the pretrained network works very well; thus, refining needs only a minimal amount of additional parameters and compute.

To demonstrate how these refinements affect the accuracy of samples and the posterior, we consider modeling strong gravitational lens systems
(Hezaveh et al., ; Cunha & Herdeiro, ; Legin et al., ; Vegetti & Koopmans, ; Vegetti et al., ),
an inverse problem in astrophysics that is challenging and requires precise posteriors for accurate modeling of observations. In galaxy-scale strong lenses, light from a source galaxy is deflected by the gravitational potential of a galaxy between the source and observer, causing multiple images of the source to be seen. Since these images and their distortions are sensitive to the distribution of matter on small scales, this can act as a probe for different dark matter models.
With upcoming and current sky surveys (Laureijs et al., ) expected to release large data catalogs in the near future, the number of known lenses will increase dramatically by several orders of magnitude. Traditional computational approaches require several minutes to many hours or days to model a single lens system. Therefore, there is an urgent need to reduce the compute and inference with learning-based methods. In this experiment, we demonstrate that using flow matching and our proposed control signals with feedback from a simulator, we obtain posterior distributions for lens modeling that are competitive with the posteriors obtained by MCMC-based methods but with much faster inference times.

Additionally, we evaluate different related variants of flow matching such as independent couplings (Tong et al., ), self-conditioning (Chen et al., ) or different loss formulations in the context of SBI using several benchmark problems. We then analyze our proposed control signals for the Lotka-Volterra model, a system of coupled ordinary differential equations (ODEs) descriping the population evolution of predators and prey over time. Our analysis underscores the essential role of simulator feedback for inference and that high accuracy is very challenging to achieve from scaling up datasets and model sizes alone.

To summarize, the main contributions of our work are:

- •
We propose a versatile strategy to improve pretrained flows with control signals based on feedback from a simulator. Control signals can be based on gradients and a cost function, if the simulator is differentiable, but they can also be learned directly from the simulator output.
- •
We assess different variants of flow matching in the context of SBI and demonstrate with the Lotka-Volterra model that performance gains due to simulator feedback are substantial and cannot be achieved by training on larger datasets alone.
- •
We demonstrate the efficacy of our proposed finetuning with control signals for inferring the parameter distributions of strong gravitational lens systems, a challenging inverse problem in astronomy that is sensitive to sample accuracy. We show that flow matching with simulator feedback is competitive with MCMC baselines and beats them significantly regarding inference time.

## 2 Related Work

#### Solving inverse problems with diffusion models

Diffusion models have been proposed to solve linear inverse problems (Kawar et al., ; ; Chung et al., ; Cardoso et al., ), as well as general inverse problems (Holzschuh et al., ; Chung et al., ; ).
In most of these works, the diffusion model learns the prior distribution and sampling from the posterior is achieved through a modified inference procedure, which guides samples via a conditioning. The conditioning can be based on a class label, text input (Song et al., ; Ho & Salimans, ; Saharia et al., ; Wu et al., ) or directly on a differentiable measurement operator (Chung et al., ; ).
In contrast to these works, we finetune a pretrained flow and learn an optimal combination of the pretrained flow and feedback from a simulator via control signals in the broader flow matching context.

#### Flow matching

Our work builds on top of prior work in flow matching (Lipman et al., ; Albergo et al., ; Pooladian et al., ; Tong et al., ; Albergo et al., ), particularly we adopt and evaluate conditional optimal transport paths (Lipman et al., ) and independent couplings or rectified flows (Liu et al., ; Tong et al., ) for simulation-based inference. We extend the existing literature by adding feedback from a simulator for scientific inverse problems.

#### Simulation-based inference

Our work directly compares to neural posterior estimation approaches for simulation-based inference (Cranmer et al., ; Lueckmann et al., , SBI). Contrary to static architectures (Dinh et al., ; Kingma & Dhariwal, ; Papamakarios et al., ; Durkan et al., ), our approach extends the continuous-time paradigm (Chen et al., ; Grathwohl et al., ). Wildberger et al. () have applied flow matching to neural posterior estimation and Sharrock et al. () have used conditional diffusion models and Langevin dynamics during sampling. In contrast to previous work, we include controls signals via problem-specific simulators and cost functions during training to significantly improve the sampling quality.

#### Strong lensing and parameter estimation

Machine learning has been successfully applied to estimate parameters of lens and source models (Hezaveh et al., ; Levasseur et al., ), however, previous methods are usually restricted to point estimates, use simple variational distributions or Bayesian Neural Networks (Schuldt et al., ; Legin et al., ; ; Poh et al., ) that are not well suited to represent more complicated high-dimensional data distributions. In this paper, we combine flow matching with problem-specific simulators to obtain highly accurate samples via feedback from control signals.

## 3 Flow Matching Theory

Continuous-time flow models transform samples $\bm{\theta}$ from a sampling distribution $p_{0}$ to samples of a target or posterior distribution $p_{1}$. This mapping can be expressed via the ODE

$$ $\displaystyle d\bm{\theta}_{t}=v_{\theta}(t,\bm{\theta}_{t})dt,$ (1) $$

where $v_{\phi}(t,\bm{\theta}_{t})$ represents a neural network with parameters $\phi$. Early works (Chen et al., ; Grathwohl et al., ) optimize $v_{\phi}(t,\bm{\theta})$ using maximum likelihood training, which is computationally demanding and difficult to scale to larger networks.
Instead, in flow matching the network $v_{\phi}(t,\bm{\theta})$ is trained by regressing a vector field $u(t,\bm{\theta})$ that generates probability paths that map from $p_{0}$ to $p_{1}$.

#### Generating probability paths

We say that a smooth vector field $u:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$, called *velocity*, generates the probability paths $p_{t}$, if it satisfies the continuity equation $\frac{\partial p}{\partial t}=-\nabla\cdot(p_{t}u_{t})$.
Informally, this means that we can sample from the distribution $p_{t}$ by sampling $\bm{\theta}_{0}\sim p_{0}$ and then solving the ODE $d\bm{\theta}=u(t,\bm{\theta})dt$ with initial condition $\bm{\theta}_{0}$. In the following, we will denote $u(t,\bm{\theta})$ by $u_{t}(\bm{\theta})$.
To regress the velocity field, we define the flow matching objective

$$ $\displaystyle\mathcal{L}_{\mathrm{FM}}(\theta):=\mathbb{E}_{t\sim\mathcal{U}(0 ,1),\bm{\theta}\sim p_{t}(\bm{\theta})}\left|\left|v_{\theta}(t,\bm{\theta})-u _{t}(\bm{\theta})\right|\right|^{2}.$ (2) $$

In order to compute this loss, we need to sample from the probability distribution $p_{t}(\bm{\theta})$ and we need to know the velocity $u_{t}(\bm{\theta})$. However, in general $u_{t}(\bm{\theta})$ is not accessible.

#### Conditioning variable

To solve this problem, we apply a trick by introducing a latent variable $\bm{z}$ distributed according to $q(\bm{z})$ and define the conditional likelihoods $p_{t}(\bm{\theta}|\bm{z})$ that depend on the latent variable so that $p_{t}(\bm{\theta})=\int p_{t}(\bm{\theta}|\bm{z})q(\bm{z})d\bm{z}$. Interestingly, if the conditional likelihoods are generated by the velocities $u_{t}(\bm{\theta}|\bm{z})$, then the velocity $u_{t}(\bm{\theta})$ can be written in terms of $u_{t}(\bm{\theta}|\bm{z})$ and $p_{t}(\bm{\theta}|\bm{z})$ with $u_{t}(\bm{\theta}):=\mathbb{E}_{q(\bm{z})}[u_{t}(\bm{\theta}|\bm{z})p_{t}(\bm{
\theta}|\bm{z})/p_{t}(\bm{\theta})]$. We can choose paths $p_{t}(\bm{\theta}|\bm{z})$ that are easy to sample from and for which we know the generating velocities $u_{t}(\bm{\theta}|\bm{z})$.
Next, we define the conditional flow matching loss

$$ $\displaystyle\mathcal{L}_{\mathrm{CFM}}(\phi):=\mathbb{E}_{t,q(\bm{z}),p_{t}( \bm{\theta}|\bm{z})}\left|\left|v_{\phi}(t,\bm{\theta})-u_{t}(\bm{\theta}|\bm{ z})\right|\right|^{2}.$ (3) $$

In contrast to the flow matching loss eq. 2, this loss is tractable and can be used for optimization. Now, one can show (Tong et al., ) that if $p_{t}(\bm{\theta})>0$ for all $\bm{\theta}\in\mathbb{R}^{d}$, then

$$ $\displaystyle\nabla_{\phi}\mathcal{L}_{\mathrm{FM}}(\phi)=\nabla_{\phi} \mathcal{L}_{\mathrm{CFM}}(\phi).$ (4) $$

This means that we can train $v_{\theta}(\bm{\theta},t)$ to regress $u_{t}(\bm{\theta})$ generating the mapping between $p_{0}$ and $p_{1}$ by optimizing the conditional flow matching loss eq. 3.

#### Couplings

The above framework allows for many degrees of freedom when specifying the mapping from $p_{0}$ to $p_{1}$ via the conditioning variable $\bm{z}$ and the conditional likelihoods $p_{t}$.
One particularly intuitive and simple choice is to consider the coupling $q(\bm{z})=p_{1}(\bm{\theta})$ (Lipman et al., ) together with conditional probability and generating velocity

$$ $\displaystyle p_{t}(\bm{\theta}|\bm{\theta}_{1})$ $\displaystyle=\mathcal{N}(\bm{\theta}|\,t\bm{\theta}_{1},(1-(1-\sigma_{\mathrm {min}})t)I)$ (5) $\displaystyle u_{t}(\bm{\theta}|\bm{\theta}_{1})$ $\displaystyle=\frac{\bm{\theta}_{1}-(1-\sigma_{\mathrm{min}})\bm{\theta}}{1-(1 -\sigma_{\mathrm{min}})t},$ (6) $$

where $\sigma_{\mathrm{min}}>0$. Conditioned on $\bm{\theta}_{1}$, this coupling transports a point $\bm{\theta}_{0}\sim\mathcal{N}(0,I)$ from the sampling distribution to the posterior distribution on the linear trajectory $t\bm{\theta}_{1}$ ending in $\bm{\theta}_{1}$ but decreasing the standard deviation from $1$ to a smoothing constant $\sigma_{\mathrm{min}}$. In this case, the transport path coincides with the optimal transport between two Gaussian distributions.

## 4 Controls for Improved Accuracy

While flow-based models $v_{\phi}(t,\bm{\theta})$ gradually transform samples from $p_{0}$ to $p_{1}$ in many steps during inference via solving the ODE eq. 1, there is no direct feedback loop between the underlying simulator, the current point on the trajectory $\bm{\theta}_{t}$, and the observation $\bm{x}_{o}$.
A central goal of our work is to reintroduce this feedback loop into inference and training by incorporating a control signal.

#### Conditioning of flows

Flows $v_{\phi}(t,\bm{\theta})$ can be conditioned on an observation $\bm{x}_{o}$ through an additional input $v_{\phi}(t,\bm{\theta},\bm{x}_{o})$, therefore modeling the conditional densities $p_{t}(\bm{\theta}|\bm{x}_{o})$ (Song et al., ). For example, in classifier free-guidance (Ho & Salimans, ), this conditioning is randomly dropped and set to 0 0 during training. The resulting models can then be used for both conditional and unconditional generation.

A fundamental problem here is that the conditioning $\bm{x}_{o}$ is static, whereas
we propose to have a dynamic control mechanism that depends on the trajectory $\bm{\theta}_{t}$, the observation, and an underlying control signal. The latter should relate $\bm{\theta}_{t}$ and observation using a physics-based model represented through a cost function $C$. As the accuracy of neural networks is inherently limited by the finite size of their weights, and smaller networks are attractive from a computational perspective, physics-based control has the potential to yield high accuracy with lean and efficient neural network models.

#### 1-step prediction

An additional issue is that the current trajectory $\bm{\theta}_{t}$ might not be close to a good estimate of a posterior sample $\bm{\theta}_{1}$, especially at the beginning of inference, where $\bm{\theta}_{0}$ is drawn from the sampling distribution.
This issue is alleviated by applying the cost function $C$ to the current estimate $\bm{\theta}_{t}$, we extrapolate $\bm{\theta}_{t}$ forward in time to obtain an estimated $\hat{\bm{\theta}}_{1}$

$$ $\displaystyle\hat{\bm{\theta}}_{1}=\bm{\theta}_{t}+(1-t)v_{\phi}(t,\bm{\theta} _{t},\bm{x}_{o}).$ (7) $$

#### Comparison with likelihood-guidance

The 1-step prediction is conceptually related to diffusion sampling using likelihood-guidance (Chung et al., ; Wu et al., ). For inference in diffusion models, sampling is based on the conditional score $\nabla_{\bm{\theta}_{t}}\log p(\bm{\theta}_{t}|\bm{x}_{o})$, which can be decomposed into

$$ $\displaystyle\nabla_{\bm{\theta}_{t}}\log p(\bm{\theta}_{t}|\bm{x}_{o})=\nabla _{\bm{\theta}_{t}}\log p(\bm{\theta}_{t})+\nabla_{\bm{\theta}_{t}}\log p(\bm{x }_{o}|\bm{\theta}_{t}).$ (8) $$

The first expression can be estimated using a pretrained diffusion model, whereas the latter is usually intractable, but can be approximated using
$p(\bm{x}_{o}|\bm{\theta}_{t})\approx p_{\bm{x}_{o}|\bm{\theta}_{0}}(\bm{x}_{o}
|\hat{\bm{\theta}}(\bm{\theta}_{t}))$,
where the denoising estimate $\hat{\bm{\theta}}(\bm{\theta}_{t}):=\mathbb{E}_{q}[\bm{\theta}_{0}|\bm{\theta}
_{t}]$ is usually obtained via Tweedie’s formula $(\mathbb{E}_{q}[\bm{\theta}_{0}|\bm{\theta}_{t}]-\bm{\theta}_{t})/t\sigma^{2}$. In practice, the estimate $\hat{\bm{\theta}}(\bm{\theta}_{t})$ is very poor when $\bm{\theta}_{t}$ is still noisy, impeding the inference in the early stages. On the contrary, flows based on linear conditional transportation paths have empirically been shown to have trajectories with less curvature (Lipman et al., ) compared to, for example, diffusion models, thus enabling inference in fewer steps and providing better estimates for $\hat{\bm{\theta}}_{1}$.

#### Controlled flow v ϕ C superscript subscript 𝑣 italic-ϕ 𝐶 v_{\phi}^{C} italic_v start_POSTSUBSCRIPT italic_ϕ end_POSTSUBSCRIPT start_POSTSUPERSCRIPT italic_C end_POSTSUPERSCRIPT

We pretrain the flow network $v_{\phi}(t,\bm{\theta},\bm{x}_{o})$ without any control signals to make sure that we can realize the best achievable performance possible based on learning alone. Then, in a second training phase, we introduce the control network $v_{\phi}^{C}(t,\bm{v},\bm{c})$ with pretrained flow $\bm{v}$ and control signal $\bm{c}$ as input. The control network is much smaller in size than the flow network, making up ca. $10\%$ of the weights $\phi$ in our large-scale experiments. We freeze the network weights of $v_{\phi}$ and train with the conditional flow matching loss eq. 3 for a small number of additional steps. This reduces training time and compute since we do not need to backpropagate gradients through $v_{\phi}(t,\bm{\theta},\bm{x}_{o})$. We did not observe that freezing the weights of $v_{\phi}$ affects the performance negatively. We include algorithms for training in appendix A.

### 4.1 Types of control signals

Figure: (a) Gradient-based control signal
Refer to caption: x2.png

Aiming for high inference accuracy, we extend self-conditioning via physics-based control signals to include an additional feedback loop between the model output and an underlying physics-based prior. We distinguish between two types of control signals.

#### Gradient-based control signal

In the first case, there is a differentiable cost function $C$ and a deterministic differentiable simulator $S$ as shown in fig. 2.
Given an observation $\bm{x}_{o}$ and the estimated prediction $\hat{\bm{\theta}}_{1}$, the control signal relates to how well $\hat{\bm{\theta}}_{1}$ explains $\bm{x}_{o}$ via some cost function $C$. The cost function can also depend directly on or be equal to the likelihood $p(\bm{x}_{o}|\hat{\bm{\theta}}_{1})$. For a differentiable cost function $C$, we define the control signal via

$$ $\displaystyle\bm{c}(\hat{\bm{\theta}}_{1},\bm{x}_{o}):=[C(S(\hat{\bm{\theta}}_ {1}),\bm{x}_{o});\nabla_{\hat{\bm{\theta}}_{1}}C(S(\hat{\bm{\theta}}_{1}),\bm{ x}_{o})].$ (9) $$

Note that while we recommend using informative control signals, we can use any control that depends on $\hat{\bm{\theta}}_{1}$ and $\bm{x}_{o}$.

#### Learning-based control signal

In the second case, the simulator is non-differentiable. To combine the simulator output with the observation $\bm{x}_{o}$, we introduce a learnable encoder model Enc with parameters $\phi_{E}$. The output of the encoder is small and of size $O(\mathrm{dim}(\bm{\theta}))$.
The control signal is then defined as

$$ $\displaystyle\bm{c}(\hat{\bm{\theta}}_{1},\bm{x}_{o}):=Enc(S(\hat{\bm{\theta}} _{1}),\bm{x}_{o}).$ (10) $$

The gradient backpropagation is stopped at the simulator output, see fig. 2.

### 4.2 Additional Considerations for Simulator Feedback

#### Stochastic simulators

Many Bayesian inference problems have a stochastic simulator. For simplicity, we assume that all stochasticity within such a simulator can be controlled via a variable $z\sim\mathcal{N}(0,I)$, which is an additional input. Motivated by the equivalence of exchanging expectation and gradient

$$ $\displaystyle\nabla_{\hat{\bm{\theta}}_{1}}\mathbb{E}_{z\sim\mathcal{N}(0,1)}[ C(S_{z}(\hat{\bm{\theta}}_{1}),\bm{x}_{o})]=\mathbb{E}_{z\sim\mathcal{N}(0,1)} [\nabla_{\hat{\bm{\theta}}_{1}}C(S_{z}(\hat{\bm{\theta}}_{1}),\bm{x}_{o})],$ (11) $$

when calling the simulator, we draw a random realization of $z$. During training, we randomly draw $z$ for each sample and step while during inference we keep the value of $z$ fixed for each trajectory.

#### Time-dependence

If the estimate $\hat{\bm{\theta}}_{1}$ is bad and the corresponding cost $C(\hat{\bm{\theta}}_{1},\bm{x}_{o})$ is high, gradients and control signals can become unreliable. In appendix B, we empirically find that the estimates $\hat{\bm{\theta}}_{1}$ become more reliable for $t\geq 0.8$. Therefore, we only train the control network $v_{\phi}^{C}$ in this range, which allows for focusing on control signals containing the most useful information. For $t<0.8$, we directly output the pretrained flow $v_{\phi}(t,\bm{\theta},\bm{x}_{o})$.

#### Theoretical correctness

Contrary to likelihood-based guidance, which uses an approximation for $\nabla_{\bm{\theta}_{t}}\log p(\bm{x}_{o}|\bm{\theta}_{t})$ as a guidance term during inference, the approximation $\hat{\bm{\theta}}_{1}$ only influences the control signal, which is an input to the controlled flow network $v_{\phi}^{C}$. In the case of a deterministic simulator, this makes the control signal a function of $\bm{\theta}_{t}$. The controlled flow network is trained with the same loss as vanilla flow matching (Lipman et al., ). Therefore all theoretical properties remain preserved.

## 5 Simulation-based Inference

This section is organized as follows. First, in section 5.1, we introduce a set of SBI benchmark tasks and provide a comparison of
popular neural posterior estimation (NPE) methods against
a baseline of flow matching without simulator feedback.
This comparison uses a similar training setup for all models and tasks.
Then, in section 5.2, we focus on an optimal task-specific network with training hyperparameters based on an extensive grid search.
We evaluate different variants of flow matching that are related to simulator feedback on the SBI tasks to push the performance as far as possible.
In section 5.3, we pick the most challenging SBI task and improve it further by introducing simulator feedback via gradient-based and learned control signals. We carefully analyze the cost-accuracy trade-off for using simulators and show that improvements from simulator feedback cannot be replicated by increasing the training dataset size alone.

### 5.1 Tasks and baselines

**Table 1: C2ST comparison with identical training setups and comparable number of network weights (ca. $300$K).**
| Method | LV | SLCP | SIR | TM |
| --- | --- | --- | --- | --- |
| CNF | 0.99 | 0.80 | 0.99 | 0.60 |
| NSF | 0.99 | - | 0.75 | 0.54 |
| FFJORD | 0.95 | 0.82 | 0.78 | 0.59 |
| Flow-Mat. | 0.93 | 0.79 | 0.79 | 0.58 |

We consider the SBI tasks Lotka Volterra LV, a coupled ODE for the population dynamics of interacting species, SIR, an epidemiological model for the spread of diseases, SLCP and Two Moons (TM), two synthetic tasks having complicated multimodal posteriors. All tasks are part of the benchmark collection from Lueckmann et al. ().
For each problem, the posterior distribution for a set of $10$ observations is known, which allows for directly comparing it with the posterior predicted by the trained model. This is measured using the C2ST score (Lopez-Paz & Oquab, ), which trains a classifier to discriminate between samples from the true posterior and samples generated from the learned model. If the classifier cannot discriminate between two sets of samples, its test accuracy will be $0.5$, whereas it increases when they become more dissimilar.

We include the following baseline methods for NPE:
Continuous normalizing flows (Chen et al., , CNF),
Neural Splince Flows (Durkan et al., , NSF),
and
FFJORD (Grathwohl et al., ).
Since we propose to include feedback from simulators, here we focus on the largest benchmark budget of $10^{5}$ simulator calls for generating the training dataset. Table 1 highlights that flow matching yields a highly competitive performance in this setting.
For details on the training setup, see appendix B.

Flow matching has also been evaluated for the SBI benchmark tasks by Wildberger et al. (), who performed an extensive hyperparamter search for each task to find optimal hyperparameters.
In the following, we focus on flow matching, and hence use the corresponding sets of optimal hyperparameters for each task.

Figure: Figure 3: Evaluation of SBI tasks using different variants of flow matching training. Lower C2ST scores are better.
Refer to caption: x4.png

### 5.2 Training variants

There are several variants of training diffusion models that can be related to simulation feedback and which we consider promising in the context of SBI. Before we go on to evaluate the simulator feedback in section 5.3, we test if we can improve the performance using any of them. In particular, we assess the following modifications:

- •
Self-conditioning: conditioning a model on something that depends on its own output can be seen as a form of self-conditioning. We evaluate an adapted version of self-conditioning (Chen et al., ). Instead of providing $\bm{\theta}_{t}$ to the flow network, the input is comprised of the concatenated vector $[\bm{\theta}_{t};\mathrm{Dropout}(\hat{\bm{\theta}}_{1})]$, where $\hat{\bm{\theta}}_{1}$ is the 1-step prediction eq. 7. For computing $\hat{\bm{\theta}}_{1}$, we require one network evaluation with the input $[\bm{\theta}_{t};0]$ and stop the gradient backpropagation at $\hat{\bm{\theta}}_{1}$. This method is similar to our simulator feedback, as it introduces a feedback loop that conditions the model on its own output, but without any simulator.
- •
Independent couplings: It is also possible to couple two non-Gaussian distributions by defining the coupling as $q(\bm{z})=p_{0}(\bm{\theta}_{0})p_{1}(\bm{\theta}_{1})$ and setting the conditional probabilities to the linear paths defined by $p_{t}(\bm{\theta}|(\bm{\theta}_{0},\bm{\theta}_{1}))=\mathcal{N}(\bm{\theta}|t
\bm{\theta}_{1}+(1-t)\bm{\theta}_{0},\sigma I)$ and $u_{t}(\bm{\theta}|(\bm{\theta}_{0},\bm{\theta}_{1}))=\bm{\theta}_{1}-\bm{
\theta}_{0}$ with bandwidth $\sigma>0$ (Liu et al., ; Tong et al., ). We can choose $p_{0}$ as the prior distribution $p(\bm{\theta})$ which we know in the SBI setting. Obtaining information in the form of an observation changes our knowledge about $\theta$ from the prior distribution to the posterior, therefore resembling a transformation similar to the noise to data transformation in diffusion models. This also suggests that the prior distribution can be closer to the posterior than a noise distribution.
- •
$x$-prediction: The reliability of the control signal depends directly on the 1-step estimate $\hat{\bm{\theta}}$. Instead of regressing the flow $u_{t}(\theta)$, we can directly predict the denoised estimate $\hat{\bm{\theta}}$ and obtain the velocity by rearranging eq. 7, giving $v_{\phi}(t,\bm{\theta}_{t},\bm{x}_{o})=\hat{\bm{\theta}}_{1}/(1-t)$. We additionally weight the $x$-prediction loss with a time-dependent weighting $w_{t}:=1/(1-t)$ to account for the scaling in eq. 7. The $x$-prediction potentially produces better estimates for $\hat{\bm{\theta}}$, thus allowing for obtaining more reliable feedback from control signals when $t<0.8$.

#### Evaluation

Figure 3 shows an evaluation of the different variants against vanilla flow matching (Gaussian sampling distribution, no self-conditioning and velocity prediction). Using task-specific priors produces outliers with better C2ST scores for SLCP but is consistently worse for LV and SIR. We conclude that normal Gaussian distributions are more suited as sampling distributions for most low-dimensional problems. Introducing self-conditioning does not show any improvements, so feedback loops without a simulator alone are not sufficient for better performance in this situation. Finally, the $x$-prediction loss consistently performs worse than the velocity prediction. Therefore, a potential improvement in the 1-step estimate is outweighed by a corresponding deterioration of the posterior correctness as indicated by the C2ST score.

### 5.3 Simulator feedback: Gradient-based and Learned

Figure: Figure 4: Evaluation of simulator feedback for LV.
Refer to caption: x8.png

In this section, we focus on the Lotka-Volterra (LV) task for a more detailed analysis. It has the highest difficulty as seen by the C2ST score, and we use it to test the different types of feedback. We reimplement the LV simulator in JAX (Bradbury et al., ) to support differentiability and evaluate the gradient-based control signal as well as the learning-based control signal, using a small multilayer perceptron (MLP). In addition, to make sure that observed improvements are not due to the increased number of network parameters and finetuning with the control network, we also evaluate a variant where we finetune with the control network but set all simulator-dependent inputs to the control network to 0 0 (Zero Controls).
We show an evaluation with C2ST in fig. 5. For both the learning and gradient-based control signals we see clear improvements with the gradient-based signal clearly ahead. The zero control signal improves only slightly, showing that the improvement can be directly attributed to the simulator.
While control signals are most useful for more high-dimensional problems with less sparse and noisy observations, this experiment demonstrates that they can also be used in low-dimensional settings. Moreover, while differentiable simulators can provide better control signals, feedback from non-differentiable simulators likewise shows clear improvements.

### 5.4 Computational efficiency

A critical issue in SBI is that calls to the simulator are potentially expensive. This imposes the question of whether compute time is better spent on extending the training dataset or training with feedback from the simulator. We empirically verify that the latter is more efficient for the LV task in this setup by comparing our method to models with an increased training dataset from a larger simulator budget. Specifically, we train with dataset sizes of $10^{6}$ and $10^{7}$. Training the gradient-based control signal took ca. $9\times 10^{6}$ simulator calls. See fig. 5 for the evaluation. There is no improvement in the C2ST for models trained without simulator feedback beyond $10^{5}$ data points, and the final train/validation loss for the $10^{7}$ model indicates that there is no more overfitting.
Nonetheless, the model trained with controls clearly outperforms the model trained with more data, indicating that the directed feedback of the simulator cannot be replaced by increased amounts of training data.

## 6 Strong Gravitational Lensing

We present our results for modeling strong gravitational lens systems, a challenging and highly relevant non-linear problem in astronomy.
Strong gravitational lensing is a physical phenomenon whereby the light rays by a distant object, such as a galaxy, are deflected by an intervening massive object, such as another galaxy or a galaxy cluster. As a result, one observes multiple distorted images of the background source. We aim to recover both the lens and source light distribution as well as the lens mass density distribution with realistic simulated observations for which we know the ground truths. We evaluate flow matching as an NPE method with gradient-based control signals from a differentiable simulator with two MCMC methods.

#### Lens modeling

The lens equation relates coordinates on the source plane $\beta$ and the observed image plane $\Theta$ via the deflection angle $\alpha$ induced by the mass profile or gravitational potential of the lens galaxy. We use a Singular Isothermal Ellipsoid (SIE) to describe this lens mass and Sérsic profiles for both the source light and light emitted from the lens galaxy (full details are provided in appendix C). There are $9$ parameters for the lens mass and $7$ parameters for each Sérsic profile, giving $23$ parameters in total. The likelihood is measured by the $\chi^{2}$-statistic, which is the modeled image plane $\Theta$ minus the observation $\bm{x}_{o}$ divided by the noise. To solve the lensing equation, we make use of the publicly available raytracing code by (Galan et al., ). We want to stress that even small perturbations of the model parameters can cause the $\chi_{2}$ to increase significantly; see fig. 13 in the appendix.

Figure: Figure 6: Reconstruction of observation. Flow matching is purely learning-based and shows noticeable residuals in the reconstruction. Including simulator feedback removes remaining residuals.
Refer to caption: x10.png

#### Datasets and pretraining

Several instrument-specific measurement effects are included when simulating the observations. We include background and Poisson noise and smoothing by a point-spread function (PSF). The pixel size corresponds to $0.04$ arc seconds. These directly affect the posterior, as more noise and a stronger PSF will widen the posterior distribution. We generate 250 000 data samples for training and 25 000 for validation. The flow network $v_{\phi}$ consists of a convolutional feature extraction neural network represented by a shallow CNN whose output is fed into a dense feed-forward neural network with residual blocks. Full details are in appendix C.

#### Finetuning with control signals

The control network $v_{\phi}^{C}$ is represented by another dense feed-forward network, which accounts for $11\%$ of all parameters in the combined model. The control signals are obtained from simulating an observation based on the predicted estimate $\hat{\bm{\theta}}_{1}$ via ray-tracing (Galan et al., ) based on the parametric models, calculating the $\chi_{2}$-statistic and computing gradients with respect to the estimate $\hat{\bm{\theta}}_{1}$. The $\chi_{2}$-statistic itself is also part of the control signal.

#### Reference posteriors

As reference posteriors, we include Hamiltonian Monte Carlo (HMC) with No-U-Turn sampler (Hoffman et al., , NUTS) and Affine-Invariant Ensemble Sampling (Goodman & Weare, , AIES), which are both two popular MCMC-methods in astronomy. We adopt implementations of both methods using numpyro (Phan et al., ; Bingham et al., ).
Details on running both methods can be found in appendix C. Additionally, we compare to diffusion posterior sampling as a learned baseline (Chung et al., , DPS), see appendix C.1 for details.
We use Euler integration for both flow matching variants.

### 6.1 Evaluation and Discussion

**Table 2: Evaluation with respect to average $\chi_{2}$ and inference time for the posterior distribution.**
| Method | Avg. $\chi_{2}\downarrow$ | Modeling Time $\downarrow$ |
| --- | --- | --- |
| NUTS | 1.83 | $\sim$ 56x (564s) |
| AIES | 1.74 | $\sim$ 67x (672s) |
| DPS | 9.98 | $\sim$ 42x (427s) |
| *Flow-Mat.* | 1.83 | 1x (10s) |
| + Simulator | 1.48 | $\sim$ 2x (19s) |

#### χ 2 subscript 𝜒 2 \chi_{2} italic_χ start_POSTSUBSCRIPT 2 end_POSTSUBSCRIPT -statistic

We show an evaluation of all methods in table 2. The average $\chi_{2}$ is computed over $1000$ randomly chosen validation systems, where for each, we draw $1000$ samples from the posterior.
If we compute the $\chi_{2}$ for the ground truth parameters, we obtain a value of $1.17$ due to the noise in the observation. Since we cannot overfit to noise with the parametric models, this represents a lower bound for $\chi_{2}$ in this experiment.
Including the physics-based control improves the $\chi_{2}$ from $1.83$ to $1.48$, representing an improvement of $53\%$ relative to the best modeling. The improved $\chi_{2}$ is even better than the best baseline method, AIES.

#### Modeling time

We define the modeling time as the average compute time required to produce $1000$ credible posterior samples. Both HMC and AIES require significant warmup times before producing the first samples from the posterior, which we include in the table. However, after warmup, it is relatively cheap to obtain new samples. On the other hand, flow matching does not require any warmup time and the modeling time increases linearly with the number of posterior samples. All methods were implemented in JAX (Bradbury et al., ) and used the same hardware.
The measurements in table 2 show that DPS is faster than the classic baselines, but yields a very sub-optimal performance in terms of its distribution.
The performance numbers also highlight that our method yields an accuracy that surpasses AIES, while being more than 30x faster.
This evaluation demonstrates that flow matching-based
methods are highly competitive even in small to moderate-sized problems
where established MCMC methods in terms of accuracy exist, clearly beating them in terms of inference time.
Flow matching with our proposed control signals is especially interesting because it is not affected as much by the curse of dimensionality as traditional inference methods and allows for having non-trivial learnable high-dimensional priors.
However, before these methods are widely trusted,
they need to demonstrate their competitiveness with classical methods.
Our results show that this is indeed the case, which opens up exciting avenues for applying and developing approaches targeting similar and adjacent inverse problems in science.

Figure: Figure 7: SBC for $x_{\mathrm{center}}$ of the source galaxy.
Refer to caption: x11.png

#### Simulation-based calibration

Acquiring truthful posterior distributions for Bayesian inference problems is difficult, which makes it hard to robustly evaluate whether the predicted posterior distribution is correct. We use simulation-based calibration (Talts et al., , SBC) as an additional evaluation tool. The data-averaged posterior obtained from averaging the posterior distribution over many problem instances has to be equal to the prior. This can be tested by considering a one-dimensional function $f:\bm{\theta}\mapsto\mathbb{R}$ and $L$ samples $\bm{\theta}^{1},...,\bm{\theta}^{L}$ drawn from an inference method. If $\bm{\theta}^{*}$ are the ground truth parameters, then the rank statistic $\sum_{l=1}^{L}\mathbf{1}_{f(\bm{\theta}^{L})<f(\bm{\theta}^{*})}$ has to be uniformly distributed over the integers $[0,L]$. If the distribution of the rank statistic is plotted as a histogram, systematic problems in the inference method can be identified visually, see fig. 7. We set $L=1000$ and plot the histograms for all $n=1000$ test problems and visualize the parameter $x_{\mathrm{center}}$, which defines the position of the source in $x$-direction. The posteriors without simulator feedback are biased, as can be seen in the deviation from uniformity in the plots. Including simulator feedback improves the distribution of the rank statistic. For an extended analysis, see appendix C.2.

### 6.2 Limitations

While introducing additional control signals increases the quality of produced samples, it comes at the cost of slower inference and training times depending on the speed of the simulator.
In general, using non-differentiable control signals is possible but removes the possibility of computing likelihoods via the instantaneous change of variables formula (Chen et al., ). Compared to MCMC approaches, inference with flow-based models requires a substantial upfront cost for training that needs to be amortized across many problems. Additionally, priors are encoded in the learned flow networks, so changing them would require retraining models with adjusted data sets.

## 7 Conclusion

We presented a method for improving flow-based models with simulator feedback using control signals.
This allows us to refine an existing flow with only a few additional weights and little training time. We thereby efficiently bridge the gap between purely learning-based methods for simulation-based inference and optimization with hand-crafted cost functions within the framework of flow matching. This improvement is critical for scientific applications where high accuracy and trustworthiness in the methods are required. Purely learning-based methods face significant difficulties in producing very accurate samples, as there is usually no feedback during inference of how good samples are.
In this paper, we demonstrated that we do not need large network sizes or tremendous amounts of data to train accurate models that are competitive with established MCMC methods if we include suitable control signals from simulators. We believe this work makes an important step towards making posterior inference in science more accurate, understandable, and reliable.

## Appendix

## Appendix A Algorithms

We include algorithms for training using flow matching and control signals, see algorithm 1. For flow matching with self-conditioning, see algorithm 2.

Figure: Algorithm 1 FM with Control Signals

## Appendix B Simulation-based Inference

#### Baselines comparison in section 5.1

For a fairer comparison, we set up all baseline methods with a similar number of network weights and available compute time.

We train all baselines and flow matching with a batch size of 512 on the largest $10^{5}$ simulation budges for all tasks. For optimization, we apply Adam (Kingma & Ba, ) with default settings and constant learning rate of $10^{-4}$ and weight decay $2\times 10^{-5}$.

All network architectures are chosen to have a similar number of ca. $3\cdot 10^{5}$ parameters. For flow matching and continuous normalizing flows (CNFs), we use the same architecture based on a dense feed-forward neural net with skip connections using 8 residual blocks with each 128 neurons and *elu* activation. As input, we concatenate time $t$ and $\bm{\theta}_{t}$.
For Neural Spline Flow (Durkan et al., ) and FFJORD (Grathwohl et al., ), we adopt the released implementation by the authors.

Depending on the time per epoch for each method, we modify the number of epochs and steps per epoch to allow all methods to train for a similar amount of time, ensuring a sufficient window for convergence. For NSF, we train for 1 000 epochs, for flow matching for 2 000 epochs, and for FFJORD and CNF 100 epochs.

#### Flow matching with optimized hyperparameters

For the experiments in section 5.2 and section 5.3, we adopt the hyperparameters and network architecture from Wildberger et al. (), which is based on a hyperparamter grid search. The hyperparameters for each task are listed in table 3. Otherwise, we follow the implementation as provided by the authors.

**Table 3: Hyperparameters for SBI from Wildberger et al. ().**
| Task | Time $\alpha$ | Batch size | Learning rate | Residual blocks |
| --- | --- | --- | --- | --- |
| LV | 1 | 32 | $10^{-3}$ | [32, 64, 128, 256, 5$\times$512, 256, 128, 64, 32] |
| SLCP | -0.5 | 256 | $5\cdot 10^{-4}$ | [32, 64, 128, 256, 5$\times$512, 256, 128, 64, 32] |
| SIR | 4 | 256 | $5\cdot 10^{-4}$ | [32, 64, 128, 256, 7$\times$512, 256, 128, 64, 32] |
| TM | 4 | 64 | $2\cdot 10^{-4}$ | [32, 64, 128, 256, 512, 3$\times$1024, 512, 128, 64, 32] |

#### Analyzing the 1-step estimate

We simulate the flow ODE from the sampling distribution at $t=0$ until $t^{*}$ ($x$-axis). Then, we compute the posterior in a single step by linearly extrapolating the flow, see eq. 7, to obtain the estimate $\hat{\bm{\theta}}_{1}$. Results are shown in fig. 8.

Figure: (a)
Refer to caption: extracted/5963824/pictures/iclr/c2st_1step.png

#### Analyzing step size

We analyze the influence of the step size of the ODE solver on the quality of the posterior distribution as shown in fig. 9.

Figure: (a)
Refer to caption: extracted/5963824/pictures/iclr/c2st_stepsize.png

#### Additional results for maximum mean discrepancy

For the evaluation in section 5.2, we show additional results for the maximum mean discrepancy (Gretton et al., , MMD) in fig. 10.

Figure: Figure 10: Evaluation of SBI tasks using different variants of flow matching training. Lower MMD scores are better.
Refer to caption: x12.png

## Appendix C Strong Gravitational Lensing

We consider the following models:

- •
For modeling the lens we use an SIE model with 6 parameters: the Einstein radius $\theta_{E}$, the ellipticities $e_{1}$ and $e_{2}$ and $x_{\mathrm{center}}$ and $y_{\mathrm{center}}$. There is shear, for which we only consider $\gamma_{1}$ and $\gamma_{2}$ as free parameters.
- •
The source is modeled by a Sersic profile with free parameters being the amplitude, the half-light radius, the Sersic index $n$, the ellipticities $e_{1}$ and $e_{2}$ as well as the positions $x_{\mathrm{center}}$ and $y_{\mathrm{center}}$.
- •
The lens light is modeled in the same way as the source, although when generating the mock data, we fix the position as well as ellipticities to be the same as the lens mass model. For training and inference, we infer positions for both lens mass and lens light model, so the model could produce different values for them. The MCMC methods use the same parameter for both lens light and lens mass position.

We list all priors in table 5, table 5 and table 6. We do not have priors on the ellipticities $e_{1}$ and $e_{2}$ directly, but we obtain them from priors on the position angle and axis ratio. Also, we obtain the shear parameters from $\gamma_{1}$ and $\gamma_{2}$ from $\phi_{\mathrm{ext}}$ and $\gamma_{\mathrm{ext}}$ by converting them polar to cartesian coordinates.
For SBI, we also include the two parameters $\mathrm{ra}_{0}$ and $\mathrm{dec}_{0}$ for the shear, which are always set to 0 when generating the training data sets, but in general our network could infer other values. Overall, there are $23$ parameters for $v_{\theta}$, which fully describe the simulation setup. However, in our dataset there are only $17$ free parameters. The MCMC methods only infer the reduced set of parameters, making use of the dependencies between them.

#### Measurement instruments

Observations have $160$ times $160$ pixels. The pixel size is $0.04$ arc seconds. We use a Gaussian points spread function (PSF) with full width at half maximum (FWHM) of 0.3.
The there is Gaussian background noise with a root mean-squared values of $0.01$ and an exposure time of 1000s.

#### Setup of MCMC-based methods

We setup both baselines methods as follows:

- 1.
Hamiltonian Monte Carlo: we use the No-U-Turn samples with a maximum tree depth of $10$ and 5 000 warmup steps.
- 2.
Affine-Invariant Ensemble Sampling: we use DEMove and StretchMove both with probability $0.5$. There are $400$ chains and we warm up for 20 000 steps before starting sampling.

Both methods are implemented in numpyro and optimized with JAX, so their runtimes are comparable with each other.

#### Network architectures and training

- •
Our flow network $v_{\phi}$ comprises a lightweight feature extraction network, representated by a CNN, which is consists of $6$ downsampling blocks with 1 layer each a $32$ channels and kernel size 3. As postprocessing of the output, we apply GroupNorm, silu and an additional 2DConv layer with kernel size $3$ and a single channel. We reshape the output and feed it through a final dense layer. The output of the feature extraction has the same dimensionality as the parameters $\bm{\theta}$.
- •
An additional dense feed-forward neural network receives the concatenated the time $t$, $\bm{\theta}_{t}$ and extracted features as input. The feed-forward neural neural networks consists of 8 residual blocks with hidden dimension 128 and elu activation.
- •
The control network $v_{\phi}^{C}$ is represented by a small feed-forward neural network, consisting of 3 residual blocks with 64 hidden layers and 3 residual blocks with 32 hidden layers. We condition each block on the time via gated linear units and use a 16 dimensional time embedding.

For training, we use a batch size of 256 for the flow network $v_{\phi}$.
When training $v_{\phi}^{C}$, we decrease the batch size to 16. We use the Adam optimizer with a learning rate of $10^{-4}$ and weight decay of $10^{-5}$. Training $v_{\phi}$ was done on a single NVIDIA Ampere A100 GPU for ca. $45$ hours. We trained $v_{\phi}^{C}$ for an additional $24$ hours. A lot of the training time was spent on running evaluation metrics, so it can be substantially improved.

**Table 4: Priors for lens mass model parameters**
| Parameter | Prior |
| --- | --- |
| $x_{\mathrm{center}}$ | $\mathcal{U}(-0.2,0.2)$ |
| $y_{\mathrm{center}}$ | $\mathcal{U}(-0.2,0.2)$ |
| position angle $\phi$ | $\mathcal{U}(0,180)$ |
| axis ratio $q$ | $\mathcal{U}(0.25,1)$ |
| external shear orientation $\phi_{\mathrm{ext}}$ | $\mathcal{U}(0,180)$ |
| external shear strength $\gamma_{\mathrm{ext}}$ | $\mathcal{U}(0,0.1)$ |
| Einstein radius $\theta_{E}$ | $\mathcal{U}(0.5,2.0)$ |

**Table 6: Priors for the lens light**
| Parameter | Prior |
| --- | --- |
| amplitude | $\mathcal{U}(5.0,10.0)$ |
| half-light radius | $\mathcal{U}(0.5,2.0)$ |
| Sersic index $n$ | $\mathcal{U}(1.5,4.0)$ |

### C.1 Diffusion Posterior Sampling (DPS)

We setup diffusion posterior sampling Chung et al. () as an additional baseline. The training dataset is the same as in 6, however since the diffusion model is unconditional, we drop any conditioning information.

#### Network architecture and training

The neural network architecture is a multilayer perceptron MLP with 8 residual blocks and 128 neurons each. The activation function is *elu*. As optimizer, we use Adam with weight decay ($10^{-5}$). We train for $2000$ epochs and for each epoch we sample $1000$ batches from the dataset using a batch size of 4. We train the network as a denoising diffusion probabilistic model (DDPM) following Ho et al. ().

#### Unconditional generation

Below, in figure 11, we visualize three samples generated by unconditionally sampling from the model. The observations are created using the lensing simulation code with the generated samples as input.

Figure: Figure 11: Visualization of unconditionally generated lensing systems.
Refer to caption: x16.png

#### Inference

We directly follow Chung et al. () Algorithm 1 for inference, where the measurement operator $\mathcal{A}$ is replaced by the lensing simulation code. The step size in the algorithm is defined via a hyperparameter $\zeta$, which needs to be finetuned depending on the problem. We empirically test different values for $\zeta$ to find an optimal choice. Our results are shown in table 7. In this evaluation, we only model a smaller number of systems ($n=25$).

**Table 7: Evaluation of DPS and choosing $\zeta$.**
| $\zeta$ | 0.0 | 0.0005 | 0.001 | 0.005 | 0.01 | 0.05 |
| --- | --- | --- | --- | --- | --- | --- |
| Avg. $\chi^{2}$ | 28.15 | 16.20 | 9.98 | 10.07 | 12.98 | 12.64 |
| Min. $\chi^{2}$ | 15.14 | 3.84 | 3.07 | 1.58 | 1.40 | 1.53 |

### C.2 Simulation-based Calibration

We use simulation-based calibration (Talts et al., , SBC) as an additional evaluation method to assess the correctness of the posterior distributions. We adopt the SBC implementation from the Python package *sbi*.
Below, in fig. 12, we show histograms for 8 parameters based on $n=1000$ lens systems with $L=1000$ posterior samples each.

Figure: (a) AIES
Refer to caption: x17.png

## Appendix D Posteriors and Reconstructions for Lens Modeling

We show how small perturbations in the lens system’s parameters affect the simulated observation in figure 13.
We show extended plots of the posteriors in fig. 14 for lens system 1 and fig. 15 for lens system 6.
Additionally, we show reconstructions based on flow matching with and without simulator feedback of lens systems 1 to 6 in fig. 16

Figure: Figure 13: We show how a small noise $\sigma$ affects the simulated observation. We add a normal Gaussian with mean 0 0 and standard deviation $\sigma$ to a lens system’s ground truth parameters $x$. Then, we plot the simulated observation based on the noised parameters and show the residuals.
Refer to caption: x20.png

Figure: Figure 14: Posterior plot for system 1.
Refer to caption: x21.png

Figure: Figure 15: Posterior plot for system 6.
Refer to caption: x22.png

Figure: Figure 16: Modeling of different lens systems: system 1 (top) to system 6 (bottom).
Refer to caption: x23.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: Please see above.

Simpler than the encoder architecture, did you consider physics-inspired ways of providing the simulator information to the drift model? For example, simply expressing the drift as (learned vector field NN) + (learned scalar or diagonal NN) x (simulator gradient), as often done in work on diffusion models for sampling of Boltzmann distributions.
soundness: 3
strengths: - The problem studied is highly relevant as foundation models (including diffusion and flow-based models) become available in various scientific domains. It is important to develop  inference methods for inverse problems that have low bias, good posterior coverage, and high efficiency of training and inferences -- this paper attempts to solve these problems.
- The proposed algorithm plausibly attacks these challenges (even if it is not very well demonstrated by the experiments, see below) and should give an asymptotically correct solution to the posterior sampling problem.
- Interesting analysis of algorithm variants in Section 5.2.
summary: The paper studies the problem of modelling the posterior $p(\theta\mid x)$ in a generative model $\theta\to x$, where $p(x\mid\theta)$ is available as a simulator, but possibly without access to exact likelihoods. When the sampling of $\theta$ given $x$ is modelled as a conditional (on $x$) neural ODE and this ODE is trained by flow matching objectives, it is proposed to place an inductive bias on the drift model: the output of simulator or its gradient, evaluated at an intermediate time point or its extrapolation to the target space, is encoded and given as an input to the drift model. Such a form of the drift model is hypothesised to improve the approximation of the target distribution by effectively guiding the drift to the modes of the posterior. Experiments are done on several low-dimensional simulation-based inference tasks, including the lens and source parameter estimation problem in strong gravitational lensing.
weaknesses: Throughout the text, there are many inaccurate or somewhat sloppy statements and references that confuse a specialist in flow matching models (and would likely impede understanding by non-specialists as well). A list follows.

- Abstract: I would suggest to revise it to explain the problem setting and approach at a higher level.
  - First sentences: "Flow-based generative modeling is a powerful tool for solving inverse problems in physical sciences" -- this is a bold claim. The use of flow-based models for inverse problems is not yet well-established; in fact, this is what this paper aims to do.
  - The next few sentences do not set up the problem well (it is not even explained that we are talking about continuous normalising flows).
  - The results at the end do not make sense without context: what does "improves the accuracy by 53%" mean?
- Introduction:
  - L046: "normalising flows transform a noise distribution to the posterior distribution" is true, but:
    - It is a statement with low specificity (VAEs and GANs also transform noise to data). 
    - These models are first introduced in the setting of training from samples, which is not what we usually have in Bayesian inference (no ground truth samples from the posterior).
    - Two of the three citations about diffusion models are actually about ODEs / flow matching. The two are of course connected, but I think it is unfair to call flow-based models an instance of "success of diffusion models [...] specifying a corruption process". For instance, flow ODEs can be learned that are not the probability flow ODEs of diffusion processes, including Liu et al.'s rectified flow (the first iteration is indeed a diffusion ODE, the later 'straightened' ones are not) and Tong et al.'s minibatch OT-based flow matching.
      - One solution could be to explicitly state the connection between FM and diffusion in the cases where it exists (e.g., for Ornstein-Uhlenbeck noise, optimal drifts of ODE and SDE are both expressed in terms of the score, so learning one is tantamount to learning the other).
  - In the paragraph starting L052, somehow we have jumped from a general distribution-matching setting (which is not how flow-based models were introduced -- they are trained from samples!) to a conditional posterior modelling setting. Please state the setting/assumptions (e.g., that we have conditional posterior samples from a simulator).
- Related work:
  - "Inverse problems with diffusion models": This seems to be better named "solving inverse problems under a diffusion model prior". Although the works that do this with Monte Carlo (e.g., Cardoso et al, Dou et al,) are mentioned (you could also consider [Song et al.](https://proceedings.mlr.press/v202/song23k.html)), there is also stochastic optimisation (e.g., [Mardani et al.](https://arxiv.org/abs/2305.04391), [Graikos et al.](https://arxiv.org/abs/2206.09012)), or amortisation by RL methods (e.g., [Black et al.](https://arxiv.org/abs/2305.13301), [Fan et al.](https://arxiv.org/abs/2305.16381), [Venkatraman et al.](https://arxiv.org/abs/2405.20971)).
  - "Flow matching": It is strange to see "optimal transport paths (Lipman et al.)" contrasted with "independent couplings or rectified flows (Liu et al., Tong et al.)". In fact, it is Rectified Flow and OT-CFM (Liu et al., Tong et al.) who consider **non-independent couplings** through rectification steps or OT couplings (thus actually approximating the dynamic OT), respectively, while Lipman et al.'s flow matching is equivalent to one using independent couplings and is OT only on the level of the conditional probability paths used for training.
- FM theory:
  - Equation (1): $\theta$ in subscript should be $\phi$.
  - L156: Because smoothness conditions are stated, they should be precise: Do you assume Bochner integrability? Continuous differentiability (how many times?) in both $\theta$ and $t$? It should also be stated that $p_t(x)=p(t,x)$ and $p$ is a function $[0,1]\times\mathbb{R}^d\to\mathbb{R}$.
  - L182 is hard to understand: what is meant by "$q(z)=p_1(\theta)$? It should be said that the conditioning variable $z$ is identified with the endpoint $\theta_1$, etc.
- Controls for improved accuracy:
  - LL201-203 do not make sense to me. The paragraph begins with conditioning of ODEs -- how is an old trick in diffusion "for example" w.r.t. such conditioning?
  - NB. Equation (7) will be the *exact* $t=1$ endpoint of integration ($\theta_1=\hat{\theta}_1$) if we have a perfectly fit OT or any model with straight integral curves (such as the converged ODE after many iterations of rectified flow)!
  - LL225-227 and later at LL352-355: Once again, I am surprised by the repeated discussion of Liu et al. and Tong et al. yet the *omission of the actual algorithms they propose* (rectification and MBOT coupling), which both produce **straighter** paths than a vanilla FM (and hence inference in fewer steps).

Experiment results are not convincing:
- Results are not consistently showing improvement (cf. Table 1) and error bars are not reported, making it impossible to assess significance.
- In Section 5, can you comment on computational efficiency in terms of wall time? 
- Lensing:
  - The evaluations do not seem to guarantee coverage ($\chi^2$ is obviously not sufficient for this, and the SBC tests use projection only on a single parameter). Have you considered coverage tests such as those used in Legin et al., Section 3? Currently it is not demonstrated that the proposed method achieves more accurate posterior sampling.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: I would be happy if authors could provide arguments against my concerns.
soundness: 1
strengths: - *Novelty*: To the best of my knowledge, the idea of including simulator feedback in flow matching for posterior estimation is novel
- *Soundness*: The idea of using simulator feedback to correct for the imprecision of score matching may indeed be helpful in certain applications. 
- *Presentation*: I found the tables and figure easy to read and informative, while being also visually appealing.
summary: This paper introduces flow matching with simulator feedback for simulation-based inference which extends offline flow matching for neural posterior estimation with an online phase that uses online simulations to improve the accuracy of the estimated posterior distribution. Indeed, the authors observe that learning perfectly the score function corresponding to the true posterior distribution is challenging and propose to access the simulator online and correct these imperfections at evaluation time. The method is supposedly much more efficient than alternative online method such as MCMC while having the potential to provide as accurate results. The paper introduces two types of control signals, a gradient-based and a learning-based control signal, which are made for differentiable (and deterministic) and non-differentiable (and potentially stochastic) simulators respectively. The method is empirically tested on 4 common SBI benchmarking tasks and a "strong gravitational lensing" problem. Results highlight that simulator feedback help improving the accuracy of the posterior distributions.
weaknesses: While I must acknowledge certain positive aspects of the paper, I have also several concerns that motivate my negative recommendation, which I am listing below. 
1. While I find the figures and tables quite enlightening, I find the presentation being a weakness. There are multiple hand-wavy explanations and claims that I find a bit confusing. See below for concrete examples.
2. Empirical validation: I was surprised by the numbers from figure 3 and 4 which seems worse than existing alternatives. In particular, the paper only compares to offline methods while it is clear from [1] that SNLE or SNPE which are sequential alternatives to NPE/NLE and perform much better than the proposed method. It is also arguable whether these benchmarks are the most relevant ones as there are quite low dimensional in term of observation dimensionality and are not fully representative of applications where SBI can shine. For instance, gravitational waveforms and spiking neurons are interesting benchmark used in many SBI papers. I also find the results of Flow matching + posterior quite bad compared to MCMC methods in figure 14 and 15 where we can clearly see an issue with bias and also variance of the predicted posteriors. 
Overall, the empirical validation seems insufficient to me and do not clearly demonstrate that the method proposed is of any real use. 
3. Relevance: Sequential SBI methods, which call the simulator online, are often complicated to motivate as they require to access the simulator while also arguing doing inference on this simulator without SBI is hard because evaluating the simulator is computationally demanding. While I agree this is not the case for all applications, I would expect the paper to clearly highlight and benchmark the method on use cases where simulating more samples offline to get a good amortised posterior is not enough and calling the simulator online does not take too much time.  

I am not confident these concerns can be addressed in the scope of the discussion (especially regarding results and presentation) but I am open to discussing with authors.

Hand-wavy explanations examples:
- l33-36: Seems an intricated way of explaining likelihood and posterior distributions which are very standard mathematical objects most reader should already know about.
- l37-43: The presentation of SBI is again a bit weird in my opinion. It does not clearly say when and why SBI may be necessary and what it solves. It may be interpreted as if SBI was only about Bayesian inference for the uninformed reader where frequentist methods exist as well.
- l48: what do you mean by "it became clear ... be specified a priori"?
- l49-50: this a vague and pretty strong statement that I would expect to be clarified and supported by reference.
- l52-54: I do not understand what you mean by saying there is "no feedback loop".
- l74-78: you mix the high level description of the method with specific implementation decisions which makes it hard to understand on what aspects the reader should focus on.
- Section 3 was presented in a way that I did not find easy to digest. 
- l205: "a fundamental problem..." why is it a fundamental problem is unclear to me.
- l254-255: What shall we conclude from that sentence?
- l277-282: It seems like a hacky solution
 - l283: While I understand you are trying to emphasise that the algorithm still learn the "true" posterior distribution, this is stated in a hand-wavy way in my opinion. 
- l345-351: This is quite confusing again. Why is there a dropout layer? It is quite complicated to grasp the details of all variants.
- 5.4: This is again a very hand-wavy claim without actually being theoretical or empirical support. 


[1]:http://proceedings.mlr.press/v130/lueckmann21a/lueckmann21a.pdf

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 2
questions: I would like to see especially the experiments and responses to the issues mentioned as weaknesses.
soundness: 1
strengths: The paper has some strengths overall, which I will outline below. 


**Strengths:**
1. Introduction of the flow matching perspective for these set of problems making a faster inference procedure.
2. Providing the solutions on using either differentiable or non-differentiable simulators.
summary: The paper considers the problem of solving inverse problems mostly in physical sciences. In particular, we are interested in posterior inference. This paper propose to use the flow matching perspective refined with additional control signals coming from a simulator. Moreover, the authors consider various scenarios depending on the differentiability of the simulator. Finally, they show the empirical results on a few simulator-based inference problems and the gravitational lensing inverse problem.
weaknesses: Despite its strengths, the paper has a few major and minor weaknesses.

**Major weaknesses:**

1. I’m really concerned about the presentation of the results and the abilities of this method, because of the lack of enough comparisons and metrics. 
2. Regarding the gravitational lensing problem (presented as the main real-world task being tackled), they are operating in a relatively small space, making the problem easy and solvable. However, they didn’t include any coverage test (e.g., TARP [1], or any other), so it’s hard to say if the posteriors are good. Moreover, as far as I understand, the evaluation test is only on their own simulations, the one that the model were being trained on. In particular, we don’t know if the model is robust to any OOD examples. Finally, the presented residuals (e.g., in Fig. 6) are looking bad.
3. Regarding the part: „(…) however, previous methods are usually restricted to point estimates, use simple variational distributions or Bayesian Neural Networks (Schuldt et al., 2021; Legin et al., 2021; 2023; Poh et al., 2022) that are not well suited to represent more complicated high-dimensional data distributions.” - Legin et al. 2021; 2023 use a likelihood-free inference (or simulation-based inference) framework to get posteriors from simple feed-forward nn (not Bayesian).
4. The authors proposed the intensive tests only on LV, which is relatively simple problem and for sure not enough for fair comparison. Moreover, the results in Tab. 1 don’t show any superiority of the proposed method - in particular, NSF is getting similar or better results.


**Minor weaknesses:**

1. The authors should include comparisons with other novel posterior sampling baseline methods than DPS, e.g., LGD−MC [2].
2. In line 319, should be „spline” 


**References:**

[1] Lemos, P., Coogan, A., Hezaveh, Y., & Perreault-Levasseur, L. (2023, July). Sampling-based accuracy testing of posterior estimators for general inference. In International Conference on Machine Learning (pp. 19256-19273). PMLR.

[2] Song, J., Zhang, Q., Yin, H., Mardani, M., Liu, M. Y., Kautz, J., ... & Vahdat, A. (2023, July). Loss-guided diffusion models for plug-and-play controllable generation. In International Conference on Machine Learning (pp. 32483-32498). PMLR.

---

confidence: 2
contribution: 2
initial_rating: 3
presentation: 2
questions: In addition to the weaknesses above: 
* What are the challenges to scaling the approach to high-dimensional spaces?
* How sensitive is the method to the quality of the simulator? What happens when the simulator contains significant approximations or errors?
* Have you explored using multiple different types of control signals simultaneously? Could this provide complementary benefits?
* Could the method be extended to handle multiple observations simultaneously in a more efficient way?
soundness: 2
strengths: * The paper tackles an interesting problem of incorporating simulator feedback into generative models trained with flow matching. With the recent interest in the flow matching for learning generative models, the problem of incorporating downstream rewards, e.g. simulator feedback, is a critical one. This paper presents one of the first efforts in that direction (based on my knowledge). 
* The application to the lensing problem is quite interesting. Generative models are well suited to such scientific inverse problems, and this method could be a useful addition to the toolbox for such problems. 
* The method supports feedback from black-box simulators as well as differentiable simulators where gradient feedback is available. 
* The method also provides considerable speed-ups in the simulation over other approaches. 
* The paper is also quite clearly written and easy to follow.
summary: This paper introduces a method to improve flow-based generative models for simulation-based inference by incorporating simulator feedback through control signals. The key idea is to refine a pretrained flow network with additional control signals based on simulator outputs, which can include gradients and problem-specific cost functions for differentiable simulators or learned signals from non-differentiable simulators. The authors demonstrate their method on several benchmark problems and show substantial improvements in accuracy (53%) while maintaining fast inference times (up to 67x faster than MCMC) on a challenging strong gravitational lensing application.
weaknesses: * The empirical evaluation seems quite a bit limited. Specifically, the paper only considers a single applicaiton on strong gravitational lensing. The other synthetic tasks are quite small and it is unclear how general the method is. Moreover, the results on the gravitational lensing experiment, the results are not convicing. The residuals indicate the samples are not capturing the posterior. Finally, the authors consider a relatively simple variant of the problem (23D). In this setting FM + simulator feedback just acts as a faster approximation to MCMC. Where FM + Simulator feedback might have an advantage is high dimensional unstructured data (e.g. images in the case of lensing)  
* Another shortcoming of the empirical evaluation is the relatively limited baselines. There are several approaches for unbiased inference with diffusion priors (like DPS) [1-4], so it would be good to add comparisons to some of these baselines.
* There is some missing discussion about related work about the guidance of flow matching models [5-6]. 
* The code to reproduce experiments is not provided though there are sufficient details in the paper. 

[1] Wu, Z., Sun, Y., Chen, Y., Zhang, B., Yue, Y., & Bouman, K. L. (2024). Principled Probabilistic Imaging using Diffusion Models as Plug-and-Play Priors. arXiv preprint arXiv:2405.18782.

[2] Wu, L., Trippe, B., Naesseth, C., Blei, D., & Cunningham, J. P. (2023). Practical and asymptotically exact conditional sampling in diffusion models. Advances in Neural Information Processing Systems, 36.

[3] Dou, Z., & Song, Y. (2024). Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In The Twelfth International Conference on Learning Representations.

[4] Chung, H., Lee, S., & Ye, J. C. (2023). Decomposed Diffusion Sampler for Accelerating Large-Scale Inverse Problems. arXiv preprint arXiv:2303.05754.

[5] Nisonoff, H., Xiong, J., Allenspach, S., & Listgarten, J. (2024). Unlocking Guidance for Discrete State-Space Diffusion and Flow Models. arXiv preprint arXiv:2406.01572.

[6] Zheng, Q., Le, M., Shaul, N., Lipman, Y., Grover, A., & Chen, R. T. (2023). Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311.13443.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: - Is there something particular about the flow matching setup (e.g. linear trajectories) that make simulator control applicable here, in contrast to a more traditional method like NPE? Does the method rely on assuming optimal transport coupling paths (eqs. 5-6), or is it more generally applicable?
- If I understand correctly, the goal of simulator control is to produce flow matching trajectories that better model the joint $(\theta, x_0)$ space by giving an additional signal that the interpolated parameter-generated samples should produce simulations consistent with the original $x_0$. This reduces the effect of simulator stochasticity. Could this partly be recreated by generating multiple samples from the same parameter point partway through training? Similarly, could producing a batch of samples for comparison to the $x_0$ further improve the flow matching control signal?
- In the strong lensing section, why was the comparison primarily made to HMC, rather than e.g. NPE?
- How critical was the $t > 0.8$ empirical threshold choice for controls, and do the authors expect this to be problem-dependent or fairly universal?
soundness: 3
strengths: - Timely paper on an important problem in the field -- using information from the simulator to speed up simulation-based inference, in order to make it more sample-efficiency and well-calibrated.
- Discussion of several control mechanisms, which could be applicable depending on the specific simulator and domain.
- Good discussion of underlying theory, both the flow matching aspect as well as simulation-based inference. Sound comparison to surrounding context literature, e.g. relevance to classifier-free guidance.
- Application to an important and challenging problem in cosmology -- strong lens modeling, and comparison with a baseline HMC approach.
summary: The paper presents a method for including simulator guidance in a flow matching-based simulation-based inference setup. The basic idea as far as I understand is to model the posterior using generative flow matching, and during training (possibly only at a fine-tuning stage), draw samples from the simulator by interpolating parameter samples from the velocity trajectory, and include an additional signal ensuring consistency between those samples and the training datapoint. The overall motivation is to increase simulation efficiency when a stochastic simulator is used for simulation inference (as is usually the case).
weaknesses: - In the strong lensing section, while a comparison with HMC is done, a comparison with "traditional" neural simulation-based inference approaches like NPE is lacking. This comparison would significantly strengthen the outcomes of this experiment.
- The pretraining/finetuning tradeoff is not comprehensively explored -- as far as I understand, a big advantage of the method is that one could just finetune using a smaller number of simulation calls. While this is mentioned briefly and tested for specific cases, a more comprehensive study of how fraction of finetuning for a fixed simulator budget affects the outcome would make the results quite a bit stronger.
- The high-level presentation could be made slightly clearer -- e.g., in Fig. 1, marking that the it is the posterior $p(\theta\mid x)$ that is being modeled/targeted, to immediately situate the reader to the problem setting.

Meta-review: The reviewers are somewhat divided (3-3-3-6-6) about the paper, but they overall lean towards rejection. The paper introduces simulator feedback as an extension to flow matching for simulation-based inference. The approach is well-motivated, but the results are not convincing. The author-reviewer discussion has been constructive and has led to a number of clarifications and improvements, with the addition of a few new results. However, the reviewers still believe substantial improvements are needed, in particular regarding the presentation, the clarity, and the evaluation of the approach. For these reasons, I recommend rejection. I encourage the authors to address the reviewers' comments and to resubmit to a future conference.

Decision: Reject