OT-Flow: Fast and Accurate Continuous Normalizing Flows via Optimal Transport

By Derek Onken, Samy Wu Fung, Lars Ruthotto. 2020

217 citations recorded by the source corpus.

A normalizing flow is an invertible mapping between an arbitrary probability distribution and a standard normal distribution; it can be used for density estimation and statistical inference. Computing the flow follows the change of variables formula and thus requires invertibility of the mapping and an efficient way to compute the determinant of its Jacobian. To satisfy these requirements, normalizing flows typically consist of carefully chosen components. Continuous normalizing flows (CNFs) are mappings obtained by solving a neural ordinary differential equation (ODE). The neural ODE's dynamics can be chosen almost arbitrarily while ensuring invertibility. Moreover, the log-determinant of the flow's Jacobian can be obtained by integrating the trace of the dynamics' Jacobian along the flow. Our proposed OT-Flow approach tackles two critical computational challenges that limit a more widespread use of CNFs. First, OT-Flow leverages optimal transport (OT) theory to regularize the CNF and enforce straight trajectories that are easier to integrate. Second, OT-Flow features exact trace computation with time complexity equal to trace estimators used in existing CNFs. On five high-dimensional density estimation and generative modeling tasks, OT-Flow performs competitively to state-of-the-art CNFs while on average requiring one-fourth of the number of weights with an 8x speedup in training time and 24x speedup in inference.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Statistical and Nonlinear Physics

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Statistical and Nonlinear Physics

Citations in the corpus

Flow Matching for Generative Modeling (2022, neutral): “…regularize the ODE to be easier to solve, e.g. , using augmentation (Dupont et al., 2019), adding regularization terms (Yang & Karniadakis, 2019; Finlay et al., 2020; Onken et al., 2021; Tong et al., 2020; Kelly et al., 2020), or stochastically sampling the integration interval (Du et al., 2022).”
Stochastic Interpolants: A Unifying Framework for Flows and Diffusions (2023, methodology): “Optimal transport perspectives in flow-based methods have primarily been employed as a regularization penalty [51, 34, 14, 46] or via imposing structure on the parameterization itself [19, 50].”
Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow (2022, background): “A number of works have been proposed to remedy this by adding regularization terms, such as these based on transport costs, to favor shorter paths; see [54, 55].”
An introduction to deep generative modeling (2021, methodology): “In Section 3, we present ﬁnite [39, 11] and continuous normalizing ﬂows [19, 48, 47, 14, 34].”
Improving and generalizing flow-based generative models with minibatch optimal transport (2023, unknown): “…distributions with neural networks; however, these require constrained architectures (Leygonie et al., 2019; Makkuva et al., 2020; Bunne et al., 2022) or use a regularized CNF, which is challenging to optimize (Tong et al., 2020; Finlay et al., 2020; Onken et al., 2021; Huguet et al., 2022a).”
JFB: Jacobian-Free Backpropagation for Implicit Networks (2021, background): “A drawback is that the adjoint equation must be solved to high-accuracy; otherwise, a descent direction is not necessarily guaranteed (Gholami, Keutzer, and Biros 2019; Onken and Ruthotto 2020; Onken et al. 2021).”
Plugin estimation of smooth optimal transport maps (2021, background): “They have also found recent uses in distributional regression (Ghodrati and Panaretos, 2021), generative modeling (Finlay et al., 2020; Onken et al., 2021), fairness in machine learning (Gordaliza et al., 2019; Black et al., 2020; de Lara et al., 2021), and in a wide range of statistical…”
Building Normalizing Flows with Stochastic Interpolants (2022, unknown): “Regulating the path can reduce the number of solver calls (Finlay et al., 2020; Onken et al., 2021), though this does not alleviate the main structural challenge of the optimization.”
TO-FLOW: Efficient Continuous Normalizing Flows with Temporal Optimization adjoint with Moving Speed (2022, methodology): “Finlay et al [37] and Onken et al [36] are dedicated to optimising trajectory spatially, thus constraining it to be straight to improve training speed.”