probabilists / lampe Goto Github PK

Description

Implementation of Balanced Neural Ratio Estimation (BNRE) as introduced in "Towards Reliable Simulation-Based Inference with Balanced Neural Ratio Estimation" (Delaunoy et al., 2022)

Implementation

An implementation of the BNRE loss similar to the NRELoss

Alternatives

None

The documentation of the build constructor argument in NRE/NPE/NSE specifies

build: Callable[[int, int], nn.Module] = MLP

and

build: The network constructor

In my opinion, this does not really help in understanding how this constructor argument should be used, as there is no mention of what the ints are referring to. It is also not clear what the nn.Module should expect as inputs and produce as outputs.

In the tutorial about embeddings, it could also be helpful to explain how the network produced with build is different from the network used to build an embedding.

Description

Given the recent popularity of score-based generative modeling, it would be great to provide a score-based inference algorithm within LAMPE.

Interface

A NSE (neural score estimation) class, a NSELoss module and a way to sample from the trained score estimator should be provided. Access to the log-density (through the probability flow ODE) is not mandatory, but would be convenient.

References

• Deep Unsupervised Learning using Nonequilibrium Thermodynamics (Sohl-Dickstein et al., 2015)
  https://arxiv.org/abs/1503.03585

• Generative Modeling by Estimating Gradients of the Data Distribution (Song et al., 2019)
  https://arxiv.org/abs/1907.05600

• Denoising Diffusion Probabilistic Models (Ho et al., 2020)
  https://arxiv.org/abs/2006.11239

• Score-Based Generative Modeling through Stochastic Differential Equations (Song et al., 2021)
  https://arxiv.org/abs/2011.13456

Description

The method described in Calibrating Neural Simulation-Based Inference with Differentiable Coverage Probability paper.

Implementation

New loss class in the inference module, similar to the NRELoss ->BNRELoss enhancement.

Two separate classes, one for NRE, and one for NPE.

Alternatives

I did not consider alternatives.

Description

The simulator tutorial references JointDataset, which does not exist.

Expectation

It should probably be updated to IterableJointDataset or JointLoader.

Description

When the domain is one-dimensional, lampe.utils.gridapply builds a grid of shape (bins,) instead of the expected (bins, 1).

Reproduce

In the following error, mat1 is x and should be of shape (128, 1).

>>> import torch
>>> import lampe
>>> A = torch.randn(1, 3)
>>> f = lambda x: x @ A
>>> domain = torch.zeros(1), torch.ones(1)
>>> lampe.utils.gridapply(f, domain)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/francois/Documents/Git/lampe/lampe/utils.py", line 104, in gridapply
    y = [f(x) for x in grid.split(batch_size)]
  File "/home/francois/Documents/Git/lampe/lampe/utils.py", line 104, in <listcomp>
    y = [f(x) for x in grid.split(batch_size)]
  File "<stdin>", line 1, in <lambda>
RuntimeError: mat1 and mat2 shapes cannot be multiplied (1x128 and 1x3)

Expected behavior

The grid should be of shape (bins, 1).

Causes and solution

The lampe.utils.gridapply function uses torch.cartesian_prod, which behaves inconsistently when given a single argument. A reshape should be enough to fix the issue.

Environment

• LAMPE version: 0.6.1
• PyTorch version: 1.12.0
• Python version: 3.9.15
• OS: Ubuntu 22.10