Mutual information maximization for amortized likelihood inference from sampled trajectories: MINIMALIST

Simulation-based inference enables learning the parameters of a model even when its likelihood cannot be computed in practice. One class of methods uses data simulated with different parameters to infer models of the likelihood-to-evidence ratio, or equivalently the posterior function. Here we frame the inference task as an estimation of an energy function parametrized with an artificial neural network. We present an intuitive approach, named MINIMALIST, in which the optimal model of the likelihood-to-evidence ratio is found by maximizing the likelihood of simulated data. Within this framework, the connection between the task of simulation-based inference and mutual information maximization is clear, and we show how several known methods of posterior estimation relate to alternative lower bounds to mutual information. These distinct objective functions aim at the same optimal energy form and therefore can be directly benchmarked. We compare their accuracy in the inference of model parameters, focusing on four dynamical systems that encompass common challenges in time series analysis: dynamics driven by multiplicative noise, nonlinear interactions, chaotic behavior, and high-dimensional parameter space.

Figure 1

(a) Distribution of energies of independent and joint pairs $(x,\theta )$. Pairs from the joint distribution have lower energy $E$, while pairs from the independent distribution have higher energy, as the majority of these independent samples are relatively unlikely under a joint model. (b) Schematic of the method. We first sample parameters $\theta$ from the prior $P\left(\theta \right)$ and then sample observations $x$ from the simulator $P\left(x\right|\theta )$ to obtain pairs $(x_i,{\theta }_i)$. To generate pairs $(x_i,{\theta }_j)$ from the independent distribution, we shuffle the two initial vectors. Both sets are used to infer a model of energy $E^{\varphi }$ by maximizing a log-likelihood in (5) with $\varphi$ parameters of the artificial neural network used for the inference.

Figure 2

Example trajectories (a,c,e,g) and posterior inference (b,d,f,h) for Ornstein-Uhlenbeck (a,b), birth-death (c,d), SIR (e,f), and Lorenz attractor (g,h). Trajectories are simulated with the parameter marked in red on the posterior plots. Circles indicate the discrete observations used for inference. Posteriors were estimated over 10 trajectories.

Figure 3

Convergence of the posteriors for Ornstein-Uhlenbeck (a,b), birth-death (c,d), SIR (e,f), and Lorenz attractor (g). To evaluate the posteriors, we first choose a reference hypothesis ${\theta }^{*}$: for Ornstein-Uhlenbeck ${\mu }^{*}=5$ and ${\sigma }^{*}=1$, for birth-death ${\beta }^{*}=10$ and ${\alpha }^{*}=0.2$, for SIR ${\beta }^{*}=0.6$ and ${\gamma }^{*}=0.2$, and for Lorenz attractor ${\rho }^{*}=35$. We show posteriors $P^{\infty }\left(\theta \right|x_{1:M})$ calculated using models trained with $N={10}^7$ trajectories. For the Ornstein-Uhlenbeck process, the exact posterior density $P\left(\theta \right|x_{1:M})$ is also shown.

Figure 4

Benchmarking of the methods. We compare the three objectives $M$, Eq. (5) (MINE), $L_f$, Eq. (7) (FDIV), and $S$, Eq. (8) (BCE) for three different metrics. We perform 10 replicates of the inference and comparison for simulation budgets $N\in \{{10}^4,{10}^5,{10}^6\}$ for four systems: Ornstein-Uhlenbeck (a,b,c), birth-death (d,e,f), SIR (g,h,i), and Lorenz attractor (j,k,l). In the first row, we compare the mutual information on held out test data using Eq. (5) with the estimated $E^{\varphi }$. For the following two metrics, we need to instead choose a hypothesis ${\theta }^{*}$; see Fig. 3 for the exact values. In the second row, we compare the Jensen-Shannon divergence $D_{\mathrm{JS}}(P^{\infty }\left(\theta \right|x_{1:M}),{\widehat{P}}_l^N\left(\theta \right|x_{1:M}))$ between the reference and inferred posteriors. In the last row, we compare the Jensen-Shannon divergence $D_{\mathrm{JS}}(P\left(x\right|{\theta }^{*}),{\widehat{P}}_l^N\left(x\right|{\theta }^{*}))$ using sampled trajectories from the simulator $P\left(x\right|{\theta }^{*})$ and the inferred distribution $\widehat{P_l}\left(x\right|{\theta }^{*})$.

Figure 5

The Ornstein-Uhlenbeck process in $d\ge 1$. We compare the three objective functions $M$ (MINE), $L_f$ (FDIV), and $S$ (BCE) for dimension $d=1,2,3,4,5$. We perform 10 replicates of the inference with simulation budgets $N={10}^4,{10}^5,{10}^6$ for changing dimension $d=1,2,3,4,5$. We compare the posterior with the analytical prediction for a hypothesis ${\left(g^{*}\right)}_{ij}=-1$. We compute the Jensen-Shannon divergence $D_{\mathrm{JS}}(P\left(\theta \right|x_{1:M}),{\widehat{P}}_l^N\left(\theta \right|x_{1:M}))$ between the true and inferred posteriors for each element of the damping matrix $\gamma$ independently. We show the divergence for diagonal (a–e) and off-diagonal terms (f–i).

Figure 6

We compare the three objective functions $M$ (MINE), $L_f$ (FDIV), and $S$ (BCE) using the area under the receiver operating characteristic (AUROC) of a classifier trained to distinguish samples from the simulator $P\left(x\right|{\theta }^{*})$ and samples from the inferred estimators $\widehat{P}\left(x\right|{\theta }^{*})$ for a specific hypothesis ${\theta }^{*}$. AUROCs closer to $1/2$ mean that the model is performing well, as its samples are indistinguishable from samples from the true distribution. This metric gives a consistent picture with respect to the $D_{\text{JS}}$ metric presented in the main text (Fig. 4, c, f, i, and l).