Improving Neural ODE Training with Temporal Adaptive Batch Normalization
Authors: Su Zheng
Zhengqi Gao
Fan-Keng Sun
Duane S Boning
Bei Yu
Martin D. Wong
Venue: NeurIPS 2024 poster (2024)
Keywords: Neural ODE
Batch Normalization
Primary area: deep_learning_architectures

Abstract: Neural ordinary differential equations (Neural ODEs) is a family of continuous-depth neural networks where the evolution of hidden states is governed by learnable temporal derivatives. We identify a significant limitation in applying traditional Batch Normalization (BN) to Neural ODEs, due to a fundamental mismatch --- BN was initially designed for discrete neural networks with no temporal dimension, whereas Neural ODEs operate continuously over time. To bridge this gap, we introduce temporal adaptive Batch Normalization (TA-BN), a novel technique that acts as the continuous-time analog to traditional BN. Our empirical findings reveal that TA-BN enables the stacking of more layers within Neural ODEs, enhancing their performance. Moreover, when confined to a model architecture consisting of a single Neural ODE followed by a linear layer, TA-BN achieves 91.1\% test accuracy on CIFAR-10 with 2.2 million parameters, making it the first \texttt{unmixed} Neural ODE architecture to approach MobileNetV2-level parameter efficiency. Extensive numerical experiments on image classification and physical system modeling substantiate the superiority of TA-BN compared to baseline methods.

