[Paper Review] Neural SDE: Stabilizing Neural ODE Networks with Stochastic Noise

Xuanqing Liu, Tesi Xiao+4 more|arXiv (Cornell University)|Jun 5, 2019

Adversarial Robustness in Machine Learning21 references89 citations

TL;DR

The paper introduces Neural SDE, a stochastic differential equation-based extension of Neural ODEs that injects noise (e.g., dropout, Gaussian) to regularize and stabilize continuous networks, improving generalization and robustness to adversarial and non-adversarial perturbations.

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.

Motivation & Objective

• Motivate and address the lack of regularization in Neural ODEs by introducing stochastic noise in a continuous framework.
• Develop a Neural SDE model that can implement common regularization techniques (dropout, Gaussian noise) within a continuous dynamics setting.
• Derive a scalable backpropagation method for training Neural SDEs, leveraging path-wise gradient and stochastic control concepts.
• Provide theoretical analysis showing that stochasticity can stabilize dynamical systems and improve robustness against perturbations.
• Empirically validate that Neural SDE improves generalization and robustness on CIFAR-10, STL-10, and Tiny-ImageNet datasets.

Proposed method

• Formulate neural dynamics as a stochastic differential equation: dh_t = f(h_t,t;w) dt + G(h_t,t;v) dB_t where B_t is Brownian motion.
• Model various noise types (additive, multiplicative, dropout-like) via appropriate choices of the diffusion term G(h_t,t;v).
• Develop a backpropagation approach based on the path-wise gradient, introducing an auxiliary variable β_t that follows an SDE to compute ∂h_{t1}/∂w efficiently.
• Provide stability analysis using stochastic Lyapunov methods to show that properly chosen diffusion can stabilize perturbations.
• Outline practical architectures where a Neural SDE module sits between feature extractor and classifier, functioning as a drop-in regularization component.

Experimental results

Research questions

• RQ1 Can injecting stochastic noise into Neural ODEs improve generalization similarly to discrete regularization techniques (dropout, Gaussian noise)?
• RQ2 How can standard regularization forms (dropout, additive/multiplicative noise) be mapped into continuous SDE dynamics for neural networks?
• RQ3 Does Neural SDE theory provide stability guarantees that explain observed improvements in robustness to perturbations and adversarial attacks?
• RQ4 Can we devise a scalable, memory-efficient gradient method to train Neural SDEs comparable to Neural ODE training?
• RQ5 Do Neural SDEs yield measurable gains in accuracy and robustness across standard vision benchmarks?

Key findings

• Neural SDEs can reproduce common regularization strategies (dropout, additive/multiplicative Gaussian noise) within a continuous-time framework.
• A path-wise gradient method yields unbiased estimators for training Neural SDEs with reduced memory (comparable to adjoint methods for Neural ODEs).
• Theoretical stability analysis shows that appropriately designed diffusion terms can make the system robust to input perturbations, and in some cases lead to almost surely exponentially stable perturbations.
• Empirical results demonstrate improved generalization over Neural ODEs across CIFAR-10, STL-10, and Tiny-ImageNet, with gains such as CIFAR-10 accuracy rising from 81.63% (ODE) to up to 84.55% (Neural SDE with TTN at testing).
• Neural SDEs also exhibit enhanced robustness to non-adversarial corruptions and adversarial perturbations, outperforming Neural ODEs under several attack and corruption settings.