ipl-uv/rbig_jax

Iterative and Parametric Gaussianization with JAX.

rbig_jax's Issues

[Algorithm][Parametric] Emerging Convolutions

Probably the only implementation that allows for invertible convolutions which spatial awareness. It is related to the autoregressive method so it will be fast for density estimates but slow for sampling.

---

Resources

• The Convolution Exponential and Generalized Sylvester Flows - Hoogeboom et al (2020)
  • akandykeller/SelfNormalizingFlows
  • ehoogeboom/convolutional_exponential_and_sylvester

  They create the exponential convolution transform. It basically uses the exponential of the convolutions. It approximates the transform via a Taylor series. It generalizes most of the householder approaches, i.e. Sylvester Flows.

Documentation.

Need a dedicated readthedocs documentation.

Toy Data for the RBIG Method

Need to do some toy experiments on the classic datasets with the Scikit-learn cluster data.

Datasets

• Classic RBIG
• Circles
• Moons
• S-Curve
• Blobs
• Swiss Roll

ToDo

• Scripts
• Loss Metrics (nll, info reduction)
• Plots (latent, original)
• Sampling
• Density Estimation
• Wandb integration (example)

[Iterative] Mutual Information Wrapper

Need a wrapper for an implementation for mutual information.

Notebooks:

• From scratch
• Wrapper function

[Algorithm][Parametric] Conditional Gaussianization

Implement conditional normalizing flows. It requires having a context argument that needs to be carried throughout all of the transformations until the base distribution.

---

Resources

• Learning Likelihoods with Conditional Normalizing Flows - Winkler et al (2019)
• FlowTorch - PyTorch Implementation
• Bayesians/nflows - Demo

[Investigate] "Squeezing Layer" for Marginal Histogram Transformation.

We still have issues with the boundaries for the marginal histogram transformation. When we do the sampling, that's when we see it the most.

Proposal

It would be interesting to see what would happen if we do a "squeezing layer" to constrain the domain space. So instead of the histogram transformation on the input domain (-inf, inf) -> [0,1], we could implement an invertible "squeezing" function to the domain, (-inf, inf) -> [0, 1], and then do the histogram transformation, [0,1]->[0,1]. At the very least, we don't have to worry about the bounds for the histogram function.

[Algorithm][Parametric] Spline Transformation

Implement the parametric spline function as an element-wise transformation, i.e. marginal Gaussianization.

---

Resources

• ChrisWaites/jax-flows - jax implemention
• Information-Fusion-Lab-Umass/NuX - jax implementation
• jfcrenshaw/pzflow - jax implementation
• bayesians/nflows - PyTorch implementation

[Algorithm][Parametric] Implement the GDN layer

A basic implementation of the GDN algorithm. It is mainly used in compression. It features a normalization which can be coupled with a convolution/linear layer.

y[i] = x[i] / sqrt(beta[i] + sum_j(gamma[j, i] * x[j]))

---

Resources

• GDN PyTorch Implementation
• Paper - Density Modeling of Images using a Generalized Normalization Transformation - Balle et al, (2016)
• Paper - Self Normalizing Flows - Keller et al (2020)

---

Potential Issues

Somehow we need to constrain a few parameters, e.g., gamma and beta. Probably a simple np.max would work.

[Iterative] KL-Divergence Wrapper

Need an implementation and demonstration of the KL-Divergence Wrapper.

Notebooks

• From Scratch
• Convenience wrapper

[Algorithm][Parametric] SVD Linear Layer

We currently have the householder transformation but we can use that for the SVD layer, X=USV^T. The rotation matrices (U,V^T) are constrained to be orthogonal and the diagonal, S, is unconstrained. It's a bit more expensive but it might help for training.

---

Resources

• bayesians/nflows

ITM Notebook

A notebook showcasing information theoretic metrics:

• Probability & Information
• Total Correlation
• Entropy
• Mutual Information
• KL-Divergence

Note: we really want to focus on the speed and simplicity. E.g. We can show how one can do it from scratch as well as the convenient wrapper functions to explain the design decisions.

[Demo] Info Loss for GaussFlow versus RBIG

It would be nice to see the differences between the information loss/reduction for the GaussFlow algorithm and the IterGauss algorithm.

Outcome: A demo notebook showcasing the info loss between layers for both algorithms.

[Iterative] Improve Total Correlation Wrapper

Need an implementation and demonstration of the Total correlation.

Notebooks

• From Scratch
• Convenience wrapper

[Parametric] Remove Dependency on objax

objax is basically a PyTorch-like version of Jax. But it is a bit limiting when trying to mix and match reviews. So I think we should remove the dependency on objax and stick with pure Jax.

---

Example

This example was taken from jax-flows.

Demo Snippet

```python def FixedInvertibleLinear(): """An implementation of an invertible linear layer from `Glow: Generative Flow with Invertible 1x1 Convolutions` (https://arxiv.org/abs/1605.08803). Returns: An ``init_fun`` mapping ``(rng, input_dim)`` to a ``(params, direct_fun, inverse_fun)`` triplet. """

def init_fun(rng, input_dim, **kwargs):
    W = orthogonal()(rng, (input_dim, input_dim))
    W_inv = linalg.inv(W)
    W_log_det = np.linalg.slogdet(W)[-1]

    def direct_fun(params, inputs, **kwargs):
        outputs = inputs @ W
        log_det_jacobian = np.full(inputs.shape[:1], W_log_det)
        return outputs, log_det_jacobian

    def inverse_fun(params, inputs, **kwargs):
        outputs = inputs @ W_inv
        log_det_jacobian = np.full(inputs.shape[:1], -W_log_det)
        return outputs, log_det_jacobian

    return (), direct_fun, inverse_fun

return init_fun

</details>

Toy Datasets for Gaussianization Flows method

Need to do some toy experiments on the classic datasets with the Scikit-learn cluster data.

Datasets

• Classic RBIG
• Circles
• Moons
• S-Curve
• Blobs
• Swiss Roll

ToDo

• Scripts
• Loss Metrics (nll, info reduction)
• Plots (latent, original)
• Sampling
• Density Estimation
• Wandb integration (example)

Deep Dive Demo Notebook

I have notebooks documenting the progression to building and RBIG model but there is no all-in-one notebook showcasing

Examples components:

• Marginal Gaussianization
• Univariate Entropy
• Information Loss
• ITMs -> TC, H

[Iterative] No convergence with with 10 layers as maxlayers.

Tried running RBIG on CIFAR data. It does not converge after a long time. Tried with 10, 20, 30, etc. Only after 50 layers does it converge.

---

Additional Information

• number of samples (50K)
• number of features (3K)
• image data (CIFAR)

Update the bisection search with max layer conditional

There is not guarantee that the bisection search will converge. We should also have a max layers stopping criteria as well.

---

Resources

• NuX

Standard Github Actions

• Markdown for Readmes
• black
• isort
• CI

[Algorithm][Iterative] Random Orthogonal Initialization

This can be one of the layers done for the iterative method. It can be done by splitting a random key with every layer and intializing the rotation matrix with every layer.

---

Source:

• Jax Documentation
• Objax Documentation