Improving Neural ODE Training with Temporal Adaptive Batch Normalization
Su Zheng': ZZengqi Gao²*, Fan-Keng Sun Duane Boning", Bei Yu', Martin Wong' *Department of CSE, CUHK
2 DDpartmentt of EECS, MIT
Abstract
Neural ordinary differential equations (Neural ODEs) is family of continuous- depth neural networks where the evolution of hidden states is governed by learnable temporal derivatives. We identify significant limitation in applying traditional Batch Normalization (BN) to Neural ODEs, due to fundamental mismatch - BN initially designed for discrete neural networks with temporal dimension, whereas Neural ODEs operate continuously over time. To bridge this gap, we introduce temporal adaptive Batch Normalization (TA-BN), novel technique that acts as the continuous-time analog to traditional BN. Our empirical findings reveal that TA-BN enables the stacking of more layers within Neural ODEs, enhancing their performance. Moreover, when confined to a model architecture consisting of single Neural ODE followed by a linear layer, TA-BN achieves 91.1% accuracy on CIFAR-10 with 2.2 million parameters, making it the first unmixed Neural ODE architecture approach MobileNerV2-Ievel parameter efficiency. Extensive numerical experiments on image classification and physical system modeling substantiate the superiority of TA-BN compared to baseline methods.
Introduction
Originally derived as the continuous limit of a residual con- nection neural ordinary differential equations (Neural ODEs)
0.95
31 is a family of continuous-depth neural networks where the evolution of hidden states is governed by Icarnable temporal derivatives. These models exhibit several intriguing features, such as the capability of temporal reversibility, which enables generative modeling 15).
37 45 g
30
18
32
19
Tast Accuracy
0.65
SGDD TA--N (Orze. lopees)
NNODE Midd-haleh layens
-NGBR IN (18
0.5
-NODE
NCDR
0.33
201
40
64
80
100
120
Previous studies on Neural ODEs parameterize the learnable temporal derivatives using shallow neural network with a lim- ited number of parameters 7 30 Without special treatment, merely stacking additional layers in the temporal derivatives does not necessarily enhance Neural ODE performance. Fur- thermore, deeper networks might increase the stiffness of the ODE system, leading to challenges with the ODE solver, such excessively small step sizes even failures due to infinite state values, as shown in Figure
Epoch
Figure 1: Test accuracies de- picted over the training epochs on CIFAR-10, 10, utilizing Neural ODEs with different numbers of convo- lution layers the backbones for learnable derivatives. Dashed hor- izontal lines denote instances of training failures.
In efforts to deepen the model and address instabilities, one might instinctively consider Batch Normalization (BN) (5] 29 16 39
as a remedy to stabilize the intermediate state values of Neural ODEs [1 35 because BN was originally proposed to accelerate
The first two authors contribute equally. Correspond to Bei Yu at byu weyuuess..uhk.cd..ht
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

the training of (discrete) neural networks and consolidate them against internal covariate shift [29]] However, preliminary works surprisingly found that the introduction of BN into Neural ODEs can even degrade performance, effect for which the underlying causes remain unclear R 35
In this paper, we demystify the aforementioned phenomenon and demonstrate that directly applying traditional BN to Neural ODEs is fundamentally flawed. The primary complication arises because the forward pass of Neural ODEs employs an adaptive step size solver that discretizes time variably. Consequently. it is not assured that two forward passes will coincide with the same discretized time grids. This variability precludes the possibility of retrieving population statisties [16] at test time. Additionally. relying on mini-batch statistics for BN renders the outputs of Neural ODEs non-deterministic and vulnerable to outliers and small batch size.
Motivated by these observations, we introduce temporal adaptive BN (TA-BN), a novel technique designed as the continuous-time counterpart to traditional BN, tailored specifically for Neural ODEs. Our empirical findings demonstrate that TA-BN facilitates the up-scaling of Neural ODE model sizes without encountering performance saturation (See Figure []]. Moreover, when confined to model architecture consisting of a single Neural ODE followed by a linear layer, TA-BN achieves 91.1% % test accuracy on CIFAR-10 10 with 2.2 million parameters, making it the first unmixed Neural ODE architecture to approach MobileNerV2-lee on image classification and physical system modeling substantiate the superiority of TA-BN compared to baseline methods.
parameter efficiency. Extensive numerical experiments
2 Preliminary
Neural ODEs Neural ODEs [4]
7
37
45 8
42
30
18
32
19
31 forge a significant linkage
between differential equations and neural networks. They model the continuous dynamics of hidden states with a learnable ODE system:
dh(t) = ()())))
(1)
dt
where (:)) RD RD is neural network parameterized by learnable parameter controlling the temporal evolution, and h(t) E R is the ODE state variable at time t. The practical usage of Neural ODEs is by solving initial value problem: given the initial state h(0) as input, the state h(T) at a later time T is computed using an ODE solver and returned as the output. The gradients required for training are computed via the adjoint method 4
18 reverse in time.
Subsequent works have expanded the capabilities and scope of Neural ODEs by introducing modifi- cations, such improving expressivity through state augmentation (7]. coupling the evolution of both weights and activations [43], accelerating convergence via semi-norm techniques (18), extend- ing to graph neural networks 37] and second-order ODEs [34], adapting irregular time-series data 19 32), and drawing connections to diffusion models 15 28). Of particular interest, two
preliminary studies [13 B5] have indicated that applying BN within Neural ODEs may compromise model performance. Our work aims to unravel this phenomenon and propose a version of BN tailored operate effectively with Neural ODEs.
21 performs a re-centering and a re-scaling operation on the given input by subtracting the mean and dividing by the standard deviation:
Batch Normalization BN 5 29 16 39
BN(x,) = BN, (I,)) = v0a + +
(2)
where E + ensuring division validity, {y E R. a E R} are learnable parameters, and represents the i-th data in one batch of size B. During training, the mini-batch statistics 1/B Σ
iz 2; and o² 1/B Σ Ei= (2; 412 are used for efficiency. In contrast, at test time, BN uses population statistics aggregated across the entire training dataset to ensure deterministic outputs [66] BN can be extended to multiple dimensions/channels by processing each one separately.
BN has been widely applied in training deep neural networks, as it can stabilize the training process and accelerate convergence. Despite its strong empirical performance, the underlying mechanisms
Throughout our paper, statistics refer mean and standard deviation (std) equivalently variance.

---

of BN have been subject to various interpretations, such as reducing internal covariate shift [16]. smoothing the optimization landscape [39] and decoupling the learning of length and direction 21 Originally, BN [6] was developed for neural networks with hidden units that do not explicitly depend
on time. However, recent advancements have extended BN to accommodate time-dependent models, including recurrent neural networks (RNNs) [3 and spiking neural networks (SNNs) [20] 44 6 17]
Unlike these models, Neural ODEs exploit adaptive time discretization, necessitating specialized modification of BN for its application.
3 Traditional Batch Normalization in Neural ODEs
Let us consider a simplified one-dimensional Neural ODE when incorporating BN as example:
dh (N
dt where A by E R. R3 are the learnable parameters. Given batch of input (0) (i.e., initial values of the ODE), the forward process employs an ODE solver to discretize Eq. B over time, visiting N sequential points T = (to 0, tj, t.. ,NN T) and computing the outputs (h,(T)) An adaptive step size solver (e.g., Runge-Kutta method) is usually preferred because of its higher efficiency compared to fixed step size solver (e.g., Forward Euler method). Thus, the differences between two consecutive time steps might not identical (e.g..
to f1 # t1 to), and the value of N also depends the batch of data. The temporal discretization indicates that we need to invoke BN at every E T for every sample:
BN(()g)) BN,, (mi.) Di.j Mj +a, where vijj he(1)) +
V03 +e
for every 1,2, NN and 1,2, where My and of are statistics associated with ty.
During training To perform the normalization shown in Eq. (4), BN applies mini-batch mean Mj 1/BBBB zi.j and variance oj 1/B under normal conditions, the statistics cannot be successfully accumulated across batches. For instance, the time grid T utilized for the first batch might differ significantly from the grid T' used for a subsequent batch. Consequently, specific time points 1 E T'' might not coincide with any time points in T, and vice versa. This might result in impractical collection of infinite statistics μi and o3. Moreover, except those associated with t and t and updated based on a limited number of batches.
E E n2)²² Although the training can proceed
T, most stored statistics will be calculated
Ratch Data
Training
to
t2
EN-1 ty
to
t Use (x")az")
to
tN-1 to Update (ee", :)
Pop BN Misalign
Pop-TI BN (")
Time Grid
t Use (r*,a²")
t
the-1 tis Use (Ha", orrit
Testing
Batch
Data
t
Figure Left: The failure of Pop BN in Neural ODEs stems from the misalignment of discretized time grids. Pop-TI BN aggregates all running mini-batch statistics into a single pair of (n, o²), implicitly assuming inneidddeenndnt population statisties. Right: Our proposed TA-BN automatically conducts temporal interpolation to accumulate statistics and update parameters during training and testing.
During inference Conventionally, BN should perform the normalization shown in Eq. using population statistics during inference. However, as shown in the left part of Figure 2 the population statistics associated with the time point l'/ E T', required by the temporal discretization during inference, might not available if the time value t' is never encountered during training. Moreover, even if the population statistics exist, they are likely to be inaccurate discussed above.

---

0.32
0.88
0.2
Too. Accurscy
0.72
Tes Accuracy
0.78
Scariaties
0.52
0.1
Misuhaack BS.
SOAEE 1X-35 10 w/ codel
0.32
NOBI
coler)
2000
304
KODE Missheac BX torlo sobee
0.12
60 80
0.58
60 Epoch
80
0.2
0.4
0.0
0.8
Epoch
Thme
Figure 3: Left: We train Neural ODE with a U-Net backbone as the learnable derivatives on CIFAR-10. Mini-batch BN shows degraded accuracies with a batch size of 2, while TA-BN can maintain high accuracies under varying batch sizes, because it uses the estimated population statistics during testing. Pop-TI BN aggregates running statistics encountered at any time points into a single pair of and o². This approach assumes time-independent statistics, leading to erroneous predictions and erratic test loss curves. Middle: When noisy data exist on average in one out of every test batch, Mini-batch BN's performance deteriorates, because the noise affects the mini-bate statisties. The backbonn for learnable derivatives in this experiment consists of 6 convolution layers. Right: We plot the output statistics from the first layer of U-Net over time; they time-dependent.
Summary seemingly straightforward remedy could be to consistently use mini-batch statistics during both training and inference, thus avoiding the need for accumulating population statistics. However, this approach has a critical limitation: the output (T) becomes highly dependent on their respective batch's ss specific characteristics. This will make the outputs inaccurate when batch size is very small (e.g..
2) or the occurrence of outliers in batch, and even fail in single-run scenarios (B = 1) because running standard deviation cannot be calculated. As shown in the left of Figure 3 when the batch size is reduced to 2, mini-batch BN displays degraded accuracies. Furthermore, when noisy data exist, mini-batch BN performance deteriorates, as shown in the middle of Figure
For clarity, we will henceforth refer the use of population statistics at test time as Pop BN, and the use of mini-batch statistics as Mini-batch BN. Notably, both methods utilize mini-batch statistics during training. It is important to emphasize that in practice, when implementing BN in Neural ODEs with popular deep learning frameworks (e.g., PyTorch with TorchDiffeq 44), actually variant of Pop BN is used, which we will refer as Pop-TI BN. Pop- BN accumulates statistics at all encountered time points during training into a single pair of and o², which implicitly assumes that all time points share the same statistics. This aggregation approach can result in outputs during test time without errors concerning missing population statistics which we would anticipate as in Pop BN. Unfortunately, these predictions are meaningless and often lead to poor performance, as illustrated in the left of Figure 3 This issue highlights the challenges identified in recent studies [13 35 nutshell, our observations can be succinctly summarized as follows:
In
Pop BN fails when Neural ODEs employ an adaptive step size solver.
Pop-TI TI BN implicitly assumes time- -meeinddeendeent statistics, leading to erroneous outputs.
Mini-batch BN fails when Neural ODEs operate with small batch sizes outliers exist.
4 Temporal Adaptive Batch Normalization
The fundamental reason why Pop-TI BN and Mini-batch BN fail in Neural ODEs is attributed to the added temporal dimension, where the distributions at different time points vary, effectively constituting a stochastic process as depicted in the right of Figure 3 This observation suggests that BN should be defined to be time- dime-dependent in Neural ODEs. Ideally, we would want BN to store population statisties at every E [0. T), with corresponding learnable parameters (t) and B(t) defined at every t. However, implementing such a model is impractical, because not every t will be encountered in training, and the update of (t), B(t), and the population statistics will be naturally sparse on f, resulting in inefficiency.
'e examined several open-source codes provided by previous works. They align with this implementation and exhibit similar erratic behavior to the Pop-TI BN's test loss curve in the left of Figure 3

---

To address the aforementioned problem. we propose an intuitive method termed temporal adaptive
BN (TA-BN) for usage in Neural ODEs. Let us illustrate with the example shown in Eq. (3). To
begin with, we evenly divide the time span [0, 7 defining (M + 1) time grids T = (4) = 0, 1
T Mr = 27
Ess = T). We associate the time grid i'm with population mean pm and population
M:
variance " as well as learnable parameters y and a for every m 0, 1, 1,2, 2.. M.. For later simplicity, we denote > = [70: 71, .74l and x²* = [a..aj, ... a Similar notations
apply to je' and o².
During training As shown in the right of Figure 2 given batch of data (h, (0)))) the forward
pass of Neural ODE might discretize the time as 7 = (to, f1. (x) which differs from the grids T* of TA-BN. However, TA-BN can calculate the temporal derivative for the i-th data at ty:
TABNy. a (2L.) = 2i,5 - Mi Migg + ay where 24, = w ho(tt) + b V0 +e
(5)
where (1) are still the mini-batch mean and variance based on {ais) Here 7/ and aj are interpolated based on 7" and a² respectively:
7 G(g,, T)), aj G a","
where Git, a. T) is a function to interpolate the value a(t) given array of values RM++ and their corresponding (M + +1) time points T. There are many choices for G(:..-), such as linear, cubic spline, and kernel smoothing. We empirically observe that linear interpolation based on two nearest neighbors suffices for our experiments, and we advocate it due to its implementation simplicity:
G(t, T) f++1 -
t -- t
a1
0141
4141
441
where is the index which makes 1 to be the largest value smaller than / in T. Using a linear G(: : implies that we approximate the underlying a(t) in a piece-wise linear manner, and as long as the number of time grids is sufficiently large, G(t, a, 7) can approximate any continuous a(t) arbitrarily well, i.e., GGtt, a, T) - a(t) + when M + co. In practice, we set M close to the number of ODE-discretized time grids. Please refer to Section 5 for ablations on the choice of G(:,.-).
During training, we also need to accumulate the encountered running mini-batch statistics (Hz, (4) to the population statistics (Hi,,, o so that later they can be used for inference. As shown in the right part of Figure 2 this requires the interpolation to be performed in the opposite direction:
Pm (1 11-1) n) pm + G( ,4, T) 0 2 (1 11-1) on +1- G( :T)).
/ [1,, 12,
|NN]"
(8)
a² lo 0z
2
where 7/ E [0, 1] is a momentum constant to perform moving averaging to update the population statistics based on current estimate of pr and t, using the new observed values G( n4, T) and G( 2², T), respectively. Note that the training does not involve the gradients with respect to
During inference After the training phase, TA-BN will have well-trained parameters Y and a², and accurately record population statistics μ* and a². During inference, when provided with a batch of data and denoting the discretized time grids as T', as depicted in the right part of Figure 2 we again proceed with the interpolation step:
Param.
J = ((t n -''
af = G(t), T")
Pop. Stat.
of2
This interpolation differs from the training phase in that now the interpolation directions are the same for parameters and the population statistics, because now we utilize them for the new time grids T'.
Implementations and summary In practice, we observe that the interpolations during training outlined in Eq. (6) have opposite directions, potentially leading to increased computational overhead. To elaborate, when encountering time point lj discretized by the ODE solver during training, the interpolation for obtaining 7j and a, in Eq. (6) using Eq. 7 can be immediately executed by identifying the index I such that t is the largest value smaller than ły. However, the interpolation
5

---

in Eq. (8) cannot be carried out at this time. It becomes feasible only after the ODE solver completes its execution and we gather all f 's into T. This sequential nature degrades computational efficiency. To mitigate this issue, we have slightly adjusted the interpolation for population statistics during training so that it can be performed concurrently with the parameter interpolation time 13. The key steps of our proposed TA-BN are summarized in Algorithm [ It will be called as subroutine for every discretized time point ł E T required by the ODE solver.
Algorithm outtines the case involving single one-dimensional neuron and one TA-BN layer within the Neural ODE, In more general cases, we may have Q TA-BN layers within a Neural ODE, and each neuron may possess D dimensions. In such cases, we can perform the normalization in dimensional manner. Thus, we must maintain 2 + 1)Q population statisties and 2 + 1)Q learnable parameters, each with D dimensions. In contrast, traditional feedforward neural network with Q BNs has 2Q population statistics and 2Q learnable parameters, each with D dimensions. Note that for image inputs, we perform channel-wise normalization, so D will be the number of channels.
Finally, since our linear interpolation in Algorithm Trelies on two nearest neighbors, it is prudent to ensure that the number of TA-BN time grids (M + 1) are roughly at the same level as the number of ODE-discretized time grids. Denser grids would be unnecessary those not adjacent to any i; would remain unused. Alternatively, as our numerical results will demonstrate, employing other interpolation methods that require all time grids may lead to increased computational overhead without significant improvements. Additionally, we employ L2 regularization l - 1ll² and || avoid significant discrepancies among the elements Y and a, making the training more stable.
Algorithm The forward pass of a TA-BN layer at time L
Input: Batched input [sa.)]_ time in
1: Get i and 1/41 such that 1f is the largest value smaller than 13 in T 5/
2: 14 W2
>> Compute the weights for linear interpolation;
3: if the model is in the training mode then
4:
n mean(x), of = =
> Compute the mini-batch statisties;
5:
= = + w27f+1 aj = =ha + +220741
6:
ut u + (1- wi wi)uf -n.wilg Pi+i + (1 (1-1) w2 Mit + W2Hi;
7:
x. x (1- 17 w1 ) +2 + w10.. j4l (1 n w2 W2) 17 w20
x.
+.2
8:
else if the model is in the evaluation mode then
9:
My = win + w27f+
w1n + w2ll41
of = w1 22 w22
+.2
10:
aj = wra + W20
11: end if
12: return x- + aaj Vazte
5 Numerical Results
Various network architecture designs have been explored in existing Neural ODE studies. For instance, the original study [4] integrates a feature extractor (e.g., CNN-based), Neural ODE module, and an MLP in sequence for classification. Another study structures neural networks with alternating Neural ODE modules and convolutional/linear layers [Σ]] While these solutions represent viable approaches for optimizing performance on given ML tasks through the synergy of conventional (discrete) neural networks and Neural ODEs, they become inappropriate for investigating the effect of specific architectural modifications (i.c., TA-BN on Neural ODEs. This is because the Neural ODE module can be entirely bypassed if its trained weights approach zero in extreme cases, thereby overshadowing the intended modification. Hence, our model comprises a Neural ODE module followed by a single linear layer as advocated by [7. which we refer to as the unmixed Neural ODE architecture. Our code is developed based on PyTorch (36). TorchDiffeq [3, and customized TA-BN layer. All experiments are run on Linux server with RTX 3090 GPUs.
5.1 Image Classification
Being consistent experimental scales of previous Neural ODE studies, we conduct image clas- sification across datasets including MNIST 26). SVHN (33, CIFAR-10 10, CIFAR-100 (22). and

---

EinyIImaeNet (24). We employ the dopri5 solver with a tolerance of 10-3 for ODE solving and adopt the Adam optimizer [27] with a learning rate of 10 -3 to train the neural networks for 128 epochs. The training batch size is 256. We set M = 100 for TA-BN.
Table T compares the test top-1 accuracy and number of parameters of various Neural ODEs. TA-BN outper- forms Mini-batch BN and Neural ODE w/o BN on SVHN,
0.9 Better
Test Accunayy
0.8
CIFAR-10, CIFAR-100, and Timy-lmageNet. Moreover, TA-BN achieves superior accuracies over Aug-NODE [7] and STEER 12] On MNIST and CIFAR-10, TA-BN has
TA-BN (ouod)
0.7
Il and
hutth
fewer parameters than Aug-NODE and STEER, which indi- cates that TA-BN can also help training in the small Neural ODE regime. We also visualize Neural ODEs' performance changes as the number of parameters varies in Figure 4 by including four additional baselines ILONO, 2nd-Ord [D]] HBNODE, and GHBNODE [41]. For reference, the popu- lar convolutional-base MobileNetV2 (38, known for its parameter efficiency, achieves approximately 94% accuracy with about 2M parameters. Our TA-BN assisted Neural ODE is the first to approach this level of performance using the ammixed architecture, while most previous Neural ODE
0.6
10'
10P
# #Parameterrs
Figure 4: Comparison between differ- ent Neural ODEs on CIFAR-10. The basclines marked by yellow triangles do not adhere to the unmixed struc- ture and are not strictly comparable to ours. It is unknown whether increasing the number of parameters inside their ODEs can lead to better accuracy. literature either falls short of this performance or/and does not follow this unmixed architecture.
Table 1: Comparison of test accuracies and number of parameters between different Neural ODEs!.
Model
MNIST
CIFARIO
SVHN
CIFAR100
Tay-Imagenee Accuracy
Accuracy
#Params
Accuracy
#Params
Accuracy
PParmms
Accuracy
&Params
#Params
Ang-NODE 2
0.982
84k
0.606
172k
0,835
172k
N/A
N/A
NIA
366k
STEER 12
0,986
84k
0.621
172k
0,841
172k
N/A
N/A
N/A
NIA
w/o BN
0.989=000001
37k
0.517=00.099
2.2M
0.096=0.025
2.2M
0,246+#0000 084
2.2M
2.2M
Pop- Tl
0.973:10.011
0.548=0.087
2.2M
0.241:0.123
0.251:0.112
2.2M
0.04420.007
2.2M
w/ Mini-batch BN
0.962:+0.013
37%
0.822=0.095
2.2M
0.906+0.001
2.2M
0.492=0.176
2.2M
0. 0.200=0.006
2.2M
w/ TA-BN (ours)
0.988=+0.001
37k
0.748=0.059
70k
0.953+00.02
220k
0,576:0.016
220k
0.436=0.013
220k
0.988-+0.001
220k
0.910-0.010
2.2M
0.958-00.00
2.2M
0.664-0.025
2.2M
0.512=00008
2.2M
1 'NIA' indicates values are not available original literature. indicates training fuiluure at the epoch. Results reported with test batch size of 256. The error bars are shown using the format meansisd.
0.64
0.52
Accaracy
0.8
Tast Accaracy
Tist. Accatacy
0.30
0.6
0.32
NO0B w/ TA-BK, Regemter NODE Mivslued IN
Pegiooki
TAA-N Rugrestx
0.4
0.16
vome wiit BeV
0.13
XO0E Miaa RVV 200
N4DD8 Pop-21 Brazncete
NEARE Pap- IS,
SODE Psp-11 Brogzente
0.2
None o/ IX
YOOEE wyj RVV
200
40
60
50
100
120
200
40
60
80
10U
120
20
40
60
80
100
120
Epach
Epoch
Epoch
Figure 5: Comparison between Neural ODEs w/ TA-BN, w/ Mini-batch BN, w/ Pop-TT BN, and w/o BN, using an 18-layer U-Net backbone for learnable derivatives. The total number of parameters including the outside linear layer is 2.2M. The datasets are SVHN (left), CIFAR-100 (middle), and Tiny- Imaeeee (right). The results on CIFAR- 10 have been shown in the left part of Figure 3
The left part of Figure 3 in Section 3 monitors the test accuracy over training epochs on CIFAR- 10. It indicates that training a large Neural ODE without special processes might fail due to numerical instability. Pop-TT TI BN performs poorly because of the ieneiindependent statistics assumption. While Mini-batch BN achieves satisfactory accuracy, our TA-BN outperforms without encountering the issues associated with Mini-batch BN, such as batch-dependent outcomes and vulnerability to outliers. Similarly, the left and middle parts of Figure dissplay test accuracy versus training epochs on the SVHN and 100 datasets. TA-BN consistently shows better and more stable accuracy than Mini-batch BN on these datasets. On the Imag- dataset (right part of Figure 5). TA-BN demonstrates superior accuracy and faster convergence.
7

---

Neural ODE Up-scaling Enabled by TA-BN As shown in the left part of Figure Neural ODEs with limited number of layers perform normally with acceptable accuracy; however, when the layer count exceeds 10, training fails due to numerical instability. In contrast, the incorporation of TA-BN enables deeper layers within Neural ODE as the learnable derivatives, scaling up the model size and enhancing accuracy. as exemplified by the middle and right sections of Figure 6 Appendix A. for architecture details.
Please see
0.92
0.92
0.92
The Accuracy
Tost: Arcuracy
Tee. Arcurwcy
0.75
0.75
0.75
WORK layees
Nouth sureer
NODDE TA
NODE wnt S0V (12 loner)
KODE 100-5 102 lenur)
NOD8 TD-BS (12 komm
0.58
NOIRE wn ligrest
0.58
NOODE TX--1X 130 Sarves
0.58
SCC000 TD B1 130 koows
WODE w/ lamie))
NODE
S0006 wypa BS (1N kmmm
lapon))
2/0000 myo lermo)
20
100
120
20
100
120
100
120
Epoch
Epoch
Epoch
Figure 6: CIFAR-I0 accuracies with increasing sizes of the backbones for learnable derivatives. These figures illustrate the scaling up of Neural ODEs without BN (left) and Neural ODEs with TA-BN (middle). We also compare the accuracies of these two settings in one figure (right).
Ablations on TA-BN We integrate TA-BN into Aug-NODE 77. HBNODE [41] and SON-
Table Comparison between Neural ODE varia- tions with and without TA-BN on CIFAR-101.
ODE (34] Table 2 reports model accuracies w/
and w/o TA-BN. It indicates that our method can also work in compatible with existing Neu-
Accuracy
Aug-NODE HBNODE
SONODE
w/o TA-BN 0.848
0.851
0.302 0.853
ral ODE variants and boost their performances.
w/ TA-BN
0.862
0.867
To show the impact of different interpolation functions G, we test the following interpolation
We use convolution layers as the backhone for learn- able derivatives.
methods: (1) Linear interpolation demonstrated
by Eq. 77). (2) Cubic interpolation: We use cubic spline interpolation [ (3) Kernel smoothing:
Given time i, we calculate G(t, a. 7) by the weighted sum of elements in a. The weight coefficient for 4 is (t, 44) = exp(-0.5b (t 4)) with b 0.I7. (4) Gaussian process: We fit a Gaussian process 10 to perform interpolation in each forward pass. As illustrated by Figure 7 in addition to linear interpolation, the cubic interpolation and kernel smoothing methods also achieve good performance. However, these methods can be much slower than linear interpolation due to more complicated interpolation mechanisms. With a larger network architecture, we observe that these methods suffer from unaffordable rumning time and instable performance. Thus, we adopt linear interpolation as our default setting. Please see Appendix for further discussion. Future work can focus on improving the interpolation strategy for TA-BN.
5.2 Physical System Modeling
0.9
We first evaluate TA-BN on physical dynamical system model- ing tasks using the Walker2d-v2 (25] and HalfCheetah-v2 [14] datasets. These datasets consist of trajectories of 3D robot sys- tems generated by the Mujoco physics engine [40]. Following
Tedt Aommrr
11.8
0.7
0.6
0.5
hrnopiation Neaad Srandling
[4), we perform temporal autoregressive prediction and use the treatments from [4) to avoid collisions. We employ TA-BN in conjunction with a 12-layer MLP serving as the backbone for learnable derivatives. TA-BN with M 100 is applied to the first 5 layers, as is done with Mini- batch BN and Pop-TI TI BN. We utilize the dopri5 solver with its default settings for solv- ing ODEs and employ the RMSprop optimizer with a learning rate of 10-3 for training the neural networks over 100 epochs. The batch size is 64 for Walker2d-v2 and 32 for HalfCheetah- v2. Figure X presents the mean absolute error (MAE) along the training epochs on these datasets. Compared with Neural ODEs w/o BN, w/ Mini-batch BN, and w/ Pop-TI TI--TIPE BN, using TA-BN can achieve lower errors and faster convergence. Moreover, Table 3 summarizes the MAE results and compares them with the results in 49
Epoch
Figure 7: Ablations on interpola- tion method using CIFAR-10 with convolution layers as the backbone for learnable derivatives. Linear in- terpolation has a gap of only 0.01 compared to the best result. How- ever, other methods have degraded accuracies and unaffordable run- time when we up-scale the model.
8

---

-10
1.86
2.8
TTABNN (0xz))
my TA-IN (CuD))
NOUE- #T TIARK (OeS))
MIDDE
Love
0.98
NODEE 20p-11 BN
Test Lasss
NCIDE my Pap-222 BSX
Lees
2.2
BV
1.56
0.96
1.8
0.91
20
40
60
80
100
1:20
20
40
60
80
100
12
14
16
Epoch
Epoch
Epori
Figure 8: Comparison of Neural ODEs with TA-BN and base- lines on Walker2d-v2 (left) and HalfCheetah--2 (right). The Neural ODE backbone for learnable derivatives an MLP with 12 layers. Please see Appendix A.2 for architecture details.
Figure 9: Comparison of dif- ferent Neural ODEs on Charge Pump eircuit modeling. The backbone for learnable deriva- tives is 8-layer MLP.
Table 3: MAE comparison between different Neu- ral ODEs on Walker2d-v2 and HallChectah-v2.
Additionally, we investigate the effectiveness of TA-BN by modeling the temporal electrical current of Charge Pump (CP) circuit [9]. The CP is an analog integrated circuit composed of MOS transistors and is commonly utilized in commercial electrical products. To simulate the semiconductor manufacturing impact on the CP, we generate 200 circuit instances based on the process design kit (PDK) file. Concretely, simulate the CP's behavior from O seconds to 200 nanoseconds using a time grid of 2 picoseconds and record the electrical currents the CP output every 200 picoseconds.
Dataset
Walker2d-v2 HalfCheetah-v2
Baseline 442
1.02
1.46
w/o BN Mini-batch BN
0.959
1.40
0.953
1.38
w/ Pop-TI BN
0.962
1.37
w/ TA-BN (ours) 0.946
1.28
In this experiment, we utilize 16 circuit features (representing the widths of the MOS transistors) and the electrical currents at the present 20 time points (with each point comprising 2 values) as the Neural ODE input, resulting in 56 input features in total. Our objective is to train the Neural ODE to predict the electrical currents the subsequent 20 time points, yielding 40 outputs in total. To construct the dataset, we randomly sample 20k data from the simulation results. The training set includes 90% of them, and the remaining data used the testing set. We use linear layers to parameterize the derivative of the Neural ODE, with each layer containing 56 neurons. Mini-batch BN, Pop-TT BN, and TA-BN are applied to the first 3 layers. We use M = 100 for TA-BN. The final linear layer outside the Neural ODE maps the 56 features to 40 prediction values. For ODE solving, we use the dopri5 method with tolerance level of 10-3 We use the SGD optimizer with learning rate of 10-2 to train the model for 16 epochs. The batch size is 200.
Figure 9 presents the mean square error (MSE) results of Neural ODEs w/o BN, w/ Mini-batch BN, w/ Pop- TI BN, and w/ TA-BN. BN can improve the prediction, and TA-BN surpasses Mini-batch BN. It indicates that TA-BN is suitable for the circuit modeling task.
6 Conclusions and Limitations
In this paper, we demystify the previously unknown reason why batch normalization (BN) may lead to performance degradation when used with Neural ODEs. The fundamental mismatch arises from BN being initially designed for discrete neural networks without temporal dimension, whereas Neural ODES operate continuously over time. To address this challenge, we propose temporal adaptive batch normalization (TA-BN), specifically tailored for Neural ODEs. Our key is associating population statistics and learnable parameters with predefined regular time grids, with temporal interpolation automatically performed during training and testing. As continuous-time counterpart to traditional BN, TA-BN ensures training stability, thereby facilitating the scaling of Neural ODE model sizes. Our extensive experiments demonstrate TA-BN's ability to enhance the performance of Neural ODEs compared to baseline methods in tasks such as image classification and physical system modeling.
The interpolation process used in TA-BN inevitably introduces runtime overhead, which slows down the execution of Neural ODEs. The parameter-free linear interpolation technique is empirically found to be stable, time-efficient, and capable of providing performance improvement. However, future work can focus on enhancing the temporal interpolation used TA-B TA-BN to further optimize its efficiency and performance.

