lampe.inference.npe

Neural posterior estimation (NPE) components.

The principle of neural posterior estimation is to train a parametric conditional distribution \(p_\phi(\theta | x)\) to approximate the posterior distribution \(p(\theta | x)\). The optimization problem is to minimize the expected Kullback-Leibler (KL) divergence between the two distributions for all observations \(x \sim p(x)\), that is,

\[\begin{split}\arg\min_\phi & ~ \mathbb{E}_{p(x)} \Big[ \text{KL} \big( p(\theta|x) \parallel p_\phi(\theta | x) \big) \Big] \\ = \arg\min_\phi & ~ \mathbb{E}_{p(x)} \, \mathbb{E}_{p(\theta | x)} \left[ \log \frac{p(\theta | x)}{p_\phi(\theta | x)} \right] \\ = \arg\min_\phi & ~ \mathbb{E}_{p(\theta, x)} \big[ -\log p_\phi(\theta | x) \big] .\end{split}\]

Normalizing flows are typically used for \(p_\phi(\theta | x)\) as they are differentiable parametric distributions enabling gradient-based optimization techniques.

Wikipedia

https://wikipedia.org/wiki/Kullback-Leibler_divergence

Classes

NPE

Creates a neural posterior estimation (NPE) normalizing flow.

NPELoss

Creates a module that calculates the negative log-likelihood (NLL) loss for a NPE density estimator.

Descriptions

class lampe.inference.npe.NPE(theta_dim, x_dim, build=<class 'zuko.flows.autoregressive.MAF'>, **kwargs)

Creates a neural posterior estimation (NPE) normalizing flow.

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], Flow]) – The flow constructor (e.g. zuko.flows.spline.NSF). It takes the number of sample and context 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-density \(\log p_\phi(\theta | x)\), with shape \((*,)\).

Return type:

Tensor

class lampe.inference.npe.NPELoss(estimator)

Creates a module that calculates the negative log-likelihood (NLL) loss for a NPE density estimator.

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

\[l = \frac{1}{N} \sum_{i = 1}^N -\log p_\phi(\theta_i | x_i) .\]
Parameters:

estimator (Module) – A log-density estimator \(\log p_\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