The vanishing gradient problem for stiff neural differential equations

Area of Science:

• Numerical Analysis
• Machine Learning
• Dynamical Systems

Background:

• Gradient-based optimization is crucial for training parameterized dynamical systems, including neural ordinary differential equations (NODEs).
• Stiff systems often exhibit vanishing gradients, impeding effective parameter learning.
• Existing research has observed this phenomenon but lacked a universal explanation.

Purpose of the Study:

• To demonstrate that vanishing gradients in stiff systems are a universal characteristic of A-stable and L-stable numerical integration schemes.
• To analyze the underlying mathematical reasons for gradient suppression in stiff regimes.
• To establish a fundamental limitation for training NODEs and other stiff parameterized models.

Main Methods:

• Analysis of the rational stability function for general stiff integration schemes.
• Derivation of explicit formulas for common stiff integration methods.
• Rigorous mathematical proof of the decay rate for parameter sensitivities.

Main Results:

• Vanishing gradients are shown to be an inherent property of all A-stable and L-stable stiff numerical integration schemes.
• Parameter sensitivities decay to zero for large stiffness, governed by the derivative of the stability function.
• The slowest possible decay rate for these sensitivities is proven to be O(|z|-1).

Conclusions:

• All A-stable and L-stable time-stepping methods inevitably suppress parameter gradients in stiff regimes.
• This poses a significant barrier for training and parameter identification in stiff neural ordinary differential equations.
• The findings reveal a fundamental limitation in optimizing stiff parameterized dynamical systems using current numerical methods.