---

References
Interpolating natural cubic splines torcheubicspline
using PyTorch.
[2] Nils Bjorck, Carla Gomes, Bart Selman, and Kilian Q Weinberger. Understanding batch normalization. Advances in Neural Information Processing Systems, 31, 2018.
[3] Ricky T. Q. Chen. torchdiffeq. httpe: http:///tcchtbcccc///c//c//////a
2018.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations, Advances in Neural Information Processing Systems, 31, 2018.
[5] Tim Cooijmans. Nicolas Ballas, César Laurent, Çaglar Gülçehre, and Aaron Courville. Recurrent batch normalization. In International Conference on Learning Representations, 2017.
[6] Chaoteng Duan, Jianhao Ding, Shiyan Chen, Zhaofei Yu, and Tiejun Huang. Temporal effective batch normalization spiking neural networks. Advances in Neural Information Processing Systems, 35:34377- 34390, 2022.
[7] Emilien Dupont, Araud Doucet, and Yee Whye Tch. Augmented neural odes. Advances neural information processing systems, 32, 2019.
[8] Chris Finlay, Jörn-Hente Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. How to train your neural ode: the world of jacobian and kinetic regularization. In International conference machine learning, pages 1154-3164. PMLR, 2020.
[9] Zhengqi Gao, Jun Tao, Fan Yang, Yangfeng Su, Dian Zhou, and Xuan Zeng. Efficient performance trade-off modeling for analog circuit based on bayesian neural network. In 2019 IEEE/ACM International Conference on Computer-Aided Design (ICCAD), pages 1-8. IEEE, 2019.
[10] R Jacob Gardner, Geoff Pleiss, Bindel David, Kilian Q Weinberger, and Andrew Gordon Wilson, GPyTorch: Blackbox matrix-matr gaussian process inference with GPU acceleration. Advances in Neural Information Processing Systems, 2018.
[11] Amir Gholami, Kurt Keutzer, and George Biros. Anode: Unconditionally accurate memory-efficient gradients for neural odes. arXiv preprint 1r02 190..1029....... 2019.
[12] Arnb Ghosh, Harkirat Behl, Emilien Dupont, Philip Tom, and Vinay Namboodiri. STEER: Simple temporal regularization for neural ode. Advances Neural Information Processing Systems, 33:14831- 14843, 2020.
[13] Julia Gusak, Larisa Markeeva, Talgat Daulbaev, Alexandr Katrutsa, Andrzej Cichocki, and Ivan Oseledets. Towards understanding normalization in neural odes. arXiv preprint arXiv:2004.0922 092222, 2020.
[14] Ramin Hasani, Mathias Lechner, Alexander Amini, Daniela Rus, and Radu Grosu. Liquid time-constant networks. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 7657-7666, 2021.
[15] Jonathan Ho, Ajay Jain, and Pieter Abheel, Denoising diffusion probabilistic models. Advances in neural information processing systems,
2020.
[16] Sergey loffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International Conference on Machine Learning, pages 448 188-45..77 PMLR, 2015.
[17] Haiyan Jiang, Vincent Zoonekynd, Giulia De Masi, Bin Gu, and Huan Xiong. Tab: Temporal accumulated batch normalization in spiking neural networks. In International Conference on Learning Representations, 2024.
[18] Patrick Kidger, Ricky TQ Chen, and Terry Lyons. hey, that's not ode": Faster ode adjoints via seminorms. ICML, pages 5443-5452, 2021.
[19] Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time series. Advances in Neural Information Processing Systems, 33:6696-6707, 2020.
[20] Youngeun Kim and Priyadarshini Panda, Revisiting batch normalization for training low-latency deep spiking neural networks from scratch. Frontiers in neuroscience, 15:773954, 2021.
10

