Improving Neural ODE Training with Temporal Adaptive Batch Normalization

Authors: su  zheng
zhengqi  gao
fan-keng  sun
duane s boning
bei  yu
martin  wong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 04:47:21

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.

Main paper: 
 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:EQUATIONwhere f θ (•, •) ∈ R D × R → R D is a neural network parameterized by learnable parameter θ controlling the temporal evolution, and h(t) ∈ R D is the ODE state variable at time t. The practical usage of Neural ODEs is by solving an 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, 31, 18] reverse in time.Subsequent works have expanded the capabilities and scope of Neural ODEs by introducing modifications, such as improving expressivity through state augmentation [7] , coupling the evolution of both weights and activations [43] , accelerating convergence via semi-norm techniques [18] , extending to graph neural networks [37] and second-order ODEs [34] , adapting to irregular time-series data [19, 32] , and drawing connections to diffusion models [15, 28] . Of particular interest, two preliminary studies [13, 35] 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 to operate effectively with Neural ODEs.Batch Normalization BN [5, 29, 16, 39, 2, 21] performs a re-centering and a re-scaling operation on the given input by subtracting the mean and dividing by the standard deviation:EQUATIONwhere ϵ ∈ R + ensuring division validity, {γ ∈ R, α ∈ R} are learnable parameters, and x i ∈ R represents the i-th data in one batch of size B. During training, the mini-batch statistics2 µ = 1/B B i=1 xi and σ 2 = 1/B B i=1 (xi -µ) 2 are used for efficiency. In contrast, at test time, BN uses population statistics aggregated across the entire training dataset to ensure deterministic outputs [16] . 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 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 [16] 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) [5, 23] and spiking neural networks (SNNs) [20, 44, 6, 17] . Unlike these models, Neural ODEs exploit adaptive time discretization, necessitating a specialized modification of BN for its application. 
