Differentiable Multiple Shooting Layers
Authors: Stefano Massaroli
Michael Poli
Sho Sonoda
Taiji Suzuki
Jinkyoo Park
Atsushi Yamashita
Hajime Asama
Venue: NeurIPS 2021 Poster (2021)
Keywords: continuous-time neural models
implicit deep learning
numerical methods
optimal control
time series

Abstract: We detail a novel class of implicit neural models. Leveraging time-parallel methods for differential equations, Multiple Shooting Layers  (MSLs) seek solutions of initial value problems via parallelizable root-finding algorithms. MSLs broadly serve as drop-in replacements for neural ordinary differential equations  (Neural ODEs) with improved efficiency in number of function evaluations (NFEs) and wall-clock inference time. We develop the algorithmic framework of MSLs, analyzing the different choices of solution methods from a theoretical and computational perspective. MSLs are showcased in long horizon optimal control of ODEs and PDEs and as latent models for sequence generation. Finally, we investigate the speedups obtained through application of MSL inference in neural controlled differential equations (Neural CDEs) for time series classification of medical data.

Differentiable Multiple Shooting Layers
Stefano Massaroli'
Michael Poli"
The University of Tokyo, DiffEqML
KAIST, DiffEqML poli_m@kaing ac kr
caasaroli@rooo.coo t. w.w.k.koot.teaeee
Sho Sonoda RIKEN
Taiji Suzuki
The University Tokyo, RIKEN
Jinkyoo Park KAIST
A Atsuhii Yamashita
Hajime Asama
The University of Tokyo
The University of Tokyo
Abstract
We detail a novel class of implicit neural models. Leveraging time-parallel methods for differential equations, Multiple Shooting Layers (MSLs) seek solutions of initial value problems via parallelizable root-finding algorithms. MSLs broadly serve as drop-in replacements for neural ordinary differential equations (Neural ODEs) with improved efficiency in number of function evaluations (NFEs) and wall- clock inference time. We develop the algorithmic framework of MSLs, analyzing the different choices of solution methods from theoretical and computational perspective. MSLs are showcased in long horizon optimal control of ODEs and PDEs and as latent models for sequence generation. Finally, we investigate the speedups obtained through application of MSL. inference neural controlled differential equations (Neural CDEs) for time series classification of medical data.
1
Introduction
For the last twenty years, has tried speed up numerical computation mainly by providing faster computers. Today, as appears that one getting closer to the maximal speed of electronic components, emphasis is put on allowing operations to be performed in parallel. In the near future, much of numerical analysis will have to be recast in more "porallel" form. Nievergelt, 1964
Discovering and exploiting parallelization opportu- nities has allowed deep learning methods to succeed across application areas, reducing iteration times for architecture search and allowing scaling to larger data sizes (Krizhevsky al., 2012; Diamos al., 2016; Vaswani et al., 2017). Inspired by multiple shooting, time-parallel methods for ODEs (Bock and Plitt. 1984; Diehl et al., 2006; Gander, 2015; Staff and Rønquist, 2005) and recent advances on the intersection of differential equations, implicit problems and deep learning, we present novel class of neural models designed to maximize paralleliza- tion across time: differentiable Multiple Shooting Layers (MSLs). MSLs seek solutions of initial value problems (IVPs) as roots of a function designed to ensure satisfaction of boundary constraints. Figure provides visual intuition the parallel nature of MSL inference.
Neural ODE
Multiple Shooting Layer by
by o
RootFind
by
Figure 1: MSLs apply parallelizable root finding
methods to obtain differential equation solutions.
Eqqual contribution. Author order was decided by flipping coin.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

