000

001
002 Generative modeling with preconditioned neural ODEs
003
004
005
006 Anonymous Authors1
007
008
009 Abstract Radford et al., 2015), and molecular design (Gilmer et al.,
010 2017; Duvenaud et al., 2015).
011 Neural ordinary differential equations (ODEs)
have shown great success in generative model- As data of interest are generally high-dimensional and
012
ing problems. However, their utilization in high- highly structured, the challenge in this domain is building
013
dimensional problems has often proved difficult models that are flexible enough to capture its complexity
014
due to expensive ODE solves. In the present work, yet still tractable to train. One such class of algorithms
015
we introduce a new architecture wherein a sys- are normalizing flows. Introduced originally by Moselhy
016
tem of decoupled ODEs are preconditioned with et. al. (El Moselhy & Marzouk, 2012), normalizing flows
017
shallow autoregressive maps. An attractive fea- construct a sequence of invertible and differentiable map-
018
ture of the proposed architecture is that it enables pings between a complex ”target” probability measure of
019
the push forward of samples from a given dis- interest and a simpler reference measure. As compared to
020
tribution by leveraging implicit ODE solvers in other generative modeling paradigms such as variational
021
an embarrassingly parallel fashion. This enables autoencoders (Kingma & Welling, 2013), generative adver-
022
training and inference in a fraction of the time re- sarial networks (Goodfellow et al., 2020), and diffusion
023
quired by standard neural ODEs, particularly for models (Sohl-Dickstein et al., 2015) normalizing flows are
024
high-dimensional problems. We also introduce a useful as both sampling and log-likelihood evaluation can
025
collocation based ODE-solver which converts the be efficient and exact.
026
027 problem of time-integration to that of nonlinear
028 root finding during training, further speeding up Normalizing flows Let x be a D-dimensional random
029 training. Experimental evaluations on standard variable. We would like to learn the underlying distribution
030 generative modeling benchmarks, such as sam- over x. The main idea of flow-based modeling is to express
031 pling for continuous normalizing flows, reveal x as a transformation T of a random vector u with a known
032 consistent performance with considerable speed reference distribution p(u):
033 up over classical neural ODEs.
x = T (u) where u ∼ p(u). (1)
034
035 Then using the change of variables formula, one can com-
036 1. Introduction pute the density, p(x), as the transformed distribution p(u)
037 denoted as T♯ u.
038 Generative modeling, that is, modeling the distribution from
039 which data are drawn, is a central task in machine learning Design goals for flow-based models can be summarized as
040 and statistics. The general problem statement is to esti- follows: (i) the transformation T must be cheaply invert-
041 mate an unknown probability density p given i.i.d samples ible, (ii) both T and T −1 must be differentiable, (iii) the
042 X = {x1 , . . . , xn } from a distribution P . By capturing the log det ∇T must be efficient to compute and, ideally, have
043 mechanisms behind the data generation process, one can an asymptotic complexity which scales linearly in the di-
044 reason about data probabilistically, analyze insightful low- mension of the random variable. Much of the normalizing
045 dimensional manifolds linked to the data, and subsequently flow literature revolves around designing transformations,
046 generate new data. Generative models have gained great mo- T , that trade-off among these criteria. El Moselhy & Mar-
047 mentum in applications such as NLP (Ramesh et al., 2021; zouk (2012) introduced triangular maps that allowed for the
048 Vaswani et al., 2017), computer vision (Sohn et al., 2015; log det ∇T component to be computed in O(D) time.
049 1
Anonymous Institution, Anonymous City, Anonymous Region, Others have demonstrated great empirical success by in-
050 Anonymous Country. Correspondence to: Anonymous Author troducing various efficient forms for T ; examples include
051 <[Link]@[Link]>.
Real NVP (Dinh et al., 2016), MAF (Papamakarios et al.,
052 2017), and GLOW (Kingma & Dhariwal, 2018). While
Preliminary work. Under review by the International Conference
053 on Machine Learning (ICML). Do not distribute. these models preserve the advantages of normalizing flows,
054
1
 Submission and Formatting Instructions for ICML 2024

055 including their ease of sampling and the straightforwardness Neural ODE (PNODE) and highlight the main contributions
056 of their training through the minimization of the forward KL below:
057 divergence, awkward restrictions on their structure, such as
058 partitioning dimensions or designing low-rank matrices (Pa- • Decoupled ODEs - We introduce a system of D-
059 pamakarios et al., 2019), severely impact their representa- decoupled ODEs that models the latent dynamics of
060 tion capability. each dimension of h(t) independently. Each member
061 of the system of ODEs can be efficiently integrated in
A common theme for normalizing flows is to compose a
062 time.
sequence of transformations T1 , T2 , . . . , Tk to increase their
063
overall expressive power. Such deep models are heavily • Preconditioners - We condition the input and output of
064
parameterized with memory cost increasing linearly with the decoupled neural ODE with shallow autoregressive
065
the number of layers. Empirically experiments have shown transformations to allow us to learn a manifold where
066
that very deep flows are required for good performance the ODE states are decoupled.
067
with high-dimensional inputs. Moreover, it was recently
068
shown that the number of layers required by some restrictive • Collocation-based ODE solvers - We introduce
069
architectures1 is expected to grow as polynomial function collocation-based ODE solvers that learn an ansatz
070
with the dimension of the input (Kong & Chaudhuri, 2020). for the state h(t) by converting the problem of time
071
This makes the prospect of deploying deep-flow models integration to that of non-linear optimization.
072
for high-dimensional, complex datasets expensive if not
073
computationally intractable. In upcoming sections, we show that these contributions al-
074
075 Neural ordinary differential equations (Neural ODEs) (Chen low us to train continuous normalizing flows at a greatly
076 et al., 2018) are promising for candidates for parameterizing reduced cost per optimization iteration than standard con-
077 flows for two reasons. First Neural ODEs define bijective tinuous normalizing flows. Importantly we accomplish this
078 transformations by construction, meaning no restriction are computational advancement without sacrificing accuracy on
079 required on their architecture. Second, Neural ODEs pro- a number of important generative modeling benchmarks.
080 vide an ”infinitely-deep” extension to residual networks
081 making well-suited for parametrizing extremely deep flows. 2. Related work
082 Neural ODEs model the differential equations governing
083 the dynamics of a latent state h using a neural network, The question of speeding up neural ODEs to extend their ap-
084 plications to high-dimensional problems has gathered much
085 dh(t) interest over the years. Researchers have explored many
= fθ (h, t), (2)
086 dt directions to speed up the training of neural ODEs. Dupont
087 et al. (2019) augment the ODE state to learn simple fields
initialized with input h(t0 ) = x and output after time
088 fθ facilitating faster integration using ODE solvers. Various
t as h(tf ) = u. The resulting map, T (x) = x +
089 Rftf regularization techniques (Finlay et al., 2020; Ghosh et al.,
t0 θ
f (h(t), t) dt, is computed using a black box ODE 2020) enable learning simple paths over which the ODE
090
solver. Note that the map is invertible as long as the network state can be efficiently integrated through.(Kidger et al.,
091
weights are finite and the nonlinearities are Lipschitz. 2021) proposed a method to accelerate the ODE integration
092
093 Neural ODEs have found abundant applications (Grathwohl process by mitigating the impact of the sensitivity of loss L
094 et al., 2018; Poli et al., 2019) and have invigorated research to the parameters θ, ∂L∂θ , on the adaptive time-stepping stage
095 in continuous deep learning (Massaroli et al., 2020; Dupont in the backward pass.R t Neural flows Biloš et al. (2021) di-
096 et al., 2019); however, these methods are held back by the rectly model “u + t0f fθ (h(t), t) dt” using neural network
097 fact that the computational cost associated with solving based maps similar to normalizing flows. For post-training
098 ODEs typically far exceeds the cost of a comparable dis- speed up, Hypersolvers (Poli et al., 2020) approximate ex-
099 crete model. For example, Grathwohl et al. (2018) report pensive trained Neural ODE solution trajectories with a
100 that the number of function evaluations (NFEs) typically neural network corrected cheaper trajectory during the in-
101 becomes prohibitively large when training continuous nor- ference stage.
102 malizing flow (CNF) models. Recent work has thus focused
Many recent works have explored collocation method in-
103 on speeding up training of Neural ODEs. In this paper, we
spired techniques to speed up neural ODEs. Daulbaev et al.
104 address the bottlenecks for training neural ODEs and pro-
(2020) interpolated the state h(t) during the backward pass
105 pose a novel, alternative strategy for learning continuous
using stored values of the state from the forward pass. This
106 normalizing flows. We name our algorithm Preconditioned
enabled them to reduce the dimension of the coupled ODE
107 1 system in the backward pass by the size of the ODE state.
The paper shows results for Sylvester and Householder flows.
108 (Koehler et al., 2021) extend the work with results for affine flows. Djeumou et al. (2022) introduced a Taylor series based
109
2
 Submission and Formatting Instructions for ICML 2024

110 ODE stepping algorithm that heuristically corrected for the the minimizer to (3) tractable. The architectures proposed
111 error introduced in each time step by the approximation. in literature actively trade-off utility of S(θ) to ease of opti-
112 The method was based on the idea that higher-order time mizing within the space.
113 derivatives of the state are easy to compute using automatic
114 differentiation (AD). 3.2. Neural ODEs
115
Recent works (Dupont et al., 2019; Zhang et al., 2020) can Chen et al. (2018) introduce a highly flexible parameteriza-
116
be interpreted to use linear preconditioning transformations tion Sθ within S(θ) described by the transformation
117
to the input u and output x of neural ODEs to augment
118 Z tf
the neural ODE dimensions. Based on our comprehension,
119 u = Sθ (x) = x + fθ (h(t), t) dt, (4)
these pieces of literature serve as initial inspirations for the t0
120
notion that effective preconditioners are useful to learn man- where h(t0 ) = x, h(tf ) = u.
121
ifolds where simplified neural ODEs can be deployed. Our
122
work is complementary to these methods, and we empiri- Here θ represents some vector of learnt parameters, fθ :
123
cally demonstrate that our method can be used to further Rd1 ,d2 ,...,dc × R → Rd1 ,d2 ,...,dc in general is any standard
124
improve the performance of these techniques. neural architecture, and h : [t0 , tf ] → Rd1 ,d2 ,...,dc is the
125
126 solution. The only requirement for the parameterized vector
127 3. Background: Normalizing flows & Neural field fθ is that it should be uniformly Lipschitz continuous
128 ODEs in h (meaning that there is a single Lipschitz constant that
129 works for all t) and continuous in t. Note that the dimension
130 To motivate the preconditioned neural ODE architecture, t can be understood as the depth of neural ODEs. For many
131 we first introduce some additional notation, followed by a applications fθ will just be a simple feedforward network.
132 brief review of pertinent normalizing flow concepts. We
Evaluating the log determinant of the Jacobian in (3) is
133 also highlight the traditional transport map structure adopted
difficult. (Grathwohl et al., 2018) introduced FFJORD for
134 along with the more flexible class of architectures used in
generative modeling using neural ODEs. Their algorithm
135 the literature. We follow with a short review of the relevant
utilized the following identity -
136 neural ODE concepts emphasizing how they are trained.
Finally, we look at collocation based ODE solvers.
 
137 ∂ ∂fθ
log det ∇Sθ = −Tr , (5)
138 ∂t ∂h
139 3.1. Normalizing flows
140 where Tr denotes the trace of an operator, to rewrite the
El Moselhy & Marzouk (2012) introduced nonlinear in- optimization problem in (3) as
141 vertible transformations T capable of mapping indepen-
142 dent samples from pu to independent samples from the M Z tf   i
143 1 Xh ∂fθ
target px and the inverse mapping S ≡ T −1 . The maps max log( pu ◦Sθ (xi ) )− Tr dt .
144 Tθ ∈T (θ) M t0 ∂h
are defined as the solution
to the optimization problem i=0
145 (6)
minS∈S DKL S♯ px || pu or equivalently -
146 In general, computing the
objective in (6) exactly costs
147 M O(D2 ) due to the Tr ∂f ∂h
θ
term. An attractive feature of
1 Xh i
148 max log( pu ◦ S(xi ) + log det ∇S(xi ) (3) neural ODEs is that given a fixed network depth, the mem-
149
S∈S M i=0 ory requirements of neural ODEs are orders of magnitude
150 s.t.∇S ≻ 0, smaller than many finite-composition flows, making them
151 ideal for generative modeling and density estimation. A fun-
152 where the constraint is necessary to ensure invertibility of S, damental drawback of neural ODEs is the long training and
153 we denote the push of pu through S by S♯ px and S is the inference time incurred by expensive integration in t. This
154 space of all diffeomorphisms. The authors further introduce cost is arguably one of the main stumbling blocks toward
155 a structure to S, known as the Knothe-Rosenblatt (KR) rear- the widespread adoption of these continuous deep models.
156 rangement, which is particularly advantageous as it leads to
Past numerical experiments show that these models appear
157 a lower triangular structured ∇S for which the determinant
too slow in their current form for general high-dimensional
158 in (3) can be cheaply computed as the product of diagonal
applications including image generation (Grathwohl et al.,
159 terms. The optimization objective in (3) effectively becomes
2018). Additionally, neural ODE implementations rely on
160 an infinite-dimensional optimization problem over the space
general-purpose ODE solvers that restrict applications to
161 of feasible KR diffeomorphisms SKR .
non-stiff differential equations in the interest of time. In
162 Much work has gone into constructing parameterized sub- general, the hidden state or the adjoint might be unstable,
163 spaces S(θ) ⊂ S or SKR (θ) ⊂ SKR to make computing exiting when stiffness is detected, and switching to methods
164
3
 Submission and Formatting Instructions for ICML 2024

165 more suitable for stiff equations (Ross, 2015). Implicit The idea of preconditioners2 is ubiquitous to machine learn-
166 schemes have much better stability properties and allow sig- ing algorithms. (Parno & Marzouk, 2018) and (Peherstorfer
167 nificantly larger time steps compared to explicit integrators, & Marzouk, 2019) introduce autoregressive preconditioners
168 at the expense of a significant increase in the number of for MCMC algorithms. (Fang et al., 2021) shows precon-
169 evaluations of fθ needed at each time iteration. ditioning GANs lead to more stable during training. In
170 this context, we introduce shallow autoregressive transfor-
171 Training Neural ODEs Chen et al. (2018) deploys an mations in SKR (θ) as conditioners to the input and output
172 optimise-then-discretise approach to train neural ODEs to of the decoupled system of neural ODEs. This effectively
173 minimize the objective function in (6). They concatenate the changes the associated initial value problem as follows:
174 first and second term in (6) solving it as aRcombined ODE
175 t dh(t)
system. We note that the second term, t0f Tr ∂f


∂h
θ
dt = fθ (h, t), (7)
176 is not actually an ODE and solving the combined system dt
177 could be disadvantageous. They subsequently produces
178 a backwards in-time differential equation, which is then with the conditioned input h(t0 ) = Sin (x), the conditioned
−1
179 solved numerically (Pontryagin et al., 1961; Chen et al., output h(tf ) = Sout (u). The parameterized transforma-
180 2018). We give more intuition to the problem in the attached tions Sin and Sout allow us to learn complex manifolds
181 appendix ??. where a decoupled system of neural ODEs can be effec-
182 tively deployed. Vanilla Neural ODEs will generally require
183 more inner layer evaluations than a fixed architecture on the
184 4. Preconditioned neural ODEs same task. What makes this even worse is that during train-
185 The central idea of flow based generative modeling is to ing, the dynamics being learned tends to become more and
186 design an efficient parameterized space S(θ) that is expres- more expensive to solve. This makes sense if we think of
187 sive enough to model the complicated probability manifolds the network as “learning to be deep” and expressing a more
188 while simultaneously being tractable under available opti- complex function. Preconditioned neural ODEs balance
189 mization routines. While normalizing flows achieve this predictive performance against compute cost, so we need
190 through a deep composition of flow layers with a huge to somehow encourage our model to be easy to solve. The
191 memory footprint, neural ODEs adopt a memory-efficient resulting composition of conditioners with the decoupled
192 continuous-time parameterization for Sθ relying on expen- neural ODEs is invertible as the composition of diffeomor-
193 sive ODE solves leading to slow training and inference. phisms is a diffeomorphism itself. Figure ?? shows the
194 Motivated by the hypothesis laid out, we introduce precondi- preconditioned neural ODE architecture in detail. We now
195 tioned neural ODEs (PNODEs): a new architecture for S(θ) describe the characteristic features of PNODEs that make
196 that brings the best of both the discrete-time and continuous- the architecture amenable to fast training while maintaining
197 time approaches to optimizing the loss function in equation comparable expressivity.
198 (3). The goal of PNODES is to be as easy to train as com-
199 mon flow architectures while retaining the expressive power 4.2. Decoupled neural ODEs
200 of neural ODEs through their efficient integration. We first
201 The preconditioned neural ODEs overcome the expensive
introduce decoupled neural ODEs and state the important cost associated with the computation of the flow objective
202 features. We then highlight preconditioners to the neural
203 and its gradients by learning a manifold where decoupled
ODEs that learn data manifolds where decoupled neural neural ODEs can efficiently push forward the individual
204 ODEs can be efficiently deployed.
205 dimensions of our latent state h(t). Our architecture allows
206 the neural ODEs to adapt to otherwise intractable high-
4.1. Preconditioned neural ODEs dimensional problems by introducing ”shallow” autoregres-
207
208 sive preconditioners for the neural ODE input and output
209 states. The transformations force the latent dimensions to
210 interact with each other during training, thereby increasing
211 the expressive power of our architecture.
212 We construct the decoupled variant of neural ODEs by en-
213 forcing a decoupled structure to the velocity field fθ . In
214 this construction, the state vector h(t) is split into D-blocks
215 {h1 , h2 , . . . , hD } of vector-valued states. Hence the dynam-
216 2
217 We understand that the word “preconditioning” may be remi-
niscent of the condition number. However, historically, “precondi-
218 Figure 1. Illustration of a preconditioned neural ODE architecture. tioning” is intended to improve the whole spectrum broadly.
219
4
 Submission and Formatting Instructions for ICML 2024

220 ics of the ith block is given by the ODE - result in an order O(CD) cost due to the decoupled nature
221 of the velocity field.
222 dhi (t)
= fθi (hi , t), ∀i ∈ {1, 2, . . . , D} (8) Calculating the gradients: The decoupled structure of
223 dt
the dynamics of each state hi extends to the dynamics of
224
under the initial condition hi (t0 ) = ui and output constraint the associated adjoint ai . Their adjoint dynamics can be
225
hi (tf ) = xi . This leads to an equivalent block expressed represented as follows:
226
for SDC as follows:
227 dai ∂fθi (hi , t)
228  Rt  = −ai (10)
x1 + t0f fθ1 (h1 , t) dt dt ∂hi
229 R tf
 x2 + t0 fθ2 (h2 , t) dt 
 
230 SDC (u) =  (9) ∀i ∈ {1, 2, . . . , D}. The decoupled structure of the state
 .. 
231  .

 dynamics, as well as the associated adjoint dynamics, has
232 Rt
xD + t0f fθD (hD , t) dt several advantages in the context of PNODEs. One advan-
233 tage is that the diagonal Jacobian structure allows PNODEs
234 As long as each fθi be uniformly Lipschitz continuous in hi to benefit from implicit ordinary differential equation (ODE)
235 and continuous in t, we can guarantee that S is invertible. solvers. Implicit solvers are known to be more effective for
236 This leads to a row-by-row assembly of h, which can be solving stiff ODEs, where the dynamics change rapidly
237 sparsified naturally while trading off the ability to quantify over small time scales. However, existing algorithms for
238 dependence between the blocks. Since decoupled neural neural ODEs often rely on specific design choices to mit-
239 ODEs are parameterized with different “weights”, we de- igate the need to solve stiff ODEs and use non-stiff ODE
240 note them θDC = {θ1 , θ2 , . . . , θD }. The Jacobian of fθ solvers requiring a large number of function evaluations in
241 simplifies to a diagonal matrix allowing for fast calculation the presence of stiff ODEs (Grathwohl et al., 2018).
242 of the Jacobian in order D.
243 Our algorithm, on the other hand, can effectively explore tra-
244 jectories governed by stiff ODEs without resorting to such
4.3. Cost of evaluating design choices. Additionally, by utilizing a collocation-
245
246 Evaluating the network: Given x, we compute all the based ODE solver, we can efficiently integrate each dimen-
247 components of the velocity field independently and in paral- sion of the decoupled ODE. This solver takes advantage
248 lel for each individual dimension hi . Consider the vanilla of automatic differentiation, enabling us to calculate the
249 neural ODE velocity field fθ where the cost of evaluating gradients while eliminating the need to solve an additional
250 fθ is of the order O(DH) where H ≈ KD is the size of ODE in reverse, which can be numerically unstable (Kim
251 the largest hidden dimension in fθ . Consider a decoupled et al., 2021).
252 neural ODE architecture with each fθi costing of the order
253 O(KD). We can potentially compute each element of fθ 5. Collocation based ODE solvers
254 for the decoupled case independently at almost identical
255 costs. Current neural ODE training routines suffer from a variety
256 of problems. Recent works (Ghosh et al., 2020; Grathwohl
Evaluating the objective: For vanilla neural ODE, the cost et al., 2018; Dupont et al., 2019) talk about the need to
257 of evaluating the normalizing flow objective in (6) reduces
258 adapt neural ODEs to practical training time. Addition-
to O(KD2 L̃), where L̃ is the number of evaluations of fθ ally (Dupont et al., 2019; Gholami et al., 2019) state empiri-
259 used by the ODE solver. For the preconditioned neural ODE
260 cal and theoretical arguments for the additional instabilities
the cost involves the total cost of applying two precondi- and errors introduced by reversing ODEs during backpropa-
261 tioners and the decoupled neural ODE while accumulating
262 gation. To this extent we introduce a collocation based ODE
the value of the log det ∇Sθ . The cost of evaluating the solver that has numerous advantages; namely -
263 two preconditioners along with their log det Jacobian val-
264 ues is O((KD2 + D)L) where L is the total number of
265 • Cost of adaptive solvers is unbounded and usually
transformations used. Note here that we assume a favorable
266 high, especially in the presence of stiffness (Wanner &
structure to our preconditioners that allow log det Jacobin
267 Hairer, 1996),
calculation in O(D). The cost of evaluating the decoupled
268 ODE is of the order O(KDL̃ + D) where the first term
269 • Each iteration of the training routine requires one ODE
results from the evaluation of decoupled fθ an L̃ number solve to calculate h(tf ) and a second ODE solve to
270 of times by the ODE solver. The cost of evaluating the
271 calculate ∂L
∂θ via the adjoint method,
log det Jacobian term is heavily dependent on the quadra-
272 ture scheme used to evaluate the integral. Utilizing a simple • The ODE solvers solve the system (??) and (??) con-
273 Monte Carlo approximate to the integral using C samples catenated together with the state h backwards in time
274
5
 Submission and Formatting Instructions for ICML 2024

275 during the gradient calculation, which is numerically 2017). We give more problem specific mdoel details in each
276 unstable (Kim et al., 2021). section below.
277
278 Collocation based ODE solver We propose the applica- 6.1. Density estimation on 2D datasets
279 tion of collocation based ODE solvers which converts the
280 problem of time-integration to that of estimating an ansatz
281 for the state h(t), usually formulated as a nonlinear root
282 finding problem. In particular, we use these solvers to in-
283 tegrate each of the D independent ODEs in our decoupled
284 ODE system.
285 PM
Let us assume an ansatz for h(t) given by i=0 λi ψi (t)
286
where ψi : R → R. On substituting this equation into (4)
287
we get the following nonlinear problem -
288
289
X X
λi ψ̇i (t) − fθ ( λi ψi (t), t) = 0, (11)
290 i i
291 X
s.t. λi ψi (t0 ) = x
292
i
293
294 The trade-off between accuracy and computational cost for
295 basis functions in the collocation method depends on two
296 main factors: the expressivity of ψ i and the choice of collo-
297 cation points. We note here that depending on the choice of
298 the basis functions and collocation points we can “mimic” Figure 2. ls on a 2-dimensional multi-modal and discontinuous
299 the behavior of various ODE solvers. For example, piece- pinwheel density.
300 wise linear basis functions with collocation points repre-
301 senting the beginning of each interval in the solver result in We first train on 2 dimensional data to visualize the learned
302 the explicit Euler ODE solver. We give more details of the model. Our PNODE model consists of identical RealNVP
303 collocation based ODE solver in the appendix. inspired autoregressive map preconditioners. Each of the
304 conditioners consist of 5 layers each with 1 hidden layer
305 of size 64. For the decoupled neural ODEs we use an ex-
6. Numerical Studies plicit ODE solver to integrate the ODEs with velocity fields
306
307 Our PNODEs algorithm was tested across various density having 1 layer and each layer with 3 hidden dimensions
308 estimation benchmarks (Chen et al., 2018; Grathwohl et al., of size 64. To compare we used the benchmarks stated
309 2018). Unlike FFJORD which solely used implicit ODE- in Grathwohl et al. (2018) where all models transform a
310 solvers, our experiments incorporated a mix of implicit reference isotropic Gaussian. Interestingly from figure 2,
311 and explicit ODE-solvers, treated as a hyperparameter and we find that PNODEs is capable of stretching the single
312 chosen as per their respective advantages. Solver tolerances mode base distribution into multiple modes, modeling the
313 followed the heuristics outlined in Grathwohl et al. (2018). areas of low probability between disconnected regions much
314 We used the exact trace during training and when reporting better than the 70-100 layer Glow model stated in the bench-
315 test results. This was done in all experiments including mark (Grathwohl et al., 2018). Although FFJORD performs
316 large image datasets as computing the exact Jacobian trace better, PNODEs is able to perform fairly close and consider-
317 was computationally feasible for PNODEs. To ensure a ably better than the deep discrete counterpart.
318 fair comparison, we run large hyperparameter searches for
319 each model and repeat each experiment 20 times to ensure Density estimation on real datasets
320 results are meaningful across initializations (see appendix
We perform density estimation on five tabular datasets pro-
321 for details).
vided in Grathwohl et al. (2018). For the PNODEs model,
322 The dynamics of PNODEs are defined by two neural net- we use identical discrete autoregressive preconditioners.
323 work based autoregressive maps and D decoupled ODEs We run two versions of PNODEs with Glow (Kingma
324 which takes as input the current state h(t) ∈ RD and the cur- & Dhariwal, 2018) inspired enhanced MAF precondition-
325 rent time t ∈ R. We simply concatenate t on to h(t) at the ers (PNODE) and MAF (Papamakarios et al., 2017) pre-
326 input to every layer and was used in all of our experiments. conditioners (PNODE (MAF)). Each of the conditioners to
327 We build autoregressive pre/post-conditioners inspired from PNODES contain 2 autoregressive layers with layer masks
328 RealNVP (Dinh et al., 2016) and MAF (Papamakarios et al., designed according to MAF. The decoupled map TDC was
329
6
 Submission and Formatting Instructions for ICML 2024

330 constructed with decoupled neural ODEs with 1 layer con-

331 taining 3 hidden dimensions of size 64. These model archi-
332 tectures were used in all the experiments for tabular datasets
333 with the reference distribution fixed to a standard Gaussian.
334 For additional details please refer to the appendix.
335
On the tabular datasets, PNODEs perform better than
336
FFJORD and it’s competitors on the Miniboone and
337
BDS300 datasets. Of the discrete normalizing flow com-
338
petitors, PNODE outperforms Glow and RealNVP on all
339
datasets. We note that PNODE in general performs better
340
than PNODE (MAF) on all datasets except the Gas dataset
341
where it performs better than PNODE and FFJORD. The Figure 3. The cost per iteration of PNODE and FFJORD as the
342
test loss is undoubtedly a function of the data, model archi- data dimension is varied for d = 1 to d = 100.
343
tecture, and model parameters and the results indicates that
344
carefully selecting the preconditioner could lead to better
345
performance than original FFJORD in certain cases. We diagonal and calculated along with the function evaluation3 .
346
find that PNODE can model the data with similar perfor-
347
mace as compared to Glow and Real NVP by enriching 7.1. Stiffness and implicit ODE solvers
348
the fewer autoregressive layers used with decoupled neu-
349 • Neural ODEs can become stiff - Steer paper
ral ODEs. We also note than using preconditioned neural
350
ODEs is able to achieve this performance while using less • people try to introduce temporal regularization or
351
than 40% as many parameters as Glow. A summary of our weight decay or augmented neural ODEs
352
results on density estimation can be found in Table 1. Full
353 • Implicit solvers have better stability
details on architectures used, our experimental procedure,
354
and additional experiments on image datasets can be found
355 • Jacobian advantageous to implicit solvers as non-linear
in Appendix B.1.
356 solvers part of implicit solvers use a lot of function
357 evaluations.
358 7. Analysis and discussion
359 7.2. Complementary to other solver
We perform a series of ablation experiments to gain a better
360
understanding of the proposed model. References
361
362 Biloš, M., Sommer, J., Rangapuram, S. S., Januschowski, T.,
363 Scalability of PNODE:
and Günnemann, S. Neural flows: Efficient alternative to
364 We plot the cost per iteration of training PNODEs compared neural odes. Advances in Neural Information Processing
365 to FFJORD using an explicit ODE solver on toy datasets Systems, 34:21325–21337, 2021.
366 of varying dimensions. The cost can be seen to grow sig-
367 nificantly around d = 50 denoting the cost of expensive Chen, R. T., Rubanova, Y., Bettencourt, J., and Duvenaud,
368 function evaluations by the neural ODEs in FFJORD. The D. K. Neural ordinary differential equations. Advances
369 highly decoupled structure of PNODEs allows us to cheaply in neural information processing systems, 31, 2018.
370 evaluate functions and save of per iteration cost. One of
371 Daulbaev, T., Katrutsa, A., Markeeva, L., Gusak, J., Ci-
the significant disadvantages of FFJORD is the expensive
372 chocki, A., and Oseledets, I. Interpolation technique
cost per iteration during training. Training FFJORD models
373 to speed up gradients propagation in neural odes. Ad-
require a huge budget in terms of time restricting abilities of
374 vances in Neural Information Processing Systems, 33:
the end user to tune hyperparameters. On the other hand we
375 16689–16700, 2020.
notice that PNODEs intuitively scale to high-dimensional
376 problems without the need to engage an encoder-decoder Dinh, L., Sohl-Dickstein, J., and Bengio, S. Density esti-
377 system as in the original FFJORD model. Additionally, the mation using real nvp. arXiv preprint arXiv:1605.08803,
378 authors in (Grathwohl et al., 2018) introduce unbiased esti- 2016.
379 mates to the Tr [] to make training faster and hyperparameter
380 tuning more efficient. For PNODEs the Jacobian matrix is Djeumou, F., Neary, C., Goubault, E., Putot, S., and Topcu,
381 U. Taylor-lagrange neural ordinary differential equations:
382 3
The latest version of PyTorch allows us to efficiently calculate
383 the Jacobian matrix (PyTorch, 2023).
384
7
 Submission and Formatting Instructions for ICML 2024

385
Table 1. Negative log-likelihood on test data for density estimation models; lower is better.
386
387
M ODEL P OWER G AS H EPMASS M INIBOONE BSDS300
388
389 R EAL NVP -0.17 -8.33 18.71 13.55 -153.28
390 G LOW -0.17 -8.15 18.92 11.35 -155.07
FFJORD -0.46 -8.59 14.92 10.43 -157.40
391 PNODE -0.14 -8.33 15.86 10.39 -158.11
392 PNODE(MAF) -0.16 -9.41 17.67 10.80 -153.56
393
394
395 Toward fast training and evaluation of neural odes. arXiv Kidger, P., Chen, R. T., and Lyons, T. J. ” hey, that’s not an
396 preprint arXiv:2201.05715, 2022. ode”: Faster ode adjoints via seminorms. In ICML, pp.
397 5443–5452, 2021.
Dupont, E., Doucet, A., and Teh, Y. W. Augmented neural
398
odes. Advances in neural information processing systems, Kim, S., Ji, W., Deng, S., Ma, Y., and Rackauckas, C. Stiff
399
32, 2019. neural ordinary differential equations. Chaos: An Inter-
400
401 Duvenaud, D. K., Maclaurin, D., Iparraguirre, J., Bom- disciplinary Journal of Nonlinear Science, 31(9):093122,
402 barell, R., Hirzel, T., Aspuru-Guzik, A., and Adams, R. P. 2021.
403 Convolutional networks on graphs for learning molecular Kingma, D. P. and Dhariwal, P. Glow: Generative flow
404 fingerprints. Advances in neural information processing with invertible 1x1 convolutions. Advances in neural
405 systems, 28, 2015. information processing systems, 31, 2018.
406
El Moselhy, T. A. and Marzouk, Y. M. Bayesian inference
407 Kingma, D. P. and Welling, M. Auto-encoding variational
with optimal maps. Journal of Computational Physics,
408 bayes. arXiv preprint arXiv:1312.6114, 2013.
231(23):7815–7850, 2012.
409
410 Fang, T., Schwing, A., and Sun, R. Precondition layer and its Koehler, F., Mehta, V., and Risteski, A. Representational
411 use for {gan}s, 2021. URL [Link] aspects of depth and conditioning in normalizing flows.
412 net/forum?id=1yXhko8GZEE. In International Conference on Machine Learning, pp.
413 5628–5636. PMLR, 2021.
414 Finlay, C., Jacobsen, J.-H., Nurbekyan, L., and Oberman,
415 A. How to train your neural ode: the world of jacobian Kong, Z. and Chaudhuri, K. The expressive power of a class
416 and kinetic regularization. In International conference on of normalizing flow models. In International Conference
417 machine learning, pp. 3154–3164. PMLR, 2020. on Artificial Intelligence and Statistics, pp. 3599–3609.
418 PMLR, 2020.
Gholami, A., Keutzer, K., and Biros, G. Anode: Uncondi-
419 tionally accurate memory-efficient gradients for neural Massaroli, S., Poli, M., Bin, M., Park, J., Yamashita, A.,
420 odes. arXiv preprint arXiv:1902.10298, 2019. and Asama, H. Stable neural flows. arXiv preprint
421 arXiv:2003.08063, 2020.
422 Ghosh, A., Behl, H., Dupont, E., Torr, P., and Namboodiri,
423 V. Steer: Simple temporal regularization for neural ode. Papamakarios, G., Pavlakou, T., and Murray, I. Masked
424 Advances in Neural Information Processing Systems, 33: autoregressive flow for density estimation. Advances in
425 14831–14843, 2020. neural information processing systems, 30, 2017.
426 Gilmer, J., Schoenholz, S. S., Riley, P. F., Vinyals, O., and
427 Dahl, G. E. Neural message passing for quantum chem- Papamakarios, G., Nalisnick, E., Rezende, D. J., Mo-
428 istry. In International conference on machine learning, hamed, S., and Lakshminarayanan, B. Normalizing flows
429 pp. 1263–1272. PMLR, 2017. for probabilistic modeling and inference, arxiv e-prints.
430 arXiv preprint arXiv:1912.02762, 2019.
431 Goodfellow, I., Pouget-Abadie, J., Mirza, M., Xu, B.,
432 Warde-Farley, D., Ozair, S., Courville, A., and Bengio, Y. Parno, M. D. and Marzouk, Y. M. Transport map acceler-
433 Generative adversarial networks. Communications of the ated markov chain monte carlo. SIAM/ASA Journal on
434 ACM, 63(11):139–144, 2020. Uncertainty Quantification, 6(2):645–682, 2018.
435 Grathwohl, W., Chen, R. T., Bettencourt, J., Sutskever, I., Peherstorfer, B. and Marzouk, Y. A transport-based mul-
436 and Duvenaud, D. Ffjord: Free-form continuous dy- tifidelity preconditioner for markov chain monte carlo.
437 namics for scalable reversible generative models. arXiv Advances in Computational Mathematics, 45(5):2321–
438 preprint arXiv:1810.01367, 2018. 2348, 2019.
439
8
 Submission and Formatting Instructions for ICML 2024

440 Poli, M., Massaroli, S., Park, J., Yamashita, A., Asama, H.,
441 and Park, J. Graph neural ordinary differential equations.
442 arXiv preprint arXiv:1911.07532, 2019.
443
444 Poli, M., Massaroli, S., Yamashita, A., Asama, H., and Park,
445 J. Hypersolvers: Toward fast continuous-depth models.
446 Advances in Neural Information Processing Systems, 33:
447 21105–21117, 2020.
448 Pontryagin, L., Boltyanskii, V., Gamkrelidze, R., and
449 Mishchenko, E. Mathematical theory of optimal pro-
450 cesses [in russian], 1961.
451
452 PyTorch. Pytorch documentation - functions, 2023. URL
453 [Link]
454 html. Accessed: 2023-05-16.
455
456 Radford, A., Metz, L., and Chintala, S. Unsupervised rep-
457 resentation learning with deep convolutional generative
458 adversarial networks. arXiv preprint arXiv:1511.06434,
459 2015.
460 Ramesh, A., Pavlov, M., Goh, G., Gray, S., Voss, C., Rad-
461 ford, A., Chen, M., and Sutskever, I. Zero-shot text-
462 to-image generation. In International Conference on
463 Machine Learning, pp. 8821–8831. PMLR, 2021.
464
465 Ross, I. M. A primer on Pontryagin’s principle in optimal
466 control. Collegiate publishers, 2015.
467
468 Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and
469 Ganguli, S. Deep unsupervised learning using nonequi-
470 librium thermodynamics. In International Conference on
471 Machine Learning, pp. 2256–2265. PMLR, 2015.
472 Sohn, K., Lee, H., and Yan, X. Learning structured output
473 representation using deep conditional generative models.
474 Advances in neural information processing systems, 28,
475 2015.
476
477 Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones,
478 L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. At-
479 tention is all you need. Advances in neural information
480 processing systems, 30, 2017.
481
482 Wanner, G. and Hairer, E. Solving ordinary differential
483 equations II, volume 375. Springer Berlin Heidelberg
484 New York, 1996.
485 Zhang, H., Gao, X., Unterman, J., and Arodz, T. Approxi-
486 mation capabilities of neural odes and invertible residual
487 networks. In International Conference on Machine Learn-
488 ing, pp. 11086–11095. PMLR, 2020.
489
490
491
492
493
494
9
 Submission and Formatting Instructions for ICML 2024

495 A. You can have an appendix here.

496
497 You can have as much text here as you want. The main body must be at most 8 pages long. For the final version, one more
498 page can be added. If you want, you can use an appendix like this one.
499 The \onecolumn command above can be kept in place if you prefer a one-column appendix, or can be removed if you
500 prefer a two-column appendix. Apart from this possible change, the style (font size, spacing, margins, page numbering, etc.)
501 should be kept the same as the main body.
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
10