Observations on Recurrent Loss in the Neural Network Model of a Partial Differential Equation: the Advection-Diffusion Equation

Technical Summary: Observations on Recurrent Loss in the Neural Network Model of a Partial Differential Equation

Problem Statement
The paper addresses a critical gap in the application of machine learning (ML) to partial differential equations (PDEs): the lack of systematic exploration regarding the stability of solutions generated by ML approaches. While existing literature utilizes ML to approximate PDE solutions—either by mimicking the method of lines or by encapsulating the entire time evolution via neural networks—these methods often lack rigorous stability guarantees. The author focuses on the advection-diffusion equation, a linear time-dependent PDE, to investigate whether recurrence in neural network training can enforce numerical stability, particularly in scenarios where traditional multistep numerical methods fail due to instability.

Methodology
The study proposes a novel framework that bridges traditional numerical analysis and machine learning by constructing a Recurrent Neural Network (RNN) that mathematically mirrors a specific numerical scheme.

1. Network Architecture as a Numerical Method: The author designs a linear neural network (with no bias and linear activation functions) that exactly replicates the evaluation of an $s$-step Adams-Bashforth multistep method paired with a collocation method for approximating spatial derivatives. The network consists of five layers:

   • Input Layer: Accepts solution vectors from previous time steps.
   • Assignment Layer: Duplicates and arranges data to match local computational node sets ($N_{k,n}$).
   • Reshape Layer: Organizes data into a matrix format.
   • Convolution Layer: Performs the spatial differentiation operation using trainable weights ($\omega$) that approximate the linear operator $L$.
   • Additive Layer: Executes the multistep time integration formula.

2. Recurrent Loss Function: Unlike standard ML approaches that minimize error at a single time step, this study employs a recurrent loss function $J(\omega)$. This function sums the squared $L_2$ errors between the network's predicted solution and the exact solution over multiple future time steps ($Q$ steps) across multiple initial condition cases ($T$). The objective is to minimize the forward error over a sequence of steps, thereby encouraging the network to learn weights that maintain stability over time.

3. Optimization: The weights $\omega$ (representing the coefficients for approximating first and second spatial derivatives) are optimized using a quasi-Newton method (specifically BFGS) to find the roots of the gradient of the loss function. The training data consists of exact solutions to the advection-diffusion equation generated via Fourier series, ensuring that errors arise solely from the numerical approximation, not data noise.

4. Stability Analysis: The paper utilizes the theory of linear stability for ODEs. Since the spatial discretization matrix $D$ is circulant, its eigenvalues can be explicitly computed. Stability is assessed by checking if the scaled eigenvalues ($\lambda_\eta h_t$) fall within the linear stability region of the multistep method.

Key Contributions and Results
The paper presents a series of computational experiments varying parameters such as the number of time steps in the loss ($Q$), the number of training cases ($T$), the smoothness of initial data ($p$), the number of nodes ($N$), and the time step size ($h_t$).

• Adaptation to Stability Regions: The primary finding is that minimizing the recurrent loss can alter the spectral properties of the discretization matrix $D$. In cases where traditional centered difference methods (paired with Adams-Bashforth) are unstable (e.g., pure advection with $s=2$), the optimization process successfully finds a set of weights such that the scaled eigenvalues of the resulting matrix fall entirely within the stability region of the multistep method.
• Success in Unstable Regimes: The method demonstrated the ability to generate stable solutions for parameter combinations where traditional methods fail, specifically when $h_t$ exceeds the critical stability limit of standard schemes.
• Unpredictability of Hyperparameters: A significant observation is the lack of a predictable pattern regarding hyperparameter selection. Increasing the number of recurrent steps ($Q$), the amount of training data ($T$), or the number of optimization iterations ($\kappa_{max}$) does not reliably lead to stable solutions. In some cases, increasing $Q$ or $\kappa_{max}$ actually degraded performance or failed to find a stable set of weights, likely due to vanishing gradients or the competition between local accuracy and global stability.
• Truncation Error: The resulting numerical methods generally exhibited truncation errors comparable to the second-order centered difference method used to initialize the optimization ($O(N^{-2})$). The optimization process did not significantly improve the order of accuracy, and in some cases, enforcing strict truncation error constraints reduced the likelihood of finding stable weights.

Significance and Claims
The paper claims that its work demonstrates the feasibility of using ML techniques to construct stable numerical methods for time-dependent PDEs, even in regimes where traditional methods fail. By framing the numerical method as a neural network, the author enables the use of traditional stability analysis tools (eigenvalue analysis) alongside efficient ML training techniques.

However, the author maintains a modest tone regarding the robustness of the approach. The paper concludes that while stable solutions can be found, the results are highly unpredictable. There is currently no reliable set of criteria for selecting hyperparameters (such as $Q$, $T$, or initial guesses) that guarantees the discovery of a stable weight set. The study highlights the difficulty of data-driven approaches in PDEs: even in a simplified linear setting with known exact solutions, the path to a stable numerical method is not systematic. The work serves as a proof-of-concept that recurrence in training can promote stability, but it underscores the need for further research into optimization strategies, network architectures (e.g., LSTMs), and initial weight selection to make such methods reliable.