Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise
Xuanqing Liu, Tesi Xiao, Si Si, Qin Cao, Sanjiv Kumar, Cho-Jui Hsieh
TL;DR
This work addresses the lack of regularization in Neural ODEs by introducing Neural SDEs that inject stochastic noise into the continuous dynamics. It provides a framework to model dropout, Gaussian, and other regularization schemes as diffusion terms, along with a memory-efficient pathwise gradient-based training method and a stability analysis showing robustness gains. Empirically, Neural SDE improves generalization on CIFAR-10 and enhances resistance to both adversarial and non-adversarial perturbations, with additional gains when using testing-time noise ensembles. Overall, the approach offers a principled, drop-in mechanism to stabilize and regularize continuous-depth networks with tangible performance and robustness benefits.
Abstract
Neural Ordinary Differential Equation (Neural ODE) has been proposed as a continuous approximation to the ResNet architecture. Some commonly used regularization mechanisms in discrete neural networks (e.g. dropout, Gaussian noise) are missing in current Neural ODE networks. In this paper, we propose a new continuous neural network framework called Neural Stochastic Differential Equation (Neural SDE) network, which naturally incorporates various commonly used regularization mechanisms based on random noise injection. Our framework can model various types of noise injection frequently used in discrete networks for regularization purpose, such as dropout and additive/multiplicative noise in each block. We provide theoretical analysis explaining the improved robustness of Neural SDE models against input perturbations/adversarial attacks. Furthermore, we demonstrate that the Neural SDE network can achieve better generalization than the Neural ODE and is more resistant to adversarial and non-adversarial input perturbations.
