lampe.inference.nre

Neural ratio estimation (NRE) components.

The principle of neural ratio estimation is to train a classifier network \(d_\phi(\theta, x)\) to discriminate between pairs \((\theta, x)\) equally sampled from the joint distribution \(p(\theta, x)\) and the product of the marginals \(p(\theta)p(x)\). Formally, the optimization problem is

\[\arg\min_\phi \frac{1}{2} \mathbb{E}_{p(\theta, x)} \big[ \ell(d_\phi(\theta, x)) \big] + \frac{1}{2} \mathbb{E}_{p(\theta)p(x)} \big[ \ell(1 - d_\phi(\theta, x)) \big]\]

where \(\ell(p) = -\log p\) is the negative log-likelihood. For this task, the decision function modeling the Bayes optimal classifier is

\[d(\theta, x) = \frac{p(\theta, x)}{p(\theta, x) + p(\theta) p(x)}\]

thereby defining the likelihood-to-evidence (LTE) ratio

\[r(\theta, x) = \frac{d(\theta, x)}{1 - d(\theta, x)} = \frac{p(\theta, x)}{p(\theta) p(x)} = \frac{p(x | \theta)}{p(x)} = \frac{p(\theta | x)}{p(\theta)} .\]

To prevent numerical stability issues when \(d_\phi(\theta, x) \to 0\), the neural network returns the logit of the class prediction \(\text{logit}(d_\phi(\theta, x)) = \log r_\phi(\theta, x)\).

References

Approximating Likelihood Ratios with Calibrated Discriminative Classifiers (Cranmer et al., 2015)

https://arxiv.org/abs/1506.02169

Likelihood-free MCMC with Amortized Approximate Ratio Estimators (Hermans et al., 2019)

https://arxiv.org/abs/1903.04057

Classes

NRE	Creates a neural ratio estimation (NRE) network.
NRELoss	Creates a module that calculates the cross-entropy loss for a NRE network.

Descriptions

class lampe.inference.nre.NRE(theta_dim, x_dim, build=<class 'zuko.nn.MLP'>, **kwargs)

Creates a neural ratio estimation (NRE) network.

Parameters:

• theta_dim (int) – The dimensionality \(D\) of the parameter space.

• x_dim (int) – The dimensionality \(L\) of the observation space.

• build (Callable[[int, int], Module]) – The network constructor (e.g. lampe.nn.ResMLP). It takes the number of input and output features as positional arguments.

• kwargs – Keyword arguments passed to the constructor.

forward(theta, x)

Parameters:

• theta (Tensor) – The parameters \(\theta\), with shape \((*, D)\).

• x (Tensor) – The observation \(x\), with shape \((*, L)\).

Returns:

The log-ratio \(\log r_\phi(\theta, x)\), with shape \((*,)\).

Return type:

Tensor

class lampe.inference.nre.NRELoss(estimator)

Creates a module that calculates the cross-entropy loss for a NRE network.

Given a batch of \(N \geq 2\) pairs \((\theta_i, x_i)\), the module returns

\[l = \frac{1}{2N} \sum_{i = 1}^N \ell(d_\phi(\theta_i, x_i)) + \ell(1 - d_\phi(\theta_{i+1}, x_i))\]

where \(\ell(p) = -\log p\) is the negative log-likelihood.

Parameters:

estimator (Module) – A log-ratio network \(\log r_\phi(\theta, x)\).

forward(theta, x)

Parameters:

• theta (Tensor) – The parameters \(\theta\), with shape \((N, D)\).

• x (Tensor) – The observation \(x\), with shape \((N, L)\).

Returns:

The scalar loss \(l\).

Return type:

Tensor