Improving Robustness and Uncertainty Modelling in Neural Ordinary Differential Equations

TL;DR

This paper introduces LT-NODE and ALT-NODE, novel Bayesian approaches to model uncertainty and improve robustness in Neural Ordinary Differential Equations by treating the end-time as a latent variable.

Contribution

The paper proposes a new Bayesian framework for NODE that models uncertainty over the end-time, enabling better robustness and model selection, with efficient inference methods.

Findings

Effective uncertainty modeling demonstrated on synthetic data.

Improved robustness in image classification tasks.

Efficient inference with a single forward pass.

Abstract

Neural ordinary differential equations (NODE) have been proposed as a continuous depth generalization to popular deep learning models such as Residual networks (ResNets). They provide parameter efficiency and automate the model selection process in deep learning models to some extent. However, they lack the much-required uncertainty modelling and robustness capabilities which are crucial for their use in several real-world applications such as autonomous driving and healthcare. We propose a novel and unique approach to model uncertainty in NODE by considering a distribution over the end-time $T$ of the ODE solver. The proposed approach, latent time NODE (LT-NODE), treats $T$ as a latent variable and apply Bayesian learning to obtain a posterior distribution over $T$ from the data. In particular, we use variational inference to learn an approximate posterior and the model parameters. Prediction is done by considering the NODE representations from different samples of the posterior and can be done efficiently using a single forward pass. As $T$ implicitly defines the depth of a NODE, posterior distribution over $T$ would also help in model selection in NODE. We also propose, adaptive latent time NODE (ALT-NODE), which allow each data point to have a distinct posterior distribution over end-times. ALT-NODE uses amortized variational inference to learn an approximate posterior using inference networks. We demonstrate the effectiveness of the proposed approaches in modelling uncertainty and robustness through experiments on synthetic and several real-world image classification data.