traditional batch normalization in neural odes: 
 Let us consider a simplified one-dimensional Neural ODE when incorporating BN as example:EQUATIONwhereθ = {γ ∈ R, α ∈ R, w ∈ R, b ∈ R} are the learnable parameters. Given a batch of input {h i (0)} B i=1 (i.e., initial values of the ODE), the forward process employs an ODE solver to discretize Eq. ( 3) over time, visiting N sequential points T = (t 0 = 0, t 1 , t 2 , • • • , t N = T ) and computing the outputs {h i (T )} Bi=1 . An adaptive step size solver (e.g., Runge-Kutta method) is usually preferred because of its higher efficiency compared to a fixed step size solver (e.g., Forward Euler method). Thus, the differences between two consecutive time steps might not be identical (e.g., t 2t 1 ̸ = t 1t 0 ), and the value of N also depends on the batch of data. The temporal discretization indicates that we need to invoke BN at every t ∈ T for every sample:EQUATIONfor everyj = 1, 2, • • • , N and i = 1, 2, • • • , B, where µ j and σ 2 j are statistics associated with t j .During training To perform the normalization shown in Eq. ( 4), BN applies mini-batch meanµj = 1/B B i=1 xi,j and variance σ 2 j = 1/B B i=1 (xi,j -µj) 2 .Although the training can proceed 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 t ′ j ∈ T ′ might not coincide with any time points in T , and vice versa. This might result in an impractical collection of infinite statistics µ j and σ 2 j . Moreover, except those associated with t = 0 and t = T , most stored statistics will be calculated and updated based on a limited number of batches. During inference Conventionally, BN should perform the normalization shown in Eq. ( 4) using population statistics during inference. However, as shown in the left part of Figure 2 , the population statistics associated with the time point t ′ j ∈ T ′ , required by the temporal discretization during inference, might not be available if the time value t ′ j is never encountered during training. Moreover, even if the population statistics exist, they are likely to be inaccurate as discussed above. Summary A 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 h i (T ) becomes highly dependent on their respective batch's specific characteristics. This will make the outputs inaccurate when batch size is very small (e.g., B = 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 3 .For clarity, we will henceforth refer to 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 [4] ), actually a variant of Pop BN is used, which we will refer to as Pop-TI BN. Pop-TI BN accumulates statistics at all encountered time points during training into a single pair of µ and σ 2 , 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] . 3 In a nutshell, our observations can be succinctly summarized as follows:• Pop BN fails when Neural ODEs employ an adaptive step size solver.• Pop-TI BN implicitly assumes time-independent statistics, leading to erroneous outputs.• Mini-batch BN fails when Neural ODEs operate with small batch sizes or outliers exist.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?• 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. Justification: There is no societal impact of the work performed. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 a 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 or 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). 
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-dependent in Neural ODEs. Ideally, we would want BN to store population statistics at every t ∈ [0, T ], with corresponding learnable parameters γ(t) and β(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), β(t), and the population statistics will be naturally sparse on t, resulting in inefficiency.To address the aforementioned problem, we propose an intuitive method termed temporal adaptive BN (TA-BN) for usage in Neural ODEs. Let us illustrate it with the example shown in Eq. (3). To begin with, we evenly divide the time span [0, T ] defining (M + 1) time gridsT ⋆ = (t ⋆ 0 = 0, t ⋆ 1 = T M , t ⋆ 2 = 2T M , • • • , t ⋆ M = T ).We associate the time grid t ⋆ m with population mean µ ⋆ m and population variance σ ⋆,2 m , as well as learnable parameters γ ⋆ m and α ⋆ m for every m = 0, 1, 2, • • • , M . For later simplicity, we denoteγ ⋆ = [γ ⋆ 0 , γ ⋆ 1 , • • • , γ ⋆ M ] T and α ⋆ = [α ⋆ 0 , α ⋆ 1 , • • • , α ⋆ M ]T . Similar notations apply to µ ⋆ and σ ⋆ .During training As shown in the right of Figure 2 , given a batch of data {h i (0)} B i=1 , the forward pass of Neural ODE might discretize the time as T = (t 0 , t 1 , • • • , t N ) which differs from the grids T ⋆ of TA-BN. However, TA-BN can calculate the temporal derivative for the i-th data at t j :EQUATIONwhere (µ j , σ 2 j ) are still the mini-batch mean and variance based on {x i,j } B i=1 . Here γ j and α j are interpolated based on γ ⋆ and α ⋆ , respectively:EQUATIONwhere G(t, a, T ) is a function to interpolate the value a(t) given an array of values a ∈ R M +1 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:EQUATIONwhere l is the index which makes t l to be the largest value smaller than t 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, T ) can approximate any continuous a(t) arbitrarily well, i.e., |G(t, a, T )a(t)| → 0 when M → ∞. 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 (µ j , σ 2 j ) to the population statistics (µ ⋆ m , σ ⋆,2 m ), 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:EQUATIONwhere η ∈ [0, 1] is a momentum constant to perform moving averaging to update the population statistics based on current estimate of µ ⋆ m and σ ⋆,2 m using the new observed values G(t ⋆ m , µ, T ) and G(t ⋆ m , σ 2 , T ), respectively. Note that the training does not involve the gradients with respect to t.During inference After the training phase, TA-BN will have well-trained parameters γ ⋆ and α ⋆ , and accurately record population statistics µ ⋆ and σ ⋆ . 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:EQUATIONThis 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 ′ .In practice, we observe that the interpolations during training outlined in Eq. ( 6)-( 8) have opposite directions, potentially leading to increased computational overhead. To elaborate, when encountering a time point t j discretized by the ODE solver during training, the interpolation for obtaining γ j and α j in Eq. ( 6) using Eq. ( 7) can be immediately executed by identifying the index l such that t ⋆ l is the largest value smaller than t j . However, the interpolation 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 t j '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 at time t j . The key steps of our proposed TA-BN are summarized in Algorithm 1. It will be called as a subroutine for every discretized time point t j ∈ T required by the ODE solver.Algorithm 1 outlines the case involving a 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 a dimensional-wise manner. Thus, we must maintain 2(M + 1)Q population statistics and 2(M + 1)Q learnable parameters, each with D dimensions. In contrast, a 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 1 relies 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 as those not adjacent to any t j 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 ∥γ ⋆ -1∥ 2 and ∥α ⋆ ∥ 2 to avoid significant discrepancies among the elements in γ ⋆ and α, making the training more stable.The forward pass of a TA-BN layer at time t j Input: Batched input x = {x i,j } B i=1 at time t j 1: Get t ⋆ l and t ⋆ l+1 such that t ⋆ l is the largest value smaller than t j in T * ; 2: ω 1 = t ⋆ l+1 -tj t ⋆ l+1 -t ⋆ l , ω 2 = t ⋆ j -t l t ⋆ l+1 -t ⋆γ j = ω 1 γ ⋆ l + ω 2 γ ⋆ l+1 , α j = ω 1 α ⋆ l + ω 2 α ⋆ l+1 ; 6: µ ⋆ l ← (1 -η • ω 1 )µ ⋆ l + η • ω 1 µ j , µ ⋆ l+1 ← (1 -η • ω 2 )µ ⋆ l+1 + η • ω 2 µ j ; 7: σ ⋆,2 l ← (1 -η • ω 1 )σ ⋆,2 l + η • ω 1 σ 2 j , σ ⋆,2 l+1 ← (1 -η • ω 2 )σ ⋆,2 l+1 + η • ω 2 σ 2 j ; 8:else if the model is in the evaluation mode then 9:µ j = ω 1 µ ⋆ l + ω 2 µ ⋆ l+1 , σ 2 j = ω 1 σ ⋆,2 l + ω 2 σ ⋆,2 l+1 ;10:γ j = ω 1 γ ⋆ l + ω 2 γ ⋆ l+1 , α j = ω 1 α ⋆ l + ω 2 α ⋆ l+1 ; 11: end if 12: return x-µj √ σ 2 j +ϵ γ j + α j ; 
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), a Neural ODE module, and an MLP in sequence for classification. Another study structures neural networks with alternating Neural ODE modules and convolutional/linear layers [11] . 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.e., 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 a customized TA-BN layer. All experiments are run on a Linux server with RTX 3090 GPUs.Being consistent in experimental scales of previous Neural ODE studies, we conduct image classification across datasets including MNIST [26] , SVHN [33] , CIFAR-10, CIFAR-100 [22] , and Tiny-ImageNet [24] . We employ the dopri5 solver with a tolerance of 10 -3 for ODE solving and adopt the AdamW 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. [7] and STEER [12] . On MNIST and CIFAR-10, TA-BN has fewer parameters than Aug-NODE and STEER, which indicates 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 IL-NODE, 2nd-Ord [30] , HBNODE, and GHBNODE [41] . For reference, the popular convolutional-based 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 unmixed architecture, while most previous Neural ODE literature either falls short of this performance or/and does not follow this unmixed architecture. Neural ODE Up-scaling Enabled by TA-BN As shown in the left part of Figure 6 , Neural ODEs with a 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 Ablations on TA-BN We integrate TA-BN into Aug-NODE [7] , HBNODE [41] , and SON-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 Neural ODE variants and boost their performances.To show the impact of different interpolation functions G, we test the following interpolation methods: (1) Linear interpolation demonstrated by Eq. ( 7). ( 2) Cubic interpolation: We use cubic spline interpolation [1] . (3) Kernel smoothing: Given time t, we calculate G(t, a, T ) by the weighted sum of elements in a. The weight coefficient for t i is K(t, t i ) = exp(-0.5b -2 (tt i ) 2 ) with b = 0.1T . (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 running time and instable performance. Thus, we adopt linear interpolation as our default setting. Please see Appendix B for further discussion. Future work can focus on improving the interpolation strategy for TA-BN. Figure 7 : Ablations on interpolation method using CIFAR-10 with 6 convolution layers as the backbone for learnable derivatives. Linear interpolation has a gap of only 0.01 compared to the best result. However, other methods have degraded accuracies and unaffordable runtime when we up-scale the model. We first evaluate TA-BN on physical dynamical system modeling tasks using the Walker2d-v2 [25] and HalfCheetah-v2 [14] datasets. These datasets consist of trajectories of 3D robot systems generated by the Mujoco physics engine [40] . Following [42] , we perform temporal autoregressive prediction and use the treatments from [42] 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 BN. We utilize the dopri5 solver with its default settings for solving 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 8 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 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 [42] . 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 in TA-BN to further optimize its efficiency and performance.Our model comprises a Neural ODE module followed by a single linear layer, which we refer to as 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 × 28 hidden features, matching the number of pixels in an MNIST image.Convolutional Neural Network (CNN): We use 4 ∼ 12 convolution layers inside the Neural ODE module. The kernel size of all layers is 3 × 3. The numbers of output channels are:4 layers: 3(input) → 32 ↓ → 64 ↓ → 32 ↑ → 3 ↑ . 6 layers: 3(input) → 32 ↓ → 64 ↓ → 128 ↓ → 64 ↑ → 32 ↑ → 3 ↑ . 7 + L c layers: 3(input) → 32 ↓ → 64 ↓ → 128 ↓ (→ 256) × L c → 128 → 64 ↑ → 32 ↑ → 3 ↑ .Note that a layer with ↓ downscales the feature maps using a stride of 2. A layer with ↑ is a transposed convolution layer that upscales the feature maps using a stride of 2.U-Net: We use a U-Net with 18 convolution layers. The kernel size of all layers is 3 × 3. The numbers of output channels are designed as3(input) → 32 ↓ → 32 → 64 ↓ → 64 → 128 ↓ → 128 → 256 → 256 → 128 ↑ → 128 → 128 → 64 ↑ → 64 → 64 → 32 ↑ → 32 → 32 → 3.The HalfCheetah-v2 and Walker2d-v2 datasets consist of trajectories of 3D robot systems generated by the Mujoco physics engine [40] . Each trajectory represents a sequence of a 17-dimensional vector describing the system'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) × 11 → 17. BN is applied to the first 5 layers. Based on the code from [42] , we do the prediction via the invariance set propagation and data-controlled neural ODE proposed in [42] . Since the input size is the same as the output size, we don't need a linear layer outside the Neural ODE module. 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 also support this point, as shown in Figure 10 (a). However, Figure 10(b) indicates that linear interpolation achieves more stable performance, and Figure 10(c ) clearly shows the superiority of linear interpolation with a larger model size. Furthermore, linear interpolation is more efficient than cubic interpolation and Gaussian process, as shown in Figure 10(d) . Therefore, we use linear interpolation by default, which can provide superior performance and stability. 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 Table 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. Guidelines:• The answer NA means that the paper has no 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 a 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 on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 of 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a 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 to usage guidelines or 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 a best faith effort.12. Licenses for existing assets Question: 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?Answer: [Yes] Justification: For the existing assets we use, we have cited the papers and repos. They are all open-source.Guidelines:• The answer NA means that the paper does not 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 a 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.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • 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 is not available online, the authors are encouraged to reach out to the asset's creators. 
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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was 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. 
crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution 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. 
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 an equivalent approval/review based on the requirements of your country or institution) were obtained?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. • Depending on 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.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.Throughout our paper, statistics refer to mean and standard deviation (std) or equivalently variance.We 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 Figure3. 
: 
 Originally derived as the continuous limit of a residual connection [4] , neural ordinary differential equations (Neural ODEs) [4, 7, 37, 45, 8, 42, 30, 18, 32, 19, 31] is a family of continuous-depth neural networks where the evolution of hidden states is governed by learnable temporal derivatives. These models exhibit several intriguing features, such as the capability of temporal reversibility, which enables generative modeling [4, 15] .Previous studies on Neural ODEs parameterize the learnable temporal derivatives using a shallow neural network with a limited number of parameters [7, 30] . Without special treatment, merely stacking additional layers in the temporal derivatives does not necessarily enhance Neural ODE performance. Furthermore, deeper networks might increase the stiffness of the ODE system, leading to challenges with the ODE solver, such as excessively small step sizes or even failures due to infinite state values, as shown in Figure 1 .In efforts to deepen the model and address instabilities, one might instinctively consider Batch Normalization (BN) [5, 29, 16, 39, 2] as a remedy to stabilize the intermediate state values of Neural ODEs [13, 35] , because BN was originally proposed to accelerate 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, an effect for which the underlying causes remain unclear [13, 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 statistics [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 1 ). 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 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. 
 
 Appendix:  


Decision: 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).

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - **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.

