Loading editor...

Train, infer, and sample from a normalizing flow: all in your browser

Imagine a random variable with probability density . A small region is mapped by a function to . This mapping does not change the total probability of this region. In other words:

A normalizing flow turns an unknown data distribution into a known one by learning an invertible map .Thus, log likelihood of an unknown distribution can be maximized by maximizing the log likelihood of a known distribution plus the log slope of the mapping function.

In 1D the cleanest such map is a cumulative distribution function CDF: it is monotonic (hence invertible) and a common target distribution to map to is uniform distribution on .

Model

In this case the CDF is generated from a mixture of three Gaussians. The parameters (to be optimized via gradient descent) are the three Gaussian components: means, standard deviations and weights.


Data

The data consists of 25 samples, clustered in two groups.

  • The data points .
  • Then we that maps these input points to points .

The flow should spread out the data points uniformly along the y-axis. Where the data is sparse, the curve should be flat, so points cluster together along the y-axis. Where the data is dense, the curve should be steep, so points spread out along the y-axis. That spreading is what makes the image uniform.


Training

  • Number of training steps is set to and learning rate is set to .
  • the loss. This might take a few seconds.

Once the gm.minimize command finishes running, the points on the y-axis should spread out a little bit to resemble a uniform distribution. To train properly, one will need to keep reducing the learning rate every few steps.


Visualizing the trained mapping

Trained for long enough, gradient descent settles on parameters close to these:

  • Means: -0.5404, -0.4750, -1.0243
  • Sigmas = 0.4519, 0.0273, 0.0498.
  • Weights = 3.79e-05, 0.71, 0.28

through this trained CDF and place the mapped points on the y-axis. They sit more evenly than the current flow's images - in other words, they are closer to uniform.


Sampling

A trained flow is also a generator. Since is invertible, drawing a sample is easy: pick a point from the known base distribution on the y-axis, then push it backward through onto the x-axis, into data space.

  • from the base distribution (uniform in our case).

Now invert each sample through the optimal mixture CDF.

Unlike the forward map, the mixture CDF has no closed-form inverse and is solved numerically giving these precomputed values:

  • : the x-position of each generated sample.

  • are our generated data. They have been drawn below the y-axis so that we can compare them with the actual data.
  • actual data.

As you can see, the and the are quite similar and come from the same distribution.