[Paper Review] Stochastic Normalizing Flows

Liam Hodgkinson, Christopher van der Heide+2 more|arXiv (Cornell University)|Feb 21, 2020

Model Reduction and Neural Networks38 references44 citations

TL;DR

Introduces stochastic normalizing flows by extending continuous normalizing flows to stochastic differential equations using rough path theory, enabling density estimation, maximum likelihood, and variational inference with SDEs and latent Brownian motion. Provides a theoretical framework and practical training method with Wong–Zakai approximations for random ODEs and Stratonovich SDEs, plus applications to stochastic MCMC hyperparameter optimization.

ABSTRACT

We introduce stochastic normalizing flows, an extension of continuous normalizing flows for maximum likelihood estimation and variational inference (VI) using stochastic differential equations (SDEs). Using the theory of rough paths, the underlying Brownian motion is treated as a latent variable and approximated, enabling efficient training of neural SDEs as random neural ordinary differential equations. These SDEs can be used for constructing efficient Markov chains to sample from the underlying distribution of a given dataset. Furthermore, by considering families of targeted SDEs with prescribed stationary distribution, we can apply VI to the optimization of hyperparameters in stochastic MCMC.

Motivation & Objective

• Provide a general framework for approximating generative models built from SDEs using continuous normalizing flows.
• Extend continuous normalizing flows to stochastic settings via rough path theory to enable density estimation, MLE, and variational approximations for SDE models.
• Offer practical training methods that integrate Stratonovich SDEs, Wong–Zakai approximations, and random ODEs within existing CNF implementations.
• Demonstrate applications to density estimation and hyperparameter optimization in stochastic MCMC.

Proposed method

• Model Z_t via Ito SDEs and convert to pathwise formulations using rough path theory.
• Use Stratonovich calculus and Wong–Zakai approximations to approximate SDEs by random ODEs that are trainable as continuous normalizing flows.
• Apply the adjoint method in the rough path setting to backpropagate through stochastic dynamics.
• Estimate densities by conditioning random ODEs on latent Brownian-path approximations and using Monte Carlo averaging over Brownian realizations.
• Propose density-estimation and VI procedures that leverage the reparameterization trick for random ODEs with SDE-driven dynamics.
• Provide algorithms and theoretical guarantees (via Theorem 2) for convergence of log-densities and gradients as Wong–Zakai approximations improve.

Experimental results

Research questions

• RQ1 Can stochastic differential equations be trained and used as latent models within a continuous normalizing flow framework?
• RQ2 How can rough path theory be used to provide a rigorous, implementable pathwise treatment of SDEs in CNF-style models?
• RQ3 Can Stratonovich SDEs be approximated by random ODEs in a way that preserves log-density and gradient computations for training?
• RQ4 How can these stochastic normalizing flows be used for variational inference and sampling, including hyperparameter optimization in stochastic MCMC?

Key findings

• A unified stochastic normalizing flows framework extends CNFs to SDEs using rough paths, enabling density estimation, MLE, and VI for SDE models.
• Stratonovich SDEs are interpreted pathwise as rough differential equations and approximated by random ODEs suitable for CNF training.
• Wong–Zakai approximations (Karhunen–Loève expansion and piecewise linear) provide practical paths to train and test SDE-based CNFs.
• The main result (Theorem 2) shows convergence of log-densities and gradients when approximating rough paths with differentiable paths.
• The framework recovers and extends stochastic adjoint methods while avoiding deficiencies of prior approaches for non-diagonal diffusions and higher-order SDE solvers.
• Numerical experiments illustrate density estimation and samplers, including a two-dimensional toy example with trainable diffusion enabling better adaptation to curvature.