- **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.

- **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.

- 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.

**References**:

[1] Chen, Ricky TQ, Brandon Amos, and Maximilian Nickel. "Learning Neural Event Functions for Ordinary Differential Equations." International Conference on Learning Representations.

[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).

[3] Kidger, Patrick, et al. "Neural controlled differential equations for irregular time series." Advances in Neural Information Processing Systems 33 (2020): 6696-6707.
questions: questions: Please see the comments above
limitations: limitations: Yes, it has discussed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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? 

Quality: 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).
questions: 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?

2) The authors claim that '“These models exhibit several intriguing features, such as the ability to compute gradients with constant memory consumption”

I 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.  

3) 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. 

Kwitt et al, A Shooting Formulation of Deep Learning, NeurIPS 2020
limitations: 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?
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.

- Lack of experiments on training efficiency.


[1] TAB: Temporal Accumulated Batch Normalization in Spiking Neural Networks;
questions: 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: limitations: See weakness. I am willing to increase my score, if my concern can be well-addressed.
paper_summary: paper_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.

The 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.

Evaluation on image classification and time-series regression tasks demonstrate that this adapted BN is effective and mitigates the issues with standard BN.
strength_weakness: strength_weakness: 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.

In terms of weaknesses that can be addressed in a reasonable rebuttal period:
- 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.
- 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.
- 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.
questions: questions: - How have hyperparameters been selected?
- 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: limitations: - The limitations are addressed in the conclusion.
- There is no broader impact statement, it is not necessary for this work but would still improve the paper.

