The Symplectic Adjoint Method: Memory-Efficient Backpropagation of Neural-Network-Based Differential Equations.

Area of Science:

• Computational mathematics
• Machine learning
• Dynamical systems modeling

Background:

• Neural networks combined with numerical integration accurately model dynamical systems and probability distributions.
• Repeated neural network use in integration deepens computation graphs, increasing memory demands for backpropagation.
• Existing methods like backpropagation and checkpointing face memory and computational challenges.

Purpose of the Study:

• To develop a memory-efficient and computationally robust method for gradient calculation in neural network-based integration.
• To address the practical limitations of backpropagation and the adjoint method in deep computation graphs.

Main Methods:

• Proposed the symplectic adjoint method, utilizing a symplectic integrator for backward-in-time gradient computation.
• Conducted theoretical analysis to compare memory consumption with existing algorithms.
• Performed experiments to validate theoretical findings and assess performance.

Main Results:

• The symplectic adjoint method achieves exact gradients (up to rounding error) with memory proportional to network size plus usage count.
• Theoretical analysis confirmed significantly lower memory usage compared to naive backpropagation and checkpointing.
• Experiments demonstrated the symplectic adjoint method is faster and more robust to rounding errors than the standard adjoint method.

Conclusions:

• The symplectic adjoint method provides an efficient solution for gradient computation in deep neural network integration tasks.
• This method overcomes the memory and computational limitations of traditional approaches.
• It offers a practical advancement for modeling complex dynamical systems and distributions using neural networks.