Continuous normalizing flow — the generative half of Chen et al., 2018.
encode integrates the augmented system forward and returns the
latent together with the flow's log-determinant; decode integrates
the same vector field with time reversed. That symmetry is the point
of the continuous formulation: a discrete flow has to be built
invertible, whereas this one is invertible because an ODE is, and the
inverse is the solver run backwards.
Parameters
configNeuralODEConfigAttributes
dynamics_CNFDynamicsNotes
Reference: Chen, Rubanova, Bettencourt, and Duvenaud, "Neural Ordinary Differential Equations", NeurIPS, 2018 (arXiv:1806.07366), §4. The Hutchinson trace option follows Grathwohl et al., "FFJORD", ICLR, 2019 (arXiv:1810.01367).
Where to run it. A solve is one small right-hand side evaluated
tens of times, so which stream wins depends entirely on how much work
each evaluation carries. Measured on an M1 Pro, one training step
(forward, loss, backward) of neural_ode_gen:
================== =========== ===========
setup "cpu" "metal"
================== =========== ===========
B=8, D=192 16.8 ms 38.1 ms
B=32, D=3072 175.0 ms 56.5 ms
B=128, D=3072 587.5 ms 70.8 ms
================== =========== ===========
Below a few hundred dimensions the tensors are too small to pay for a
GPU dispatch and Accelerate wins; at anything resembling a real
training batch the ratio reverses and keeps widening — 8.3x at the
bottom row, where the CPU cost grows with the batch and the GPU's
barely does. Train on "metal"; keep the paper's two-dimensional
density experiments on "cpu".
Examples
>>> import lucid
>>> from lucid.models.generative.neural_ode import (
... NeuralODEConfig, NeuralODEModel,
... )
>>> cfg = NeuralODEConfig(sample_size=2, in_channels=1, out_channels=1,
... hidden_dim=8, num_blocks=1)
>>> model = NeuralODEModel(cfg).eval()
>>> z, log_det = model.encode(lucid.rand((2, 1, 2, 2)))
>>> z.shape, log_det.shape
((2, 4), (2,))
>>> model.decode(z).shape
(2, 1, 2, 2)Used by 2
Constructors
1Properties
4int: Flattened dimensionality D the flow acts on.
int: Vector-field evaluations spent by the most recent solve.
A continuous flow has no layer count to report, so this is the closest thing to one — and it grows as the field stiffens during training, which the paper documents.
This is the forward count, snapshotted when the solve returns.
Reading the live counter instead let the adjoint's backward sweep
keep incrementing it after the fact, so the number changed depending
on whether .backward() had run yet — and §3 (Fig. 3c/3d) treats
forward and backward NFE as two separate quantities.
str: Name of the factorised latent prior.
str: Whether the divergence is summed exactly or estimated.
Instance methods
5Per-sample negative log-likelihood in bits per dimension.
Invert the flow — latent (B, D) back to samples (B, C, H, W).
The same vector field, integrated from t = 1 to t = 0. No
separate inverse exists to get wrong, so the round trip is exact up
to the tolerances the solver was given; tighten rtol / atol
and it tightens with them.
Map samples to the latent space, accumulating the log-determinant.
Integrates from t = 0 to t = 1
with , so that
Parameters
(B, C, H, W).Returns
Notes
When trace_method resolves to "hutchinson" the returned
log_det is an unbiased estimate, not the exact value; its
expectation is right, which is what a maximum-likelihood gradient
needs, but two calls on the same input will not agree.
Examples
>>> import math
>>> import lucid
>>> from lucid.models.generative.neural_ode import (
... NeuralODEConfig, NeuralODEModel,
... )
>>> cfg = NeuralODEConfig(sample_size=2, in_channels=1, out_channels=1,
... hidden_dim=8, num_blocks=1)
>>> model = NeuralODEModel(cfg).eval()
>>> x = lucid.rand((2, 1, 2, 2))
>>> z, log_det = model.encode(x)
>>> z.shape, log_det.shape
((2, 4), (2,))
>>> model.trace_method # D = 4 is small enough for the exact trace
'exact'
>>> const = 0.5 * model.input_dim * math.log(2 * math.pi)
>>> log_pz = -0.5 * (z**2).sum(dim=-1) - const # the Gaussian prior
>>> bool(lucid.allclose(model.log_prob(x), log_pz + log_det))
True
>>> bool(lucid.allclose(model.decode(z), x, atol=1e-4)) # solve it backwards
Trueforward(x: Tensor)Encode x and report the density along with it.
Parameters
(B, C, H, W) data samples.Returns
NormalizingFlowOutputlatent , log_det_jacobian accumulated over the
solve, and log_prob in nats. No
loss — that is the task wrapper's to define.
Notes
Same computation as log_prob, returning the intermediates
instead of discarding them. With trace_method resolving to
"hutchinson" the log-determinant is an unbiased estimate, so
two calls on the same input will not agree.
Per-sample log-likelihood in nats.