## Example: Using Planar Flow

Here we provide a minimal demonstration of learning a synthetic 2d banana distribution using planar flows (Renzende et al. 2015) by maximizing the Evidence Lower Bound (ELBO). To complete this task, the two key inputs are:

- the log-density function of the target distribution,
- the planar flow.

#### The Target Distribution

The `Banana` object is defined in `example/targets/banana.jl`, see the source code for details.

    p = Banana(2, 1.0f-1, 100.0f0)
    logp = Base.Fix1(logpdf, p)

Visualize the contour of the log-density and the sample scatters of the target distribution:

#### The Planar Flow

The planar flow is defined by repeated applying a sequence of invertible transformations to a base distribution $q_0$. The building blocks for a planar flow of length $N$ are the following invertible transformations, called planar layers:

\[\text{planar layers}: T_{n, \theta_n}(x)=x+u_n \cdot \tanh \left(w_n^T x+b_n\right), \quad n=1, \ldots, N, \]

where $\theta_n = (u_n, w_n, b_n), n=1, \dots, N$ are the parameters to be learned. Thankfully, `Bijectors.jl` provides a nice framework to define a normalizing flow. Here we used the `PlanarLayer()` from `Bijectors.jl` to construct a 20-layer planar flow, of which the base distribution is a 2d standard Gaussian distribution.

    using Bijectors, FunctionChains

    function create_planar_flow(n_layers::Int, q₀)
        d = length(q₀)
        Ls = [f32(PlanarLayer(d)) for _ in 1:n_layers]
        ts = fchain(Ls)
        return transformed(q₀, ts)
    end

    # create a 20-layer planar flow
    flow = create_planar_flow(20, MvNormal(zeros(Float32, 2), I))
    flow_untrained = deepcopy(flow) # keep a copy of the untrained flow for comparison

Notice that here the flow layers are chained together using `fchain` function from `FunctionChains.jl`. Alternatively, one can do

    ts = reduce(∘, [f32(PlanarLayer(d)) for i in 1:20])

However, we recommend using `fchain` to reduce the compilation time when the number of layers is large. See this comment for how the compilation time might be a concern.

#### Flow Training

Then we can train the flow by maximizing the ELBO using the `train_flow` function as follows:

    using NormalizingFlows
    using ADTypes
    using Optimisers

    sample_per_iter = 10
    # callback function to track the number of samples used per iteration
    cb(iter, opt_stats, re, θ) = (sample_per_iter=sample_per_iter,)
    # defined stopping criteria when the gradient norm is less than 1e-3
    checkconv(iter, stat, re, θ, st) = stat.gradient_norm < 1e-3
    flow_trained, stats, _ = train_flow(
        elbo,
        flow,
        logp,
        sample_per_iter;
        max_iters=200_00,
        optimiser=Optimisers.ADAM(),
        callback=cb,
        hasconverged=checkconv,
        ADbackend=AutoZygote(), # using Zygote as the AD backend
    )

Examine the loss values during training:

    using Plots

    losses = map(x -> x.loss, stats)
    plot(losses; xlabel = "#iteration", ylabel= "negative ELBO", label="", linewidth=2)

## Evaluating Trained Flow

Finally, we can evaluate the trained flow by sampling from it and compare it with the target distribution. Since the flow is defined as a `Bijectors.TransformedDistribution`, one can easily sample from it using `rand` function, or examine the density using `logpdf` function. See documentation of `Bijectors.jl` for details.

    using Random, Distributions

    nsample = 1000
    samples_trained = rand(flow_trained, n_samples) # 1000 iid samples from the trained flow
    samples_untrained = rand(flow_untrained, n_samples) # 1000 iid samples from the untrained flow
    samples_true = rand(p, n_samples) # 1000 iid samples from the target

    # plot
    scatter(samples_true[1, :], samples_true[2, :]; label="True Distribution", color=:blue, markersize=2, alpha=0.5)
    scatter!(samples_untrained[1, :], samples_untrained[2, :]; label="Untrained Flow", color=:red, markersize=2, alpha=0.5)
    scatter!(samples_trained[1, :], samples_trained[2, :]; label="Trained Flow", color=:green, markersize=2, alpha=0.5)
    plot!(title = "Comparison of Trained and Untrained Flow", xlabel = "X", ylabel= "Y", legend=:topleft)

## Reference

- Rezende, D. and Mohamed, S., 2015. Variational inference with normalizing flows. International Conference on Machine Learning

Theme

This document was generated with Documenter.jl version 0.27.25 on Saturday 19 August 2023. Using Julia version 1.9.2.