Differentiable Multiple Shooting Layers

1 Introduction . 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 . 2 Multiple Shooting Layers . 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 ) 3 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 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 . 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 γθ ( 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

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 -  * 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.