lampe.inference.dcp¶
Differentiable Coverage Probability (DCP) components.
This module implements Differentiable Coverage Probability (DCP) regularizer for training Neural Ratio Estimator (NRE) and Neural Posterior Estimator (NPE), please refer to the documentation of the respective modules to learn about the methods.
DCP regularizer relies on the Expected Coverage Probability (ECP) of Highest Posterior Density Region (HPDR)
where \(\Theta_{\hat{p}(\theta|x)}(1-\alpha)\) is the HPDR of the approximate posterior model \(\hat{p}(\theta|x)\) at level \(1-\alpha\), and \(\mathbb{I}[\cdot]\) is an indicator function.
It can be shown that
where \(\alpha(\hat{p},\theta_j,x_j) := \int \hat{p}(\theta|x_j) \mathbb{I}[\hat{p}(\theta|x_j) < \hat{p}(\theta_j|x_j)]\textrm{d}\theta\).
A conservative (calibrated) posterior model is one that has \(\text{ECP}(1 - \alpha)\) above (at) level \(1-\alpha\) for all \(\alpha \in (0,1)\). DCP regularizer turns this condition into a differentiable optimization objective minimized in parallel to the main optimization objective of NRE or NPE.
References
Installation¶
This module relies on differentiable sorting package
torchsort which is not a default
dependency of lampe. torchsort supports both CPU and GPU (CUDA) computation.
In order to install the package with CUDA support, you need to have C++ and CUDA
compiler installed. For more information please refer to the installation section
in the official repository.
Please keep in mind, that compiler’s CUDA version has to match that of PyTorch!
For certain Python + CUDA + pytorch versions combinations there are pre-build wheels available.
Below is an example of how to install torchsort in a conda environment with
Python 3.11 + CUDA 12.1.1 + PyTorch 2.0.0:
$ conda create -y --name lampe python=3.11
$ conda activate lampe
$ unset -v CUDA_PATH
$ conda install -y -c "nvidia/label/cuda-12.1.1" cuda
$ conda install -y -c conda-forge gxx_linux-64=11.4.0
$ conda install -y pip
$ pip install --upgrade pip
$ pip install torch==2.0.0 --index-url https://download.pytorch.org/whl/cu121
$ TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing;Ampere" pip install --no-cache-dir torchsort
Classes¶
Creates a module that calculates cross-entropy loss and the regularization loss for an NRE network. |
|
Creates a module that calculates the negative log-likelihood loss and the regularization loss for a NPE normalizing flow. |
Descriptions¶
- class lampe.inference.dcp.DCPNRELoss(estimator, prior, proposal=None, lmbda=5.0, n_samples=16, calibration=False, sort_kwargs=None)¶
Creates a module that calculates cross-entropy loss and the regularization loss for an 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 \frac{1}{M} \sum_{j=1}^M (\text{ECP}(1 - \alpha_j) - (1 - \alpha_j))^2\end{split}\]where \(\ell(p) = -\log p\) is the negative log-likelihood and \(\text{ECP}(1 - \alpha_j)\) is the Expected Coverage Probability at credibility level \(1 - \alpha_j\).
References
Calibrating Neural Simulation-Based Inference with Differentiable Coverage Probability (Falkiewicz et al., 2023)- Parameters:
estimator (Module) – A log-ratio network \(\log r_\phi(\theta, x)\).
prior (Distribution) – Prior distribution \(p(\theta)\).
proposal (Distribution) – If given, proposal distribution module for Importance Sampling \(I(\theta)\), otherwise prior is used for IS.
lmbda (float) – The weight \(\lambda\) controlling the strength of the regularizer.
n_samples (int) – Number of samples in MC estimate of rank statistic
calibration (bool) – Boolean flag of calibration objective (default: False)
sort_kwargs (dict) – Arguments of differentiable sorting algorithm, see doc (default: None)
- class lampe.inference.dcp.DCPNPELoss(estimator, proposal=None, lmbda=5.0, n_samples=16, calibration=False, sort_kwargs=None)¶
Creates a module that calculates the negative log-likelihood loss and the regularization loss for a NPE normalizing flow.
Given a batch of \(N \geq 2\) pairs \((\theta_i, x_i)\), the module returns
\[\begin{split}l & = \frac{1}{N} \sum_{i = 1}^N -\log p_\phi(\theta_i | x_i) \\ & + \lambda \frac{1}{M} \sum_{j=1}^M (\text{ECP}(1 - \alpha_j) - (1 - \alpha_j))^2\end{split}\]where \(\ell(p) = -\log p\) is the negative log-likelihood and \(\text{ECP}(1 - \alpha_j)\) is the Expected Coverage Probability at credibility level \(1 - \alpha_j\).
References
Calibrating Neural Simulation-Based Inference with Differentiable Coverage Probability (Falkiewicz et al., 2023)- Parameters:
estimator (Module) – A normalizing flow \(p_\phi(\theta | x)\).
proposal (Distribution) – If given, proposal distribution module for Importance Sampling \(I(\theta)\), otherwise regularizer is estimated by directly sampling from the model (default: None).
lmbda (float) – The weight \(\lambda\) controlling the strength of the regularizer.
n_samples (int) – Number of samples in MC estimate of rank statistic
calibration (bool) – Boolean flag of calibration objective (default: False)
sort_kwargs (dict) –
Arguments of differentiable sorting algorithm, see doc (default: None)
- rsample_and_log_prob(x, shape=())¶