Stateful ODE-Nets using Basis Function Expansions

1 Introduction . The interpretation of neural networks ( NNs ) as discretizations of differential equations [ 7 , 17 , 31 , 43 ] has recently unlocked a fruitful link between deep learning and dynamical systems . The strengths of so-called ordinary differential equation networks ( ODE-Nets ) are that they are well suited for modeling time series [ 11 , 39 ] and smooth density estimation [ 14 ] . They are also able to learn representations that preserve the topology of the input space [ 9 ] ( which may be seen as a “ feature ” or as a “ bug ” ) . Further , they can be designed to be highly memory efficient [ 13 , 51 ] . However , one major drawback of current ODE-Nets is that the predictive accuracy for tasks such as image classification is often inferior as compared to other state-of-the-art NN architectures . The reason for the poor performance is two-fold : ( a ) ODE-Nets have shallow parameterizations ( albeit long computational graphs ) , and ( b ) ODE-Nets do not include a mechanism for handling layers with internal state ( i.e. , stateful layers , or stateful modules ) , and thus can not leverage batch normalization layers which are standard in image classification problems . That is , in modern deep learning environments the running mean and variance statistics of a batch normalization module are not trainable parameters , but are instead part of the module ’ s state , which does not fit into the ODE framework . Further , traditional techniques such as stochastic depth [ 21 ] , layer dropout [ 12 ] and adaptive depth [ 10 , 28 ] , which are useful for regularization and compressing traditional NNs , do not necessarily apply in a meaningful manner to ODE-Nets . That is , these methods do not provide a systematic scheme to derive smaller networks from a single deep ODE-Net source model . This limitation prohibits quick and rigorous adaptation of current ODE-Nets across different computational environments . ∗Equal contributions . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Re Ne ANODEV1 Ne a ODE ; FFJORD C Ne S cOdeB c OdeB c Re U /B c Re U /B c OdeB c OdeB c OdeB c S c S c S c FC Se A e MLP BN C ReLU BN C ReLU Te [ ] E bed OdeB c POS [ ] Se A e MLP N a ODE : linea eigh , ada i e e ing OdeB c OdeB c ANODEV1 : con an eigh , man fi ed e OdeB c OdeB c R N : con an eigh , one E le e Re B c Re B c C N V2 : a bi a ba i f nc ion OdeB c P [ ] De e P Te [ 0 ] Te [ 1 ] Te [ 2 ] ... Te [ N ] E bed S a O B Se A e MLP Se A e MLP POS [ 0 ] POS [ 1 ] POS [ 2 ] ... POS [ N ] Se A e MLP I : O : Figure 1 : Sketch of a continuous-in-depth transformer-encoder . The model architecture consists of a sparse embedding layer followed by an OdeBlock that integrates the encoder to feed into a classification layer to determine parts of speech . The model graph is generated by nesting the residual R on the right into a time-integration scheme in the OdeBlock . The weights for each call toR are determined by evaluating the basis expansion . Internal hidden states evolve continuously during the forward pass . This is illustrated by a smoothly varying attention matrix . Hence , there is a need for a general but effective way of reducing the number of parameters of deep ODE-Nets to reduce both inference time and the memory footprint . To address these limitations , we propose a novel stateful ODE-Net model that is built upon the theoretical underpinnings of numerical analysis . Following recent work by [ 34 , 37 ] , we express the weights of an ODE-Net as continuous-in-depth functions using linear combinations of basis functions . However , unlike prior works , we consider parameter transformations through a change of basis functions , i.e. , adding basis functions , decreasing the number of basis functions , and function projections . In turn , we are able to ( a ) design deep ODE-Nets that leverage meaningful continuousin-depth batch normalization layers to boost the predictive accuracy of ODE-Nets , e.g. , we achieve 94.4 % test accuracy on CIFAR-10 , and 79.9 % on CIFAR-100 ; and ( b ) we introduce a methodology for compressing ODE-Nets , e.g , we are able to reduce the number of parameters by a factor of two , without retraining or revisiting data , while nearly maintaining the accuracy of the source model . In the following , we refer to our model as Stateful ODE-Net . Here are our main contributions . • Stateful ODE-Block . We introduce a novel stateful ODE-Block that enables the integration of stateful network layers ( Sec . 4.1 ) . To do so , we view continuous-in-depth weight functions through the lens of basis function expansions and leverage basis transformations . An example of our continuous-in-depth model is shown in Figure 1 . • Stateful Normalization Layer . We introduce stateful modules using function projections . This enables us to introduce continuous-in-depth batch normalization layers for ODE-Nets ( Sec . 4.2 ) . In our experiments , we show that such stateful normalization layers are crucial for ODE-Nets to achieve state-of-the-art performance on mainstream image classification problems ( Sec . 7 ) . • A Posteriori Compression Methodology . We introduce a methodology to compress ODE-Nets without retraining or revisiting any data , based on parameter interpolation and projection ( Sec . 5 ) . That is , we can systematically derive smaller networks from a single deep ODE-Net source model through a change of basis . We demonstrate the accuracy and compression performance for various image classification tasks using both shallow and deep ODE-Nets ( Sec . 7 ) . • Advantages of Higher-order Integrators for Compression . We examine the effects of training continuous-in-depth weight functions through their discretizations ( Sec . 6 ) . Our key insight in Theorem 1 is that higher-order integrators introduce additional implicit regularization which is crucial for obtaining good compression performance in practice . Proofs are provided in Appendix A . 2 Related Work . The “ formulate in continuous time and then discretize ” approach [ 7 , 17 ] has recently attracted attention both in the machine learning and applied dynamical systems community . This approach considers a differential equation that defines the continuous evolution from an initial condition x ( 0 ) = xin to the final output x ( T ) = xout as ẋ ( t ) = F ( θ̂ , x ( t ) , t ) . ( 1 ) Here , the functionF can be any NN that is parameterized by θ̂ , with two inputs x and t. The parameter t ∈ [ 0 , T ] in this ODE represents time , analogous to the depth of classical network architectures . Using a finer temporal discretization ( with a smaller ∆t ) or a larger terminal time T , corresponds to deeper network architectures . Feed-forward evaluation of the network is performed by numerically integrating Eq . ( 1 ) : xout = xin + ∫ T 0 F ( θ̂ , x ( t ) , t ) dt = OdeBlock [ F , θ̂ , scheme , ∆t , t ∈ [ 0 , T ] ] ( xin ) . ( 2 ) Inspired by this view , numerous ODE and PDE-based network architectures [ 1 , 8 , 9 , 15 , 32 , 42 , 48 , 49 ] , and continuous-time recurrent units [ 4 , 29 , 30 , 40 , 41 ] have been proposed . Recently , the idea of representing a continuous-time weight function as a linear combination of basis functions has been proposed by [ 34 ] and [ 37 ] . This involves using the following formulation : ẋ ( t ) = R ( θ ( t ; θ̂ ) , x ( t ) ) , ( 3 ) whereR is now parameterized by a continuously-varying weight function θ ( t ; θ̂ ) . In turn , this weight function is parameterized by a countable tensor of trainable parameters θ̂ . Both of these works noted that piecewise constant basis functions algebraically resemble residual networks and stacked ODE-Nets . A similar concept was used by [ 6 ] to inspire a multi-level refinement training scheme for discrete ResNets . In addition , [ 34 ] uses orthogonal basis sets to formulate a Galërkin Neural ODE . In this work , we take advantage of basis transformations to introduce stateful normalization layers as well as a methodology for compressing ODE-Nets , thus improving on ContinuousNet and other prior work [ 37 ] . Although basis elements are often chosen to be orthogonal , inspired by multi-level refinement training , we shall consider non-orthogonal basis sets in our experiments . Table 1 highlights the advantages of our model compared to other related models . 3 Basis Function View of Continuous-in-depth Weight Functions . Let θ ( t ; θ̂ ) be an arbitrary weight function that depends on depth ( or time ) t , which takes as argument a vector of real-valued weights θ̂ . Given a basis φ , we can represent θ as a linear combination of K ( continuous-time ) basis functions φk ( t ) : θ ( t ; θ̂ ) = K∑ k=1 φk ( t ) θ̂k . ( 4 ) The basis sets which we consider have two parameters that specify the family of functions φ and cardinality K ( i.e. , the number of functions ) to be used . Hence , we represent a basis set by ( φ , K ) . While our methodology can be applied to any basis set , we restrict our attention to piecewise constant and piecewise linear basis functions , common in finite volume and finite element methods [ 44 , §4 ] . • Piecewise constant basis . This orthogonal basis consist of basis functions that assume a constant value for each of the elements of width ∆t = T/K , where T is the “ time ” at the end of the ODENet . The summation in Eq . ( 4 ) involves piecewise-constant indicator functions φk ( t ) satisfying φk ( t ) = { 1 , t ∈ [ ( k − 1 ) ∆t , k∆t ] 0 , otherwise . ( 5 ) • Piecewise linear basis functions . This basis consists of evenly spaced elements where each parameter θ̂k corresponds to the value of θ ( tk ) at element boundaries tk = T ( k − 1 ) / ( K − 1 ) . Each basis function is a “ hat ” function around the point tk , φk ( t ) = ( t− k∆t ) /∆t , t ∈ [ ( k − 1 ) ∆t , k∆t ] 1− ( t− k∆t ) /∆t , t ∈ [ k∆t , ( k + 1 ) ∆t ] 0 , otherwise . ( 6 ) Piecewise linear basis functions have compact and overlapping supports ; i.e. , they are not orthogonal , unlike piecewise constant basis functions . 3.1 Basis Transformations . We consider two avenues of basis transformation to change function representation : interpolation and projection . Note that these transformations will often introduce approximation error , particularly when transforming to a basis with a smaller size . Furthermore , interpolation and projection do not necessarily give the same result for the same function . Interpolation . Some basis functions use control points for which parameters correspond to values at different tk such that θ ( tk , θ̂ ) = θ̂k . Given θ1 , the parameter coefficients for θ2 can thus be calculated by evaluating the θ1 at the control points t2k : θ̂2k = θ 1 ( t2k , θ̂ 1 ) = K1∑ b=1 φ1b ( t 2 k ) θ̂ 1 b for k = 1 , ... , K2 . ( 7 ) Interpolation only works with basis functions where the parameters correspond to control point locations . For piecewise constant basis functions , tk corresponds to the cell centers , and for piecewise linear basis functions , tk corresponds to the endpoints at element boundaries . Projection . Function projection can be used with any basis expansion . Given the function θ1 ( t ) , the coefficients θ̂2k are solved by a minimization problem : min θ̂2k ∫ T 0 ( θ1 ( t , θ̂1 ) − θ2 ( t , θ̂2 ) ) 2 dt = min θ̂2k ∫ T 0 ( K1∑ a=1 θ̂1aφ 1 a ( t ) − K2∑ k=1 θ̂2kφ 2 k ( t ) ) 2 dt . ( 8 ) Appendix C.1 includes the details of the numerical approximation of the loss and its solution . The integral is evaluated using Gaussian quadrature over sub-cells . The overall calculation solves the same linear problem repeatedly for each coordinate of θ used by the call toR .

Summary: The paper introduces a new neural ODE with time-continuos depth and time-dependent parameterization of the network. The time-dependent parameterization is composed of a linear combination of time-dependent basis function and learnable “base parameters”.   By relying on this continuous reparameterization using basis function the authors highlight the close connection with the compressibility of the network by doing a basis change during inference time to a lower-dimensional space of basis function, where the “base parameters” in the lower dimensional space can be computed either via projection or linear interpolation from the higher-dimensional basis space.   Additionally, the authors introduces a novel notion of time-continuous batch normalization and they present theoretical results pertaining to the implicit regularization potential of their model when solved via RK-methods and empirical results for image classification and time-continuous analogues of transformers.