An implicit neural differential equation MSL inference is built on the interplay of numerical methods for root finding problems and differential equations. This property reveals the proposed method as a missing link between implicit-depti architectures such as Deep Equilibrium Newtorks (DEQs) (Bai al., 2019) and continuous-depth models (Weinan, 2017; Chen et al., 2018; Massaroli et al., 2020; Kidger et al., 2020b; Li et al., 2020). Indeed, MSLs can be broadly applied as drop-in replacements for Neural ODEs, with the advantage of often requiring a smaller number of function evaluations (NFEs) for neural networks parametrizing the vector field. MSL variants and their computational signature are taxonomized on the basis of the particular solution algorithm employed. such Newton and parareal (Maday and Turinici, 2002) methods.
Faster inference and fixed point tracking Differently from classical multiple shooting methods, MSLs operate in regimes where function evaluations of the vector field be significantly more expensive than surrounding operations. For this reason, the reduction in NFEs obtained through time-parallelization leads significant inference speedups. In full-batch training regimes, MSLs provably enable tracking of fixed points across training iterations, leading to drastic acceleration of forward passes (often the cost of single root finding step). We apply the tracking technique to optimal control of ODEs and PDEs, with speedups in the order of several times over Neural ODEs. MSLs are further evaluated in sequence generation via latent variant, and as a faster alternative to neural controlled differential equations (Neural CDE) (Kidger et al., 2020b) in long-horizz time series classification.
2 Multiple Shooting Layers
Consider the initial value problem (IVP)
(t) = fo(t, =(t)
L LE [0,
(2.1)
2 200 2g
with state E Z C R parameters 8 W for space W of functions [0, -> R and
smooth vector field fo [0,T] x 2 x W - Z. For all 2 E Z, s, E [0, T]: < t we denote with (2(2,,, t, the solution of (2.1) at time t starting from time i.e. 0n s, (2, s, z(t).
The crux behind multiple-shoote
methods for differential equations to turn the initial value problem (2.1) into a boundary value problem (BVP) We split the the time interval [0, () in N sub-intervals [tn, +1] with 0 = to < t1 < ty = T and define left boundary subproblems
zn b, and (t) (t, 2n, (t)), E [L.z, In-
(2.2)
where b, are denoted as shooting parameters. At each time / E [o, T]. the solution of (2.2) matches the one of (2.1) iff all the shooting parameters b, are identical to =(fn), bn do (zo, to, ln). Using =ffnn) = do (fo(2(1n11)))) bn fn). we obtain the equivalent conditions
bo (0 la, lo) 20
b, de (bo, to 11) 20 20(t1)
by 20 (bx IN,,, = nN 2
Let B = (bo. bj. ... by) and 70 B, 20) == (zo. de thus turn the IVP (2.1) 1) into the roots-finding finding problem the of function go defined
(e(( (bo. to t1). 2e (bN tx))). We can
(B, 2o) = B No (B, z0)
Definition (Multiple Shooting Layer (MSL)). With L. X + Z and ly N+1 + y nwo affine maps, multiple shooting layer is defined as the implicit input-output mapping + y:
20 (x(
B go go(B" B", =0)
(2.3)
(B()))
2

---

Realization of Multiple Shooting Layers
The remarkable property of MSL is the possibility of computing the solutions of all the N IVPs (2.2) in parallel from the shooting parameters in B with any standard ODE solver. This allows for drastic reduction in the number of vector field evaluations at the cost of higher memory requirement for the parallelization take place. Nonetheless, the forward pass of MSLs requires the shooting parameters B to satisfy the nonlinear algebraic matching condition go B, =o) 0, which has also to be solved numerically.
3.1 Forward Model
The forward MSL model involves the synergistic combination of two main classes of numerical methods: ODE solvers and root finding algorithms, to compute "oo B, 20) and B", respectively. There exists a hierarchy between the two classes of methods: the ODE solver will be invoked at each step k of the root finding algorithm compute "io (*, 20) and evaluate the matching condition go (B*, 20).
Newton methods for root finding Let us denote with B the solution of the root finding problem
at the k-th step of the Newton method and let Dge B* z0) be the Jacobian of ge computed in B* The solution 8" ge(B", 20) 0 can be obtained by iterating the Newton Nowton-Raphson fixed point iteration
Bk+1 = B* a [IN @ hn - (y 20)] [B* - -70(B" 20)]
(3.1)
which converges quadratically to B" (Nocedal and Wright, 2006). The exact Newton iteration theoretically (3.1) requires the inverse of the Jacobian ly @ l,, Doj ( (B+,2))). Without the special structure of the MSL problem, the Jacobian would have had be the computed in full, in the case of DEQs (Bai et al., 2019). Being the Jacobian of dimension Nn, x xNn: its computation with reverse-mode automatic differentiation (AD) tools scales poorly with state dimensions and number of shooting parameter (cubically in both n, and N).
Instead, the special structure of the MSL matching function ge (B, 20) B - 70(B, (o(B,)o) zo) and its Jacobian, opens up application of direct updates where inversion is not required.
Direct multiple shooting Following the treatment of Chartier and Philippe (1993), we can obtain direct formulation of the Newton iteration which does not require the composition of the whole Jacobian nor its inversion. The direct multiple shooting iteration is derived by setting = 1 and multiplying the Jacobian on both sides of (3.1) yielding
[On lng (B* =0)] (B*+1 - - (B* z0) B*
which leads the following update rule for the individual shooting parameters by (see Fig. 2):
g++1 00 (b%) ) + Deo (b+t) (6/+1 bt)) 6
zg
(3.2)
where D (bt) ) = do (bt:) (db) the sensitivity of each individual flow to its initial condition.
Due to the dependence of k+1 on k++1 complete Newton iteration theoretically requires N sequential stages.
Finite-step convergence Iteration (3.2) exhibits convergence to the exact solution of the IVP (2.1) in N steps (Gander, 2018, Theorem 2.3). In par- ticular, given perfect integration of the sub-IVPs. 6% coincides with the exact solution @0(20, to, (x) from iteration index & = onward, i.e. at iteration & only the last N shooting parameters are actually up- dated. Thus, the computational and memory footprint of the method diminishes with the number of iterations. representation of iteration (3.2) in Figure 3 while further details are discussed in Appendix B.I.
Single Stage of Newton Iteration
00.m
D
b++1
Figure 2: One stage Newton iteration (3.2).
This result can be visualized in the graphical
3

---

Time-Iteration Propagation of Newton Scheme
Numerical implementation Practical imple- mentation of the Newton iteration (3.2) re- quires an ODE solver to approximate the flows Φ (bA) and an algorithm to compute their sen- sitivities w.r.t. by Besides direct application of AD, we show efficient alternative to obtain D in parallel alongside the flows, with single call of the ODE solver.
7
8
8
B
k
b
20
B
20
B
Efficient exact sensitivities Differentiating through the steps of the forward numerical ODE solver using reverse-mode AD straight- forward, but incurs in high memory cost, additional computation to unroll the solver steps and introduces further numerical error on Døg Even though the memory footprint might be mitigated by applying the adjoint method (Pontryagin et al., 1962), this still re- quires solve backward the N - le adjoint ODEs and sub-IVPs (2.2), at each iteration k. We leverage forward sensitivity analysis to compute Døa, n alongside Φ8,n in single call of the ODE solver. This approach, which might be considered as the continuous variant of forward- mode AD, scales quadratically with ny, has low memory cost, and explicitly controls numerical error.
B Figure 3: Propagation & n of the Newton iteration (3.2). The intertwining between the updates in r. and & leads to the finite step convergence result. In fact, by setting 8g 20, see how the correcting term b 6% multiplying the flow sensitivity D progressively nullifies at the same rate n and k. As a result, the exact sequential solution of the IVP (2.1) unfolds on the diagonal & = and the only active part of the algorithm is the above the diagonal (highlighted in yellow).
Proposition 1 (Forward Sensitivity (Khalil, 2002)). Let (0(2, be the solution of (2.1). Then, v(l) = (2 8, I) satisfies the linear matrix-valued differential equations
itt) Dfo(t, (t))o(t, v(s)
ln_ where Dfo denotes ff/
Therefore, at iteration 1: D (b)) can be computed in parallel while performing the forward integration of the N - k IVPs (2.2) and their forward sensitivities, i.c
ForwardSensitivityy {b (6 Déo REDEN
which enables full vectorization of Jacobian matrix products between 0fe/Oz and v well as maximizing re-utilization of vector field evaluations. Detailed derivations are provided in Appendix B.2. Appendix C.1 analyzes practical considerations and software implementation of the algorithm.
Zero-order approximate iteration In high state dimension regimes, the quadratic memory scaling of the forward sensitivity method might be infeasible. If this is the case, zero zero-orda order approximation of the Newton iteration preserving the finite-step converge property can employed: the parareal method Lions et al. (2001). From the Taylor expansion of 60,- (bx+1) around by;
(bx+t) = Φe (b,.) Dé (b*) (b++1 bt)) + (116 b%
we have the following approximant for the correction term of (3.2)
Dés.n (b,,) (ba+1 bt:) ~ (o, - Φe.n (btt)
(3.3)
Parareal computes the RHS of (3.3) by coarse? numerical solutions woon (b%).. (bhht!) 60,n (b%), (b, leading the forward iteration,
of
b++ (b)) + the (b(+1) vo.n (V%))
Pe.g. few steps of a loweorder ODE solver

---

Forward
Backward
B
sensitivty
reverse autodiff.
fx
B
Lo
V Lo
heroth-order
sequential adjoint
Figure 4: Scheme of the
pass of MSLs. After applying the input map la to the input and choosing initial shooting parameters B", the forward pass is iteratively computed with one of the numerical schemes described in Sec. 3.1 which, in turn, makes use of some ODE solver to compute du., in parallel, each step. Once the output y and the loss computed computed by applying and Lg B", the loss gradients can computed by standard adjoint methods reverse mode automatic differentiation.
3.2 Properties of MSLs
Differentiating through MSL Computing loss gradients through MSLs can be performed by directly back-propagating through the steps of the numerical solver via reverse-mode AD.
A memory efficient alternative is to apply the sequential adjoint method to the underlying Neural ODE. In particular, consider a loss functions computed independently with the values of different shooting parameters, L(x, 8 0) = Σ co(x, b;)) The adjoint gradient for the MSL is then given by
VOl. = L/ (T
where the Lagrange multiplier X(t) satisfies backward piecewise soncewise-contiuuuus linear ODE
X(t) (t))A(()
if (tn.tat)) X(T) (I,,b(x,b)))
A (t,,, = A(t,,)))))))))() + V co(2,, b,.)
The adjoint method typically requires the IVP (2.1) be solved backward alongside À to retrieve the value of z()) needed to compute the Jacobians Df and Vale. This step introduces additional errors on the final gradients: numerical errors accumulated on b% = x(T) during forward pass, propagate to the gradients and sum up with errors on the backward integration of (2.1).
Here take different, more robust direction by interpolating the shooting parameters and drop the integration of (2.1) () during the backward pass. The values of the shooting parameters retrieved by the forward pass of MSLs solution points of the IVP (2.1), i.e. by o( 20, fo, k,,)) (up the forward numerical solver tolerances). On this assumption, we construct a cubic spl ine interpolation {t) of the shooting parameters 6", and we query during the integration of À compute the Jacobians of fa. Further results on back-propagation of MSLs are provided in Appendix B.3. Appendix C.4 practical aspects of the backward model alongside software implementation of the interpolated adjoint. The entire scheme of forward backward pass of an MSL is shown in Figure 4.
One-step inference: fixed point tracking Consider training MSL. minimize twice- differentiable loss function L(x, B" 0) with Lipschitz constant m2 through the gradient descent iteration
0p+1 0p Trp VOL(I, B,, 0p)
where 7p is a positive learning rate and B is the exact root of the matching function go computed with parameters Op (i.e. the exact solution of the IVP 12.1) at the boundary points). Due to Lipschitzness of L, we have the following uniform bound on the variation the parameters across training iterations
10 - 0p l S mp
we also assume "e to Lipschitz continuous w.r.t = and with constants m m and differentiable w.r.t. 0 we can obtain the variation of the fixed point 3 to small changes in the model parameters by linearizing solutions around 0p
Tne, (B =o)
B B = 10p+1 - 0p]
+ o(110p+1 0,.
00
to obtain the uniform bound
|B Bill 5 mp f o(110 12))
5

