lampe.inference.bnre¶

Balanced neural ratio estimation (BNRE) components.

References

Towards Reliable Simulation-Based Inference with Balanced Neural Ratio Estimation (Delaunoy et al., 2022)

https://arxiv.org/abs/2208.13624

Classes¶

BNRELoss

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

Descriptions¶

class lampe.inference.bnre.BNRELoss(estimator, lmbda=100.0)¶

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

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

\[\begin{split}l & = \frac{1}{2N} \sum_{i = 1}^N \ell(d_\phi(\theta_i, x_i)) + \ell(1 - d_\phi(\theta_{i+1}, x_i)) \\ & + \lambda \left(1 - \frac{1}{N} \sum_{i = 1}^N d_\phi(\theta_i, x_i) + d_\phi(\theta_{i+1}, x_i) \right)^2\end{split}\]

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

Parameters:

estimator (Module) – A log-ratio network \(\log r_\phi(\theta, x)\).
lmbda (float) – The weight \(\lambda\) controlling the strength of the balancing condition.

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