Comments: title: Thanks for the clarification and our further thoughts
Thanks for your response and clarification on our contribution
Thanks for your response-Raise the score
Thank you for the response, my score remains the same
comment: Thank you very much for taking the time to review our work. Your insightful feedback and constructive suggestions are invaluable and greatly appreciated. Your support and encouragement mean a lot to us. We will incorporate our discussions into the updated manuscript.
After reading the discussion with all the reviewers, I would like to thank the authors for addressing the outstanding concerns. New results look reasonable, confirm the hypothesis, and improve the value of the paper. 

I think the paper has scope and value of both presenting the reasons behind the failure of BN in Neural ODEs and the remedy for it and deserves to be accepted.
Many thanks, it answers my question. I think it would be actually good to add this discussion, for example, into the appendix. 

Meanwhile, I'm looking at the rest of your rebuttal and will revise the score after.
Thanks so much for the swift response and the clarification! We really appreciate your time involved in the discussion. Now we understand the question better.


To summarize, hypernet and shooting methods (referred to as HS below) will replace the parameter $\theta$ with a time-varying parameter $\theta(t)$ in the Neural ODE equation, yielding $\frac{dx}{dt}=f(x(t),\theta(t))$, This approach decouples the layers and learnable parameters along the time axis.


Traditional BN involves not only learnable parameters but also the crucial aspect of obtaining batch population statistics during training to reuse during inference. However, HS methods, in their original formulation, do not define how to acquire these batch population statistics because they do not consider BN.


