Neural Ratio Estimation (NRE)

Introduction

As we have seen, the output of prior + simulator is the array of pairs $({\mathit{x}}_i,{\mathit{\theta }}_i)$ is drawn from the joint distribution

$$({\mathit{x}}_i,{\mathit{\theta }}_i)\sim p(\mathit{x},\mathit{\theta })=p(\mathit{x}|\mathit{\theta })p(\mathit{\theta })$$

We now consider the shuffled pairs $({\mathit{x}}_i,{\mathit{\theta }}_j)$, where ${\mathit{x}}_i$ is the output of the forward-modeled input ${\mathit{\theta }}_i,i\ne j$. These pairs are sampled from the product distribution

$$({\mathit{x}}_i,{\mathit{\theta }}_j)\sim p(\mathit{x})p(\mathit{\theta })$$

The idea of NRE is to train a classifier to learn the ratio

$$r(\mathit{x},\mathit{\theta })\equiv \frac{p(\mathit{x},\mathit{\theta })}{p(\mathit{x})p(\mathit{\theta })}=\frac{p(\mathit{x}|\mathit{\theta })}{p(\mathit{x})},$$

which is equal to the likelihood-to-evidence ratio. The application of Bayes' theorem makes the connection between $r(\mathit{x},\mathit{\theta })$ and the Bayesian inverse problem:

$$r(\mathit{x},\mathit{\theta })=\frac{p(\mathit{x},\mathit{\theta })}{p(\mathit{x})}=\frac{p(\mathit{\theta }|\mathit{x})}{p(\mathit{\theta })}.$$

In other words, $r(\mathit{x},\mathit{\theta })$ equals the posterior-to-prior ratio. Therefore, one can get samples from the posterior distribution of $\mathit{\theta }$ from the approximate knowledge of $r(\mathit{x},\mathit{\theta })$ and prior samples from $\mathit{\theta }$.

More specifically, the binary classifier $d_{\mathit{\varphi }}(\mathit{x},\mathit{\theta })$ with learnable parameters $\mathit{\varphi }$ is trained to distinguish the $({\mathit{x}}_i,{\mathit{\theta }}_i)$ pairs sampled from the joint distribution from their shuffled counterparts. We label pairs with a variable $y$, such that $y=1$ refers to joint pairs, and $y=0$ to shuffled pairs. The classifier is trained to approximate

$$\begin{array}{rl}{d}_{\mathit{\varphi }}(\mathit{x},\mathit{\theta })& \approx p(y=1|\mathit{x},\mathit{\theta })\\ & =\frac{p(\mathit{x},\mathit{\theta }|y=1)p(y=1)}{p(\mathit{x},\mathit{\theta }|y=0)p(y=0)+p(\mathit{x},\mathit{\theta }|y=1)p(y=1)}\\ & =\frac{p(\mathit{x},\mathit{\theta })}{p(\mathit{x})\mathit{\theta }+p(\mathit{x},\mathit{\theta })}\\ & =\frac{r(\mathit{x},\mathit{\theta }\mathbf{)}}{1+r(\mathit{x},\mathit{\theta }\mathbf{)}},\end{array}$$

where we used $p(y=0)=p(y=1)=0.5$.

The classifier learns the parameters $\mathit{\varphi }$ by minimizing the binary-cross entropy, defined as

$$L(d_{\mathit{\varphi }})=-\int d\mathit{\theta }\int d\mathit{x}p(\mathit{x},\mathit{\theta })\log d_{\mathit{\varphi }}(\mathit{x},\mathit{\theta })-p(\mathit{x})\mathit{\theta }\log (1-d_{\mathit{\varphi }}(\mathit{x},\mathit{\theta }))$$

References

[1]: Hermans, Joeri, Volodimir Begy, and Gilles Louppe. "Likelihood-free mcmc with amortized approximate ratio estimators." International conference on machine learning. PMLR, 2020.

[2]: Miller, Benjamin K., et al. "Truncated marginal neural ratio estimation." Advances in Neural Information Processing Systems 34 (2021): 129-143.

[3]: Anau Montel, Noemi, James Alvey, and Christoph Weniger. "Scalable inference with autoregressive neural ratio estimation." Monthly Notices of the Royal Astronomical Society 530.4 (2024): 4107-4124.