---

Solve IVP
Sensitivities D
Compute Step
2n falt, 2nd,
2n (fm)) b
Up (fm 1
in
(t Zn Un,
Up (tn) =
B
(1+ jap(nxt)
MFE
D/ a
O(N A)
B
Soll, zn):
Zn (ln) b%
do (b
Iv sensitivty zeroth-order
MFE
XFE
Figure 5: Single iteration computational span (McCool et al., 2012) in MSL. We normalize to the cost of evaluating fo. The N - & sub IVPs solved in parallel in their sub intervals, thus requiring minimum span NFE,, Forward sensitivity introduces Jacobian matrix products costs amounting jmp, which can be further parallelized into jvps.
Having bounded variation the solutions of the MSL for small changes of the model parameters 8, we might think of recycling the previous shooting parameters B as an initial guess for the direct Newton algorithms in the forward pass succeeding the gradient descent update. We show that by choosing sufficiently small learning rate dp one Newton iteration can be sufficient to track the true value of B" during training. In partieular, the following bound can be obtained.
Theorem (Quadratic fixed-point tracking). If Se is twice continuously differentiable in z then
| - S Mn
(3.4)
for some M > O. Br is the result of one Newton iteration applied to B
The proof, reported in Appendix A.I, relies on the quadratic converge of Newton method. The quadratic dependence of the tracking error bound Yl allows use of typical learning rates for standard gradient based optimizers to keep the error under control. this way, we can turn the implicit learning problem into explicit one where the implicit inference pass reduces to one Newton iteration. This approach leads to the following training dynamics:
0 - 0 VLL ,) B" a3.2y((3.2), B*)
We note that the main limitation of this method is the assumption on input x to be constant across training iterations (i.c. the initial condition zo is constant as well). If the input changes during the training (e.g. under mini-batch SGD training regime), the solutions of the IVP (2.1) and thus its corresponding shooting parameters may drastically change with even for small learning rates.
Numerical scaling Each class of MSL outlined in Section 3.1 is equipped with unique computa- tional scaling properties. In the following, we denote with NFE, the total number of vector field so evaluations done, in parallel across shooting parameters, a single sub-interval [l,, /n]] Similarly, NFE indicates the function evaluations required by the coarse solver used for parareal approxima- tions. Here, we set out to investigate the computational signature of MSLs as parallel algorithms. To this end, we decompose single MSL iteration into two core steps: solving the IVPs across sub-intervals and computing sensitivities DØ (or their approximation). Figure 5 provides a summary of the algorithmic span' of a single MSL iteration as a function of number of Jacobian vector products (jvp), Jacobian matrix products (jmp), and vector field evaluations (NFE).
Fu sensitivity MSL. frontloads the cost of computing Déoo., by solving the forward sensitivity ODEs of Proposition alongside the evaluation of 7/0- Forward sensitivity equations involve a jmp, which can be optionally further parallelized n. jvps by paying memory overhead. Once sensitivities have been obtained along with 79. no additional computation needs to take place until application of the shooting parameter update formula. The forward sensitivity approach thus enjoys the highest degree of time-paralle time-parlllelizabiity at a larger memory cost.
Zeroth-orde MSL computes T via total of NFE, evaluations fe. parallelized across nub-intervals. The cheaper IVP solution in both memory and compute is however counterbalanced during calculation
'Longest sequential cost, in terms of computational primitives, that is not parallelizable due to problem- specific dependencies. A specific example for MSLs are the sequential NFE, calls required by the sequential ODE solver for each mub-interval.

---

of the sensitivities, as this MSL approach approximates the sensitivities Dōs,, by a saroth-order update requiring N k sequential calls coarse solver.
The analysis of MSL. backpropagation sealing is straightforward, as sequential adjoints for MSLs mirror standard sensitivity techniques for Neural ODEs in both compute and memory footprints. Alternativeyy, AD can be utilized to backpropagate through the operations of the forward pass methods in use. This approach introduces a non-constant memory footprint which scales in the number of forward iterations and thus depth of computational graph.
4 Applications
4.1 Variational MSL
Let R R",, be an observable of some continuous- time process and let X {xr M,, To. KO....,,NN E R(M+++1)x*% be a sequence of observations of xr(() at time instants L-M
< (.y. We seek a model able to predict T1,,.. , .N given past observations
I E-M...... Tu, equivalent to approximating the conditional distribution p(Z1.N (x--M.)). To this end we introduce variational MSLs (uMSLs) as the following latent variable model:
(1, = (z=
Encoder &w Approx. Posterior
B' : go (B', (B", 20)
Decoder Da
(200/==-MM0 = NN(,,)) 2)
21...., IN {(B")
Readout l
2g qu (20l x M.0)
Reparametrization
Once trained, such model can be also used generate new realistic sequences of the observable a=()) by querying the decoder network at a desired 2u- vMSLs are designed to scale data generation to longer sequences, exploiting wherever possible parallel computation in time in both encoder well as decoder modules. The structure of & is designed to leverage modern advances in representation learning for time-series via temporal convolutions (TCNs) attention operators (Vaswani al., 2017) to offer a higher degree of parallelizability compared to sequential encoders e.g RNNs, ODE- RNNS (Rubanova et al., 2019) Neural CDEs (Kidger al., 2020b). This, in turn, allows the encoder to match the decoder in efficiency, avoiding unnecessary bottlenecks. The decoder Da is composed of a MSL which is tasked unroll the generated trajectory in latent space. vMSLs are trained via traditional likelihood methods. The iterative optimization problem can be cast as the maximization of an evidence lower bound (ELBO):
min Eso M.a) Σ log (L(g.||W(O.0) 1))
(0,w)
with Pr. (f..) N (#n, σn) and the standard deviations on are left as a hyperparameters.
Sequence generation under noisy observations We apply vMSLs on a sequence generation task given tra- jectories corrupted by correlaore noise. We con- sider a baseline latent model sharing the same over- all architecture vMSLs, with Neural ODE as de- coder. In particular, the Neural ODE is solved via the dopri5 solver with absolute and relative tolerances set to 10 whereas the vMSL decoder is an instance of fw sensitivity MSL. The encoder for both models is comprised of two layers of temporal convolutions (TCNs). All decoders unroll the trajectories directly in output space without additional readout networks. The validation on sample quality then performed by in- spection of the learned vector fields and the error against the nominal across the entire state-space. The proposed model obtains equivalent results as the baseline at a significantly cheaper computational cost. As shown in Figure 6, eMSLs require 60% less NFEs for a single training iteration as well as for sample generation, achieving results comparable to standard Latent Neural ODEs (Rubanova et al., 2019). We report further details and results Appendix E.
Variational MSLs & Latent ODEs
NFE Analysis
50
40
NFF
30
20
0
500
1,000 Training Step
uMSL - - Latent ODE Figure 6: Mean and standard deviation of NFEs during vMSL and Latent Neural ODE across training trials. vMSLs require 60% less NFEs during both training and inference.
7

---

4.2 Neural Optimal Control
Neural Optimal Policy via MSLs Vector Field
Learned Controller
Beyond sequence generation, the proposed framework can be applied to optimal control. Here we can fully exploit the drastic computa- tional advantages of MSLs. In fact, we leverage on the natural assumption of finiteness of initial conditions where the controlled system (or plant) is initialized verify the result of Th. Let us consider controlled dynamical system
(T
a(1) [(t, ((), Te(L, 2)), z(0) = 20 (4.1)
2
with parametrized policy I,, - xo (1, and initial conditions 2g ranging finite set Zo = {= We consider the problems of stabilizing low-dimensional limit cycle and deriving an optimal boundary policy for linear 300-dimeniional ODE.
0
-2
PDE discretized
Figure 7: [Left] Closed-loop vector fields and trajec- tories corresponding to the ug-controlled MSI.. [Right] Learned controller u=(=) (z E R²) for the two different desired limit cycles. Although, the inference of the all trajectory performed with just steps of RK4 (8 NFE). the initial accuracy dopri5 (> 3000 NFE) is preserved throughout training.
Limit cycle stabilization We consider a sta- bilization task where, given desired closed curve S = {= E s2(2) = 0}, we mini- mize the 1-norm between the given curve and the MSL solution of (4.1) across the timestamp- well as the control effort [Tel- We verify the approach system, aiming with closed curves Sa of various shapes. Following the assumptions on fixed point tracking and slow woww-varynn--foee flows of Th. we initialize BD by dopri5 adaptive-step solver with tolerances 10 Then, at each training iteration, we perform inference with single parallel rk4 step for each sub-interval [t,, 11+1] followed by single Newton update. Figure 7 shows the learned vector fields and controller, confirming successful system stabilization of the system to different types of closed curves. We compare with range of baseline Neural ODEs, solved via rk4 and dopri5. Training of the controller via MSL achieved with orders of magnitude less wall-clock time and NFEs. Figure 8 shows the difference in NFEs w.r.t. dopri5 while Fig. compares wall-clock times per training iteration of MSL, dopri5 and rk4.
degree of freedom mechanical
We further provide Symmetric Mean Average Percentage Error (SMAPE) measurements be- tween trajectories obtained via MSLs and an adaptive-Ste step solver. MSLs initialized with re- cycled solutions are able to track the nominal trajectories across the entire training process. Additional details on the experimental setup, including wall-clock time comparisons with rk4 and dopri5 bascline Neural ODEs is are provided in Appendix E.2.
Neural Optimal Control via MSLs
-10-4
SMAPE of MSL
0.5
500
1,000
1,500
2,000
2,500
Training Step
NFE Comparison aguinst Neural ODE
NPB
102
MSL [2 steps of RK4]
Neural Boundary Control of the Timo- shenko Beam We further show how MSLs can be scaled to high-dimensional regimes by tackling the optimal boundary control prob- lem for linear partial differential equation. In particular, we consider the Timoshenko beam (Macchelli and Melchiorri, 2004) model in Hamiltonian form. We derive formalize the boundary control problem and obtain structure structure-preserving spectral discretization yielding a 660-iimesio Hamiltonian ODE.
Neural ODE [dopris]
10
500
1,000
1,500
2,000
2,500
Training Step
Figure 8: Symmetric Mean Average Percentage Error (SMAPE) between solutions of the controlled systems obtained by MSLs and nominal. Compared to Neural ODEs, MSLs solve the optimal control problem with NFE savings of several orders of magnitude by carrying forward their solution across training iterations.
We parameterize the boundary control policy with multi-layer perceptron taking input (control feedback) the dimennnnall discrete state. We train the model in similar setting to the previous example having the MSL equipped with fw sensitivity and one step of rk4 for the parallel

---

2D Limit Cycle (Circle) Control
Timoshenko Beam
wall-clock time [8]
10"
wall-clock time
1.5
0.5
10-1
10
10
10a
10
10a
10a
training iteration
training iteration
MSL
dopri5 rk4
MSL rk4
Figure 9: Mean and standard deviation of wall- clock time of complete training iteration (for- ward/backward passes GD update) for differ- ent solvers on the circle experiment.
Figure 10: Mean and standard deviation of wall-clock time per training iteration for MSL and Neural ODE across training trials. MSLs require are three times faster than the sequential rk4 with same accuracy (step size).
integration. We compare the training wall-clock time with Neural ODE solved by sequential rk4 in Fig. 10. The resulting speed up of MSL with forward sensitivity is three time faster than the baseline Neural ODE proving that the proposed method able to scale to high-dimensional regimes. We include a formal treatment of the boundary control problem in Appendix D while further experimental details are provided in Appendix E.3.
4.3 Fast Neural CDEs for Time Series Classification
Training Curve sepsis Pred.s
To further verify the broad range of applicability of MSLs, apply them to time series classification as faster alternatives to neural controlled differen- rial equations (Neural CDEs) (Kidger et al., 2020b).
AUROC
0.8
0.6
MSL
rk4
0.4
Here, MSLs remain applicable since Neural CDEs practically solved ODEs with special struc- ture, as described in Appendix E.4. We tackle the PhysioNet 2019 challenge (Goldberger et al., 2000) on sepsis prediction, following the exact experi- mental procedure described by Kidger al. (2020b),
200
400
600
800
1,000
wall-clock time [s]
Figure 11: Mean and standard deviation of AUROC during training MSLs and baseline Neural CDEs on sepsis prediction.
including hyperparameters and Neural CDE architectures. However, we train all models on the full dataset enable application of the fixed point tracking technique for MSLs Figure 11 visualized training convergence of zeroth- order MSL Neural CDEs and the baseline Neural CDE solved with rk4 as in the original paper. Everything else being equal, including architecture and backpropagation sequential adjoints, MSL Neural CDEs converge with total wall-clock time one order of magnitude smaller than the baseline.
Related Work
Parallel-in-time integration & multiple shooting MSLs belong to the framework of time- parallel integration algorithms. The study these methods relatively recent, with seminal work in e.g.
the 60s (Nievergelt, 1964). The multiple shooting formulation of time-parallel integration, (Bellen and Zennaro, 1989) or (Chartier and Philippe, 1993), finally lead to the modern algorithmic form using the Newton iteration reported in (3.2). Parareal (Lions al., 2001) has been successively introduced as a cheaper approximated solution of multiple shooting problem, rapidly spreading across application domains, e.g. optimal control of partial differential equations (Maday and Turinici, 2002). We refer to (Gander, 2015, 2018) as excellent introduction to the topic. We also note recent work (Vialard et al., 2020) introducing single shooting terminology for Neural ODEs (Chen al., 2018), albeit in the unrelated context of learning time-varying parameters.
WWe note that the training for all models has been performed on single NVIDIA RTX A6000 with 48Gb of GPU memory.

---

Time parrraaaaallliiaiioon in neural models In the pursuit for increased efficiency, several works have proposed approaches parallelization across time or depth in neural networks. (Gunther al., 2020; Kirby et al., 2020; Sun et al., 2020) use multigrid and penalty methods to achieve speedups in ResNets. Meng al. (2020) proposed a parareal variant Physics-informet neural networks (PINNs) for PDEs. Zhuang et al. (2021) uses a penalty-variant of multiple shooting with adjoint sensitivity for parameter estimation in the medical domain. Solving the boundary value problems with regularization term, however, is not guaranteed to converge continuous solution. The method of Zhuang et al. (2021) further optimizes its parameters of (1) achieves drastic speedups while preserving convergence guarantees. Recent theoretical work (Lorin, 2020) has applied parareal methods to Neural ODEs. However, their analysis is limited to the theoretical computational complexity setting and does not involve multiple shooting methods nor derives its implicit differentiation formula.
full-batch regime, where application
In contrast our objective is to introduce novel class of implicit time-parallel models, and to validate their computational efficiency across settings.
6 Conclusion
This work introduces differentiable Multiple Shooting Layers (MSLs), a parallel-in-time alternative neural differential equations. MSLs seek solutions of differential equations via parallel application of root finding methods across solution subintervals. Here, we analyze several model variants, further proving a fixed point tracking property that introduces drastic speedups in full-batch training. The proposed approach is validated on different tasks: as generative models, MSLs are shown to achieve same task performance as Neural ODE baselines with 60% less NFEs, whereas they are shown to offer several orders of magnitude faster in optimal control tasks.
Remarkably few methods have been proposed for parallel integration of ODEs. In part this is because the problems do not have much natural parallelism. (Gear, 1988)
Funding Statement
This work financially supported by The University of Tokyo, KAIST and RIKEN AIP. All experiments were run on GPUs provided by The University of Tokyo and KAIST.
References
M. Alnæs, J. Blechta, J. Hake, A. Johansson, B. Kehlet, A. Logg. C. Richardson, J. Ring, M. E. Rognes, and G. N. Wells. The fenics project version 1.5. Archive of Numerical Software, 3(100), 2015.
S. Bai, J. Z. Kolter, and V. Koltun. Deep equilibrium models. In Advances in Neural Information Processing Systems, pages 690-701, 2019.
A. Bellen and M. Zennaro. Parallel algorithms for initial-value problems for difference and differential equations. Journal of Computational and applied mathematics, 25(3):341-350 1989.
H. G. Bock and K.-J Plitt. A multiple shooting algorithm for direct solution of optimal control problems. IFAC Proceedings Volumes, 17(2):1603-166
1984.
J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, VanderPlas, Wanderman-
and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: httt::///ithull //ggtttttb.... google/g/jeaae
C. G. Broyden. A class of methods for solving nonlinear simultaneous equations. Mathematics of computation, 19(92):577-593, 1965.
P. Chartier and B. Philippe. A parallel shooting technique for solving dissipative ode's. Computing. 51(3-4):209-255 1993.
R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud. Neural ordinary differential equations. In S. Bengio, H. Wallach, Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 31, pages 6571-6583. Curran Associates, Inc., 2018. URL https: http:///e//rogooppe//eoeoo//ee http://w.....riis.aaa/////////// cc/ppaper/201//filee 69386f6bb1 Ldw. Paper pdf.
10

---

G. D. Clifford, I. Silva, B. Moody, Li, D. Kella, A. Shahin, T. Kooistra, D. Perry, and R. G. Mark. The physionet/ompputng in cardiology challenge 2015: reducing false arrhythmia alarms in the icu. In 2015 Computing in Cardiology Conference (CinC), pages 273-276. IEER, 2015.
G. Diamos, S. Sengupta, B. Catanzaro, M. Chrzanowski, A. Coates, E. Elsen, J. Engel, A. Hannun, and Satheesh. Persistent mns:: Stashing weights on chip. 2016.
M. Diehl, Bock, H. Diedam, and P.-B. Wieber. Fast direct multiple shooting algorithms for optimal robot control. In Fast motions in biomechanics and robotics, pages 65-93. Springer, 2006.
M. J. Gander. 50 years of time parallel time integration. In Multiple shooting and time domain decomposition methods, pages 69-113. Springer, 2015.
M. Gander. Time parallel time integration. 2018.
C. W. Gear. Parallel methods for ordinary differential equations. Calcolo, 25(1-2):1 1988.
A. Goldberger, A. Amaral, Glass, J. Hausdorff, Ivanov, R. Mark, J. E. Mietus, G. Moody, C.-K. Peng, and H. E. Stanley. Physiobank, physiotoolkit, and physionet: components of new research resource for complex physiologic signals. circulation, 101(23):e215 e220, 2000.
Gunther, L. Ruthotto, J. B. Schroder, E. C. Cyr, and N. R. Gauger. Layer-parallel training of deep residual neural networks. SIAM Journal on Mathematics of Data Science, 2(1):1-23, 2020.
J. Jia and A. R. Benson. Neural jump stochastic differential equations. 190.v:1905.1044 2019.
arXiv preprint
H. K. Khalil, Nonlinear systems, volume 3. Prentice Hall, 2002.
Kidger, T. Chen, and T. Lyons. hey, that's not an ode": Faster ode adjoints with 12 lines of code. arXiv preprint arXiv:2009. 2XXYv:2009009457, 2020a.
P. Kidger, J. Morrril,, J. Foster, and T. Lyons. Neural controlled differential equations for irregular time series. arXiv preprint arXiv:2005.0m 000.:20055.092.. 08926. 2020b.
D. P. Kingma and J. Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.698,, 2014.
A. Kirby, S. Samsi, M. Jones, A. Reuther, J. Kepner, and V. Gadepally. Layer-parallel training with gpu concurrency of deep residual neural networks via nonlinear multigrid. In 2020 IEEE High Performance Extreme Computing Conference (HPEC), pages 1-7. IEEE, 2020.
A. Krizhevsky, I. Sutskever, and G. F. Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25:1097-1105, 2012.
X. Li, T.-K. L. Wong, R. T. Chen, and D. Duvenaud, Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, pages 3870-3882. PMLR, 2020.
J. Lions, Y. Maday, and Turinici. Résolution d'edp par un schéma en temps pararéel. Comptes Rendus de I' l'Académie des Sciences-Series I-Mathematics, 332(():661-663) 2001.
E. Lorin. Derivation and analysis of parallel-in-time neural ordinary differential equations. Annals of Mathematics and Artificial Intelligence, 1035-199,,,, 2020.
I. Loshehilov and F. Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05100 2017.
A. Macchelli and C. Melchiorri. Modeling and control of the timoshenko beam. the distributed port hamiltonian approach. SIAM Journal on Control and Optimization, 43(2):743-767, 2004.
A. Macchelli, A. J. Van Der Schaft, and C. Melchiorri. Port hamiltonian formulation of infinite dimensional systems i. modeling. In 2004 43rd IEEE Conference Decision and Control (CDC)IEEE Cat. No. 04CH37601), volume 4, pages 3762-3767. IEEE, 2004.
Y. Maday and G. Turinici. A parareal in time procedure for the control of partial differential equations. Comptes Rendus Muthematique, 335(4))387-392 2002.
S. Massaroli, M. Poli, J. Park, A. Yamashita, and H. Asama. Dissecting neural odes. arXiv preprint arXiv:2002.08077 2020.
M. McCool, Reinders, and A. Robison. Structured parallel programming: patterns for efficient computation. Elsevier, 2012.

---

X. Meng, Z. Li, D. Zhang, and G. E. Karniadakis. Ppinn: Parareal informed neural network for umeedeppndentt pdes. Computer Methods in Applied Mechanics and Engineering, 370:1 13250, 2020.
Nievergelt. Parallel methods for integrating ordinary differential equations. Communications of the ACM, 7(12):731-733 1964.
J. Nocedal and S. Wright. Numerical optimization. Springer Science & Business Media, 2006.
A. Pal. Y. Ma, V. Shah, and C. Rackauckas, Opening the blackbox: Accelerating neural differential equations by regularizing internal solver heuristics. arXiv preprint arXiv:2105.
2021.
A. Paszke, S. Gross, Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, high-perfone -igh-pefformance deep learning library. arXiv 2019.
preprint 19Iiv:1912.011
M. Poli, S. Massaroli, A. Yamashita, H. Asama, and J. Park. Hypersolvers: Toward fast continuous- depth models. arXiv preprint arXiv:200700960. 2020a.
M. Poli, Massaroli, A. Yamashita, Asama, and J. Park. Torchdyn: A neural differential equations library. arXiv preprint arXiv:2009.0990 09346, 2020b.
L. S. Pontryagin, E. Mishchenko, V. Boltyanskii, and R. Gamkrelidze. The mathematical theory of optimal processes. 1962.
C. Rackauckas, M. Innes, Y. Ma, J. Bettencourt, L. White, and V. Dixit. Diffeqflux, jl-a julia library for neural differential equations. arXiv preprint arXiv: 2019.
Y. Rubanova, R. T. Q. Chen, and D. K. Duvenaud. Latent ordinary differential equations for arregularly-samplee time series. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché-Bue, E. Fox, and R. Garnett, editors, Advances Neural Information Processing Systems, volume 32, pages 5320-5330. Curran Associates, Inc., 2019. URL https: /ttprc/eeooo......///o// / www.eeeednnnn..neiiii..ee. cc/ Paper pdf.
L. N. Smith and N. Topin. Super-convergence: Very fast training of neural networks using large learning rates. In Artificial Intelligence and Machine Learning for Aolti- Operations Applications, volume 11006, page 1100612. International Society for Optics and Photonics, 2019.
G. A. Staff and E. M. Rønquist. Stability of the parareal algorithm. In Domain decomposition methods in science and engineering, pages 449 456. Springer, 2005.
Q. Sun, H. Dong, Z. Chen, W. Dian, J. Sun, Y. Sun, Z. Li, and B. Dong. Penalty and augmented la- grangian methods for layer-parallel training of residual networks. arXiv preprint arXi::2009.01462 2020.
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. arXiv preprint 1706.0::7000376::: 2017.
E-X. Vialard, R. Kwitt, S. Wei, and M. Niethammer. A shooting formulation of deep learning. Advances in Neural Information Processing Systems, 33, 2020.
E. Weinan. A proposal machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
Zhuang, N. Dvornek, S. Tatikonda, X. Papademetris, Ventola, and Duncan, Multiple-shooting adjoint method for whole-brain dynamic causal modeling. arXiv preprint arXiv:2 10/3, 2021.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] See Section 3.2 for numerical
scaling of the model, and the Appendix for additional information.
(c) Did you discuss any potential negative societal impacts of your work? [Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you ran experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, cloud provider)? [Yes] See the Appendix for details hardware used.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent obtained from people whose data you're you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "laTCYYm0t51", "reviewer": "Reviewer_cwaG", "summary": "This paper presents a plug-in replacement for vanilla neural ODEs. The methodology relies on the multiple shooting idea, which divides a long sequence into subsequences and solves them concurrently.  The shooting variables are computed by Newton's root finding method (which requires one step per subsequence), which is then approximated by the \"fixed point tracking\" idea. The experiments show that in certain cases, the resulting technique leads to a reduction in the number of function evaluations and wall-clock time while achieving the same level of performance as neural ODEs.", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work brings an old but important idea into today's relevant literature. To the best of my knowledge, Newton's method for root finding has never been considered in the context of ODEs - so this is indeed a nice contribution. Utilizing forward sensitivities in the forward pass is also an interesting idea, which could impact future research as adjoints have been much more popular than sensitivities so far. \n\nThat being said, I believe the paper has certain clarity issues. Also, more (ablation style) experiments are needed to support/illustrate the claims. Finally, (although I might be mistaken in saying that), I believe the forward sensitivity derivation is wrong and should be revisited. I would be happy to increase my score if the authors address these points. Below are more detailed comments.\n\n- The introduction section has clarity issues. I believe what the proposed framework does and why it is important are mixed. It would have been better to first clearly describe the problem, motivate why the problem is significant, and finally how the method addresses the described problem. \n- Is proposition 1 correct? A modern derivation of sensitivities is given in [1]. If you look at eq. (8), you will see the term $\\partial f/ \\partial\\theta$, which is missing in the submitted work (this term is missing because one needs to compute the total derivative whereas A.10 only computes the partial derivative by the chain rule.).\n- Since the ODE solutions are computed numerically, this might impact the behavior of Newton's method. This should be demonstrated experimentally.\n- Is the \"one-step inference\" method used in all experiments? If so, this should be clearly stated and a comparison between this approximation and the exact Newton solution should be given. Furthermore, in addition to the overall performance (demonstrated in 4.2), a simple experimental illustration showing the behavior of the bound is needed. \n- Why are the NODE solver tolerances so high in sequence generation tasks? Comparison with default (if not tuned) values would be better. Also, I am not sure how to interpret Figure 6. Although NFE drops, the Jacobian computation might take significant time, making the method slower compared to latent ODE.\n- The idea seems highly related to DEQs, which requires more explanation and possibly a conceptual comparison.\n- More explanation/motivation for the \"interpolated adjoints\" is needed.\n\nMinor comments:\n- Figure 1 can be improved. In particular, the \"RootFind\" and b's can be shortly explained.\n- Line 53: Is W a function of time? If so, this should be clearly explained. Otherwise, I'm having difficulty understanding the mapping $[0,T] \\rightarrow R^{n_\\theta}$. \n- Formally defining $\\phi_\\theta(z,s,t)$ would help to quickly understand.\n- l_x and l_y need clarification. If they refer to previous/next layers, a diagram would be better for illustrative purposes.\n- Does the cross symbol denote \n- Line 77: (3.1) before \"theoretically\".\n- In line 55, $\\phi$ has three arguments while it has only one in (3.2).\n- Showing the derivation of (3.2) explicitly would be nice.\n- Important references to sensitivities [2,3] are missing.\n- Double \"computed\" in Figure 4 caption.\n- Line 290: [CITE]\n\n[1] Fröhlich, Fabian, et al. \"Scalable parameter estimation for genome-scale biochemical reaction networks.\" PLoS computational biology 13.1 (2017): e1005331.\n[2] Kokotovic, Petar, and James Heller. \"Direct and adjoint sensitivity equations for parameter optimization.\" IEEE Transactions on Automatic Control 12.5 (1967): 609-610.\n[3] Heinonen, Markus, et al. \"Learning unknown ODE models with Gaussian processes.\" International Conference on Machine Learning. PMLR, 2018.\n\n**Post-rebuttal comment:** Thanks for the very detailed response. The authors nicely addressed almost all the points I raised. I will raise my score to 6 \"marginally above\" since I believe the manuscript still suffers from certain issues:\n- Reading the introduction, it is very difficult to clearly isolate the research problem, the gap in the literature, and the summary of the proposed approach.\n- The following sections do not clearly describe the research problem either. What is the problem/posterior/likelihood/optimization objective? - - Obviously, the approach advertises itself as a replacement for neural ODEs. This is indeed a nice idea; however, for an outsider that quickly looks at the paper, immediately seeing the addressed problem is not really possible.\n- Fig4 is definitely a much higher level explanation of the methodology and should precede all the figures except Fig1. So, I recommend moving the figure upwards and also refer it in the text.\n- As a final suggestion, some of the technical material can be moved to the appendix and more verbal explanations of the concepts can be inserted.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors addressed both the limitations and societal impact. However, I do believe the limitations should be experimentally illustrated. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "A8xuC1DAMAV", "reviewer": "Reviewer_1tje", "summary": "This paper proposes multiple shooting layers that solves a neural ODE dynamical system in parallel by reformulating IVPs as BVPs. The authors discussed the basic update rules and compared potential forward/backward modeling schemes, such as forward sensitivity, parareal method, and the potential recycling of fixed points (i.e., shooting parameters). The experiments show the benefits of the MSL formulation of ODE flows, which may significantly help the scalability of these (good but incredibly slow) implicit networks.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think the paper is clearly written and tackles an important issue (parallelization of ODE flows; bridging DEQs and NODEs). The idea of applying direct multiple shooting to model neural ODEs is also quite novel. To summarize:\n- Good presentation and clear intuition explanation. The paper analyzed and compared different forward/backward dynamics for training an MSL and provided both theoretical and empirical analysis for the efficiency. I checked most of the proof, which seem correct.\n- Applying the idea of direct multiple shooting to accelerate neural ODE solving seems novel to me, and bridges the equilibrium networks with the ODE-flow-based methods.\n- Relatively thorough set of experiments showing the advantage of MSLs. However, considering that neural ODEs have also been applied in other more conventional domains (e.g., image classification), it would be interesting to see how well and efficiently the MSLs can perform in those settings.\n\n-----------------------------\n\nI also have a few questions/feedbacks:\n\n1. Although the Newton-based update rule (3.2) is straightforward, I feel the second half of Sec. 3.1 is a bit redundant. For example, if I simply change (3.2) to $b_{n+1}^{k+1} = \\phi_{\\theta, n}(b_n^k)$, I would've arrived at the same conclusion in terms of \"finite-step convergence\". This will in practice make a difference because of 1) parallelism; and 2) the parareal method with coarse solvers; but at least theoretically I feel the time-iteration propagation discussion is a bit useless. Moreover, I don't think the notation $\\phi_{\\theta,n}$ is formally defined in the paper (and it took me a while to understand).\n\n2. One discussion that is interesting but missing in the main text is the initialization $B^0$. How are these initializations picked, and how sensitive are the MSL convergences to them?\n\n3. The authors suggest that unlike in deep equilibrium networks where one has to rely on either full vJP or low-rank quasi-Newton approximation to compute $(I-\\mathsf{D}\\gamma_\\theta)^{-1}$, we can exploit the known structure of $D_\\gamma$ and explicitly write out the inverse by a finite expansion $\\left(I - \\mathsf{D}\\gamma_\\theta(B^*)\\right)^{-1} = I + \\sum_{n=1}^N [\\mathsf{D}\\gamma_\\theta(B^*)]^n$. However, besides the nilpotency, such usage based on Neumann series also require that the spectral radius of $\\gamma_\\theta$ to be $<1$. Is this generally a problem in practice?\n\n4. The fixed-point tracking is indeed a great and interesting benefit of the fixed-point-based formulation (which applies also to works like DEQs, I presume?), but as the authors mentioned, they are also limited to the full-batch training mode. But in practice, unless we are dealing with an extremely small dataset, I don't feel this can be generally applied?\n\n5. Is there any caveat that the authors find for using the zero-order (parareal) method coarse solver? I'm curious if this actually creates a problem when the underlying dynamical systems is quite involuted.\n\n6. A major disadvantage of the approach is that, although the current experiments suggest acceleration, it is mainly because Neural ODEs are too slow and too small, not because MSLs are fast. For example, although Fig. 6 shows a significant reduction in NFEs, this is mainly a blessing of parallelization and we still have to undergo several root-solving iterations. If I'm understanding correctly, this means the actual FLOPs used for computation is actually larger, not smaller, than Neural ODEs. Therefore, will this acceleration benefit gradually diminish as we compare MSLs and NODEs on very large-scale and high-dim problems where parallelization yields only little gain?\n\n**Minor points**:\n- L137: $b_{n+1}^{b+1}$ should be $b_{n+1}^{k+1}$.\n- L175: $\\eta_k$ should be $\\eta_p$.\n- L290: There is a \"[CITE]\".\n\n--------------------\n--------------------\n\n#### Post-rebuttal\n\nI'd like to thank the authors for the response and agree that fixed-point tracking is an important direction to study for the implicit models. I can imagine how this benefit could enable implicit models to surpass/replace explicit models in many settings (e.g., where adjacent batches are correlated, such as in temporal sequences). I also appreciate the clarification on parallelism and FLOPs. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes, the authors have adequately addressed the limitations.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "A8NRyVlqL3z", "reviewer": "Reviewer_7kd1", "summary": "This work proposes a novel neural model (Multiple Shooting Layer) that is a significant improvement over Neural ODEs. MSLs seek solutions of initial value problems via parallelizable root-finding algorithms which lead to - \n* significant speedups in wall-clock inference time;\n* significantly less Number of Function Evaluations (NFE);\n\nThe authors further used their algorithm for -\n* Neural Optimal Control - In this experiment they demonstrate that MSLs lead to orders of magnitude of NFE savings as compared to Neural ODEs.\n* Time series classification - In this experiment they demonstrate that MSLs converge an order of magnitude faster than Neural ODEs.", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality:\nThis work uses Multiple Shooting Methods (a well known technique to solve PDEs in parallel) to propose a novel architecture that can serve as a direct replacement for Neural ODEs. MSLs come with several advantages - \n* Require significantly smaller number of NFEs;\n* Faster wall-clock time;\nwhich are clear improvements over Neural ODEs.\n\nQuality:\nThe authors have demonstrated the effectiveness of MSLs by conducting the following experiments -\n* Neural Optimal Control - In this experiment they demonstrate that MSLs lead to orders of magnitude of NFE savings as compared to Neural ODEs.\n* Time series classification - In this experiment they demonstrate that MSLs converge an order of magnitude faster than Neural ODEs.\n\nClarity:\nThe submission is well written.\n\nSignificance:\nThe effectiveness of MSLs over Neural ODEs is noteworthy as demonstrated in Neural Optimal Control and Time Series Classification experiments.\n\n-------\n# Post Rebuttal Comments\nI've read the concerns raised by the reviewers and the authors have responded to them convincingly. For this reason I'd like to stick to my original score of 7.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Nil", "review_text": "", "extra_scores": {"ethical_concerns": "Nil", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "2 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: Bringing multiple shooting, time–parallel methods for ODEs to bear on improving inference and training of deep architectures based on Neural-ODE is a strong technical and conceptually novel contribution. As such, the paper may be expected to generate interest in both ML and Numerical methods (Differential Equations) communities. The reviews do ask for improved clarity in presentation and problem formulation in the final version. 
Decision: Accept (Poster)