---

[21] Jonas Kohler, Hadi Daneshmand. Aurelien Lucchi, Thomas Hofmann, Ming Zhou, and Klaus Neymeyr. Exponential convergence rates for batch normalization: The power of Length-direction decoupling in non-convex optimization. In The 22nd International Conference on Artificial Intelligence and Statistics. pages 806-815 815. PMLR, 2019.
[22] Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images. 2009.
[23] César Laurent, Gabriel Pereyra, Philémon Brakel, Ying Zhang, and Yoshua Bengio. Batch normalized recurrent neural networks. In IEEE International Conference Acoustics, Speech and Signal Processing (ICASSP), pages 2657-2660 2016.
[24] Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. 2015.
[25] Mathias Lechner and Ramin Hasani. Mixed-memory RNNs for learning long-tern dependencies in irregularly sampled time series. 2022.
[26] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 (6(12):2278 23244, 1998.
[27] Ilya Loshchilov and Frank Hutter. arXiv:1711.05101, 2017.
Decoupled
weight decay regularization. arXiv preprint
[28] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems, 5577775-5787, 2022.
[29] Ping Luo, Xinjiang Wang. Wenqi Shao, and Zhanglin Peng. Towards understanding regularization in batch normalization. arXiv preprint arXiv: (ee::0000000440 00846, 2018.
[30] Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissecting neural odes. Advances in Neural Information Processing Systems, 33:3952 3963, 2020.
[31] Takashi Matsubara, Yuto Miyatake, and Takaharu Yaguchi. Symplectic adjoint method for exact gradient of neural ode with minimal memory. Advances in Neural Information Processing Systems, 34:20772-20784, 2021.
[32] James Morrill, Cristopher Salvi, Patrick Kidger, and James Foster. Neural rough differential equations for long time series. In International Conference Machine Learning, pages 7829-7838. PMLR, 2021.
[33] Yuval Netzer. Tao Wang. Adam Coates, Alessandro Bissacco, Baolin Wu, Andrew Y Ng. et al. Reading digits in natural images with unsupervised feature learning. In NIPS workshop on deep learning and unsupervised feature learning. volume 2011. page 7, Granada, Spain, 2011.
[34] Alexander Norcliffe, Cristian Bodnar, Ben Day, Nikola Simidjievski, and Pietro Liò. On second order behaviour in augmented neural odes, Advances neural information processing systems, 33:5911-599 2020.
[35] Viktor Oganesyan, Alexandra Volokhova, and Dmitry Vetrov. Stochasticity in neural odes: An empirical study, arXiv preprint XXix:2002.09779 2020.
[36] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high-performance deep lcarning library. Advances in neural information processing systems, 32, 2019.
[37] Michael Poli, Stefano Massaroli, Junyoung Park, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Graph neural ordinary differrnntial equations. arXiv preprint 191 1.07532, 2019.
[38] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang Chich Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 4510-4520, 2018
[39] Shibani Santurkar, Dimitris Tsipras, Andrew Ilyas, and Aleksander Madry. How does batch normalization help optimization? Advances in Neural Information Processing Systems, 31, 2018.
[40] Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In JEEP/RSJ international conference on intelligent robots and systems, pages 5026-5033. IERF, 2012.
[41] Hedi Xia, Vai Suliafu, Hangjie Ji, Tan Nguyen, Andrea Bertozzi, Stanley Osher, and Bao Wang. Heavy ball neural ordinary differential equations. Advances in Neural Information Processing Systems, 34: 34:16646 18659, 2021.
11

---

[42] Wei Xiao, Tsun-Ilsuan Wang, Ramin Hasani, Mathias Lechner. Yuiong Ban, Chuang Gan, and Daniela Rus. On the forward invariance of neural odes. In International conference machine learning, pages 38100-38124, PMLR, 2023.
[43] Tianjun Zhang, Zhewei Yao, Amir Gholami, Joseph E Gonzalez, Kurt Keutzer, Michael W Mahoney, and George Biros. Anodev2: coupled neural ode framework. Advances in Neural Information Processing Systems, 32, 2019.
[44] Hanle Zheng, Yujie Wu, Deng, Yifan Hu, and Guoqi Li. Going deeper with sirectly-trainee larger spiking neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 11062-1111 11070, 2021.
[45] Juntang Zhuang. Nicha Dvomek, Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Duncan. Adaptive checkpoint adjoint method for gradient estimation in neural ode. In International Conference on Machine Learning, pages 11639-
1639-11644 11649, PMLR, 2020.
12

---

A Implementation Details
A.1 Image Classification
Our model comprises a Neural ODE module followed by a single linear layer, which we refer to the unmixed Neural ODE architecture. The single linear layer maps the activations provided by Neural ODE to the final class probability (or logit). We have 3 settings for the Neural ODE module.
Multilayer Perceptron (MLP): We use 2 linear layers inside the Neural ODE module. On MNIST, each linear layer has 28 x 28 hidden features, matching the number of pixels an MNIST image.
Convolutional Neural Network (CNN): We use 4 ~ 12 convolution layers inside the Neural ODE module. The kemel size of all layers 3 x 3. The numbers of output channels are:
4 layers: (imput) + 32, + 64; + 32, + 3,.
6 layers: (input) + 32, + 64. 128 + 64, 32, 37.
7 - Le layers: ((inuut) + 32, + 64, - 128, (- 256) x L - 128 - 64, - 32, - 3
Note that layer with J downscales the feature maps using stride of 2. A layer with is transposed convolution layer that upscales the feature maps using stride of 2.
U-Net: We use U-Net with 18 convolution layers. The kernel size of all layers is 3 x 3. The numbers output channels are designed as (imput + 32 + 32 + 64 + 64
1284 + 128
256 + 256 128- + 128 + 128 64, - 64 64 + 32, 32
32
+ 3.
A.2 Physical Dynamical Systems Modeling
The HalfCheetah-v2 and Walker2d-v2 datasets consist of trajectories of 3D robot systems generated by the Mujoco physics engine [40]. Each trajectory represents sequence of 17-dimensional vector describing the system's 's state, such as the robot's joint angles and poses. The Neural ODEs are required to predict the trajectories in an autoregressive manner, as in 25 42).
In these tasks, the input size is 17, and the Neural ODE module is designed as 17(input) (- 64) x 11 - 17. BN is applied to the first 5 layers. Based on the code from (42, we do the prediction via the invariance size is the same as the output size, we don't 't need a linear layer outside the Neural ODE module.
propagation and data-controlled neural ODE proposed in 42 Since the input
0.64
0.64
Muth
Anmuuy
The: Avuuaay
0.32
0.32
Interpoiatioon
Linou Interpolation
0.16
Cubic Interpelation
Gausitan Proceus
Rovel Secodbiing
00
T00
Linear Interpolation
Knree Smoothing
40
Epoch
Epoch
per Step (Tast Fgunc)
(d) Figure 10: (a) Kernel smoothing can achieve slightly better accuracy on CIFAR-100 with 6 convolu-
(a)
(b)
(c)
tion layers for learnable derivatives. (b) Linear interpolation achieves more stable performance across
runs. (c) Linear interpolation performs better with a larger model size. (d) We compare different
methods on the runtime per step at the last epoch show the efficiency of linear interpolation.
B Discussion on Interpolation Methods
In Figure 7 we compare the performance of four interpolation methods, including linear interpolation. cubic interpolation, kernel smoothing, and Gaussian process. The cubic interpolation and kernel smoothing methods achieve good performance with 6 convolution layers for learnable derivatives.
The results on CIFAR-100 100 also support this point, as shown in Figure a0 However, Figure [0] b) indicates that linear interpolation achieves more stable performance, and Figure [0(c) clearly shows the superiority of linear interpolation with larger model size. Furthermore, linear interpolation is more efficient than cubic interpolation and Gaussian process, as shown in Figure 10 d). Therefore, use linear interpolation by default, which can provide superior performance and stability.
13

---

Additional Ablation Studies
Table 4: Ablation study on ODE solvers.
C.1 Ablation Study on ODE Solvers
Method
ODE Solver
Accuracy
We have conducted extra experiments on CIFAR-10 using fixed-step solvers like Euler method, with the 8-layer backbone. The results are shown in 5 Regardless of the solver, TA-BN achieves the best performance among the techniques. We also explored midpoint and rk4 solvers; however, they are much slower and haven't finished in the limited time constraint.
w/o BN w/ Pop-TI BN
Euler
0.839=0.002
Euler
0.631=0.203
w/ Mini-batch BN
Euler
0.8641:0.002
w/ TA-BN (ours)
Euler
0.872=000003
w/o BN w/ Pop-TT TI BN
Dopri5
0.843=0.004
Dopri5
0.392=0.000
w/ Mini-batch BN
Dopri5 0.005
w/ TA-BN (ours)
Dopri5
0.87410.001
C.2 Ablation Study on Time Grids
Table 5: Ablation study on ODE solvers.
Regarding the grid size hyperparameter M, we ran experiments without BN and found that the number of function evaluations (NFE) is around
Method
Time Grids M
Accuracy
w/ TA-BN (ours)
10
0.851=00015
hundreds. Thus, we set M 100 in our paper. We have performed extra ablation studies it, as reported in the following table. Using
w/ TA-BN (ours)
50
0.851=0.019
w/ TA-BN (ours)
100
0.874+100001
w/ TA-BN (ours)
500
0.870
M > 100 brings no improvement but too much
runtime overhead. We don't t have the confidence interval M 500 due to the time limit.
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: We clearly describe our main claims such as the findings, contributions, and results in the abstract and introduction. They match our idea, methodology, and experimental results that we present in the paper.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We describe the limitations of our method in Section 5 "Conclusion and Limitations".
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
15

---

Answer:
[NA]
Justification: Our method is based on batch normalization. We do not introduce additional assumptions and theories.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversell, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have fully disclosed the details of neural ODE architectures, ODE solvers, and optimizers in the paper, along with the used Python packages and hardware platform. Combined with the detailed experimental settings our paper, the provided information can be used to reproduce the main experimental results.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
16

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, described in supplemental material?
Question:
Answer: [Yes]
Justification: We put part of the code for reproduciblity in supplementary. It will be released upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nnis.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/nips. wttpcc/pcccccccccccooooooop httpa//gggudeee/dduuuuubbiiiiiii wetpamssssonPoe. swiniiossiiiiccyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We have specified all the training and test details the experiment section and the appendix. They follow existing papers and open-source code.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification: We report error bars in the major result table. We also show box plots for reporting error bars in some experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We have described the hardware platform in the experiment section.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer:
[Yes]
Justification:
We strictly conform the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: There is no societal impact of the work performed.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The paper poses risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
For the existing assets we use, have cited the papers and repos. They all open-source.
Justification:
Guidelines:
The answer NA means that the paper does use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
19

---

If assets are released, the license, copyright information, and terms of use in the thcode has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssi
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
20

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
21


Reviews:
[{"review_id": "dwXjybB06A", "reviewer": "Reviewer_6xoP", "summary": "This work identified the fundamental mismatch between Neural ODEs and traditional batch normalization (BN) techniques. To address this issue, the authors introduced Temporal Adaptive Batch Normalization (TA-BN), incorporating temporal interpolation to accumulate mini-batch statistics during training and use them as population statistics during inference. Neural ODEs with TA-BN outperformed those with traditional BN techniques and without BN on image classification and physical system modeling tasks. Additionally, TA-BN enhanced the performance of some other Neural ODE variants.", "questions": "Please see the comments above", "limitations": "Yes, it has discussed the limitations.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The work provided a thorough analysis of why traditional BN fails in Neural ODEs, which adds depth to the understanding of the problem.\n\nThe proposed TA-BN is well-motivated and has the potential to improve the performance of Neural ODEs in multiple fields.\n\nThis paper is well-organized and well-written.", "weaknesses": "- **Overclaim of experimental results**. The work claims that Neural ODEs can approach MobileNetV2-level efficiency in the abstract and introduction. However, training Neural ODEs can be slow, and TA-BN may exacerbate this problem due to increased computational complexity. Even if Neural ODEs with TA-BN achieve similar parameter efficiency to MobileNetV2, time efficiency should also be considered.\n\n- **Lack of theoretical analysis**. The work focused heavily on empirical results but could benefit from a more in-depth theoretical analysis of why TA-BN works so well, potentially providing insights that could lead to further improvements.\n\n- **Inappropriate experiment design**.  Vanilla Neural ODEs should not be used for Walker2d and HalfCheetah tasks due to potential collisions in such robotic systems [1]. Dynamical systems such as predator-prey equations and double pendulums are recommended.\n\n- More ablation studies should be performed, such as testing Neural ODEs with traditional BN, with TA-BN, and without BN using fixed-step-size ODE solvers to support the claim that adaptive step-size ODE solvers cause the failure of traditional BN in Neural ODEs. Additionally, more NDE models with TA-BN, such as ODE-RNNs [2], Latent ODEs [2], and Neural CDEs [3], should be tested to demonstrate TA-BN's performance.\n\n**References**:\n\n[1] Chen, Ricky TQ, Brandon Amos, and Maximilian Nickel. \"Learning Neural Event Functions for Ordinary Differential Equations.\" International Conference on Learning Representations.\n\n[2] Rubanova, Yulia, Ricky TQ Chen, and David K. Duvenaud. \"Latent ordinary differential equations for irregularly-sampled time series.\" *Advances in neural information processing systems* 32 (2019).\n\n[3] Kidger, Patrick, et al. \"Neural controlled differential equations for irregular time series.\" Advances in Neural Information Processing Systems 33 (2020): 6696-6707.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate Reviewer 6xoP's thorough review and valuable insights. Below we respond to each raised concern.\n\n---\n**Q1: Overclaim of experimental results about \"MobileNetV2-level efficiency\"**\n\nWe will clarify in our updated manuscript that \"efficiency\" refers to the accuracy a model can achieve given the number of learnable parameters, i.e., **parameter efficiency**—a concept commonly used in tiny/efficient ML [1], [2]—**not training efficiency** mentioned by the reviewer. Our TA-BN can elevate an unmixed Neural ODE to MobileNetV2-level parameter efficiency. Alternatively, we agree with the reviewer that Neural ODEs, w/ or w/o TA-BN, is relatively slow due to ODE solving.\n\n---\n**Q2: Lack of theoretical analysis on why TA-BN works.**\n\nWe respectfully highlight that our explanation (e.g., Figure 2 and right of Figure 3) sufficiently reveals the inadequacy of traditional BN for Neural ODE, and why TA-BN works. We agree that a theoretical analysis would make our argument more convincing, which is provided below:\n\nThe ideal normalization has mean $\\\\mu_k$ and variance $\\\\sigma^2_k$ at every time step $t_k$. Regardless of time, traditional BN uses the mean $\\\\bar{\\\\mu}$ and variance $\\\\bar{\\\\sigma}$ obtained by averaging the statistics over all time steps. We can express accurate ODE solving (e.g. Euler method) as: \n$$\n\\\\small\nh(t_n) = b \\\\sum_{k=0}^{n-1} ( \\\\frac{f_{\\\\theta} (h(t_k), t_k) - \\\\mu_k}{\\\\sigma_k} ), \\\\; i \\\\in [1, 2, \\\\cdots, N], \\\\; t_0 = 0, t_N = T. \n$$\nwhere $b$ is the time step, the input and output are $h(0)$ and $h(T)$, respectively. In traditional BN, we use a similar equation but replace $h(t_k)$, $\\\\mu_k$, and $\\\\sigma_k$ with $h_e(t_k)$, $\\\\bar{\\\\mu}$, and $\\\\bar{\\\\sigma}$, respectively. \n\nThe difference between ideal normalization and traditional BN is $\\\\Delta_n = h(t_n) - h_e(t_n) = b \\\\sum_{i=0}^{n-1} \\\\delta_k$, where\n$\n\\\\delta_k = ( \\\\frac{f_{\\\\theta} (h(t_k), t_k)}{\\\\sigma_k} - \\\\frac{f_{\\\\theta} (h_e(t_k), t_k)}{\\\\bar{\\\\sigma}} ) - ( \\\\frac{\\\\mu_k}{\\\\sigma_k} - \\\\frac{\\\\bar{\\\\mu}}{\\\\bar{\\\\sigma}} ). \n$\nWith some assumptions on $\\\\mu_k$, $\\\\sigma_k$ and $f_{\\\\theta}$, we can roughly prove $|\\\\Delta_n|$ has a lower bound grows asymptotically as $\\\\Omega (n)$.\n\nThe above implies that traditional BN fails because time-independent statistics assumption will lead to increasing error. However, TA-BN does not suffer from it due to the time-dependent statistics, and the linear interpolation in TA-BN brings little errors to statistic estimation. Considering the interpolation of $\\\\mu$ as an example, to estimate mean $\\\\mu^{(t)}$ at time $t$, we use $G(t,\\\\mathbf{\\\\mu},\\\\mathcal{T})$ (Eq. (7) in our paper). Assuming that $\\\\mu$ is Lipschitz-continuous in the range $[t_l, t_{l+1}]$ with a Lipschitz constant $k_l$, we can derive the upper bound of error: \n$$\n\\\\Vert G(t,\\\\mathbf{\\\\mu},\\\\mathcal{T}) - \\\\mu^{(t)} \\\\Vert = \\\\Vert \\\\frac{t_{l+1}-t}{t_{l+1}-t_l} (\\\\mu_l - \\\\mu^{(t)}) + \\\\frac{t-t_l}{t_{l+1}-t_l} (\\\\mu_{l+1} - \\\\mu^{(t)}) \\\\Vert \\\\leq \\\\frac{t_{l+1}-t}{t_{l+1}-t_l} k_l (t - t_l) + \\\\frac{t-t_l}{t_{l+1}-t_l} k_l (t_{l+1} - t) \\\\leq \\\\frac{k_l}{2} (t_{l+1} - t_l). \n$$\nWe can make the time interval $[t_l, t_{l+1}]$ small enough to have a small slope of $\\\\mu$, and thus a small $k_l$.\n\nDue to the time limit of response, we will leave a comprehensive detailed proof for future work.\n\n---\n**Q3: Inappropriate experiment design. Vanilla Neural ODEs are not for Walker2d and HalfCheetah.**\n\nFirst, we would like to clarify that the Neural ODEs for Walker2d and HalfCheetah (Table 3), are based on Reference [42] in our paper. They **incorporate the special treatments proposed by [42] to avoid collisions, rather than employing vanilla Neural ODEs.** Concretely, [42] prevents the system from entering unsafe regions by enforcing constraints via invariance propagation. Based on [42], we employ a larger Neural ODE and test various BN techniques. We appreciate the reviewer's insight regarding predator-prey and double pendulums, which are valuable and will be considered in future research.\n\n---\n**Q4: Ablation studies on fixed-step size solver and other Neural ODE models.**\n\nWe have conducted extra experiments using fixed-step solvers like Euler method, with the 8-layer backbone introduced in Appendix A.1. The results are shown in following table. Regardless of the solver, TA-BN achieves the best performance among the techniques. We also explored midpoint and rk4 solvers; however, they are much slower and haven't finished in the limited time constraint.\n\n|Method|Model|ODE Solver|Accuracy|\n|-|-|-|-|\n|TA-BN|8-layer|dopri5|0.874±0.001|\n|Mini-batch BN|8-layer|dopri5|0.865±0.004|\n|Pop-TI BN|8-layer|dopri5|0.332±0.090|\n|w/o BN|8-layer|dopri5|0.843±0.004|\n|TA-BN|8-layer|euler|0.872±0.003|\n|Mini-batch BN|8-layer|euler|0.864±0.002|\n|Pop-TI BN|8-layer|euler|0.631±0.203|\n|w/o BN|8-layer|euler|0.839±0.002| \n\nWe have considered using TA-BN with ODE-RNNs, Latent ODEs, and Neural CDEs. However, we realized it was not appropriate or out of our context because (i) These models are for irregular time series, while BN is mostly used for image problems. Variants of normalizations can be applied to time series (e.g., layer norm), but these are already distant from BN. As an extension of BN, TA-BN may not be suitable for time series. (ii) Neural ODE applications in time series often use shallow models with a small number of parameters (e.g., see Supplementary 5 of [3]). These models perform adequately in related tasks, and TA-BN does not bring significant improvement.\n\n---\n**References:** \n\n[1] S. Han et al., 'Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding,' ICLR 2016.\n\n[2] H. Mostafa et al., 'Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization,' ICML 2019.\n\n[3] Y. Rubanova et al., 'Latent ordinary differential equations for irregularly-sampled time series,' NeurIPS 2019.", "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": "", "extra_scores": {}}, {"review_id": "eZQvvuX8bI", "reviewer": "Reviewer_EgcE", "summary": "The paper proposes a remedy for batch normalisation in NeuralODE training, which proposes using a time grid for estimating the depth-dependent statistics, which leads to improved accuracy of the model.", "questions": "1) Piece-wise linear interpolation of the batch parameters, as outlined in Algorithm 1, would not be differentiable at the grid points (see line 2 of the algorithm); does it cause any problems for the convergence at the ODE integration?\n\n2) The authors claim that '“These models exhibit several intriguing features, such as the ability to compute gradients with constant memory consumption”\n\nI don’t think it would be an intriguing feature of Neural ODEs as  it has been achieved by coupling of NeuralODE parameters throughout the depths, which also could be done with the standard ResNets as well. Furthermore, the authors themselves define a time-dependent batch normalisation parameters grid which essentially highlights the trade-offs between the number of parameters and memory consumption vs accuracy.  \n\n3) There has been a line of work on how to parametrise NeuralODEs in a way that allows for continual change of parameters, which includes, e.g. shooting methods (Kwitt et al, 2020) and hypernetworks (the original paper Chen et al, 2018 itself, page 5, last paragraph of the intro in the section 4 called Time-dependent dynamics). I guess such line of work should be referenced as well, perhaps even showing whether such an alternative parameterisation through shooting methods/hypernetworks can equally address the problem. \n\nKwitt et al, A Shooting Formulation of Deep Learning, NeurIPS 2020", "limitations": "More empirical analysis would be necessary to highlight the failure modes: is such observed improvement of efficacy of batch normalisation observed for different ODE solvers? How does it compare with the hyper networks approach in the original NeuralODE paper?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "Clarity: the paper is clearly written, with good motivation\n\nQuality: the quality of the paper is good in general, however, see the weaknesses section\n\nSignificance: training NeuralODEs in a more efficient way, which is also practical and computationally efficient, would be a significant insight for the community. However, this idea lacks back up for the given insights: please refer to the originality section of the weaknesses for more discussion.", "weaknesses": "Originality: the authors need to confirm the originality of the paper: it seems like as the idea of using time grid for batch normalisation parameters is intuitive, there needs to be more empirical/theoretical insight drawn from this observation to make it a full NeurIPS paper. Now, such insight seems to be limited and boils down to the statement that grid-wise batch normalisation parameters improve when they are selected using grid. How would such model behave in case of different ODE solvers, would it influence the efficacy of batch normalisation? There is also a need to explore the relationship of such model with the ResNet with coupled layer parameters, which is an Euler discretisation of the NeuralODE (see Chen et al, 2018) : does the model still give advantages over such model? Does the same effect on batch normalisation  repeat in such scenario? \n\nQuality: the experimental results seems to not show any confidence intervals. They also do not seem to provide the answer on how to select the grid size hyperparameter (although I may have misunderstood it, so the author's clarification would be much appreciated).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We extend our sincere gratitude to Reviewer EgcE for the constructive feedback. We will include the following discussions in the updated manuscript.\n\n---\n**Q0.1 Originality: more insights and TA-BN with other ODE solvers**\n\nWe have conducted extra experiments using fixed-step solvers like the Euler method, with the 8-layer backbone introduced in Appendix A.1. The results are shown in the following table. Regardless of the solver, TA-BN achieves the best performance compared to baselines. We also explored midpoint and rk4 solvers; however, they were too slow to finish within the limited time.\n\n|Method|Model|ODE Solver|Accuracy|\n|-|-|-|-|\n|TA-BN|8-layer|dopri5|0.874±0.001|\n|Mini-batch BN|8-layer|dopri5|0.865±0.004|\n|Pop-TI BN|8-layer|dopri5|0.332±0.090|\n|w/o BN|8-layer|dopri5|0.843±0.004|\n|TA-BN|8-layer|euler|0.872±0.003|\n|Mini-batch BN|8-layer|euler|0.864±0.002|\n|Pop-TI BN|8-layer|euler|0.631±0.203|\n|w/o BN|8-layer|euler|0.839±0.002|\n\n\nRegarding our originality, with all due respect, we would like to justify our work from the following perspectives:\n\n+ Although it is natural to consider applying BN to Neural ODEs, all preliminary efforts fail to obtain stable and superior results. We for the first time thoroughly demystify the underlying reasons.\n\n+ We proposed TA-BN to resolve the above problem. Its time-dependent statistics and interpolation method can accurately normalize Neural ODEs' continuous-time dynamics. \n\n+ Although prior works using *unmixed* structure (a Neural ODE module followed by a linear layer) enable the clear analysis of Neural ODE architectures, they show poor performance. We achieve scalable and performant *unmixed* structures for the first time. \n\n---\n**Q0.2 Quality: confidence interval and grid size hyperparameter**\n\nIn our original manuscript, we included a set of results on confidence intervals (e.g., Fig. 10(b)) in the Appendix. We omitted others in the main text, as we empirically observed our method to be stable across independent runs. To justify it, we have added the repeated results in the following table, with the backbones 8-layer and UNet introduced in Appendix A.1. According to the results, TA-BN consistently outperforms other methods in accuracy and stability. We will update **the whole Table I** in our revised manuscript using the format of 'mean±std'. \n\n|Method|Dataset|Model|Accuracy|\n|-|-|-|-|\n|TA-BN|CIFAR10|8-layer|0.874±0.001|\n|TA-BN|CIFAR10|UNet|0.910±0.010|\n|Mini-batch BN|CIFAR10|UNet|0.822±0.095|\n|Pop-TI BN|CIFAR10|UNet|0.548±0.087|\n|w/o BN|CIFAR10|UNet|0.517±0.049|\n|TA-BN|SVHN|UNet|0.958±0.004|\n|Mini-batch BN|SVHN|UNet|0.906±0.310|\n|Pop-TI BN|SVHN|UNet|0.241±0.123|\n|w/o BN|SVHN|UNet|0.096±0.025|\n\nRegarding the grid size hyperparameter $M$, we ran experiments without BN and found that the number of function evaluations (NFE) is around hundreds. Thus, we set $M=100$ in our paper. We have performed extra ablation studies on it, as reported in the following table. Using $M>100$ brings no improvement but too much runtime overhead. We don't have the confidence interval of $M=500$ due to the time limit. \n\n|Method|Model|Time Grids|Accuracy|\n|-|-|-|-|\n|TA-BN|8-layer|10|0.851±0.015|\n|TA-BN|8-layer|50|0.851±0.019|\n|TA-BN|8-layer|100|0.874±0.001|\n|TA-BN|8-layer|500|0.870|\n\n---\n**Q1: Non-differentiability of piecewise linear interpolation at grid points**\n\nIf a variable $a=a(t)$ is defined piecewise linearly on time, then $\\\\frac{da}{dt}$ at defining time grids will be non-differentiable. However, our algorithm's backward propagation during training **does not involve such time-based gradients.** Specifically, using the symbols in Algorithm 1 of our manuscript, the time variables occur only in $\\\\omega_1$ and $\\\\omega_2$. According to the chain rule, we have $\\\\frac{dL}{d\\\\gamma_l^\\\\star}=\\\\frac{dL}{d\\\\gamma_j}\\\\frac{d\\\\gamma_j}{d\\\\gamma_l^\\\\star}=w_1\\\\frac{dL}{d\\\\gamma_j}$ for line 5, which doesn't involve any gradients with respect to time, and thus fully differentiable.\n\n---\n**Q2: Constant memory consumption claimed by Neural ODE**\n\nWe would like to clarify that the statement mentioned by the reviewer was intended as an introduction to Neural ODEs. This statement is **not related to**, nor a new property from our TA-BN, but is **directly referenced from** the original Neural ODE paper [1]. In that context, constant memory consumption means no intermediate results need to be stored during training because gradients are calculated by solving ODE (i.e., adjoint method) instead of backward propagation. We agree with the reviewer that this statement is confusing and will remove it in the updated manuscript.\n\n---\n**Q3: Comparison with hypernetwork and shooting methods**\n\nThanks for highlighting these methods. The shooting method [2] provides a time-varying weight trajectory. Hypernetwork [1] uses time-dependent model weights $\\\\theta(t)$. Both methods make model parameters dependent on time, enhancing flexibility. Alternatively, TA-BN addresses a different aspect, using time-dependent statistics to normalize the time-dependent outputs of each layer, stabilizing training and enabling deeper architectures. Note that while [1] and [2] can improve flexibility, TA-BN is still necessary to address training instabilities in deep Neural ODEs. In essence, the shooting method and hypernetwork (i.e., time-varying model parameters) address a different issue from TA-BN (i.e., time-dependent statistics).\n\n---\n**Q4: Limitations: More empirical analysis on TA-BN with other ODE solvers, and compare it with hypernetworks**\n\nWe kindly refer the reviewer to our response to your Q0.1 for the question of TA-BN with other ODE solvers, and our response to your Q3 for the comparison with hypernetworks.\n\n---\n**References:**\n\n[1] Qicky T.Q. Chen et al., 'Neural Ordinary Differential Equations', NeurIPS 2018.\n\n[2] Kwitt et al, 'A Shooting Formulation of Deep Learning', NeurIPS 2020.", "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": "", "extra_scores": {}}, {"review_id": "467sOTZob0", "reviewer": "Reviewer_97fY", "summary": "The paper presents Temporal Adaptive Batch Normalization (TA-BN) which is tailored for Neural Ordinary Differential Equations (Neural ODEs). This method addresses the limitation of applying traditional Batch Normalization to Neural ODEs by acting as a continuous-time analog. The use of TA-BN in Neural ODEs is shown to allow for deeper architectures, which in turn improves performance. The paper demonstrates the effectiveness of TA-BN in image classification tasks, achieving a high test accuracy on CIFAR-10 comparable to more complex models, and also shows its advantage in physical system modeling.", "questions": "- What is the impact of TA-BN on the training efficiency? I observe this paper selected STEER as a baseline, however, STEER is designed to accelerate the Neural ODEs instead of enhancing its performance.", "limitations": "See weakness. I am willing to increase my score, if my concern can be well-addressed.", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The writing and presentation of this paper are good.\n- This paper studies the issue of using BN in Neural ODEs, which is an interesting topic.\n- The experiment results are good.", "weaknesses": "- The comparison to related works in the context of SNN is not enough. Especially, I found TAB is highly related to this paper (even though the name is highly similar). There should be a comprehensive comparison to clarify the novelty and contribution of this paper.\n\n- Lack of experiments on training efficiency.\n\n\n[1] TAB: Temporal Accumulated Batch Normalization in Spiking Neural Networks;", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate reviewer 97fY for dedicating time to review our paper. The thoughtful feedback and constructive comments have been invaluable in improving the quality of our work. \n\n---\n**Q1: Comparison to related works in the context of SNN and clarifications on novelty.**\n\nThank you for the constructive feedback. In the submitted manuscript, we cited several papers on BN applied in SNN, including the TAB paper mentioned by the reviewer. Both TAB and TA-BN design BN by considering the temporal characteristics of models, but they differ significantly in the following aspects. The proposed TA-BN employs a temporal adaptive scheme through interpolation and is specifically designed for Neural ODEs with any type of solver (i.e., fixed step-size or adaptive step-size). In contrast, TAB [1] is designed for an SNN using fixed step-size discretization on the LIF neuron model, accumulating statistics based on all time steps visited up to the current time point: $\\\\mu_{1:t} = \\\\frac{1}{t} \\\\sum_{s=1}^{t} \\\\mu[s]$ and $\\\\sigma^2_{1:t} = \\\\frac{1}{t} \\\\sum_{s=1}^{t} \\\\sigma^2[s]$. Namely, TAB only records $\\\\{\\\\mu[1],\\\\mu[2],\\\\cdots,\\\\mu[t]\\\\}$ on uniform time grids, and thus **TAB cannot be direcly applied** in the Neural ODE context when adaptive solver is used, since we will need values like $\\\\mu[1.3]$, which is not available. Moreover, the learnable parameters in the BN layers at different time points are independent in TAB ($\\\\gamma[t]$ and $\\\\beta[t]$ in their symbols), while our proposed TA-BN also applies the temporal interpolation for them ($\\\\gamma$ and $\\\\alpha$ in our symbols). \n\nRegarding our contributions, Reviewer 6xoP appreciates our \"thorough analysis of why traditional BN fails in Neural ODEs\" and our \"well-motivated\" methodology. Reviewer XMSL concurs that we address a \"meaningful problem\" with a solution that is \"elegant in its simplicity.\" With all due respect, we would like to justify our work from the following three main perspectives:\n\n+ Although it is natural to consider applying BN to Neural ODEs, all preliminary efforts unfortunately fail to obtain stable and superior results. We for the first time thoroughly demystify the underlying reasons.\n\n+ We proposed temporal accumulated BN (TA-BN) to resolve the above problem. TA-BN's time-dependent statistics and interpolation method can accurately normalize Neural ODEs' continuous-time dynamics. \n\n+ Most Neural ODE studies use *mixed* structure. For instance, they insert a Neural ODE module into a middle layer of a CNN for image classification. However, the specific contribution of the Neural ODE module in such a setting remains unclear. Alternatively, prior works using *unmixed* structure (a pure Neural ODE followed by a learnable linear projection), such as Augmented Neural ODE, show poor performance. We achieve scalable and performant *unmixed* structures for the first time. \n\n---\n**Q2: Lack of experiments on training efficiency.**\n\nThank you for the constructive feedback. To demonstrate the impact of TA-BN on training efficiency, we compare the convergence speed and accuracy in the following table. When compared to Mini-batch BN, TA-BN can converge faster (CIFAR10 and CIFAR100) or with significantly higher accuracy (SVHN). As shown in the paper (e.g. Figures 3 and 5), Pop-TI BN and w/o BN are usually unstable, and some of them cannot converge. Therefore, TA-BN exhibits superior training efficiency. Moreover, although TA-BN inevitably introduces a small computational overhead, the overall training time may not increase due to the faster convergence speed. To better address the reviewer's concern, we will add the above discussion in the updated manuscript.\n\nTable. Convergence (defined as no improvement in 10 epochs) comparison \n\n| Dataset  | Model   | Method        | Epoch | Accuracy |\n|--------|-------|-------------|-----|--------|\n| MNIST    | 8-layer | TA-BN         | 40    | 0.984    |\n| MNIST    | 8-layer | Mini-batch BN | 40    | 0.984    |\n| CIFAR10  | 8-layer | TA-BN         | 89    | 0.868    |\n| CIFAR10  | 8-layer | Mini-batch BN | 92    | 0.864    |\n| SVHN     | UNet    | TA-BN         | 47    | 0.956    |\n| SVHN     | UNet    | Mini-batch BN | 33    | 0.920    |\n| CIFAR100 | UNet    | TA-BN         | 88    | 0.588    |\n| CIFAR100 | UNet    | Mini-batch BN | 100   | 0.584    |\n\n---\n**Q3: Is STEER a proper baseline?**\n\nThank you for your constructive feedback on our experimental details. STEER [2] employs random sampling of the end time of the ODE during training, which is a regularization technique that, as the reviewer mentioned, can accelerate training. However, as also stated in their abstract, and here we quote, '...the proposed regularization can significantly decrease training time and **even improve performance over baseline models**.' Therefore, we respectfully point out that STEER is a proper baseline. From a high-level perspective, both STEER and our proposed TA-BN are techniques that can improve Neural ODE performance, and thus they should be compared.\n\n---\n**Reference:**\n\n[1] H. Jiang et al. 'TAB: Temporal Accumulated Batch Normalization in Spiking Neural Networks,' ICLR 2024.\n\n[2] A. Ghosh et al. 'STEER: Simple temporal regularization for neural ode.' NeurIPS, 2020.", "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": "", "extra_scores": {}}, {"review_id": "D7d55Ry35Q", "reviewer": "Reviewer_XMSL", "summary": "This paper focusses on adapting standard Batch Normalization to be applied to Neural ODEs. The paper proposes the reason for standard BN failing on Neural ODEs is that the population statistics cannot be meaningfully tracked for continuous $t$. That is, since Neural ODEs can be viewed as having continuous depth, the population statistics should also be time-dependent.\n\nThe proposed solution is to track the population statistics on a predetermined uniform time grid. During an ODE solve, for arbitrary $t$, interpolation on this grid is carried out to get $\\mu(t)$ and $\\sigma(t)$. This is also done for the learnable scale and locations of each Temporal BN layer. The interpolation is also used to update the population statistics on the grid based on batch statistics calculated at arbitrary times during training.\n\nEvaluation on image classification and time-series regression tasks demonstrate that this adapted BN is effective and mitigates the issues with standard BN.", "questions": "- How have hyperparameters been selected?\n- I'm interested to hear a deeper explanation about why standard BN fails. The dynamics function of the Neural ODE is $f_\\theta(x, t)$. Assume this is an MLP with one hidden layer, and we only apply BN at the hidden layer. Aren't the batchnorm statistics more about tracking the mean and variance of the hidden state of this dynamics function, across the distribution of all possible hidden states given $x$ and $t$. And so, since this applies to the dynamics function, as long as many possible $x(t)$ and $t$ are seen during training BN is still valid? If so couldn't the training issues be avoided with lower learning rate or gradient clipping. My question is essentially: is this an engineering trick to improve training, or a solution to a deeper fundamental problem with applying BN to Neural ODEs? Either answer is fine, I'm just curious.", "limitations": "- The limitations are addressed in the conclusion.\n- There is no broader impact statement, it is not necessary for this work but would still improve the paper.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is nicely written.\n- A meaningful problem has been identified, explored and a solution has been proposed. The solution is elegant in its simplicity.\n- The evaluation is extensive and convincing.", "weaknesses": "There are very few weaknesses of this paper. My general opinion is that it should be accepted, ultimately the main weakness is that the work is \"incremental\". This is not a major weakness in my view it only limits the paper attaining the top scores.\n\nIn terms of weaknesses that can be addressed in a reasonable rebuttal period:\n- As given in part 7 of the checklist, there are only error bars included for some of the experiments, but not all. Are the numbers available? More detail needs to be given about how many repeats are carried out, and if repeats have not been carried out for some experiments, this should be justified. This applies to all results given in the main paper.\n- It would also be good to see how this method of BN can be applied to Neural CDEs and Continuous Normalizing Flows. This is not necessary but would still improve the paper.\n- It seems like this method would not be necessary if we train and carry out inference with fixed solver step size. Since here the population statistics can be stored and are meaningful at these points. This will negatively impact the accuracy of the solve, however it would be good to see this tradeoff in another experiment if possible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the insightful feedback from Reviewer XMSL. Below we respond to each raised concern.\n\n---\n**Q0.1: Error bar.**\n\nWe had some error bar results in Fig. 10 in the Appendix. We omitted others in the main text, as we observed TA-BN's stability across independent runs. To justify it, we add the repeated results in the following table. We agree with the reviewer on the importance of error bars. In the revised manuscript, we will update **the whole Table I** using the format of 'mean±std'. \n\n|Method|Dataset|Model|Accuracy|\n|-|-|-|-|\n|TA-BN|CIFAR10|8-layer|0.874±0.001|\n|TA-BN|CIFAR10|UNet|0.910±0.010|\n|Mini-batch BN|CIFAR10|UNet|0.822±0.095|\n|Pop-TI BN|CIFAR10|UNet|0.548±0.087|\n|w/o BN|CIFAR10|UNet|0.517±0.049|\n|TA-BN|SVHN|UNet|0.958±0.004|\n|Mini-batch BN|SVHN|UNet|0.906±0.310|\n|Pop-TI BN|SVHN|UNet|0.241±0.123|\n|w/o BN|SVHN|UNet|0.096±0.025|\n\n---\n**Q0.2: How TA-BN can be applied to Neural CDEs and CNFs?**\n\nNeural CDE is for time series, and Continuous Normalizing Flow (CNF) is for density matching and generation tasks. We considered integrating TA-BN with them initially but realized it was improper for three reasons: (i) BN is mostly used for image problems. Other normalizations can be applied (e.g., layer norm), but these are already distant from BN. As an extension of BN, TA-BN may not be suitable for time series and density matching. (ii) Neural CDE and CNF often use shallow models that perform adequately, and adding TA-BN does not yield significant improvement. (iii) The need for integrating TA-BN into high-dimensional CNF has been eliminated by flow matching [1], where CNF no longer needs ODE solving. \n\n---\n**Q0.3: TA-BN is not necessary with a fixed step size.**\n\nWe view TA-BN as being simplified in a fixed step-size solver rather than unnecessary. In such solvers, TA-BN degenerates to a simpler version without interpolation, aligning with the reviewer's point on storing statistics at fixed time grids. However, standard BN will still malfunction with a fixed step-size solver. For detailed explanations of each BN technique in the Neural ODE context, please refer to our response to Q2.\n\nWe have tried different ODE solvers, including fixed step-size ones, and added the results in the table below. Regardless of the ODE solver, TA-BN is the best one compared to other techniques. We also tested midpoint and rk4 solvers, but they are much slower and didn’t finish within our time constraint.\n\n|Method|Model|ODE Solver|Accuracy|\n|-|-|-|-|\n|TA-BN|8-layer|dopri5|0.874±0.001|\n|Mini-batch BN|8-layer|dopri5|0.865±0.004|\n|Pop-TI BN|8-layer|dopri5|0.332±0.090|\n|w/o BN|8-layer|dopri5|0.843±0.004|\n|TA-BN|8-layer|euler|0.872±0.003|\n|Mini-batch BN|8-layer|euler|0.864±0.002|\n|Pop-TI BN|8-layer|euler|0.631±0.203|\n|w/o BN|8-layer|euler|0.839±0.002|\n\n---\n**Q1: How to select hyperparameters?**\n\nBelow we briefly show our hyperparameter setting criteria and extra ablation studies on hyperparameters. They will be added to the updated manuscript.\n\n+ We chose the popular dopri5 ODE solver. Since the ODE tolerance impacts the accuracy and runtime, we tested tolerances in $[10^{-5},10^{-1}]$ and chose $10^{-3}$, which has decent accuracy and affordable solving time. Additional experiments on different solvers can be found in our response to Q0.3 above.\n\n+ Regarding the number of TA-BN time grids $M$, we ran experiments without BN and found that the number of function evaluations (NFE) is usually around hundreds. Thus, we used $M=100$ in our manuscript. Ablation studies on the value of $M$ for TA-BN are shown in the following table. The confidence interval of $M=500$ is absent due to the time limit. \n\n|Model|Time Grids|Accuracy with TA-BN|\n|-|-|-|\n|8-layer|10|0.851±0.015|\n|8-layer|50|0.851±0.019|\n|8-layer|100|0.874±0.001|\n|8-layer|500|0.870|\n\n+ In image classification, we used typical hyperparameters (e.g., AdamW optimizer with a learning rate of 1e-3). We train models for 128 epochs to ensure convergence. The learning rate is decreased by a factor of 0.1 at the 64th epoch.\n\n+ In physical systems modeling, we empirically selected the number of MLP layers to make the models deep enough so that the effect of BN can be observed. We tried different solvers like AdamW, RMSprop, and SGD with typical learning rates and selected the best one.\n\n---\n**Q2: Deeper explanation on why standard BN fails.**\n\nHere, we present clear pseudocode and deep illustrations. Considering the Neural ODE mentioned by the reviewer: $\\\\frac{dx}{dt}=f_\\\\theta(x,t)$, where $f_\\\\theta$ is an MLP with one hidden layer and standard BN, its `forward` function is:\n\n```python\ndef forward(x,t):\n    a1 = ReLU(self.linear1(x))\n    a2 = self.bn(a1)\n    return self.linear2(a2)\n```\nWhen solving from the input $x(0)$ to the output $x(T)$, this model will be called $N$ times at different time points $\\\\lbrace t_n \\\\rbrace_{n=1}^N$. This forces `self.bn` to update its statistics based on all $\\\\lbrace t_n \\\\rbrace_{n=1}^N$. Essentially, it assumes $a_1$ has the same distribution regardless of $t$, which is not true. This is the problem of applying standard BN to Neural ODE in any DL frameworks, which we call **Pop-TI BN** in our manuscript. The issue is more profound than an engineering problem, and simple variants of BN also fail. First, **Mini-batch BN** uses $a_1$'s mini-batch statistics at each time step. It suffers from outliers and small batches. Second, **Pop BN** uses $N$ statistics for $\\\\lbrace t_n \\\\rbrace_{n=1}^N$ separately. It doesn't work because adaptive solvers change $\\\\lbrace t_n \\\\rbrace_{n=1}^N$ for every batch. For fixed-step-size solvers, Pop BN can work, and our TA-BN is equivalent to Pop BN in this case.\n\n---\n**Q3: Broader impact statement.**\n\nWe have reached our response length limit but will include a broader impact statement in the updated manuscript.\n\n---\n**References:** \n\n[1] Yaron Lipman et al., 'Flow Matching for Generative Modeling', ICLR 2023.", "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": "", "extra_scores": {}}]

Author rebuttal:
We sincerely thank all the reviewers for their thoughtful and constructive feedback on our manuscript. We are encouraged by their positive remarks, noting that our proposed TA-BN is well-motivated (Reviewer 6xoP, Reviewer EgcE), addresses an interesting and important problem (all reviewers), and that our experiments are robust and convincing (Reviewer 97fY, Reviewer XMSL). We also appreciate the comments highlighting that the paper is well-organized and well-written (all reviewers). We diligently addressed all the concerns raised by providing ample evidence and the requested results. The raised points will be thoughtfully considered and integrated into the revised manuscript. Here is the summary of our responses to the raised questions:

1. **Extra numerical results (e.g., confidence interval, other ODE solvers)**: We have conducted additional experiments to analyze the confidence intervals of various methods, the effects of different ODE solvers, and the results of hyperparameter ablation studies, among others. Our TA-BN exhibits better accuracy and stability under different ODE solvers compared to baselines. The selected number of time grids makes a superior trade-off between accuracy and speed. 

2. **TA-BN with other Neural ODE models (e.g., ODE-RNN, Neural CDE, CNF)**: Regarding the application of TA-BN with methods like ODE-RNNs [1], Neural CDEs [2], and continuous normalizing flow (CNF) [3], we did perform a quick, small-scale experiment. However, we realized it was not appropriate or out of our context due to the following reasons: (i) BN was originally proposed and is mostly used for image problems. As an extension of BN, TA-BN may not be suitable for time series or density matching. (ii) ODE-RNN, Neural CDE, and CNF often use shallow models with a small number of parameters as the trainable dynamics. These shallow models perform adequately well in related tasks, and adding TA-BN to these shallow neural networks does not yield significant improvement.

3. **Novelty and contribution:** Our TA-BN contributes to the field of Neural ODEs in the following three aspects: (1) We for the first time thoroughly demystify the underlying reasons why preliminary efforts of applying BN to Neural ODEs failed to obtain stable and superior results. (2) TA-BN's time-dependent statistics and interpolation method can accurately normalize Neural ODEs' continuous-time dynamics, resolving the above problem. (3) Although prior works using *unmixed* structure (a Neural ODE module followed by a linear layer) enable the clear analysis of Neural ODE architectures, they show poor performance. We achieve scalable and performant *unmixed* structures for the first time. 

4. **Experimental setup (e.g., hyperparameters, baseline choices)**: We select hyperparameters based on prior research and our empirical experience. For physical system modeling baselines, we meticulously choose a base method that satisfies the system constraints before modifying the architecture and integrating different BN methods. These efforts ensure the effectiveness and validity of our experiments. Please refer to individual responses for details.

5. **Theoretical analysis**: We prove that traditional BN fails to normalize the activations with appropriate statistics, while TA-BN avoids this problem by accurately estimating the statistics at each time step. Furthermore, we prove that TA-BN's linear interpolation incurs negligible estimation error in statistics.

6. **Clarifications (e.g., efficiency, differentiability)**: To avoid misunderstanding, we clarify that "efficiency" in our discussion refers to the parameter efficiency in efficient/tiny-ML. For certain non-differentiable operations in TA-BN, we emphasize that training does not involve their gradients. We have also addressed other potential points of confusion concerning our baselines, Neural ODEs' memory consumption, and so on.

We would again like to thank all reviewers for their valuable time and feedback, and we hope that our changes adequately address all concerns. Any further questions are highly welcomed. Below, we will provide individual responses to address each reviewer’s concerns.


**References:**

[1] Y. Rubanova et al., 'Latent ordinary differential equations for irregularly-sampled time series,' NeurIPS 2019.

[2] P. Kidger et al. 'Neural controlled differential equations for irregular time series,' NeurIPS 2019.

[3] Qicky T.Q. Chen et al., 'Neural Ordinary Differential Equations', NeurIPS 2018.

Meta-review: The reviewers all agreed on the premise of the paper: an algorithm for tracking Batch Normalization within the temporal field of Neural ODEs is introduced that uses interpolation. The paper explains the problems with naive implementation to BN and provides experiments comparing against other Neural ODEs benchmarks.

The reviewer scores average to a weak accept, but reviewer 97fY gave a borderline reject recommendation. This reviewer had a lower confidence score than the others. The reviewer opined that the work was not very novel compared to another published method (TAB) but the authors highlighted an additional insight from their work that did not appear in the TAB paper. 
The existence of similar algorithms was also addressed by reviewer EgcE; the AC is also aware of similar algorithms for BN in Neural ODEs; reviewer XMSL also thought the paper could be considered incremental. These reviewers did not think that novelty / incrementality were blockers for publication.
 Additionally, after the discussion, reviewer 97fY insisted that additional comparisons on training efficiency would be necessary, but the authors did not think it would be possible to do so in the time allotted.

New experimental results were added during the discussion. These results were required for two of the reviewers to increase their score to 6. Therefore, it is necessary for the authors to include these results in their revision. The new experimental results are three distinct tables, so it should be possible to fit the results and a short discussion into the revision. 

The authors also stated that a new proof in discussion would not be possible to include within time constraints. Including the proof was not required by the reviewer (6xoP) to increase the score.

Overall, the arguments to accept the paper outweigh the limitations noted by the reviewers. As such, the paper can be accepted (with the note that the reviewers did require the additional tables described above).
Decision: Accept (poster)