Stateful ODE-Nets using Basis Function Expansions
This work addresses efficiency and performance limitations in ODE-Nets for machine learning practitioners, though it appears incremental as it builds on existing ODE-Net frameworks.
The authors tackled the problem of improving ordinary differential equation networks (ODE-Nets) by formulating a stateful ODE-Block using basis function expansions, which enabled continuous-in-depth batch normalization for state-of-the-art performance and weight compression without retraining to reduce inference time and memory footprint while maintaining near state-of-the-art performance.
The recently-introduced class of ordinary differential equation networks (ODE-Nets) establishes a fruitful connection between deep learning and dynamical systems. In this work, we reconsider formulations of the weights as continuous-in-depth functions using linear combinations of basis functions which enables us to leverage parameter transformations such as function projections. In turn, this view allows us to formulate a novel stateful ODE-Block that handles stateful layers. The benefits of this new ODE-Block are twofold: first, it enables incorporating meaningful continuous-in-depth batch normalization layers to achieve state-of-the-art performance; second, it enables compressing the weights through a change of basis, without retraining, while maintaining near state-of-the-art performance and reducing both inference time and memory footprint. Performance is demonstrated by applying our stateful ODE-Block to (a) image classification tasks using convolutional units and (b) sentence-tagging tasks using transformer encoder units.