Entropy estimation with non-invertible generative models

[austin-hoover]

This PR adds the class NNDist to the gpsr.beams module. With a non-invertible transformation from the base distribution (NNTransform), we can't evaluate the probability density $p(x)$ at an arbitrary point $x$. But we can still evaluate the density at a batch of sampled points $\{ x^{(i)} \} \sim p(x)$ by computing the Jacobian matrix of the transformation. This can provide differentiable estimates of the KL divergence or other integrals involving the distribution function.

Changes

• Added method compute_batched_jacobian(x: torch.Tensor, function: Callable) -> torch.Tensor: to gpsr.beams module. This method computes the Jacobian matrix of callable function for a batch of particles x, following the example here.
• Added class gpsr.beams.NNDist. This is a subclass of GenModel that uses the NNTransform as the transformation from the base distribution. It has methods sample and sample_and_log_prob but not log_prob.
• Renamed class gpsr.beams.NSF to gpsr.beams.NSFDist for consistency. Also changed arguments in NSFDist to match NNDist (width, depth).
• Added method entropy(n: int, prior: Any) -> torch.Tensor to the gpsr.beams.GenModel class. This is a convenience method to estimate the entropy or KL divergence from a batch of n samples, using the sample_and_log_prob method.
• Added examples which train the model by minimizing the KL divergence.

[roussel-ryan]

Really cool stuff @austin-hoover . Is there a reference for the batched jacobian method of estimating the entropy using this technique?

[roussel-ryan]

Can we move this to a unit test instead of having a python notebook?

[austin-hoover]

I added methods to test_beams.py here: https://github.com/austin-hoover/gpsr/blob/913c247477011841ba6d1bf3a47c162b9956a63a/gpsr/tests/test_beams.py#L40-L71. The first one tests that Jacobian of torch.sin is correct. The second one computes the Jacobian of the NNTransform forward method.

[roussel-ryan]

this is closer to a usable example, so we should rename to kl_nn.ipynb. A separate unit test would be useful for validating changes to beams.py

[austin-hoover]

Renamed kl_nn.ipynb

[austin-hoover]

For references: the formula for KL divergence is the same as for normalizing flows (Section 2.3.2 of https://arxiv.org/pdf/2404.00502). There's some discussion about computing Jacobian determinant for non-invertible models here.

[austin-hoover]

I guess I could say "vectorized" instead of "batched"?

[austin-hoover]

Updated KL minimization examples.

[austin-hoover]

Marking this as a draft for now.

• The KL minimization looks like it's working in 2D examples, but convergence is slower than the normalizing flow and sometimes fails to match the target distribution, depending on initialization. So far I can't reproduce correct MaxEnt solutions for 2D reconstruction.
• I also tried to train the NN using a sliced Wasserstein loss function to regress to a target distribution. In this case, the NN has an even harder time converging. So maybe the log-prob calculation is fine and these issues have to do with the specific NN architecture used.

[austin-hoover]

Added test under /gpsr/tests/test_beams.py that verifies the log_prob calculation is the same as the exact result from the normalizing flow.

[roussel-ryan]

Thanks @austin-hoover. I'm looking into the testing errors now since its effecting the main branch