🤖 AI Summary
In neural ODE training, automatic differentiation (AD) with high-order ODE solvers—such as Leapfrog or linear multistep methods (LMMs)—often induces gradient distortion, leading to non-convergent oscillations that degrade training stability and generalization. This work identifies, for the first time, that standard AD breaks consistency between discrete gradients and the underlying continuous flow updates when applied to reversible or symmetric high-order discretizations—constituting the primary cause of oscillatory behavior. To address this, we propose a lightweight, solver-agnostic post-processing gradient correction method that enforces geometric consistency of discrete gradients with the flow map, thereby restoring gradient convergence without re-implementing the solver. Theoretical analysis establishes convergence guarantees under mild assumptions, and extensive numerical experiments demonstrate substantial suppression of training oscillations. Our approach consistently improves training stability and generalization accuracy across multiple benchmark tasks.
📝 Abstract
Does the use of auto-differentiation yield reasonable updates to deep neural networks that represent neural ODEs? Through mathematical analysis and numerical evidence, we find that when the neural network employs high-order forms to approximate the underlying ODE flows (such as the Linear Multistep Method (LMM)), brute-force computation using auto-differentiation often produces non-converging artificial oscillations. In the case of Leapfrog, we propose a straightforward post-processing technique that effectively eliminates these oscillations, rectifies the gradient computation and thus respects the updates of the underlying flow.