Going one step further, if we extend HS to make the population statistics of BN also time-dependent $\mu(t)$ and $\sigma(t)$ in $[0,T]$,  then it will encounter the same problem as POP BN we discussed in our manuscript, since not every $t$ in $[0,T]$ will be visited by the adaptive ODE solver during training, and at inference time, the required $\mu$ and $\sigma$ at $\\{t_1,t_2,…,t_N\\}$ might not be available. Thus, a temporal interpolation technique is inevitable for recording the population statistics $\mu(t)$ and $\sigma(t)$ correctly, which is exactly the TA-BN approach. 

The above discussion reflects our understanding on this question, and we hope it addresses the concern raised. We are keen to hear the reviewer’s further thoughts and insights on this matter, as your perspective is invaluable to us. We remain open to continued discussions to refine our explanations or clarify any remaining questions.
Many thanks for a swift response! Checking the other concerns in the meantime, and I'll update my score accordingly as soon as I finish (no questions on them yet, just need to read it carefully). Many thanks for preparing  a really thorough rebuttal.

Just to elaborate on the question about TA-BN vs hypernets: I  understand that the hypernets and shooting methods are different from TA-BN. There is no batch normalisation in hypernet-parameterised neural ODEs,  it is instead a way to parameterise layers dynamically. My question is different: fundamentally, batch normalisation does not work with standard neural ODE with coupled layers, as the authors show. But if the problem with NeuralODEs and difference with resnets were solely in coupling the layers (which may or may not be true), then other methods which decouple the layers such as hypernets and shooting methods may potentially be used jointly with batch normalisation similarly to the proposed approach and deliver the same effect. Therefore, my question is whether the authors have any idea how does TA-BN compare over hypernets+BN or shooting+BN.
Thanks for acknowledging our response.

