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 proposes weight compression technique for learnt ODE-Nets without retraining the model while maintaining near state-of-the-art performance. Their proposal includes treating the weight matrices as the linear combination of continuous-in-time basis functions, where the linear combination coefficients are learnt by optimizing the loss function between the learnt weights and the approximate weights (computed by combining basis functions). Empirical evidence suggests that such a scheme compresses ODE-Nets on image classification and sentence-tagging tasks.