[Paper Review] Interpolated Adjoint Method for Neural ODEs.

Talgat Daulbaev, Alexandr Katrutsa+4 more|arXiv (Cornell University)|Mar 11, 2020

Model Reduction and Neural Networks29 references7 citations

TL;DR

This paper introduces the Interpolated Adjoint Method, a novel approach to stabilize training in Neural ODEs by reconstructing the forward trajectory during the backward pass using smooth function interpolation. By avoiding direct adjoint integration, the method mitigates stiffness and numerical instability, significantly improving training stability and convergence in practice.

ABSTRACT

In this paper, we propose a method, which allows us to alleviate or completely avoid the notorious problem of numerical instability and stiffness of the adjoint method for training neural ODE. On the backward pass, we propose to use the machinery of smooth function interpolation to restore the trajectory obtained during the forward integration. We show the viability of our approach, both in theory and practice.

Motivation & Objective

• To address the numerical instability and stiffness issues inherent in the standard adjoint method for training Neural ODEs.
• To develop a backward pass mechanism that avoids direct integration of the adjoint system, which is prone to instability.
• To enable stable and accurate backpropagation in Neural ODEs by reconstructing the forward trajectory using interpolation.
• To provide a theoretically sound and practically effective alternative to traditional adjoint computation in Neural ODEs.

Proposed method

• The method reconstructs the forward trajectory during the backward pass using smooth function interpolation of the states recorded during forward integration.
• Instead of solving the adjoint ODE, the method interpolates the solution path from saved intermediate states, reducing reliance on numerically sensitive adjoint integration.
• It leverages high-order interpolation schemes to ensure smooth and accurate reconstruction of the state trajectory.
• The interpolated trajectory is then used to compute gradients via the chain rule, avoiding the need to integrate the adjoint system forward in time.
• The approach is designed to be compatible with standard ODE solvers and can be integrated into existing Neural ODE frameworks with minimal modification.
• Theoretical analysis shows that the method maintains consistency and convergence under appropriate interpolation conditions.

Experimental results

Research questions

• RQ1 Can trajectory reconstruction via interpolation effectively replace adjoint ODE integration in Neural ODEs without sacrificing gradient accuracy?
• RQ2 Does the interpolated approach reduce numerical stiffness and improve training stability compared to standard adjoint methods?
• RQ3 What level of interpolation accuracy is required to maintain convergence and generalization in Neural ODE training?
• RQ4 How does the method perform across different ODE solvers and neural network architectures?
• RQ5 Can the method be implemented efficiently without significant computational overhead?

Key findings

• The Interpolated Adjoint Method successfully mitigates numerical instability and stiffness in Neural ODE training by avoiding direct adjoint integration.
• The method achieves stable training on stiff ODE problems where standard adjoint methods fail or diverge.
• Gradient computation via interpolation maintains sufficient accuracy for effective model optimization, as demonstrated in empirical evaluations.
• The approach enables reliable training with higher-order ODE solvers that are typically unstable under standard adjoint backpropagation.
• The method provides a practical and robust alternative to the standard adjoint method, particularly in settings with stiff dynamics.