Stateful ODE-Nets using Basis Function Expansions
Year: 2022

Abstract:
The recently-introduced class of ordinary differential equation networks (ODENets) establishes a fruitful connection between deep learning and dynamical systems. In this work, we reconsider formulations of the weights as continuous-indepth functions using linear combinations of basis functions which enables us to leverage parameter transformations such as function projections. In turn, this view allows us to formulate a novel stateful ODE-Block that handles stateful layers. The benefits of this new ODE-Block are twofold: first, it enables incorporating meaningful continuous-in-depth batch normalization layers to achieve state-of-theart performance; second, it enables compressing the weights through a change of basis, without retraining, while maintaining near state-of-the-art performance and reducing both inference time and memory footprint. Performance is demonstrated by applying our stateful ODE-Block to (a) image classification tasks using convolutional units and (b) sentence-tagging tasks using transformer encoder units.

heading: 1 Introduction
text: 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 cannot 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.

heading: 2 Related Work
text: 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.

heading: 3 Basis Function View of Continuous-in-depth Weight Functions
text: 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.

heading: 3.1 Basis Transformations
text: 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.

heading: 4 Stateful ODE-Nets
text: Modern neural network architectures leverage stateful modules such as BatchNorms [22] to achieve state-of-the-art performance, and recent work has demonstrated that ODE-Nets also can benefit from normalization layers [16]. However, incorporating normalization layers into ODE-Nets is challenging; indeed, recent work [45] acknowledges that it was not obvious how to include BatchNorm layers in ODE-Nets. The reason for this challenge is that normalization layers have internal state parameters with specific update rules, in addition to trainable parameters.

