Stateful ODE-Nets using Basis Function Expansions
Authors: Alejandro Francisco Queiruga
N. Benjamin Erichson
Liam Hodgkinson
Michael W. Mahoney
Venue: NeurIPS 2021 Poster (2021)
Keywords: Neural ODEs
Dynamical Systems
Differential Equations

Abstract: The recently-introduced class of ordinary differential equation networks (ODE-Nets) establishes a fruitful connection between deep learning and dynamical systems. In this work, we reconsider formulations of the weights as continuous-in-depth 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-the-art 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.

Stateful ODE-Nets using Basis Function Expansions
Alejandro Queiruga' Google Research
N. Benjamin Erichson" University of Pittsburgh erichson&pitt.co edu
Liam Hodgkinson ICSI and UC Berkeley
afq@google.co com
liam. www..hh.kkinnoombbreelel.eeeeeee edu
Michael W. Mahoney
ICSI and UC Berkeley
mmahoneydstat. umahoneeyssttt.eerkee www.eeeetttatt...eleeeee..ee.e.e edu
Abstract
The recentl--inrooduce class of ordinary differential equation networks (ODE- Nets) establishes fruitful connection between deep learning and dynamical sys- tems. In this work, we reconsider formulations of the weights continuous-in- depth 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 meaningful pantinuous-in-depe batch normalization layers to achieve state-of of art performance; second, it enables compressing the weights through change of basis, without retraining, while maintaining near state-of-the-ar performance and reducing both inference time and memory footprint. Performance is demon- strated by applying our stateful ODE-Block to (a) image classification tasks using convolutional units and (b) sentence-tagging tasks using transformer encoder units.
twofold: first, it enables incorporating
Introduction
The interpretation of neural networks (NNs) as discretizations of differential equations 7
has recently unlocked 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 [1 399 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 "bug"). Further, they can be designed to be highly memory efficient [3 [1). However, one major drawback of current ODE-Nets is that the predictive accuracy for tasks such as image classification is often inferior compared to other state-of-the-ar 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 mechanism for handling layers with internal state (i.c., 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 221, layer dropout [2] and adaptive depth [10 28 which are useful for regularization and compressing traditional NNs, do not necessarily apply meaningful manner to ODt-Nets. That is, these methods do not provide systematic scheme to derive smaller networks from 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 (NcurIPS 2021).

---

InRut:
StatefulDdeBBlock
Outrut:
R(((t).
Text[1]
Tex][o]
Embed
((L)
Σ )c())a
POS[O]
POS[1]
x(1)
Tex
Dense
POS[2]
Tex][N]
Position
POS[N]
Self Attention
c(O)
x(t)
(
MLP
dar
A(t 0)
A(t) = (=())(())()=(()))))
A(t T)
Figure I: Sketch of a continoous-in-depth transformer-encoder. The model architecture consists of a sparse embedding layer followed by OdeBlock that integrates the encoder feed into classification layer to determine parts of speech. The model graph is generated by nesting the residual R on the right into time-integration scheme in the OdeBlock. The weights for each call to R are determined by evaluating the basis expansion. Internal hidden states evolve continuously during the forward pass. This is illustrated by smoothly varying attention matrix.
Hence, there 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 built upon the theoretical underpinnings of numerical analysis. Following recent work by 34 37], express the weights of an ODE-Net as continuous- continuous-in-dppti functions using linear combinations functions. However, unlike prior works, consider parameter transformations through change of basis functions, i.., adding basis functions, decreasing the number of basis functions, and function projections. In turn, we able (a) design deep ODE-Nets that leverage meaningful continuous- in-depth batch normalization layers to boost the predictive accuracy of ODE-Nets, e.g., we achieve 94.4% test accuracy CIFAR-IO, and 79.9% CIFAR-100: and (b) introduce methodology for compressing ODE-Nets, c.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.
basis
In the following, we refer to our model as Stateful ODE-Net. Here are our main contributions.
Stateful ODE-Block. We introduce novel stateful ODE-Block that enables the integration of stateful network layers (Sec.[4.]. To do so, we view continuous-in-deph 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 [
Stateful Normalization Layer. We introduce stateful modules using function projections. This enables us to introduce continuous-in-depti batch normalization layers for ODE-Nets (Sec. 4.2). our experiments, we show that such stateful normalization layers are crucial for ODE-Nets to achieve thte-offthhe-at performance on mainstream image classification problems (Sec.
A Posteriori Compression Methodology. We introduce a methodology to compress ODE-Nets without retraining revisiting any data, based parameter interpolation and projection (Sec. 5) That is, we can systematically derive smaller networks from single deep ODE-Net source model through 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 Higher-order Integrators for Compression. We examine effects of training continuous-in-denn weight functions through their discretizations (Sec. 6 Our key insight in Theorem is that higher-order integrators introduce additional implicit regularization which is crucial for obtaining good compression performance in practice. Proof's are provided in Appendix
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 xx0)= Zin to the final output (TT) Fout as
#(t) = F(0, (t),
(I)
Here, the function F can be any NN that is parameterized by 0, with two inputs and f. The parameter t E [0, I"] in this ODE represents time, analogous the depth of classical network architectures. Using finer temporal discretization (with a smaller At) or a larger terminal time T, corresponds to deeper network architectures. Feed-forward evaluation of the network is performed by numerically integrating Eq. (
Tout Jen + L. F(0,x(1),)) =ddeBlock [F. 0, scheme, Ar, E [0, (Fan).
Inspired by this view, numerous ODE and FDE-based network architectures [ 8 49], and continuous-time recurrent units 4
15
29 30
40
41 have been proposed.
Recently, the idea of representing continuous- time weight function linear combination of basis functions has been proposed by [34] and 33.. This involves using the following formulation:
*(t) (R(((t .()).
where R is now parameterized by continuously-varying weight function ((t; 0). In turn, this weight function is parameterized by 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 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 often chosen to be orthogonal, inspired by multi-level refinement training, we shall consider non-orthogonal basis sets in our experiments. Table[ the advantages of our model compared other related models.
Table 1: Comparison of our Stateful ODE-Net to other dynamical system inspired models.
Model
Multi-level Compression
Basis Function View
Stateful Layers
NODE
Multi-level ResNet 6
Galerkin ODF-Net
X
ContinuousNet [37
X
X
Stateful ODE-Net (ours)
3 Basis Function View of Continuous-in-depth Weight Functions
Let (t;; be an arbitrary weight funetion that depends on depth (or time) t, which takes as argument a vector of real-valued weights é. Given a basis o, we can represent as a linear combination of K (continuous-time) basis functions (t):
((t:0) Σ (t) 6g.
(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 A! = T/K, where T is the "time" at the end of the ODE- Net. The summation Eq. involves piesewise-constant indicator functions &b (f) satisfying
11 te [ 1)At, kAt
(x(()=
(5)
0, otherwise.
3

---

Piecewise linear basis functiins. This basis consists of evenly spaced elements where each parameter 0g corresponds to the value of f(t2) at element boundaries Ea T(k h: 1)/(K 1). Each basis function is a "hat" function around the point bK.
tt Att
E [b - 1)2t, kAt
(x(1)
- (t kAt)/At, E [kAt, (h + 1)4t)
(6)
0,
otherwise.
Piecewise linear basis functions have compact and overlapping supports; i.c,, they are not orthogo- nal, unlike piecewise constant basis functions.
3.1 Basis Transformations
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 necessarily give the same result for the same function.
Interpolation. Some basis functions use control points for which parameters correspond values at different ty such that ((tx. Q) = 4.. Given 8', the parameter coefficients for 8 can thus be calculated by evaluating the 0¹ at the control points th:
Ki
8 = 0 0 01)
(1 for k 1, K2.
Interpolation only works with basis functions where the parameters correspond to control point locations. For piecewise constant basis functions, lg corresponds to the cell centers, and for piecewise linear basis functions, 4 corresponds the endpoints at element boundaries.
Function projection can be used with any basis expansion. Given the function (t), the coefficients 0% are solved by minimization problem:
Projection.
min (' (0'(tt)))) 62
K. ( - Σ 0 A=1
K₂
ott, dt = min 0%
dr.
(8)
Appendix C.I lincludes 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 to R.
Stateful ODE-Nets
Modern neural network architectures leverage stateful modules such as BatchNorms [22] achieve state-of-the-ar artt performance, and recent work has demonstrated that ODE-Nets also can benefit from normalization layers [66. 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.
4.1 Stateful ODE-Block
To formulate stateful ODE- Block, consider the following differential equation:
(9) which parameterized by two continuous-in-depth weight functions (1, 0s) and 0* 0²). re-
*(t) R (t), xt))
spectively. For example, the continuous-in-deptt BatchNorm function has two gradient-paramete functions in 09 (t) scale x()) and bias b(t) and two state-parameter functions in (*(() - mean p(l)
and variance ot). Using functional programming paradigm inspired by Flax [20), replace the

---

internal state update of the module with secondary output. The continuous update function R is split into two components: the forward-pass output i R2 and the state update output 6** R Then, we solve è using methods for numerical integration (denoted by scheme), followed by basis projection of (s*,, We obtain the following imput-output relation for the forward pass during training:
StatefulOdeBlock, [Rz, scheme, At, Φ) (09, 0= ,in) StatefulDdeBlock, [R,, scheme, At, Ø) (09, 0- Tim,
Bout 0
(10)
which we jointly optimize for (t) and (t) during training. Optimizing for 0 and 0 involves two coupled equations, updating the gradient with respect to the loss L and a fixed-point iteration on 0",
0 (t) = 8 (t) 30P(C)
aL (00 (t), 0 (t), (t),
(11)
0 (f) R. (1),,0((),)(())
(12)
The updates (1) are computed by backpropagation of the loss through the ODF-Block with respect its basis coefficients Smaue Optimizing (t) involves projecting the update rule back onto the basis during every training step.
80,
4.2 Numerical Solution of Stateful Normalization Layers
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 to R. at times t; E [0, during the forward pass outputs an updated state 0: = (t,).. This generates a point cloud (tz, ) which be projected back onto the basis set by minimizing the point-wise least- squared-error
Natate
arg min ga
Σ
0 P du (t.)))
(13)
Algorithm Min the appendix describes the calculation of Eq. (T)). fused into one loop. When using forward Euler integration and piecewise constant basis functions, this algorithm reduces to the forward pass and update rule a ResNet with BatchNorm layers. This theoretical formalism and algorithm generalizes any combination of stateful layer, basis set, and integration scheme.
5 A Posteriori Compression Methodology
Basis transformations of continuous functions are natural choice for compressing ODE-Nets because they couple well with continuous- time operations. Interpolation and projection can be applied to coefficients 09 and ja to change the number of parameters needed to represent the continuous functions. Given the coefficients g to a function 0 (t.) on a basis (Ø) K1), ) we can determine new coefficients 8 on a different space (d² K2). Changing the basis size can reduce the total number of parameters (to K₂ < K1). and hence the total storage requirement. This representation enables compression in a systematic way in particular, we can consider change of basis to transform
0.2
yu
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Network depth t
Network depth
Network depth
Figure 2: Example of projecting a component of the query kernel from a continuous-in in-puu-----ppppte depth transformer. The model was trained with K 64 piecewise constant basis functions (left) and then projected to: K = 16 piecewise constant basis functions (middle), and K = 16 piecewise linear basis functions (left). Circles denote the control points (knots) corresponding to parameters 0,.
5

---

learned parameter coefficients d' to other coefficients 0² To illustrate this, Figure shows different basis representations of 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.
Advantages of Higher-order Integrators for Compression
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, 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 0.
examine the advantages of certain integrators on compression
To quantify this effect, derive 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 use h to denote step size in place of At. A p-stage explicit RK method for an ODE (t) = f(yo(t). O(t)) provides an approximation yh, of yo for a given step size h on discrete grid of points {0, h, 2h....} (which can then be joined together through linear interpolation). As the order 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 implicil regularization, where the choice of integrator impacts the trained parameters 0.
Implicit regularization is of interest in machine learning very generally [3], and in recent years it has received attention for NNs [25] 35). In typical theory-centric approach to describing implicit regularization, one derives an approximately equivalent explicis regularizer, relative to a more familiar loss function. In Theorem 1 we demonstrate that for any scalar loss function L, using 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
defined by AR [xx, 1) A 1x-110 [xx ( + h)) - l), with J [(x () = [xx, (). Also, for any E [0, T), 1
we let Ln (t) = [t/h] denote the nearest point on the grid {0, h, 2h..... to
Theorem (Implicit regularization) There exists polynomial P depending only on the Runge- Kutta scheme satisfying P(O)= = O, and a smooth function To (t) depending on f, 6, and the scheme, such that for any E [0, as h- - 0+,
(1)) = ((To())) h²V (t) + O(htt1),
(14)
where (t,
of By
P (t)) and
0' Si Fyt f, (,),,,(t))000 ((n(())))
(h)(()= h- ---ma"" Ar
(15)
L+m d+mpp.
The proof is provided in Appendix A| Theorem demonstrates that Runge- Kutta integrators implicitly regularize toward smaller derivatives/fnite differences of f of orders k S p 1. To demonstrate the effect this has on compression, we consider the sensitivity of the solution y to O(t), and hence to the choice of basis functions, using Gateaux derivatives. A smaller sensitivity to 0(t) allows for coarser approximations before encountering a significant reduction in accuracy. The sensitivity of the solution in (() is contained in the following lemma.
Lemma There exists a smooth,function 0 depending only on 0 such that for any smooth function (p(().
D(ph, = d de vn vh, = S "Phholae)) 01 00 (s), (()))((ss)))ss))) OOPP)
where t)) =ft oy (fn, ((z)))d. (()))(

---

1.0
1.0 3.2M
3.2M
2,80
0.9
2.8M
2.4M
Accuracy
2.4M
2,0M
2.0M DM
Test
6M
0.4
baseline
baseline
RK4
RK4
0.6
Buler
Euler
1.0
1.2
1.4
1.8
2.0
1.0
1.2
1.6 Compression Ratio
1.8
2.0
Compression Ratio
(a) Compression using interpolation.
(b) Compression using projections.
Figure 3: Higher-order integrators introduce additional implicit regularization into learned continuous weight functions. By experimenting CIFAR-IN 10, it is empirically observed that training with the RK4 scheme improves the compression-performaan compared to using the forward Euler scheme.
Simply put, Lemma shows that the sensitivity of (tt in decreases monotonically with decreasing derivatives of the integrand f in each of its arguments. These are precisely the objects that appear in Theorem [
Therefore, a reduced sensitivity can occur in one of two ways:
(I) A higher-order integrator used, causing higher-order derivatives to appear the implicit regularization term 6n,0- By the Landau-Kolmoge Kolmogorov inequality [26), for any differential and integer m 2 1, || 2 Cm f[(( If for some Cy. > O. Hence, by implicitly regularizing towards smaller higher-order derivatives/finite differences, should expect more dramatic reductions in the first-order derivative/finite difference as well.
(II) A larger step size used. Since doing so can lead stability concerns during training, we later consider refinement training scheme 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 Ist-order Euler scheme. Unfortunately, higher- order integrators typically increase the runtime on the order of Opp). Therefore, some of our later experiments will focus on strategy (II), which avoids this issue.
Empirical results
We present empirical results to demonstrate the predictive accuracy and compression performance of Stateful ODE-Nets for both image classification and NL.P tasks. Each experiment repeated with eight different seeds, and the figures report mean and standard deviations. Details about the training provided in Appendix H
process, and different model configurations part of the following GitHub repository: httpps://itte www.b...om/afuueeeeeee wttp///fguuiiirug/ttttttuuuuuode Our models implemented in Jax using Flax 20
Research code is provided as
7.1 Compressing Shallow ODE-Nets for Image Classification Tasks
First, we consider shallow ODF-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 10 (24). Here, we train Stateful ODE-Net using the classic Runge- Kutte- (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 single OdeBlock with units, and each unit has 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 bDE-d networks on this task, as shown Table 2 We also show an ablation model by training our ODE-Net without continuous-in-deph batch normalization layers while keeping everything else fixed. The ablation experiment shows that normalization provides only slight accuracy improvements; this is to be expected as MNIST is relatively easy problem.
7

---

Table 2: Compression performance and test accuracy of shallow ODE-Nets on MNIST.
Model
Best
Average
Min
# Parameters
Compression
Inference
NODE
96.4%
84K
ANODE
98.2%
84K
2nd-Order
99.2%
20K
A4+NODE+NDDE 50
98.5%
84K
Ablation Model
99.3%
99. %
98.9%
18K
Stateful ODE-Net (ours)
99.6 %
99,4 %
99.3% %
18K
baseline
1.7 (s)
+ (compressed)
99.4%
99.2%
99.1%
10K
45%
1.2 (8)
(compressed)
97.8%
96.9%
95.7%
7K
%
1.1 (s)
Table 3: Compression performance and test accuracy of shallow ODE-Nets on CIFAR-10.
Model
Best
Average
Min
# Parameters
Compression
Inference
HyperODENel 45
87.9%
460K
SDE BNN (+ STL)
89.1% 1%
460K
Hamiltonian 42
89.3%
264K
NODE
53.7%
172K
ANODE
60.6%
172K
A4+NODE+NDDE (50
59.9%
107K
Ablation Model
88.9%
88.4%
88.1% 1%
204K
Stateful ODE-Net (ours)
90.7% %
90,4 %
90.1 %
207K
baseline 45%
2.1 (s)
- (compressed)
90.3%
89.9%
89.6%
114K
1.6
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 accept 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 units, and the units 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 ResNet-20 with 270K parameters (a ResNet-20 yields about 91 25% accuracy [9)). In Table we see that our model outperforms all other ODE-Nets on average, while having a comparable parameter count. The ablation model 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 down to 4. The resulting model has 114K parameters and achieves about 89 89.9% accuracy. Despite the compression ratio of nearly 2, the performance of our compressed model still better as compared to other ODE-Nets. Again, it can be seen that the inference time on the test set is greatly reduced.
7.2 Compressing Deep ODE-Nets for Image Classification Tasks
Next, we demonstrate that we can also compress high-capacity deep ODE-Nets trained on CIFAR- We consider model that has 3 stateful ODE-blocks. The units within the blocks have an increasing number channels: 16,32,64. Additional results for CIFAR- 100 are provided in Appendix D
Table 4 shows results for CIFAR-10 10. Here we consider two configurations: (c1) is model trained with refinement training. which has piecewise linear basis functions; (c2) is model trained without refinement training, which has piecewise constant basis functions. It can be seen that model (c2) achieves high predictive accuracy, yet yields poor compression performance.
contrast, model (cl) is about 1.5% less accurate, but it shows to be highly compressible. We can compress the param-

---

Table 4: Compression performance and test accuracy of deep ODE-Nets on CIFAR-10.
Model
Best
Average
Min
# Parameters
Compression
Inference
ResNet- 110
[9
93.4%
1.73M
ResNNet-
93.8%
1.92M
MidPoint-6l [5]
92.8%
1.78M
Continuous.Net 37
94.0%
93.8%
93.5%
3.19M
Ablation Model
10.0%
10.0%
10.0%
1.62K
Stateful ODE-Net (cl)
93.0%
92,4%
92.1%
1.63M
baseline 48%
3.4 (s)
(compressed)
92.2%
91.8%
91.1%
0.85M
2.3 (s)
Stateful ODB-Net (c2)
94.4%
94.1% %
93.8%
1.63M
baseline 48%
3.4
(compressed)
69.9%
60.5%
52.7%
0.85M
2.3 (S)
1.0
1.0
1.6M
1.4M
1.2M
1.1
0.9M
0.7M
K=16 K-14 K-12
K-10
K=8
K-6
Test Accuracy
0.9
0.5M
Test Accuracy
0.9
0.8
0.8
model
model
N: 16
0.7
Plecewise Constant
0.7
N; 12
Plecewise Linear
N:
Piecewise Linear Piecewise Const.
N:
10
20
0.6
10
20
30
Effective Compression %
Effective Compression %
(a) Compressing basis coefficients with fixed Ny.
(b) Decreasing Nr using the red line in (a).
Figure 4: In (a) we show the prediction accuracy on CIFAR-10 10 as function of weight compression for models that are trained with different basis functions. can be seen that there advantage of using piecewise linear basis functions during training. In (b) we compare models with different numbers of time steps. Reducing the number of time-steps reduces the the number of FLOPs.
eters by nearly 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 continuousin--depl batch normalization layers does not converge and achieves performance similar to tossing 10-sided dice.
In Figure 4a we show that the compression performance of model (cl) 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 Ny during inference time, which in turn leads to reduction of the number of FL.PPs during inference time. Recall, Ny refers to the number of time steps, which in turn determines the depth of the model graph. For instance, reducing Nr from 16 to 8, reduces the inference time by about a factor of 2, while nearly maintaining the predictive accuracy.
7.3 Compressing Continuous Transformers
A discrete transformer-based encoder can be written
Ft++1 It + T(q 2x) + M (p, I, + T()xx)))
(16)
where T'g, .)) is self attention (SA) with parameters q (appearing twice due the internal skip connection) and M is a multi-layer perceptron (MLP) with parameters p. Repeated transformer-bast encoder blocks can phrased as an ODE with the equation
t ==q((t), + M(p(t), + T))(
(17)
Recognizing ((1) (q()),((1), this formula can be directly plugged into the basis-functions and OdeBlocks to create continuous-in-depth transformer, illustrated in Figure [ Note that the

---

0.90
0.90
64
24K
0.88
6.6M
5.0M
0.88
Test Accuracy
Test Accuracy
0.86
MMM
0.86
K=4
2
0.84
PC Piecewise Constant
2.0M
0.84
My 64
K=2
PC Piecewise Linear
My
N; 16
0.82
PC- Discontinuous Linear
Source Model
1.9M
0.82
Source Model
Nr -
0.80
0.80
20
40
60
20
40
60
80
Effective Compression
Effective Compression
(a) Compressing basis coefficients with fixed Nr.
(b) Decreasing Nr using the red line in (a).
Figure 5: In (a) we show the accuracy of a transformer on a part-of-speech tagging problem as a function of weight compression. The transformer model can be compressed by factor of 2, while nearly maintaining its predictive accuracy. Discounting the 1.8M embedding parameters, the smallest model achieves 98% compression, albeit losing accuracy. In (b) we show models with different number of time steps N₂ (i.e., models with with shorter computational graphs).
OdeBlock is continuous along the depth of the computation; the sequence still uses discrete tokens and positional embeddings. With a forward Fuler integrator, Át 1, a piecewise constant basis, and K = KN Eq. 177 generates an algebraically equivalent graph to the discrete transformer.
We apply the encoder model to 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/Vaits and MLP dimensions of 128 with a single attention head. The final OdeBlock has K = 64 piecewise constant basis functions and takes Nr 64 steps.
present the compression performance for different basis sets. Staying piecewise constant basis functions yields the best performance during testing (red line). The performance slightly drops when project the piecewise constant basis onto piecewise linear basis (blue line). Projecting to a discontinuous linear basis (green line) performs approximately as well as the piecewise constant basis. In Figure 5bb] we show the effect of decreasing the number of time steps Nr during inference time on the predictive accuracy. In the regime where NT < K, there is no significant loss obtained by reducing K. Note that there a divergence when Nr > K, where the integration method will skip over parameters. However, projection to Nr < K incorporates information across larger window. Thus, projection improves graph shortening as compared to previous ODE-Nets.
In Figure 5a
8 Conclusion
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 also allows us to implement meaningful stateful normalization layers that are crucial for achieving predictive accuracy comparable to ResNets. Indeed, experiments showcase that Stateful ODE- Nets outperform other recently proposed ODE-Nets, achieving 94.1% accuracy on CIFAR-10.
take advantage of parameter transformations. In addition, this formulation
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 small amount of accuracy. On 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, demonstrate that our compression method reliably generates consistent models without requiring retraining and without needing to revisit any data. However, 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 further improve the compression performance. We can also explore smarter strategies for compression and computational graph shortening by pursuing hp-adaptivity algorithms (38).
10

---

Acknowledgments
We are grateful for the generous support from Amazon AWS and Google Cloud. NBE and MWM would like to acknowledge IARPA (contract W91 INF20C0033) 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.
References
[1] Bai Z. Kolter, and V. Koltun. Deep equilibrium models. In Advances in Neural Information
Processing Systems, pages 690-701, 2019.
[2] E. Borges Völker, M. Wendt, F. Hennig, and A. Köhn. HDT-UD: A very large Universal Dependencies treebank for German. Workshop on Universal Dependencies, pages 46-57, 2019.
[3] Bradbbry, R. Frostig. P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke,
J. VanderPlas, S. Wanderman-Milne, and Q Zhang. JAX: composable transformations of
Python+NumPy programs, 2018.
[4] B. Chang, M. Chen, E. Haber, and E. H. Chi. AntisymmerrieNN:: A dynamical system view on recurrent neural networks. In International Conference Learning Representations, 2019. [5] B. Chang. L. Meng. E. Haber, Ruthotto, D. Begert, and E. Holtham. Reversible architectures for arbitrarily deep residual neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, 2018.
[6] B. Chang, I.. Meng, E. Haber, F. Tung, and D. Begert. Multi-level residual networks from dynamical systems view. In International Conference on Learning Representations, 2018.
[7] T. Q. Chen, Y. Rubanova, J. Bettencourt, and K. Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, pages 6571-6583, 2018. [8] M. Cranmer, S. Greydanus, S. Hoyer, P. Battaglia, D. Spergel, and S. Ho. Lagrangian neural networks. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential
Equations, 2020.
[9] E. Dupont, A. Doucet, and Y. W. Teh. Augmented neural odes. In Advances in Neural Information Processing Systems, pages 3134-3 4144, 2019.
[10] M. Elbayad, J. Gu, H. Grave, and M. Auli. Depth-adaptive transformer. In International Conference on Learning Representations, 2020.
[11] N. B. Erichson, O. Azencot, A. Queiruga, and M. W. Mahoney. Lipschitz recurrent neural
networks. arXiv preprint arXiv:2006. 20700, 2020.
[12] A. Fan, E. Grave, and A. Joulin. Reducing transformer depth on demand with structured dropout.
In International Conference on Learning Representations, 2020.
[13] A. Gholami, K. Keutzer, and Biros. ANODE: unconditionally accurate memory-efficient gradients for neural ODEs. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 730-736, 2019.
[14] W. Grathwohl, R. T. Chen, J. Bettencourt, Sutskever, and D. Duvenaud. FFJORD: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.01367. 2018.
[15] S. Greydanus, M. Dzamba, and J. Yosinski. Hamiltonian neural networks. In Advances in Neural Information Processing Systems, pages 15353-15363, 2019.
[16] J. Gusak, L. Markeeva, T. Daulbaev, A. Katrutsa, A. Cichocki, and I. Oseledets. Towards understanding normalization in neural ODEs. In ICLR 2020 Workshop on Integration of Deep Neural Models and Differential Equations, 2020.
[17] E. Haber and L. Ruthotto. Stable architectures for deep neural networks. Inverse Problems,
3401):014004, 2017.
[18] E. Hairer, P. Nørsett, and G. Wanner. Solving ordinary differential equations. I. Nonstiff problems. Springer-Verlat Verlag, 1993.
[19] K. He, X. Zhang,, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the Conference on Computer Vision and Pattern Recognition, pages 770-778, 2016.
[20] Heek, A. Levskaya, A. Oliver, M. Ritter, Rondepierre, A. Steiner, and M. van Zee. Flax: neural network library and ecosystem for JAX, 2020.
[21] G. Huang, Y. Sun, Z. Liu, D. Sedra, and K. Q. Weinberger. Deep networks with stochastic depth.
In Proceedings of the European Conference on Computer Vision, pages 646-661. Springer,
2016.
[22] S. loffe and C. Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pages 448-456, 2015.
11

---

[23] C. Jordan and K. Jordán. Calculus of finite differences, volume 33, American Mathematical
Soc., 1965.
[24] A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. 2009.
[25] J. Kukacka, V. Golkov, and D. Cremers. Regularization for deep learning: A taxonomy.
Technical Report Preprint: arXiv:1710. 10686, 2017.
[26] M. K. Kwong and A. Zettl. Norm inequalities for derivatives and differences. Springer, 2006,
[27] Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner. Gradient-based learning applied to document
recognition. Proceedings of the IEEE,
1998.
[28] X. Li, A. Cooper Stickland, Y. Tang, and X. Kong. Deep transformers with latent depth.
Advances in Neurul Information Processing Systems, 33, 2020.
[29] S. H. Lim. Understanding recurrent neural networks using nonequilibrium response theory. arXiv preprini arXiv:2006. 1052, 2020.
[30] S. H. Lim, N. B. Erichson, L. Hodgkinson, and M. W. Mahoney. Noisy recurrent neural
networks. arXiv preprint arXiv:2 102. 2021.
[31] Y. Lu, A. Zhong, Q. Li, and B. Dong. Beyond finite layer neural networks:
Bridging deep architectures and numerical differential equations. In International Conference on Machine
Learning, pages 5181-5190, 2018.
[32] Y. Lu, A. Zhong. Q. Li, and B. Dong. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International Conference on Machine Learning, pages 3276-3285. PMLR, 2018.
[33] M. W. Mahoney. Approximate computation and implicit regularization for very large-scale data analysis. In ACM Symposium on Principles of Database Systems, pages 143-154, 2012. [34] S. Massaroli, M. Poli, J. Park. A. Yamashita, and H. Asma. Dissecting neural odes. In Advances
in Neural Information Processing Systems, 2020.
[35] B. Neyshabur. Implicit regularization in deep learning.
Technical Report Preprint:
arXiv:1 709.01953, 2017.
[36] Nivre, M.-C. de Marneffe, F. Ginter, J. Hajië, D. Manning, S. Pyysalo, S. Schuster, F. Tyers, and D. Zeman. Universal dependencies v2: An evergrowing multilingual treebank collection.
arXiv preprint arXiv:2004. 10643, 2020.
[37] A. F. Queiruga, N. B. Erichson, D. Taylor, and M. W. Mahoney. Continuous -nn-ptt neural
networks. arXiv preprint arXiv:2008. 02389, 2020.
[38] W. Rachowiez, D. Pardo, and L. Demkowicz. Fully automatic hp-adaptivity in three dimensions. Computer methods in applied mechanics and engineering,
(1177-40):4816448444 2006. [39] Y. Rubanova, R. T. Chen, and D. Duvenaud. Latent ODEs for irreuularl--sapple time series. In International Conference on Neural Information Processing Systems, pages 5320-5330, 2019.
[40] T. K. Rusch and S. Mishra. Coupled oscillatory recurrent neural network (coRNN): An accurate and (gradient) stable architecture for learning long time dependencies. In International
Conference on Learning Representations, 2021.
[41] T. K. Rusch and S. Mishra. Unicornn: A recurrent model for learning very long time dependen-
cies. arXiv preprint arXiv:2 2021.
[42] L. Ruthotto and E. Haber. Deep neural networks motivated by partial differential equations. Journal of Mathematical Imaging and Vision, pages 1-13, 2019.
[43] E. Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
[44] P. Wriggers. Nonlinear finite element methods. Springer Science & Business Media, 2008
[45] W. Xu, R. T. Chen, X. Li, and D. Duvenaud. Infinitely deep bayesian neural networks with
stochastic differential equations, arXiv preprini arXiv:2102 06559, 2021.
[46] S. Zagoruyko and N.
Wide residual networks. arviv preprint arXiv:
2016.
[47] A. Zeldes. The GUM corpus: Creating multilayer resources in the classroom. Language
Resources and Evaluation, 5 51(3):581-61) 2017.
[48] T. Zhang, Z. Yao, A. Gholami, J. E. Gonzalez, K. Keutzer, M. W. Mahoney, and G. Biros.
ANODEV2: A coupled neural ODE framework. In Advances in Neural Information Processing
Systems, pages 1152-51622 2019.
[49] Y. D. Zhong, B. Dey, and A. Chakraborty. Symplectic ODE-Net: Learning Hamiltonian
dynamics with control. In International Conference on Learning Representations, 2019. [50] Q. Zhu, Y. Guo, and W. Lin. Neural delay differential equations. In International Conference
on Learning Representations, 2021.
[51] J. Zhuang, N. Dvorekk X. Li, S. Tatikonda, X. Papademetris, and J. Duncan. Adaptive checkpoint adjoint method for gradient estimation in neural ODE, In International Conference on Machine Learning, volume 119, pages 1 11639-11649 2020.
12


Reviews:
[{"review_id": "SF9HasEIdAi", "reviewer": "Reviewer_FQch", "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nStrengths:\n---------------\n\n- Compression of the weight matrices is done post training and does not require re-training or any feed-forward passes through training data. \n\n\nWeakness:\n---------------\n\n- Novelty of the compression scheme seems incremental.\n\n- Model compression using the proposed scheme only guarantees drop in model size. Paper lacks computational cost of performing inference ( Algorithm~1 suggests that the inference cost will be higher than forward pass of the uncompressed model ). \n\nQuestions for Authors:\n---------------\n\n- It would have been better to showcase the inference time comparison between the compressed and uncompressed scheme. If for each input, the weights are constructed on the fly with the basis function, this additional cost is non-trivial and needs to be accounted for. Instead if the weights are simply constructed from the compressed model, then the inference proceeds as in the uncompressed scheme, then we did not gain anything from the computational standpoint. \n\n\n- Could authors elaborate a bit as to why is it hard to incorporate BatchNorm as one of the layers in the general function $f$ for the ODE $\\frac{dx}{dt} = f(x)$?\n\n- Is there any ablative study as to what benefits do you get by incorporating CT-Batch Norm in your proposed scheme?\n\n- Do you count the basis function parameters as well as linear coefficients in the number of parameters shown in the tables?\n\n- Did you run any of the baselines learning schemes with the architectures used in the experiments section? Its not clear from the paper if the baselines used the same architecture design as ContinuousNetV2 (barring the continuous time batch-norm).\n\nWriting Clarity:\n---------------\n\n- Line 33: Clearly state what are gradient parameters and how they differ from state parameters.\n\n- Missing related works \n\n\tODE/PDEs:\n\tContinuous Time RNN (by Rosenblatt),\n\tNeural Controlled Differential Equations for Irregular Time Series (https://proceedings.neurips.cc//paper/2020/file/4a5876b450b45371f6cfe5047ac8cd45-Paper.pdf), \n\tIncremental RNNs (https://openreview.net/forum?id=HylpqA4FwS), \n\tNeuPDE ( http://proceedings.mlr.press/v107/sun20a.html),\n\tTime Adaptive RNNs (https://openaccess.thecvf.com/content/CVPR2021/papers/Kag_Time_Adaptive_Recurrent_Neural_Network_CVPR_2021_paper.pdf)\n\t\n\tRelated works on the deep network compression is missing.\n\t", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "KNwtzDfU60O", "reviewer": "Reviewer_1fyo", "summary": "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.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, this paper uses concepts from continuous-time/depth neural networks to improve their discretized variants, e.g. a ResNet. I like that the paper shows models with higher-order discretizations and continuous basis functions behave more similarly to continuous-depth neural networks, whereas the piecewise-constant (ResNet) counterpart behaves worse, especially during basis function compression.\n\nHowever, the scope of the paper seems to be limited to learning simple mappings from input to output, and does not discuss applications that ODEs are more useful at, e.g. irregularly-sampled time series prediction and continuous normalizing flows. These applications either require more complex trajectories or directly make use of the right-hand-side of ODEs, and both benefit significantly from adaptive step size solvers in practice.\n\nWith the use of fixed step size solvers, it’s not clear whether the discretized model is still an approximation to the true ODE solution. The observation that piecewise-constant basis functions do not compress well may simply be due to them not being a close approximation of an ODE. It’d be more convincing to see experiments on the aforementioned applications.\n\nAnother concern with the use of piecewise-constant/linear basis functions is the lack of smoothness for higher-order solvers. When using a method such as RK4, the step sizes must correspond to the basis function’s Delta t. This again goes against the use of adaptive step sizes. This property makes it difficult to perform long-range predictions, e.g. for irregularly-sampled time series predictions. Why not explore more complex or smoother basis functions?\n\nRegarding continuous-time batch normalization, it doesn’t seem clear to me that using this is beneficial in practice. While the paper discusses efficiently storing a time-dependent mean/var/bias/scale, it’s not clear if this actually improves convergence during training. Normalization seems like a weird operation to have in an ODE since it forces the model to predict outputs based on relative (and stochastic due to mini-batching) differences instead of absolute position. This dependence makes its interpretation as a dynamical system following a vector field less intuitive. In any case, it would be more convincing if the authors can run ablation experiments to discuss the merits of using continuous-time batch normalization.\n\nFurthermore, the lack of a Lipschitz bound for normalization is concerning. If a hidden unit is the same value (e.g. zero) for all samples in a mini-batch, then there is a division by zero. Likewise, it was shown that multihead attention layers are not Lipschitz [1]. Though I suspect these aren’t major concerns unless one tries to solve the ODEs extremely accurately.\n\n[1] “The Lipschitz Constant of Self-Attention” Kim et al.\n\n----\n\nUpdate after author response: Thanks for the clarifications. I generally agree with the response; however, the paper likely should be rewritten to make it clear that this is regarding discrete-time networks and to re-iterate on the scope and expectations of this work. The naming convention for continuous-time BN can also be misleading, since this is not applied in a continuous-time setting. Just my two cents, but answering these questions may be helpful: why is compressing an ODE-net an important task? What if you simply said you're compressing ResNets (and mention extending to networks that are parameterized by more sophisticated ODE solver schemes)? This may help increase the number of readers while also clarifying scope.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The scope is limited to discretized networks that are simply expected to behave like ODE solutions. This expectation is a small leap of faith and should be discussed more in the paper. It’d be more convincing to see experiments making use of adaptive step size solvers and/or applications involving strong reliance on accurate ODE solutions. The batch normalization method should be ablated against as it's not clear if this leads to faster training or performance improvements. Furthermore, the numerics involved with solving an ODE with potentially very large or unbounded Lipschitz is not discussed sufficiently.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "BNMr6aFFG6", "reviewer": "Reviewer_KZk9", "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”. \n\nBy 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. \n\nAdditionally, 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. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "## Score\n\nOverall, I definitely appreciated the novelty and the rigor of the proposed time-continuous model and I believe there is scientific merit and the potential to influence the broader community. However, I am little concerned about the current presentation of the results. Specifically, they are covering lots of grounds from introducing a new type of model to theoretical results to compression to empirical results. Consequently, some of the presentation is very superficial and it’s hard to understand the important details that might help other researcher pick up on this work and build upon it. \n\nSo in short, I imagine the core contribution to be quite interesting but this has to be backed by the corresponding depth and detailed discussions of the method in the paper. Right now, there is lots of breadth and little depth across the sections.\n\n\n## Ways to Improve My Score\n\nPlease address the weaknesses I am outlining below. I would specifically focus on trying to clean up some of the writing and re-structuring it to be able to go into depth about individual aspects of the work while potentially removing other aspects.\n\n## Strengths\n\n* Intuitive and powerful model that promises a lot of flexibility down the road. \n\n* Interesting experimental results pertaining both to image classification and NLP/transformers. \n\n* Good introduction and explanation about the concept of basis functions. \n\n* I enjoyed learning about the combination of using a NN parameterization based on basis functions  + time-continuous models. \n\n## Weaknesses\n\n* By far the biggest weakness to me is that there is so much going on that it is very hard to formulate a take-away message at the end of the paper. Sure, I might be able to download the code, re-run the experiments, and modify the models a little bit to adapt them to my task. But I don’t think the papers contributes a deep understanding of the model that someone could leverage in their own research in order to come up with new and innovative solutions.\n\n* Introduction: I don’t think it’s really useful to touch upon so many concepts in the introduction. Try re-structuring the core concepts, simplifying them, shortening the intro, and then go into more details about each aspect later on in the paper. \n\n* The time-continuous batch normalization: This concept has been mentioned many times throughout the sections and I suppose I can somewhat imagine how it works but overall the language of the paper never goes beyond the conceptual idea of what the authors want to achieve. For example, looking at Algorithm 1 I cannot figure out where $\\hat\\theta^{s*}$ is being used? It seems to be an output of the algorithm but what are we doing with it? Is it part of the model as a batch norm layer? Where would this layer be? It’s confusing but still seems to be very important so it would be crucial to understand how this fits into the model. \n\n* I cannot really find a take-away message from Section 5. It seems to be a somewhat interesting result from the space of dynamical systems and uses a lot of very atypical concepts that I cannot imagine are common knowledge among the broader NeurIPS community (or even the more narrow neural ODE community). If you really want to include these theoretical results, please simplify the section in the main paper and put a much more in-depth discussion into the appendix. In the current shape, I am left to simply take the authors’ word for granted about the connection between their theorem and the implicit regularization potential without really being able to back it up myself. \n\n* The experiments are good but I would ask the authors to significantly expand upon the concept of (multi-level) refinement. This seems to be a central empirical “trick” in order to make their model work in practice. What exactly is the concept behind it? How to tune the different levels? What is the crucial insight there that makes the difference in the empirical performance? \n\n* Choice of basis functions: The basis functions seem overly simplistic (either convex combination or piecewise). What about other basis functions and what constitutes the motivation for these basis functions? The concept and the notation are very heavy in order to introduce these basis functions. So is there any benefit beyond these simple basis functions? \n\n* Related to my question above: It appears to me that the authors rely on a fixed-step solver like RK4. In addition, they often use piecewise-constant basis functions. My question is how this model even differs from a regular NN that explicitly models multiple layers? Specifically, if for each time step in RK4 there is a different “piecewise-constant basis function” that is activated this implies that at any given depth you call the _same layer architecture_ with a _distinct set of parameters_ much like you would simply repeat the same layer architecture with separately trained parameters in a regular NN with residual connections? And even if there is more time steps in RK4 than there are basis functions it very much still seems to boil down to a regular NN with a little bit of parameter sharing, right? And same goes for the piecewise linear basis functions, right? I am not saying that the proposed formulation isn’t more expressive in theory but it seems that all instantiations of the model framework are essentially very close to  a regular network with multiple layers of the same layer architecture stacked on top of each other. Could the authors comment on this observation? \n\n* I would also appreciate if the authors could spend more time in the introduction and the related work on how their work differs from prior work. It seems basis functions have been explored before from what I understand from their paper and also of course neural ODEs are not a new concept. So it would be very helpful to contextualize their work within existing work more clearly and delineate their contributions. \n\n--- \n\n## Update after rebuttal \n\nI am raising my score to 6 notwithstanding some of the open issues that I feel need to be addressed, whether during a potential camera-ready version of the paper or a potential re-submission depending on the outcome of the review process. More details can be found below in my discussion with the authors. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See above weaknesses especially with regards to the refinement in the training process. \n\nAdditionally, I glanced over Section 4.1 in the appendix regarding the implementational details of the projection step. It seems to require various derivatives including a Hessian computation. How does the projection step scale with the model? Is that something that’s feasible for larger models and datasets? What is the practical runtime for this step? It would be good to understand that. ", "review_text": "", "extra_scores": {"ethical_concerns": "No concerns. ", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a compressed representation of ODE-based neural networks, and analyzes its computational and implicit regularization properties. Experiments cover a good range of architectures, including both convolutional and transformer-based.

The reviewers believe there are worthwhile ideas in the paper and haven't identified anything that looks like a critical flaw. They raised some concerns about clarity and identified numerous ways in which the work could be extended. Reviewers also expressed skepticism of the usefulness, either because compression might not translate into practical gains, or because this architecture might lose various benefits associated with neural ODEs. 

My own impression is that the submission is pretty well written by the standards of a NeurIPS paper, and the methods are described clearly even if the motivation could sometimes be made more explicit. The reviewers' suggestions seem to me (as a non-expert on this topic) to be sensible ones which would indeed improve the paper, but any paper that tried to do all of them would be unreasonably expansive. Regarding significance, the proposed architecture seems like a natural one to explore, and a paper which does it carefully (as this one does) is making a worthwhile contribution. Therefore, I recommend acceptance despite the slightly-low scores.

Decision: Accept (Poster)