Go with the flow: Adaptive control for Neural ODEs

1 INTRODUCTION . The interpretation of artificial neural networks as continuous-time dynamical systems has led to both theoretical and practical advances in representation learning . According to this interpretation , the separate layers of a deep neural network are understood to be a discretization of a continuous-time operator so that , in effect , the net is infinitely deep . One important class of continuous-time models , neural ordinary differential equations ( NODEs ) ( Chen et al. , 2018 ) , have found natural applications in generative variational inference ( Grathwohl et al. , 2019 ) and physical modeling ( Köhler et al. , 2019 ; Ruthotto et al. , 2020 ) because of their ability to take advantage of black-box differential equation solvers and correspondence to dynamical systems in nature . Nevertheless , NODEs suffer from known representational limitations , which researchers have tried to alleviate either by lifting the NODE activation space to higher dimensions or by allowing the transition operator to change in time , making the system non-autonomous ( Dupont et al. , 2019 ) . For example , Zhang et al . ( 2020 ) showed that NODEs can arbitrarily approximate maps from Rd to R if NODE dynamics operate with an additional time dimension in Rd+1 and the system is affixed with an additional linear layer . The same authors showed that NODEs could approximate homeomorphisms from Rd to itself if the dynamics were lifted to R2d . Yet , the set of homeomorphisms from Rd to itself is in fact quite a conservative function space from the perspective of representation learning , since these mappings preserve topological invariants of the data space , preventing them from “ disentangling ” data classes like those of the annulus data in Fig . 1 ( lower left panel ) . In general , much remains to be understood about the continuous-time framework and its expressive capabilities . In this paper , we propose a new approach that we call neurally-controlled ODEs ( N-CODE ) designed to increase the expressivity of continuous-time neural nets by using tools from control theory . Whereas previous continuous-time methods learn a single , time-varying vector field for the whole input space , our system learns a family of vector fields parameterized by data . We do so by mapping the input space to a collection of control weights which interact with neural activity to optimally steer model dynamics . The implications of this new formulation are critical for model expressivity . In particular , the transformation of the input space is no longer constrained to be a homeomorphism , since the flows associated with each datum are specifically adapted to that point . Consequently , our system can easily “ tear ” apart the two annulus classes in Fig . 1 ( lower right panel ) without directly lifting the data space to a higher dimension . Moreover , when control weights are allowed to vary in time , they can play the role of fast , plastic synapses which can adapt to dynamic model states and inputs . The rest of the paper proceeds as follows . First , we will lay out the background for N-CODE and its technical formulation . Then , we will demonstrate its efficacy for supervised and unsupervised learning . In the supervised case , we show how NCODE can classify data by learning to bifurcate its dynamics along class boundaries as well as memorize high-dimensional patterns in real-time using fast synapses . Then , we show how the flows learned by N-CODE can be used as latent representations in an unsupervised autoencoder , improving image generation over a base model . 2 BACKGROUND . Neural ODEs ( NODEs ) ( Chen et al. , 2018 ) are dynamical systems of the form dx dt = f ( x , θ , t ) , ( 1 ) whereX is a space of features , θ ∈ Θ is a collection of learnable parameters , and f : X×Θ×R 7→ X is an equation of motion which we take to be differentiable on its whole domain . f defines a flow , i.e . a triple ( X , R , Φθ ) with Φθ : X × R 7→ X defined by Φθ ( x ( 0 ) , T ) = x ( 0 ) + T∫ 0 f ( x ( t ) , θ , t ) dt ( 2 ) which relates an initial point x ( 0 ) to an orbit of points { x ( t ) = Φθ ( x ( 0 ) , t ) , t ∈ R } . For a fixed T , the map x 7→ Φθ ( x , T ) is a homeomorphism from X to itelf parametrized by θ . Several properties of such flows make them appealing for machine learning . For example , the ODEs that govern such such flows can be solved with off-the shelf solvers and they can potentially model data irregularly sampled in time . Moreover , such flows are reversible maps by construction whose inverse is just the system integrated backward in time , Φθ ( . , t ) −1 = Φθ ( . , −t ) . This property enables depth-constant memory cost of training thanks to the adjoint sensitivity method ( Pontryagin Lev Semyonovich ; Boltyanskii V G & F , 1962 ) and the modeling of continuous-time generative normalizing flow algorithms ( ? ) . Interestingly , discretizing Eq . 2 yields the recursive formulation of a residual network ( He et al. , 2015 ) with a single residual operator fθ : ΦResNet ( x ( 0 ) , T ) = x ( 0 ) + T∑ t=1 fθ ( xt−1 ) ( 3 ) In this sense , NODEs with a time-independent ( autonomous ) equation of motion , f , are the infinitely-deep limit of weight-tied residual networks . Relying on the fact that every nonautonomous dynamical system with state x ∈ Rd is equivalent to an autonomous system on the extended state ( x , t ) ∈ Rd+1 , Eq . 2 can also be used to model general , weight-untied residual networks . However it remains unclear how dependence of f in time should be modeled in practice and how their dynamics relate to their discrete counterparts with weights evolving freely across blocks through gradient descent . 3 N-CODE : LEARNING TO CONTROL DATA-DEPENDENT FLOWS . General formulation - The main idea of N-CODE is to consider the parameters , θ ( t ) , in Eq . 1 as control variables for the dynamical state , x ( t ) . Model dynamics are then governed by a coupled system of equations on the extended state z ( t ) = ( x ( t ) , θ ( t ) ) . The initial value of the control weights , θ ( 0 ) , is given by a mapping γ : X → Θ . Throughout , we assume that the initial time point is t = 0 . The full trajectory of control weights , θ ( t ) , is then output by a controller , g , given by another differentiable equation of motion g : Θ× X × R 7→ Θ with initial condition γ ( x0 ) . Given an initial point , x ( 0 ) , we can solve the initial value problem ( IVP ) dz dt = h ( z , t ) z ( 0 ) = z0 = ( dx dt , dθ dt ) = ( f ( x , θ , t ) , g ( θ , x , t ) ) ( x ( 0 ) , θ ( 0 ) ) = ( x0 , γ ( x0 ) ) , ( 4 ) where h = ( f , g ) . We may think of g and γ as a controller of the dynamical system with equation of motion , f . We model g and γ as neural networks parameterized by µ ∈ Rnµ and use gradient descent techniques where the gradient can be computed by solving an adjoint problem ( Pontryagin Lev Semyonovich ; Boltyanskii V G & F , 1962 ) that we describe in the next section . Our goal here is to use the metaparameterization in the space Θ to create richer dynamic behavior for a given f than directly optimizing fixed weights θ. Well-posedness - If f and g are continuously differentiable with respect to x and θ and continuous with respect to t , then , for all initial conditions ( x ( 0 ) , θ ( 0 ) ) , there exists a unique solution z for Eq . 4 by the Cauchy-Lipschitz theorem . This result leads to the existence and uniqueness of the augmented flow ( X × Θ , R , Φµ ) with Φµ : ( X ×Θ ) × R 7→ X ×Θ . Moreover , considering the restriction of such a flow on X , we are now endowed with a universal approximator for at least the set of homeomorphisms on X given that this restriction constitutes a non-autonomous system . We discuss now how g and γ affect the evolution of the variable x ( t ) , exhibiting two forms of control and noting how they relate to previous extensions of NODEs . 3.1 OPEN AND CLOSED-LOOP CONTROLLERS . If the controller outputs control weights as a function of the current state , x ( t ) , then we say it is a closed-loop controller . Otherwise , it is an open-loop controller . Open-loop control : First , we consider the effect of using only the mapping γ in Eq . 4 as a controller . Here , γ maps the input space X to Θ so that f is conditioned on x ( 0 ) but not necessarily on x ( t ) for t > 0 . In other words , each initial value x ( 0 ) evolves according to its own learned flow ( X , R , Φγ ( x ( 0 ) ) ) . This allows for trajectories to evolve more freely than within a single flow that must account for the whole data distribution and resolves the problem of non-intersecting orbits ( see Figure 4 ) . Recently , ( Massaroli et al. , 2020b ) proposed a similar form of data-conditioned openloop control with extended state ( x ( t ) , x ( 0 ) ) ) . This is a version of our method in which γ is of the form γ ( x ) = θ ( 0 ) = [ C : id ] with C a constant vector , id is the identity function , and : denotes concatenation . Our open-loop formulation makes an architectural distinction between controller and dynamics and is consequently generalizable to the following closed-loop formulation . stationary stimulus . We evaluate this hypothesis in section . 5.3 . The expressivity of N-CODE compared to other continuous-time neural networks is encapsulated in the following proposition . The result , proven in appendix , shows that both openloop and closed-loop control systems overcome NODEs ’ expressivity constraint with two distinct strategies , data-conditioning and statespace augmentation . Proposition 1 - There exists a transformation φ : Rd → Rd which can be expressed by N-CODE but not by NODEs . In particular , φ is not a homeomorphism . 3.2 TRAINING . Loss function : Dynamics are evaluated according to a generalized loss function that integrates a cost over some interval [ 0 , T ] : l ( z ) : = T∫ 0 ` ( z ( t ) , t ) dt = T∫ 0 ` ( x ( t ) , θ ( t ) , t ) dt ( 5 ) The loss in Eq . 5 is more general than in Chen et al . ( 2018 ) since ` can be any Lebesgue-measurable function of both states , x , and control parameters , θ . In particular , this includes penalties at discrete time points or over the entire trajectory ( Massaroli et al. , 2020b ) but also regularizations on the weights or activations over the entire trajectory rather than the final state z ( T ) of the system . In order to estimate a control function θ ( t ) that is optimal with respect to Eq . 5 , we invoke Pontryagin ’ s maximum principle ( PMP ) ( Pontryagin Lev Semyonovich ; Boltyanskii V G & F , 1962 ) , which only requires mild assumptions on the functional control space Θ and applies to functions f that are non-smooth in θ . The PMP gives necessary conditions on θ ( t ) at optimality via the augmented adjoint variable a ( t ) . This quantity is the Jacobian of ` with respect to both x ( t ) and θ ( t ) . In the case of θ ( t ) being differentiable with respect to the meta-parameters µ , solving for the augmented adjoint state a ( t ) as in Chen et al . ( 2018 ) allows us to compute the gradient of the loss with respect to µ thanks to Theorem 1 . Theorem 1 - Augmented adjoint method : Given the IVP of equation 4 and for ` defined in equation 5 , we have : ∂l ∂µ = T∫ 0 a ( t ) T ∂h ∂µ dt , such that a satisfies da dt = −aT . ∂h ∂z − ∂ ` ∂z a ( T ) = 0 ( 6 ) where ∂h ∂z is the Jacobian of h with respect to z : ∂h∂z = ∂f∂x ∂f∂θ ∂g ∂x ∂g ∂θ . In practice , we compute the Jacobian for this augmented dynamics with open source automatic differentiation libraries using Pytorch ( Paszke et al. , 2019 ) , enabling seamless integration of NCODE modules in bigger architectures . We show an example of such modules in section B of Appendix .

Summary: Comment: Summary: This paper presents a technique for more expressive neural ordinary differential equations (NODE) flows. Instead of learning a fixed set of parameters $\theta$ that governs the ODE dynamic, the proposed approach learns dynamic parameters that evolve over time. The authors propose two variants of the methods: $\textit{open-loop}$ and $\textit{closed-loop}$. The former only maps the initial observation as the controller whereas $\theta$ in the later model follows another NODE $g$. Model performance is demonstrated on several tasks. The model is shown to solve the well-known "crossing curves" problem, on which NODE fails. Also, it's demonstrated that the presented technique yields sharper images compared to VAEs and also better at classification and interpolation.