heading: 4.1 Stateful ODE-Block
text: To formulate a stateful ODE-Block, we consider the following differential equation:
ẋ(t) = R (θg(t), θs(t), x(t)) , (9)
which is parameterized by two continuous-in-depth weight functions θg(t, θ̂g) and θs(t, θ̂s), respectively. For example, the continuous-in-depth BatchNorm function has two gradient-parameter functions in θg(t) – scale s(t) and bias b(t) – and two state-parameter functions in θs(t) – mean µ(t) and variance σ(t). Using a functional programming paradigm inspired by Flax [20], we replace the
internal state update of the module with a secondary output. The continuous update function R is split into two components: the forward-pass output ẋ = Rx and the state update output θs∗ = Rs. Then, we solve ẋ using methods for numerical integration (denoted by scheme), followed by a basis projection of θ̂s∗. We obtain the following input-output relation for the forward pass during training:{
xout = StatefulOdeBlockx [Rx, scheme,∆t, φ] (θ̂g, θ̂s, xin) θ̂s∗ = StatefulOdeBlocks [Rs, scheme,∆t, φ] (θ̂g, θ̂s, xin),
(10)
which we jointly optimize for θg(t) and θs(t) during training. Optimizing for θg and θs involves two coupled equations, updating the gradient with respect to the loss L and a fixed-point iteration on θs,
θg∗(t) = θg(t)− ∂L∂θg(t) (θ g(t), θs(t), x0), (11) θs∗(t) = Rs (θg(t), θs(t), x(t)) . (12)
The updates to θg(t) are computed by backpropagation of the loss through the ODE-Block with respect to its basis coefficients ∂xout
∂θ̂g . Optimizing θ̂s(t) involves projecting the update rule back onto
the basis during every training step.

heading: 4.2 Numerical Solution of Stateful Normalization Layers
text: While Eq. (12) can be computed given a forward pass solution x(t), this naive approach requires implementing an additional numerical discretization. Instead, consider that each call toR at times ti ∈ [0, T ] during the forward pass outputs an updated state θ̄si = Rs(ti). This generates a point cloud {ti, θ̄si } which can be projected back onto the basis set by minimizing the point-wise leastsquared-error
θ̂s∗ = arg min θ̂s∗ Nstate∑ i=1 ( θ̄si − K∑ k=1 φk(ti)θ̂ s∗ k )2 . (13)
Algorithm 1 in the appendix describes the calculation of Eq. (10), fused into one loop. When using forward Euler integration and piecewise constant basis functions, this algorithm reduces to the forward pass and update rule in a ResNet with BatchNorm layers. This theoretical formalism and algorithm generalizes to any combination of stateful layer, basis set, and integration scheme.
5 A Posteriori Compression Methodology
Basis transformations of continuous functions are a natural choice for compressing ODE-Nets because they couple well with continuous-time operations. Interpolation and projection can be applied to the coefficients θ̂g and θ̂s to change the number of parameters needed to represent the continuous functions. Given the coefficients θ̂1 to a function θ1(t) on a basis (φ1,K1), we can determine new coefficients θ̂2 on a different space (φ2,K2). Changing the basis size can reduce the total number of parameters (to K2 < K1), and hence the total storage requirement. This representation enables compression in a systematic way — in particular, we can consider a change of basis to transform
learned parameter coefficients θ̂1 to other coefficients θ̂2. To illustrate this, Figure 2 shows different basis representations of a weight function.
Note that this compression methodology also applies to basis function ODE-Nets that are trained without stateful normalization layers. Further, continuous-time formulations can also decrease the number of model steps, increasing inference speed, in addition to decreasing the model size.

heading: 6 Advantages of Higher-order Integrators for Compression
text: To implement any ODE-Net, it is necessary to approximately solve the corresponding ODE using a numerical integrator. There are many possible integrators one can use, so to take full advantage of our proposed methodology, we examine the advantages of certain integrators on compression from a theoretical point of view. In short, for the same step size, higher-order integrators exhibit increasing numerical instability if the underlying solution changes too rapidly. Because these large errors will propagate through the loss function, minimizing any loss computed on the numerically integrated ODE-Net will avoid choices of the weights that result in rapid dynamics. This is beneficial for compression, as slower, smoother dynamics in the ODE allow for coarser approximations in θ.
To quantify this effect, we derive an asymptotic global error term for small step sizes. Our analysis focuses on explicit Runge–Kutta (RK) methods, although similar treatments are possible for most other integrators. For brevity, here we use h to denote step size in place of ∆t. A p-stage explicit RK method for an ODE ẏθ(t) = f(yθ(t), θ(t)) provides an approximation yh,θ(t) of yθ(t) for a given step size h on a discrete grid of points {0, h, 2h, . . . } (which can then be joined together through linear interpolation). As the order p increases, the error in the RK approximation generally decreases more rapidly in the step size for smooth integrands. However, a tradeoff arises through an increased sensitivity to any irregularities. We consider these properties from the perspective of implicit regularization, where the choice of integrator impacts the trained parameters θ.
Implicit regularization is of interest in machine learning very generally [33], and in recent years it has received attention for NNs [25, 35]. In a typical theory-centric approach to describing implicit regularization, one derives an approximately equivalent explicit regularizer, relative to a more familiar loss function. In Theorem 1, we demonstrate that for any scalar loss function L, using a RK integrator implicitly regularizes towards smaller derivatives in f . Since θ can be arbitrary, we consider finite differences in time. To this effect, recall that the m-th order forward differences are defined by ∆mh f(x, t) = ∆ m−1 h f(x, t+ h)−∆ m−1 h f(x, t), with ∆ 0 hf(x, t) = f(x, t). Also, for any t ∈ [0, T ], we let ιh(t) = bt/hc · h denote the nearest point on the grid {0, h, 2h, . . . } to t.
Theorem 1 (Implicit regularization) There exists a polynomial P depending only on the RungeKutta scheme satisfying P (0) = 0, and a smooth function ȳθ(t) depending on f, θ, and the scheme, such that for any t ∈ [0, T ], as h→ 0+,
L(yh,θ(t)) = L(ȳθ(t)) + h p∇L(ȳθ(t)) · eh,θ(t) +O(hp+1), (14)
where ėh,θ(t) = ∂f∂y (yh,θ(t), θ(ιh(t)))eh,θ(t) + P ◦ D p h,θ(t) and Dph,θ(t) = { h−m∆mh ∂lfj ∂yl fj(yh,θ(t), θ(ιh(t))) } j=1,...,d, l+m≤p+1 . (15)
The proof is provided in Appendix A. Theorem 1 demonstrates that Runge-Kutta integrators implicitly regularize toward smaller derivatives/finite differences of f of orders k ≤ p + 1. To demonstrate the effect this has on compression, we consider the sensitivity of the solution y to θ(t), and hence to the choice of basis functions, using Gateaux derivatives. A smaller sensitivity to θ(t) allows for coarser approximations before encountering a significant reduction in accuracy. The sensitivity of the solution in θ(t) is contained in the following lemma.
Lemma 1 There exists a smooth function θ̄ depending only on θ such that for any smooth function ϕ(t),
Dϕyh,θ(t) := d
d yh,θ+ ϕ(t) ∣∣∣∣ =0 = ∫ t 0 eFh,θ(s,t) ∂f ∂θ (yh,θ(s), θ̄(s))ϕ(s)ds+O(hp),
where Fh,θ(s, t) = ∫ t s ∂f ∂y (yh,θ(u), θ̄(u))du.
Simply put, Lemma 1 shows that the sensitivity of yh,θ(t) in θ decreases monotonically with decreasing derivatives of the integrand f in each of its arguments. These are precisely the objects that appear in Theorem 1. Therefore, a reduced sensitivity can occur in one of two ways:
(I) A higher-order integrator is used, causing higher-order derivatives to appear in the implicit regularization term eh,θ. By the Landau-Kolmogorov inequality [26], for any differential D and integer m ≥ 1, ‖Dmf‖ ≥ cm‖f‖(‖Df‖/‖f‖)m for some cm > 0. Hence, by implicitly regularizing towards smaller higher-order derivatives/finite differences, we should expect more dramatic reductions in the first-order derivative/finite difference as well.
(II) A larger step size h is used. Since doing so can lead to stability concerns during training, we later consider a refinement training scheme [6, 37] where h is slowly reduced.
In Figure 3, we verify strategy (I), showing that the 4th-order RK4 scheme exhibits improved test accuracy for higher compression ratios over the 1st-order Euler scheme. Unfortunately, higherorder integrators typically increase the runtime on the order of O(p). Therefore, some of our later experiments will focus on strategy (II), which avoids this issue.

heading: 7 Empirical results
text: We present empirical results to demonstrate the predictive accuracy and compression performance of Stateful ODE-Nets for both image classification and NLP tasks. Each experiment was repeated with eight different seeds, and the figures report mean and standard deviations. Details about the training process, and different model configurations are provided in Appendix E. Research code is provided as part of the following GitHub repository: https://github.com/afqueiruga/StatefulOdeNets. Our models are implemented in Jax [3], using Flax [20].

heading: 7.1 Compressing Shallow ODE-Nets for Image Classification Tasks
text: First, we consider shallow ODE-Nets, which have low parameter counts, in order to compare our models to meaningful baselines from the literature. We evaluate the performance on both MNIST [27] and CIFAR-10 [24]. Here, we train our Stateful ODE-Net using the classic Runge-Kutta (RK4) scheme and the multi-level refinement method proposed by [6, 37].
Results for MNIST. Due to the simplicity of this problem, we consider for this experiment a model that has only a single OdeBlock with 8 units, and each unit has 12 channels. Our model has about 18K parameters and achieves about 99.4% test accuracy on average. Despite the lower parameter count, we outperform all other ODE-based networks on this task, as shown in Table 2. We also show an ablation model by training our ODE-Net without continuous-in-depth batch normalization layers while keeping everything else fixed. The ablation experiment shows that normalization provides only a slight accuracy improvements; this is to be expected as MNIST is a relatively easy problem.
Next, we compress our model by reducing the number of basis functions and timesteps from 8 down to 4. We do this without retraining, fine-tuning, or revisiting any data. The resulting model has approximately 10K parameters (a 45% reduction), while achieving about 99.2% test accuracy. We can compress the model even further to 7K parameters, if we are willing to accept a 2.6% drop in accuracy. Despite this drop, we still outperform a simple NODE model with 84K parameters. Further, we can see that the inference time (evaluated over the whole test set) is significantly reduced.
Results for CIFAR-10. Next, we demonstrate the compression performance on CIFAR-10. To establish a baseline that is comparable to other ODE-Nets, we consider a model that has two OdeBlocks with 8 units, and the units in the first block have 16 channels, while the units in the second block have 32 channels. Our model has about 207K parameters and achieves about 90.4% accuracy. Note, that our model has a similar test accuracy to that of a ResNet-20 with 270K parameters (a ResNet-20 yields about 91.25% accuracy [19]). In Table 3 we see that our model outperforms all other ODE-Nets on average, while having a comparable parameter count. The ablation model is not able to achieve state-of-the-art performance here, indicating the importance of normalization layers.
As before, we compress our model by reducing the number of basis functions and timesteps from 8 down to 4. The resulting model has 114K parameters and achieves about 89.9% accuracy. Despite the compression ratio of nearly 2, the performance of our compressed model is still better as compared to other ODE-Nets. Again, it can be seen that the inference time on the test set is greatly reduced.

heading: 7.2 Compressing Deep ODE-Nets for Image Classification Tasks
text: Next, we demonstrate that we can also compress high-capacity deep ODE-Nets trained on CIFAR-10. We consider a model that has 3 stateful ODE-blocks. The units within the 3 blocks have an increasing number of channels: 16,32,64. Additional results for CIFAR-100 are provided in Appendix D.
Table 4 shows results for CIFAR-10. Here we consider two configurations: (c1) is a model trained with refinement training, which has piecewise linear basis functions; (c2) is a model trained without refinement training, which has piecewise constant basis functions. It can be seen that model (c2) achieves high predictive accuracy, yet it yields a poor compression performance. In contrast, model (c1) is about 1.5% less accurate, but it shows to be highly compressible. We can compress the param-
eters by nearly a factor of 2, while increasing the test error only by 0.6%. Further, the ablation model shows that normalization layers are crucial for training deep ODE-Nets on challenging problems. Here, the ablation model that is trained without our continuous-in-depth batch normalization layers does not converge and achieves a performance similar to tossing a 10-sided dice.
In Figure 4a, we show that the compression performance of model (c1) depends on the particular choice of basis set that is used during training and inference time. Using piecewise constant basis functions during training yields models that are slightly less compressible (green line), as compared to models that are using piecewise linear basis functions (blue line). Interestingly, the performance can be even further improved by projecting the piecewise linear basis functions onto piecewise constant basis functions during inference time (red line). In Figure 4b, we show that we can also decrease the number of time steps NT during inference time, which in turn leads to a reduction of the number of FLOPs during inference time. Recall, NT refers to the number of time steps, which in turn determines the depth of the model graph. For instance, reducing NT from 16 to 8, reduces the inference time by about a factor of 2, while nearly maintaining the predictive accuracy.

heading: 7.3 Compressing Continuous Transformers
text: A discrete transformer-based encoder can be written as xt+1 = xt + T (q, xt) +M(ρ, xt + T (q, xt)) (16)
where T (q, x) is self attention (SA) with parameters q (appearing twice due to the internal skip connection) andM is a multi-layer perceptron (MLP) with parameters ρ. Repeated transformer-based encoder blocks can be phrased as an ODE with the equation
ẋ = T (q(t), x) +M(ρ(t), x+ T (q(t), x)). (17) Recognizing θ(t) = {q(t), ρ(t)}, this formula can be directly plugged into the basis-functions and OdeBlocks to create a continuous-in-depth transformer, illustrated in Figure 1. Note that the
OdeBlock is continuous along the depth of the computation; the sequence still uses discrete tokens and positional embeddings. With a forward Euler integrator, ∆t = 1, a piecewise constant basis, and K = NT , Eq. (17) generates an algebraically equivalent graph to the discrete transformer.
We apply the encoder model to a part-of-speech tagging task, using the English-GUM treebank [47] from the Universal Dependencies treebanks [36]. Our model uses an embedding size of 128, with Key/Query/Value and MLP dimensions of 128 with a single attention head. The final OdeBlock has K = 64 piecewise constant basis functions and takes NT = 64 steps.
In Figure 5a, we present the compression performance for different basis sets. Staying on piecewise constant basis functions yields the best performance during testing (red line). The performance slightly drops when we project the piecewise constant basis onto a piecewise linear basis (blue line). Projecting to a discontinuous linear basis (green line) performs approximately as well as the piecewise constant basis. In Figure 5b, we show the effect of decreasing the number of time steps NT during inference time on the predictive accuracy. In the regime where NT ≤ K, there is no significant loss obtained by reducingK. Note that there is a divergence whenNT > K, where the integration method will skip over parameters. However, projection to NT ≤ K incorporates information across a larger depth-window. Thus, projection improves graph shortening as compared to previous ODE-Nets.

heading: 8 Conclusion
text: We introduced a Stateful ODE-based model that can be compressed without retraining, and which can thus quickly and seamlessly be adopted to different computational environments. This is achieved by formulating the weight function of our model as linear combinations of basis functions, which in turn allows us to take advantage of parameter transformations. In addition, this formulation also allows us to implement meaningful stateful normalization layers that are crucial for achieving predictive accuracy comparable to ResNets. Indeed, our experiments showcase that Stateful ODENets outperform other recently proposed ODE-Nets, achieving 94.1% accuracy on CIFAR-10.
When using the multi-level refinement training scheme in combination with piecewise linear basis functions, we are able to compress nearly 50% of the parameters while sacrificing only a small amount of accuracy. On a natural language modeling task, we are able to compress nearly 98% of the parameters, while still achieving good predictive accuracy. Building upon the theoretical underpinnings of numerical analysis, we demonstrate that our compression method reliably generates consistent models without requiring retraining and without needing to revisit any data. However, a limitation of our approach is that the implicit regularization effect introduced by the multi-level refinement training scheme can potentially reduce the accuracy of the source model.
Hence, future work should investigate improved training strategies, and explore alternative basis function sets to further improve the compression performance. We can also explore smarter strategies for compression and computational graph shortening by pursuing hp-adaptivity algorithms [38].

heading: Acknowledgments
text: We are grateful for the generous support from Amazon AWS and Google Cloud. NBE and MWM would like to acknowledge IARPA (contract W911NF20C0035), NSF, and ONR for providing partial support of this work. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.

Review summaries:
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. 

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.

This paper discusses using piecewise constant or piecewise linear basis functions to represent time-dependent parameters in an ODE-net. This can also be used to represent parameters in a time-dependent BatchNorm layer. A change of basis can be performed to reduce the number of parameters or increase the discretization step sizes.