[Paper Review] Leveraging Hamilton-Jacobi PDEs with time-dependent Hamiltonians for continual scientific machine learning

Paula Chen, Tingwei Meng+3 more|arXiv (Cornell University)|Nov 13, 2023

Model Reduction and Neural Networks4 citations

TL;DR

This paper establishes a theoretical link between regularized scientific machine learning (SciML) with integral losses and Hamilton-Jacobi (HJ) PDEs with time-dependent Hamiltonians, enabling a new Riccati-based training method that naturally supports continual learning. By interpreting model updates as solutions to a time-evolving HJ PDE, the approach avoids catastrophic forgetting and achieves significant computational and memory efficiency, particularly in streaming data regimes, as demonstrated in 1D and 2D regression tasks with low relative $L_2$ errors and reduced memory usage compared to standard least squares.

ABSTRACT

We address two major challenges in scientific machine learning (SciML): interpretability and computational efficiency. We increase the interpretability of certain learning processes by establishing a new theoretical connection between optimization problems arising from SciML and a generalized Hopf formula, which represents the viscosity solution to a Hamilton-Jacobi partial differential equation (HJ PDE) with time-dependent Hamiltonian. Namely, we show that when we solve certain regularized learning problems with integral-type losses, we actually solve an optimal control problem and its associated HJ PDE with time-dependent Hamiltonian. This connection allows us to reinterpret incremental updates to learned models as the evolution of an associated HJ PDE and optimal control problem in time, where all of the previous information is intrinsically encoded in the solution to the HJ PDE. As a result, existing HJ PDE solvers and optimal control algorithms can be reused to design new efficient training approaches for SciML that naturally coincide with the continual learning framework, while avoiding catastrophic forgetting. As a first exploration of this connection, we consider the special case of linear regression and leverage our connection to develop a new Riccati-based methodology for solving these learning problems that is amenable to continual learning applications. We also provide some corresponding numerical examples that demonstrate the potential computational and memory advantages our Riccati-based approach can provide.

Motivation & Objective

• To address interpretability and computational efficiency in scientific machine learning (SciML) by linking regularized learning problems to Hamilton-Jacobi PDEs with time-dependent Hamiltonians.
• To develop a new training methodology for SciML that inherently supports continual learning and avoids catastrophic forgetting by leveraging optimal control and viscosity solution theory.
• To demonstrate the computational and memory advantages of the proposed Riccati-based approach over standard least squares in streaming data settings.
• To extend prior work on multi-time HJ PDEs to the infinite-time, time-continuous case for broader applicability in continual learning.
• To explore the potential of reusing established HJ PDE and optimal control solvers for efficient, interpretable SciML training.

Proposed method

• The paper establishes a generalized Hopf formula connection between regularized learning problems with integral-type losses and the viscosity solution of a time-dependent HJ PDE.
• It reinterprets incremental model updates as the time evolution of a viscosity solution to an HJ PDE, with all prior data intrinsically encoded in the solution.
• For linear regression, the method derives a Riccati differential equation that governs the evolution of the model parameters, enabling efficient, incremental updates.
• The approach uses a time-continuous formulation where data is processed sequentially along a propagation direction (e.g., along the y-axis in 2D), avoiding storage of past data.
• The HJ PDE solution is computed via numerical integration of the generalized Hopf formula, with time evolution handled by standard ODE solvers such as RK4.
• The method is contrasted with standard least squares (LSE) using Monte Carlo integration, highlighting memory and scalability advantages.

Experimental results

Research questions

• RQ1 Can regularized learning problems with integral losses be reinterpreted as solutions to time-dependent Hamilton-Jacobi PDEs?
• RQ2 How can the connection between SciML and HJ PDEs be leveraged to design memory-efficient, continual learning algorithms?
• RQ3 What are the computational and memory advantages of the proposed Riccati-based approach compared to standard least squares in streaming data scenarios?
• RQ4 Can the HJ PDE framework be extended to nonconvex or discontinuous Hamiltonians to broaden applicability to general learning problems?
• RQ5 How does the choice of propagation direction affect the memory and accuracy performance in higher-dimensional problems?

Key findings

• The Riccati-based approach achieved relative $L_2$ errors of 0.07% for $u$ and 0.03% for $f$ after processing all data, outperforming standard least squares in accuracy despite using only 1D data slices at a time.
• The standard least squares method with $10^4$ Monte Carlo points had 0.30% and 0.09% relative $L_2$ errors for $u$ and $f$, respectively, but could not scale to $10^5$ points due to memory constraints on a standard laptop.
• The proposed method required only $N$ data points at a time, while the LSE required an $N \times 10^5$ grid, demonstrating a significant memory advantage in big data regimes.
• The method inherently avoids catastrophic forgetting by encoding all historical data in the HJ PDE solution, eliminating the need to store past data.
• The approach maintained high accuracy even with noisy data, achieving 34.00% and 30.30% relative $L_2$ errors for $u$ and $f$ after processing only 25% of the domain.
• Theoretical analysis confirms that solving the regularized learning problem is equivalent to solving an optimal control problem and its associated HJ PDE with time-dependent Hamiltonian.