Research map / Architectures and efficiency / RNNs and normalization
Improving Neural ODE Training with Temporal Adaptive Batch Normalization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes temporal adaptive batch normalization (TA-BN), which stores population statistics and affine parameters on a fixed time grid and linearly interpolates them, enabling stable BN in Neural ODEs and deeper Neural ODE stacks.
Problem
Batch normalization degrades or fails when applied inside Neural ODEs because adaptive ODE solvers produce variable time grids, so population statistics cannot be accumulated and test-time outputs become erroneous or batch-dependent.
Approach
Associate BN population statistics and learnable scale/shift parameters with a predefined regular grid of time points, interpolate them linearly to the solver's time points during training and testing, and accumulate running statistics through the same interpolation weights (the continuous-time analog of BN).
Key results
- On CIFAR-10 with a single unmixed Neural ODE plus linear layer and 2.2M parameters, TA-BN reaches 0.910±0.010 test accuracy, versus 0.822±0.095 for Mini-batch BN and 0.517±0.049 without BN.
- On SVHN, CIFAR-100 and Tiny-ImageNet (2.2M params), TA-BN achieves 0.958, 0.664 and 0.512 accuracy respectively, outperforming Mini-batch BN (0.906, 0.492, 0.200).
- On Walker2d-v2 and HalfCheetah-v2 MAE, TA-BN achieves 0.946 and 1.28, improving over the prior baseline [42] (1.02 and 1.46).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace no normalization with TA-BN | Neural ODE · CIFAR-10 | accuracy: 0.517±0.049 -> 0.910±0.010 | Helped |
| replace Mini-batch BN with TA-BN | Neural ODE · CIFAR-10 | accuracy: 0.822±0.095 -> 0.910±0.010 | Helped |
| replace Mini-batch BN with TA-BN | Neural ODE · SVHN | accuracy: 0.906±0.031 -> 0.958±0.004 | Helped |
| replace Mini-batch BN with TA-BN | Neural ODE · CIFAR-100 | accuracy: 0.492±0.176 -> 0.664±0.025 | Helped |
| replace Mini-batch BN with TA-BN | Neural ODE · Tiny-ImageNet | accuracy: 0.200±0.006 -> 0.512±0.008 | Helped |
| add TA-BN (ablation) | Aug-NODE · CIFAR-10 | accuracy: 0.848 -> 0.862 | Helped |
| add TA-BN (ablation) | HBNODE · CIFAR-10 | accuracy: 0.851 -> 0.867 | Helped |
| add TA-BN (ablation) | SONODE · CIFAR-10 | accuracy: 0.302 -> 0.853 | Helped |
| replace linear interpolation with cubic interpolation (ablation) | TA-BN · CIFAR-10 | accuracy | No clear effect |
| replace linear interpolation with kernel smoothing (ablation) | TA-BN · CIFAR-10 | accuracy | Mixed |
| replace linear interpolation with Gaussian process interpolation (ablation) | TA-BN · CIFAR-10 | accuracy | Hurt |
| use Euler instead of Dopri5 solver (ablation) | TA-BN Neural ODE · CIFAR-10 | accuracy: 0.872±0.003 vs 0.874±0.001 | No clear effect |
Surprising findings
- Pop-TI BN, the variant used in common PyTorch/TorchDiffeq implementations, performs worse than no BN on several tasks (e.g., 0.332 accuracy with dopri5 on CIFAR-10 vs 0.843 without BN).
- Mini-batch BN fails at batch size 2 and under noisy test batches, even though it works well at standard batch sizes.
Limitations
- Interpolation adds runtime overhead that slows Neural ODE execution.
- Linear interpolation is heuristic; cubic, kernel, and Gaussian process alternatives were slower or less stable at larger scale.
- No theoretical analysis is provided; claims rest on empirical results.
Experimental setup
Datasets: MNIST SVHN CIFAR-10 CIFAR-100 Tiny-ImageNet Walker2d-v2 HalfCheetah-v2 Charge Pump circuit simulation data
Benchmarks: MNIST SVHN CIFAR-10 CIFAR-100 Tiny-ImageNet Walker2d-v2 HalfCheetah-v2 Charge Pump
Metrics: test top-1 accuracy mean absolute error mean squared error parameter count
Baselines: Aug-NODE STEER IL-NODE 2nd-Ord (SONODE) HBNODE GHBNODE Mini-batch BN Pop-TI BN w/o BN MobileNetV2
Builds on: Neural ODE Batch Normalization Augmented Neural ODE TorchDiffeq
Closest papers by meaning
- Neural Ordinary Differential Equations (NeurIPS 2018 · similarity 0.90)
- STEER : Simple Temporal Regularization For Neural ODE (NeurIPS 2020 · similarity 0.89)
- How to Train Your Neural ODE: the World of Jacobian and Kinetic Regularization (ICML 2020 · similarity 0.88)
- ANODEV2: A Coupled Neural ODE Framework (NeurIPS 2019 · similarity 0.88)
- How Deep Do We Need: Accelerating Training and Inference of Neural ODEs via Control Perspective (ICML 2024 · similarity 0.87)
- Implicit regularization of deep residual networks towards neural ODEs (ICLR 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.