CNF Layer Functions

CNFLayer

Abstract interface for continuous normalizing flow layers implemented by DiffEqFlux.

Interface

Concrete subtypes are Lux wrapper layers and are callable as:

(logpx, λ₁, λ₂), new_state = layer(x, ps, st)

where x is a batch of samples, ps are Lux parameters, and st is Lux state. The first returned value contains the log density estimate and regularization terms used by continuous normalizing flows.

Rules

• The wrapped Lux model must be stored in a field named model.
• The state must contain regularize and monte_carlo flags.
• Solver keyword arguments supplied to the constructor are forwarded to solve.

Implementations

FFJORD is the public implementation of this interface.

FFJORD(model, tspan, input_dims, args...; ad = nothing, basedist = nothing, kwargs...)

Constructs a continuous-time recurrent neural network, also known as a neural ordinary differential equation (neural ODE), with fast gradient calculation via adjoints [1] and specialized for density estimation based on continuous normalizing flows (CNF) [2] with a stochastic approach [2] for the computation of the trace of the dynamics' jacobian. At a high level this corresponds to the following steps:

1. Parameterize the variable of interest x(t) as a function f(z, θ, t) of a base variable z(t) with known density p_z.
2. Use the transformation of variables formula to predict the density p_x as a function of the density p_z and the trace of the Jacobian of f.
3. Choose the parameter θ to minimize a loss function of p_x (usually the negative likelihood of the data).

After these steps one may use the NN model and the learned θ to predict the density p_x for new values of x.

Arguments

• model: A Flux.Chain or Lux.AbstractLuxLayer neural network that defines the dynamics of the model.
• basedist: Distribution of the base variable. Set to the unit normal by default.
• input_dims: Input Dimensions of the model.
• tspan: The timespan to be solved on.
• args: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.
• ad: The automatic differentiation method to use for the internal jacobian trace. Defaults to AutoForwardDiff() if full jacobian needs to be computed, i.e. monte_carlo = false. Else we use AutoZygote().
• kwargs: Additional arguments splatted to the ODE solver. See the Common Solver Arguments documentation for more details.

Fields

• model: Lux layer used for the CNF dynamics.
• basedist: Optional base distribution. If nothing, a standard normal base density is used.
• ad: ADTypes.jl automatic differentiation backend for the Jacobian trace estimate.
• input_dims: Dimensions of one sample, excluding the batch dimension.
• tspan: Integration time span.
• args: Positional solver arguments, usually including the ODE algorithm.
• kwargs: Keyword solver arguments forwarded to solve.

Returns

A CNFLayer. Calling the layer as ffjord(x, ps, st) returns ((logpx, λ₁, λ₂), new_state).

Examples

using DiffEqFlux, Lux, Random

rng = Random.default_rng()
model = Lux.Chain(Lux.Dense(2 => 8, tanh), Lux.Dense(8 => 2))
ffjord = FFJORD(model, (0.0f0, 1.0f0), (2,))
ps, st = Lux.setup(rng, ffjord)
(logpx, λ₁, λ₂), st = ffjord(rand(Float32, 2, 4), ps, st)

References

[1] Pontryagin, Lev Semenovich. Mathematical theory of optimal processes. CRC press, 1987.

[2] Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. "Neural ordinary differential equations." In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 6572-6583. 2018.

[3] Grathwohl, Will, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. "Ffjord: Free-form continuous dynamics for scalable reversible generative models." arXiv preprint arXiv:1810.01367 (2018).

FFJORDDistribution(model, ps, st)

Wrap an FFJORD layer as a Distributions.jlContinuousMultivariateDistribution.

Arguments

• model: The FFJORD layer used for density evaluation and sampling.
• ps: Lux parameters for model.
• st: Lux state for model, including regularize and monte_carlo.

Fields

• model: Stored FFJORD layer.
• ps: Stored Lux parameters.
• st: Stored Lux state.

Returns

A distribution that supports length, eltype, logpdf, pdf, and rand.

Examples

using DiffEqFlux, Distributions, Lux, Random

rng = Random.default_rng()
model = Lux.Chain(Lux.Dense(2 => 8, tanh), Lux.Dense(8 => 2))
ffjord = FFJORD(model, (0.0f0, 1.0f0), (2,))
ps, st = Lux.setup(rng, ffjord)
d = FFJORDDistribution(ffjord, ps, st)
logp = logpdf(d, rand(Float32, 2))