[Paper Review] Augmented Neural ODEs

Emilien Dupont, Randal Douc+1 more|arXiv (Cornell University)|Apr 2, 2019

Model Reduction and Neural Networks18 references57 citations

TL;DR

The paper identifies topology-preserving limitations of Neural ODEs and introduces Augmented Neural ODEs (ANODEs), which lift data into higher dimensions to learn simpler, more expressive flows with lower computational cost, better generalization, and improved stability.

ABSTRACT

We show that Neural Ordinary Differential Equations (ODEs) learn representations that preserve the topology of the input space and prove that this implies the existence of functions Neural ODEs cannot represent. To address these limitations, we introduce Augmented Neural ODEs which, in addition to being more expressive models, are empirically more stable, generalize better and have a lower computational cost than Neural ODEs.

Motivation & Objective

• Demonstrate that Neural ODEs preserve input-space topology and cannot represent certain functions.
• Propose Augmented Neural ODEs to overcome representational limits and improve training stability and efficiency.
• Empirically compare NODEs and ANODEs on toy tasks and image datasets to assess expressivity, NFEs, generalization, and scalability.

Proposed method

• Define Neural ODEs as continuous-time limits of ResNets and formalize the NODE flow as φ_t mapping inputs via an IVP with h(0)=x.
• Prove that NODEs cannot represent specific simple and higher-dimensional functions due to topology preservation (flows are homeomorphisms).
• Introduce Augmented Neural ODEs by augmenting the state from R^d to R^{d+p} and solving the ODE on the augmented space with initial condition [x; 0].
• Hypothesize that augmentation yields smoother, simpler flows requiring fewer function evaluations and enabling representation of previously inexpressible functions.
• Empirically compare NODEs and ANODEs on toy functions and image datasets, examining training loss, NFEs, accuracy, stability, and generalization.

Experimental results

Research questions

• RQ1 What classes of functions can Neural ODEs not represent due to topology preservation?
• RQ2 Can augmenting the latent space enable NODEs to represent more complex functions with simpler flows?
• RQ3 Do ANODEs offer lower computational cost, better generalization, and greater stability compared to NODEs on image data?

Key findings

• NODEs cannot represent certain 1D and higher-dimensional functions that require crossing trajectories, due to flow being a homeomorphism.
• ANODEs learn simpler, smoother flows in augmented space and require significantly fewer NFEs than NODEs for similar tasks.
• ANODEs achieve lower training losses and better generalization on MNIST, CIFAR-10, SVHN, and 64×64 ImageNet with comparable parameter counts.
• Augmentation improves training stability, reduces NFEs, and yields higher accuracy on image datasets.
• NODEs exhibit instability and high NFEs when fitting certain functions, while ANODEs remain stable and efficient.