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: The paper proposes a new class of neural ODE based models, called neurally-controlled ODEs (N-CODE). Instead of directly learning the weights of the neural network parameterizing the vector field f of the ODE, the paper instead proposes to learn a controller that takes as input the initial state and outputs the initial weights of f and optionally also updates the weights of f as the ODE is solved. This allows the model to parameterize a family of (potentially time varying) vector fields instead of a single vector field. The authors go on to show that their model can be trained using an augmented adjoint method (similar to the adjoint method used in the original neural ODE paper).