Improving Neural ODE Training with Temporal Adaptive Batch Normalization

We appreciate the insightful feedback from Reviewer XMSL. Below we respond to each raised concern.

---

Q0.1: Error bar.

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

---

Q0.2: How TA-BN can be applied to Neural CDEs and CNFs?

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

---

Q0.3: TA-BN is not necessary with a fixed step size.

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

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

---

Q1: How to select hyperparameters?

Below we briefly show our hyperparameter setting criteria and extra ablation studies on hyperparameters. They will be added to the updated manuscript.

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

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

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

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

---

Q2: Deeper explanation on why standard BN fails.

Here, 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:

def forward(x,t):
    a1 = ReLU(self.linear1(x))
    a2 = self.bn(a1)
    return self.linear2(a2)

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

---

Q3: Broader impact statement.

We have reached our response length limit but will include a broader impact statement in the updated manuscript.

---

References:

[1] Yaron Lipman et al., 'Flow Matching for Generative Modeling', ICLR 2023.

We extend our sincere gratitude to Reviewer EgcE for the constructive feedback. We will include the following discussions in the updated manuscript.

---

Q0.1 Originality: more insights and TA-BN with other ODE solvers

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

Regarding our originality, with all due respect, we would like to justify our work from the following perspectives:

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

• We proposed TA-BN to resolve the above problem. Its time-dependent statistics and interpolation method can accurately normalize Neural ODEs' continuous-time dynamics.

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

---

Q0.2 Quality: confidence interval and grid size hyperparameter

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

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

---

Q1: Non-differentiability of piecewise linear interpolation at grid points

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

---

Q2: Constant memory consumption claimed by Neural ODE

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

---

Q3: Comparison with hypernetwork and shooting methods

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

---

Q4: Limitations: More empirical analysis on TA-BN with other ODE solvers, and compare it with hypernetworks

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

---

References:

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

[2] Kwitt et al, 'A Shooting Formulation of Deep Learning', NeurIPS 2020.

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.

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.

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.

---

Q1: Comparison to related works in the context of SNN and clarifications on novelty.

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

Regarding 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:

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

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

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

---

Q2: Lack of experiments on training efficiency.

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

Table. Convergence (defined as no improvement in 10 epochs) comparison

---

Q3: Is STEER a proper baseline?

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

---

Reference:

[1] H. Jiang et al. 'TAB: Temporal Accumulated Batch Normalization in Spiking Neural Networks,' ICLR 2024.

[2] A. Ghosh et al. 'STEER: Simple temporal regularization for neural ode.' NeurIPS, 2020.

We appreciate Reviewer 6xoP's thorough review and valuable insights. Below we respond to each raised concern.

---

Q1: Overclaim of experimental results about "MobileNetV2-level efficiency"

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

---

Q2: Lack of theoretical analysis on why TA-BN works.

We 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:

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

$$\\small h(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.$$

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

The 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 $\\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}} ).$ With 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)$.

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

$$\\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).$$

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

Due to the time limit of response, we will leave a comprehensive detailed proof for future work.

---

Q3: Inappropriate experiment design. Vanilla Neural ODEs are not for Walker2d and HalfCheetah.

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

---

Q4: Ablation studies on fixed-step size solver and other Neural ODE models.

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

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

---

References:

[1] S. Han et al., 'Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding,' ICLR 2016.

[2] H. Mostafa et al., 'Parameter efficient training of deep convolutional neural networks by dynamic sparse reparameterization,' ICML 2019.

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