To the best of our knowledge, the hypernet and shooting methods are different from TA-BN and they cannot solve the problem TA-BN addresses. TA-BN normalizes the output from the previous layer using time-dependent statistics at each time grid, resulting in a final distribution with a mean of zero and a standard deviation of one. By maintaining consistent distributions across layers, it can mitigate the issue of internal covariate shift, where significant changes in distributions can make training unstable. However, hypernet and shooting methods do not explicitly normalize the output distributions. Consequently, models using hypernet or shooting methods still suffer from internal covariate shift and require TA-BN for normalization.

We would be grateful if the reviewer could kindly elaborate on how they believe the hypernet and shooting methods could be applied in our case. This would enable us to respond more effectively and precisely. We are open to further discussion.

In the meantime, we hope that our other responses have addressed your concerns and that they lead to a more positive assessment.
Many thanks for the insightful rebuttal! I am going through all the responses to all the authors and will follow up with the message after.

I have one follow-up question though.

"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)."

My question was mainly: the proposed method uses grid to mitigate the issues of time-dependent statistics. Although shooting methods/hypernetworks address a different problem, which is time-varying model parameters, there is a possibility that addressing the time-varying model parameters problem as per, e.g., [1-2] may (or may not) help enable using the batch normalisation techniques and equally solve the problem of time-dependent statistics. It would be good if the authors clarify upon it.
Thank you for your thoughtful feedback. 

We acknowledge the reviewer's observation that our proposed TA-BN shares several similarities with TAB (as we detailed in our response to Q1). However, we would like to respectfully emphasize that TA-BN is not our sole contribution. Another significant contribution is our demonstration of **why standard BN fails in Neural ODEs. This failure explanation can not be inferred from the TAB work in SNNs**, and we kindly ask the reviewer to consider this key aspect.

Regarding training efficiency, since no specific metric was mentioned in the rebuttal, we opted to use the converged epoch, as it is a commonly accepted standard in traditional deep learning research. We sincerely apologize for any confusion this may have caused, as it differs from the NFE metric the reviewer had in mind. We will conduct experiments reporting the NFE metric, although it may not be possible to meet the discussion deadline. We will promptly update the reviewer with the results once available.
Thank you for your response. I have read your response and other reviews. I think it is natural to extend TAB via interpolation techniques. This to some extent reduces the novelty and impact of this work. Moreover, for the training efficiency experiments, what I like to see is the comparison regarding NFE, since this is actually the key fact of the training efficiency for Nerual ODE. Hence, I would like to keep my score.
We sincerely appreciate your reply. Your review is immensely valuable and deeply appreciated.
Thank you very much for your response. You have addressed most of my concerns, so I raise the score to 6.
Thank you for your reply. We will consider applying TA-BN to time series and generation in future work. Your suggestion is valuable and sincerely appreciated. In addition to the reasons provided in the reply, we would like to note that we did conduct quick, small-scale experiments on TA-BN with the mentioned models, and our preliminary results did not exhibit significant improvements.
Thank you for the detailed response. I have now read all reviews and responses.

I maintain my view that this is a strong paper and maintain my score.

In particular I am grateful for the explanation of why we need Temporal Batchnorm rather than just normal batchnorm.

However, I am not convinced by the argument that TBN may not be suitable for time-series and CNFs. SInce the exact same argument can be made that continuous depth models are not suitable for image recognition.