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