Differentiable Multiple Shooting Layers

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.

heading: 1 Introduction
text: For the last twenty years, one has tried to speed up numerical computation mainly by providing ever faster computers. Today, as it appears that one is 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 a more “parallel” form. Nievergelt, 1964
Discovering and exploiting parallelization opportunities has allowed deep learning methods to succeed across application areas, reducing iteration times for architecture search and allowing scaling to larger data sizes (Krizhevsky et al., 2012; Diamos et 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 a 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 1 provides visual intuition of the parallel nature of MSL inference. ∗Equal contribution. Author order was decided by flipping a coin.
35th Conference on Neural Information Processing Systems (NeurIPS 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–depth architectures such as Deep Equilibrium Newtorks (DEQs) (Bai et 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 as 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 can be significantly more expensive than surrounding operations. For this reason, the reduction in NFEs obtained through time–parallelization leads to 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 a 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 a latent variant, and as a faster alternative to neural controlled differential equations (Neural CDE) (Kidger et al., 2020b) in long–horizon time series classification.

heading: 2 Multiple Shooting Layers
text: Consider the initial value problem (IVP)
ż(t) = fθ(t, z(t)) z(0) = z0 , t ∈ [0, T ]. (2.1)
with state z ∈ Z ⊂ Rnz , parameters θ ∈ W for some spaceW of functions [0, T ] → Rnθ and a smooth vector field fθ : [0, T ]× Z ×W → Z . For all z ∈ Z, s, t ∈ [0, T ]; s < t we denote with φθ(z, s, t) the solution of (2.1) at time t starting from z at time s , i.e. φθ(z, s, t) : (z, s, t) 7→ z(t). The crux behind multiple–shooting methods for differential equations is to turn the initial value problem (2.1) into a boundary value problem (BVP). We split the the time interval [0, T ] in N sub–intervals [tn, tn+1] with 0 = t0 < t1 < · · · < tN = T and define N left boundary subproblems
zn(tn) = bn and żn(t) = fθ(t, zn(t)), t ∈ [tn, tn+1] (2.2) where bn are denoted as shooting parameters. At each time t ∈ [0, T ], the solution of (2.2) matches the one of (2.1) iff all the shooting parameters bn are identical to z(tn), bn = φθ(z0, t0, tn). Using z(tn) = φθ(z(tn−1), tn−1, tn), we obtain the equivalent conditions
b0 = φθ(z0, t0, t0) = z0 b1 = φθ(b0, t0, t1) = z0(t1)
... bN = φθ(bN−1, tN−1, tN ) = zN−1(tN )
Let B := (b0, b1, · · · , bN ) and γθ(B, z0) := (z0, φθ(b0, t0, t1), · · · , φθ(bN−1, tN−1, tN )). We can thus turn the IVP (2.1) into the roots–finding problem the of a function gθ defined as
gθ(B, z0) = B − γθ(B, z0)
Definition 1 (Multiple Shooting Layer (MSL)). With `x : X → Z and `y : ZN+1 → Y two affine maps, a multiple shooting layer is defined as the implicit input–output mapping x 7→ y:
z0 = `x(x)
B∗ : gθ(B ∗, z0) = 0
y = `y(B ∗)
(2.3)

heading: 3 Realization of Multiple Shooting Layers
text: 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 a drastic reduction in the number of vector field evaluations at the cost of a higher memory requirement for the parallelization to take place. Nonetheless, the forward pass of MSLs requires the shooting parameters B to satisfy the nonlinear algebraic matching condition gθ(B, z0) = 0, which has also to be solved numerically.

heading: 3.1 Forward Model
text: The forward MSL model involves the synergistic combination of two main classes of numerical methods: ODE solvers and root finding algorithms, to compute γθ(B, z0) andB∗, 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 to compute γθ(Bk, z0) and evaluate the matching condition gθ(Bk, z0).
Newton methods for root finding Let us denote with Bk the solution of the root finding problem at the k-th step of the Newton method and let Dgθ(Bk, z0) be the Jacobian of gθ computed in Bk. The solution B∗ : gθ(B∗, z0) = 0 can be obtained by iterating the Newton–Raphson fixed point iteration
Bk+1 = Bk − α [ IN ⊗ Inz − Dγθ(Bk, z0) ]−1 [ Bk − γθ(Bk, z0) ] (3.1)
which converges quadratically to B∗ (Nocedal and Wright, 2006). The exact Newton iteration theoretically (3.1) requires the inverse of the Jacobian IN ⊗ Inz − Dγθ(Bk, z0). Without the special structure of the MSL problem, the Jacobian would have had to be the computed in full, as in the case of DEQs (Bai et al., 2019). Being the Jacobian of dimension RNnz×Nnz , its computation with reverse–mode automatic differentiation (AD) tools scales poorly with state dimensions and number of shooting parameter (cubically in both nz and N ).
Instead, the special structure of the MSL matching function gθ(B, z0) = B − γθ(B, z0) 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 a 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
[ IN ⊗ Inz − Dγθ(Bk, z0) ] (Bk+1 −Bk) = γθ(Bk, z0)−Bk
which leads to the following update rule for the individual shooting parameters bkn (see Fig. 2):
bk+1n+1 = φθ,n(b k n) + Dφθ,n(b k n) ( bk+1n − bkn ) , bk+10 = z0 (3.2)
where Dφθ,n(bkn) = dφθ,n(b k n)/dbn is the sensitivity of each individual flow to its initial condition. Due to the dependence of bk+1n+1 on b k+1 n , a complete Newton iteration theoretically requires N − 1 sequential stages.
Finite-step convergence Iteration (3.2) exhibits convergence to the exact solution of the IVP (2.1) in N − 1 steps (Gander, 2018, Theorem 2.3). In particular, given perfect integration of the sub–IVPs, bkn coincides with the exact solution φθ(z0, t0, tn) from iteration index k = n onward, i.e. at iteration k only the last N − k shooting parameters are actually updated. Thus, the computational and memory footprint
of the method diminishes with the number of iterations. This result can be visualized in the graphical representation of iteration (3.2) in Figure 3 while further details are discussed in Appendix B.1.
Time–Iteration Propagation of Newton Scheme n
k
l
l
l
l
l
l
l
l
l z0 b
0 1 b 0 2
z0 b 1 1 b 1 2
z0 b 2 1 b 2 2
φθ
φθ
φθ
φθ φθ
φθ
l
l l
bkn
bk+1n b k+1 n+1
B2
B1
B0
Bk+1
of the ODE solver. This approach, which might be considered as the continuous variant of forward– mode AD, scales quadratically with nz , has low memory cost, and explicitly controls numerical error. Proposition 1 (Forward Sensitivity (Khalil, 2002)). Let φθ(z, s, t) be the solution of (2.1). Then, v(t) = Dφθ(z, s, t) satisfies the linear matrix–valued differential equations
v̇(t) = Dfθ(t, z(t))v(t), v(s) = Inz where Dfθ denotes ∂fθ/∂z.
Therefore, at iteration k all Dφθ,n(bkn) can be computed in parallel while performing the forward integration of the N − k IVPs (2.2) and their forward sensitivities, i.e.
ForwardSensitivity : { bkn 7→ (φθ,n,Dφθ,n) } k<n≤N
which enables full vectorization of Jacobian–matrix products between ∂fθ/∂z and v as 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, a zero–order approximation of the Newton iteration preserving the finite–step converge property can be employed: the parareal method Lions et al. (2001). From the Taylor expansion of φθ,n(bk+1n ) around b k n
φθ,n(b k+1 n ) = φθ,n(b k n) + Dφθ,n(b k n) ( bk+1n − bkn ) + o ( ‖bk+1n − bkn‖22 ) ,
we have the following approximant for the correction term of (3.2)
Dφθ,n(b k n) ( bk+1n − bkn ) ≈ φθ,n(bk+1n )− φθ,n(bkn). (3.3)
Parareal computes the RHS of (3.3) by coarse2 numerical solutions ψθ,n(bkn), ψθ,n(b k+1 n ) of φθ,n(b k n), φθ,n(b k+1 n ), leading to the forward iteration,
bb+1n+1 = φθ,n(b k n) + ψθ,n(b k+1 n )− ψθ,n(bkn).
2e.g. few steps of a low–order ODE solver

heading: 3.2 Properties of MSLs
text: 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,B∗, θ) = ∑N n=1 cθ(x, b ∗ n). The adjoint gradient for the MSL is then given by
∇θL = ∫ T
0
λ>(t)∇θfθ(t, z(t))dt
where the Lagrange multiplier λ(t) satisfies a backward piecewise–continuous linear ODE
λ̇(t) = −Dfθ(t, z(t))λ(t) if t ∈ [tn, tn+1) λ−(tn) = λ(tn) +∇>b cθ(x, bn) λ(T ) = ∇>b cθ(x, bN )
The adjoint method typically requires the IVP (2.1) to be solved backward alongside λ to retrieve the value of z(t) needed to compute the Jacobians Dfθ and ∇θfθ. This step introduces additional errors on the final gradients: numerical errors accumulated on b∗N ≈ z(T ) during forward pass, propagate to the gradients and sum up with errors on the backward integration of (2.1).
Here we take a 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 are solution points of the IVP (2.1), i.e. b∗n = φ(z0, t0, tn) (up to the forward numerical solver tolerances). On this assumption, we construct a cubic spline interpolation ẑ(t) of the shooting parameters b∗n and we query it during the integration of λ to compute the Jacobians of fθ. 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 a forward–backward pass of an MSL is shown in Figure 4.
One-step inference: fixed point tracking Consider training a MSL to minimize a twice– differentiable loss function L(x,B∗, θ) with Lipschitz constant mθL through the gradient descent iteration θp+1 = θp − ηp∇θL(x,B∗p , θp) where ηp is a positive learning rate andB∗p is the exact root of the matching function gθ computed with parameters θp (i.e. the exact solution of the IVP (2.1) at the boundary points). Due to Lipschitzness of L, we have the following uniform bound on the variation of the parameters across training iterations
‖θp+1 − θp‖2 ≤ ηpmθL. If we also assume γθ to be Lipschitz continuous w.r.t z and θ with constantsmθγ , m z γ and differentiable w.r.t. θ we can obtain the variation of the fixed point B∗ to small changes in the model parameters by linearizing solutions around θp
B∗p+1 −B∗p = [θp+1 − θp] ∂γθp(B
∗ p , z0)
∂θ + o(‖θp+1 − θp‖22)
to obtain the uniform bound
‖B∗p+1 −B∗p‖2 ≤ ηpmθLmθγ + o(‖θp+1 − θp‖22).
Having a bounded variation on the solutions of the MSL for small changes of the model parameters θ, we might think of recycling the previous shooting parameters B∗p as an initial guess for the direct Newton algorithms in the forward pass succeeding the gradient descent update. We show that by choosing a sufficiently small learning rate ηp one Newton iteration can be sufficient to track the true value of B∗ during training. In particular, the following bound can be obtained. Theorem 1 (Quadratic fixed-point tracking). If fθ is twice continuously differentiable in z then
‖B∗p+1 − B̄∗p‖2 ≤Mη2p (3.4) for some M > 0. B̄∗p is the result of one Newton iteration applied to B ∗ p .
The proof, reported in Appendix A.1, relies on the quadratic converge of Newton method. The quadratic dependence of the tracking error bound on ηk allows use of typical learning rates for standard gradient based optimizers to keep the error under control. In this way, we can turn the implicit learning problem into an explicit one where the implicit inference pass reduces to one Newton iteration. This approach leads to the following training dynamics:
θ ← θ − η∇θL(x,B∗, θ) B∗ ← apply{(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.e. the initial condition z0 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 x even for small learning rates.
Numerical scaling Each class of MSL outlined in Section 3.1 is equipped with unique computational scaling properties. In the following, we denote with NFEφ the total number of vector field fθ evaluations done, in parallel across shooting parameters, in a single sub–interval [tn, tn+1]. Similarly, NFEψ indicates the function evaluations required by the coarse solver used for parareal approximations. Here, we set out to investigate the computational signature of MSLs as parallel algorithms. To this end, we decompose a 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 span3 of a single MSL iteration as a function of number of Jacobian vector products (jvp), Jacobian matrix products (jmp), and vector field evaluations (NFE).
Fw sensitivity MSL frontloads the cost of computing Dφθ,n by solving the forward sensitivity ODEs of Proposition 1 alongside the evaluation of γθ. Forward sensitivity equations involve a jmp, which can be optionally further parallelized as nz jvps by paying a memory overhead. Once sensitivities have been obtained along with γθ, 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–parallelizability at a larger memory cost.
Zeroth–order MSL computes γθ via a total of NFEφ evaluations fθ, parallelized across sub–intervals. The cheaper IVP solution in both memory and compute is however counterbalanced during calculation
3Longest 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 sub–interval.
of the sensitivities, as this MSL approach approximates the sensitivities Dφθ,n by a zeroth–order update requiring N − k sequential calls to a coarse solver. The analysis of MSL backpropagation scaling is straightforward, as sequential adjoints for MSLs mirror standard sensitivity techniques for Neural ODEs in both compute and memory footprints. Alternatively, 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.

heading: 4 Applications

heading: 4.1 Variational MSL
text: Let x : R → Rnx , be an observable of some continuous–time process and let X = {x−M , . . . , x0, . . . , xN} ∈ R(M+N+1)×nx be a sequence of observations of x(t) at time instants t−M < · · · < t0 < · · · < tN . We seek a model able to predict x1, . . . , xN given past observations x−M , . . . , x0, equivalent to approximating the conditional distribution p(x1:N |x−M :0). To this end we introduce variational MSLs (vMSLs) as the following latent variable model:
(µ,Σ) = Eω(x−M :0) Encoder Eω qω(z0|x−M :0) = N (µ,Σ) Approx. Posterior
z0 ∼ qω(z0|x−M :0) Reparametrization
B∗ : gθ(B ∗, z0) = 0 Decoder Dθ
x̂1, . . . , x̂N = `(B ∗) Readout `
Once trained, such model can be also used to generate new realistic sequences of the observable x(t) by querying the decoder network at a desired z0. vMSLs are designed to scale data generation to longer sequences, exploiting wherever possible parallel computation in time in both encoder as well as decoder modules. The structure of Eω is designed to leverage modern advances in representation learning for time–series via temporal convolutions (TCNs) or attention operators (Vaswani et al., 2017) to offer a higher degree of parallelizability compared to sequential encoders e.g RNNs, ODERNNs (Rubanova et al., 2019) or Neural CDEs (Kidger et al., 2020b). This, in turn, allows the encoder to match the decoder in efficiency, avoiding unnecessary bottlenecks. The decoder Dθ is composed of a MSL which is tasked to 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 (θ,ω) Ez0∼qω(z0|x−M:0)
N∑
n=1
log pn(x̂n)− KL(qω||N (0, I))
with pn(x̂n) = N (xn, σn) and the standard deviations σn are left as a hyperparameters.
0 500 1,000
20 30 40 50
Training Step
N F E
Variational MSLs & Latent ODEs NFE Analysis
vMSL 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.
Sequence generation under noisy observations We apply vMSLs on a sequence generation task given trajectories corrupted by state–correlated noise. We consider a baseline latent model sharing the same overall architecture as vMSLs, with a Neural ODE as decoder. In particular, the Neural ODE is solved via the dopri5 solver with absolute and relative tolerances set to 10−4, 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 is then performed by inspection 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, vMSLs 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 in Appendix E.1.
4.2 Neural Optimal Control Neural Optimal Policy via MSLs Vector Field Learned Controller
−5
0
5
πθ(z)
−2 0
2
Figure 7: [Left] Closed–loop vector fields and trajectories corresponding to the uθ-controlled MSL. [Right] Learned controller uθ(z) (z ∈ R2) for the two different desired limit cycles. Although, the inference of the all trajectory is performed with just two steps of RK4 (8 NFE), the initial accuracy of dopri5 (> 3000 NFE) is preserved throughout training.
Beyond sequence generation, the proposed framework can be applied to optimal control. Here we can fully exploit the drastic computational advantages of MSLs. In fact, we leverage on the natural assumption of finiteness of initial conditions z0 where the controlled system (or plant) is initialized to verify the result of Th. 1. Let us consider a controlled dynamical system
ż(t) = f(t, z(t), πθ(t, z)), z(0) = z0 (4.1)
with a parametrized policy πθ : t, z 7→ πθ(t, z) and initial conditions z0 ranging in a finite set Z0 = {zj0}j . We consider the problems of stabilizing a low–dimensional limit cycle and deriving an optimal boundary policy for a linear PDE discretized as a 200–dimensional ODE.
Limit cycle stabilization We consider a stabilization task where, given a desired closed curve Sd = {z ∈ Z : sd(z) = 0}, we minimize the 1-norm between the given curve and the MSL solution of (4.1) across the timestamp-
sas well as the control effort |πθ|. We verify the approach on a one degree–of–freedom mechanical system, aiming with closed curves Sd of various shapes. Following the assumptions on fixed point tracking and slow–varying–flows of Th. 1, we initialize B0j by dopri5 adaptive–step solver set with tolerances 10−8. Then, at each training iteration, we perform inference with a single parallel rk4 step for each sub–interval [tn, tn+1] followed by a single Newton update. Figure 7 shows the learned vector fields and controller, confirming a successful system stabilization of the system to different types of closed curves. We compare with a range of baseline Neural ODEs, solved via rk4 and dopri5. Training of the controller via MSLs is achieved with orders of magnitude less wall–clock time and NFEs. Figure 8 shows the difference in NFEs w.r.t. dopri5 while Fig. 9 compares wall–clock times per training iteration of MSL, dopri5 and rk4.
We further provide Symmetric Mean Average Percentage Error (SMAPE) measurements between trajectories obtained via MSLs and an adaptive–step solver. MSLs initialized with recycled 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 baseline Neural ODEs is are provided in Appendix E.2.
Neural Boundary Control of the Timoshenko Beam We further show how MSLs can be scaled to high–dimensional regimes by tackling the optimal boundary control problem for a 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 a
structure–preserving spectral discretization yielding a 160–dimensional Hamiltonian ODE.
We parameterize the boundary control policy with a multi–layer perceptron taking as input (control feedback) the 160–dimensional 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
integration. We compare the training wall–clock time with a 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 is 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.

heading: 4.3 Fast Neural CDEs for Time Series Classification
text: To further verify the broad range of applicability of MSLs, we apply them to time series classification as faster alternatives to neural controlled differential equations (Neural CDEs) (Kidger et al., 2020b). Here, MSLs remain applicable since Neural CDEs are practically solved as ODEs with a special structure, as described in Appendix E.4. We tackle the PhysioNet 2019 challenge (Goldberger et al., 2000) on sepsis prediction, following the exact experimental procedure described by Kidger et al. (2020b),
including hyperparameters and Neural CDE architectures. However, we train all models on the full dataset to enable application of the fixed point tracking technique for MSLs4. 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 via sequential adjoints, MSL Neural CDEs converge with total wall–clock time one order of magnitude smaller than the baseline.

heading: 5 Related Work
text: Parallel–in–time integration & multiple shooting MSLs belong to the framework of time– parallel integration algorithms. The study of these methods is relatively recent, with seminal work in the 60s (Nievergelt, 1964). The multiple shooting formulation of time–parallel integration, see e.g. (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 et 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 et al., 2018), albeit in the unrelated context of learning time–varying parameters.
4We note that the training for all models has been performed on a single NVIDIA RTX A6000 with 48Gb of GPU memory.
Time–parallelization in neural models In the pursuit for increased efficiency, several works have proposed approaches to parallelization across time or depth in neural networks. (Gunther et al., 2020; Kirby et al., 2020; Sun et al., 2020) use multigrid and penalty methods to achieve speedups in ResNets. Meng et al. (2020) proposed a parareal variant of Physics–informed 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 a regularization term, however, is not guaranteed to converge to a continuous solution. The method of Zhuang et al. (2021) further optimizes its parameters in a full–batch regime, where application 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.
In contrast our objective is to introduce a novel class of implicit time–parallel models, and to validate their computational efficiency across settings.

heading: 6 Conclusion
text: This work introduces differentiable Multiple Shooting Layers (MSLs), a parallel–in–time alternative to 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 was 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.

Review summaries:
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 -  * significant speedups in wall-clock inference time; * significantly less Number of Function Evaluations (NFE);  The authors further used their algorithm for - * Neural Optimal Control - In this experiment they demonstrate that MSLs lead to orders of magnitude of NFE savings as compared to Neural ODEs. * Time series classification - In this experiment they demonstrate that MSLs converge an order of magnitude faster than Neural ODEs.

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.

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.