Mixture of Neural Operators:  Incorporating Historical Information for Longer Rollouts

Authors: harris abdul majid
francesco  tudisco
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-10 18:02:46

Abstract: Traditional numerical solvers for time-dependent partial differential equations (PDEs) notoriously require high computational resources and necessitate recomputation when faced with new problem parameters. In recent years, neural surrogates have shown great potential to overcome these limitations. However, it has been paradoxically observed that incorporating historical information into neural surrogates worsens their rollout performance. Drawing inspiration from multistep methods that use historical information from previous steps to obtain higher-order accuracy, we introduce the Mixture of Neural Operators (MoNO) framework; a collection of neural operators, each dedicated to processing information from a distinct previous step. We validate MoNO on the Kuramoto-Sivashinsky equation, demonstrating enhanced accuracy and stability of longer rollouts, greatly outperforming neural operators that discard historical information.

Main paper: 
 introduction: 
 Traditional numerical methods, such as finite difference, finite element, and spectral methods, are the standard go-to methods for solving partial differential equations (PDEs) (Thomas, 2013) ; however, these methods may require prohibitive computational costs, especially for high-dimensional timedependent problems. Furthermore, traditional methods may struggle with complex geometries, and suffer from accuracy and stability issues when dealing with multiscale or nonlinear phenomena (Tadmor, 2012) . Given these challenges, there has been a growing interest in the use of machine learning for solving PDEs by directly learning the solution operator (Thuerey et al., 2021; Brunton & Kutz, 2023) . This paradigm shift is driven by the ability of machine learning to learn from data, identify patterns, and enable fast predictions, offering new perspectives in scientific computing.The use of machine learning for predicting PDE trajectories can be broadly divided into two main categories. The first category involves using initial conditions to predict an entire trajectory up to a pre-specified time step (Raissi et al., 2019; Kovachki et al., 2021) ; however, this approach often struggles with generalizing beyond the training data. The second category adopts a sequential strategy (Stachenfeld et al., 2021; Brandstetter et al., 2022b ), akin to traditional numerical methods. Here, the solution at the current time step is used to predict the solution at the next time step, which is then used as a stepping stone to make further predictions or unroll the trajectory. While this approach has been quite successful at predicting accurate and stable trajectories over extended periods, it is not without shortcomings. A key challenge is the accumulation of small one-step prediction errors, which compound over time, leading to significant deviations from the true trajectory. The one-step prediction errors are exacerbated at time steps that extend beyond the training range, highlighting an area of improvement in the application of autoregressive or rollout methods.Intuition suggests that incorporating historical information in the training data, i.e. solutions at several previous time steps, will reduce one-step prediction errors and long-term predictions as a consequence, akin to traditional multistep approaches in scientific computing. However, current approaches tend to show the opposite behaviour: while incorporating historical information does reduce one-step prediction errors, it paradoxically worsens the accuracy and stability of long-term predictions (Lippe et al., 2023) . As one typically trains the network to learn residuals (the difference between the solution at the next time step and the current time step), the problem arising when incor-porating history is that the difference between the inputs is highly correlated with the model's target, the residual at the next time step. This leads the neural operator to focus on modeling second-order differences, deteriorating rollout performance.Inspired by linear multistep methods in numerical analysis (Wanner & Hairer, 1996) , which use a combination of solution values at multiple time steps, in this work we propose to use a combination of neural operators, each dedicated to processing the solution from a distinct previous time step. This collection of operators is coupled with a gating network that acts as a dynamic selector, adapting the contribution of each neural operator based on the solution at the current time step. The resulting model is an end-to-end, highly parallelizable approach that leverages historical information to, not only reduce one-step prediction errors but also, improve the accuracy and stability of long-term predictions, largely outperforming standard single-step neural operator baselines. 
problem setting: 
 In this paper, we focus on time-dependent PDEs of the formu t + N (t, x, u, u x , u xx , . . . ) = 0,where t ∈ [0, T ] represents the temporal dimension, x ∈ X represents the (possibly multiple) spatial dimension(s), and u(t, x) : [0, T ] × X → R n represents the solution. Here, N is a nonlinear operator that governs the dynamics of the system, describing how the different variables and their derivatives interact. Further, we focus on initial conditions given by u(0, x) = u 0 (x), along with periodic boundary conditions. Discretizing the temporal dimension with step size ∆t transforms the continuous PDE into a discrete PDE, yielding a sequence of solutions at discrete time steps U 0 , U 1 , . . . , U N , where N = T /∆t is the number of steps. This discretization introduces an evolution operator G, which maps the solution at any given time step to the solution at the subsequent time step:G (U n ) = U n+1 .In the context of operator learning (Lu et al., 2019; 2021; Li et al., 2020; Kovachki et al., 2021) , the objective is to approximate the evolution operator G with a learned neural operator G θ that captures the dynamics of the system. A key consideration is the choice between directly predicting the solution at subsequent time steps, or predicting the differences between the solutions at two successive time steps:G (U n ) = U n+1 -U n .Empirical observations suggest that, for small step sizes, predicting residuals is advantageous for maintaining accuracy and stability over extended periods (Li et al., 2021; Lippe et al., 2023) . See Appendix A for further details.In both traditional numerical methods and data-driven methods, the choice of step size ∆t is a key hyperparameter with its set of trade-offs (de Hoop et al., 2022) . While a larger temporal step size decreases computational costs, it leads to larger one-step prediction errors and greater divergence from the true trajectory. In autoregressive neural methods, a larger temporal step size can significantly lower computational costs; for example, a doubling of the temporal step size halves the number of discretization points, thereby halving training and inference costs. However, if the chosen step size ∆t is too small, the trajectory can suffer from an accumulation of rounding errors.Incorporating historical information into data-driven models is intuitively appealing, as it potentially provides the models with additional context of the underlying dynamics, thereby improving the accuracy and stability of long-term predictions. The most straightforward strategy for incorporating historical information into data-driven methods involves the introduction of additional channels, where each m-dimensional solution at a previous time step U n-1 , U n-2 , . . . is assigned a separate channel, analogous to the RGB channels in color images. Using this representation as input, and mirroring the findings by Lippe et al. (2023) , we observe a reduction of one-step prediction errors but a worsened high-correlation time (see Figure 1 ). These effects are exacerbated with smaller step sizes and longer histories. Lippe et al. (2023) explain that the problem arises because the difference between the current and previous solution U n -U n-1 is highly correlated with the residual of the next time step U n+1 -U n -this leads the learned neural operator to focus on second-order differences which is known to deteriorate performance in explicit schemes. To mitigate the aforementioned challenges, we propose training a collection of neural operators, each dedicated to processing the solution from a distinct previous time step, before weighting and aggregating their outputs. This framework prevents the model from focusing on higher-order differences, allowing effective use of historical information. 
mixture of neural operators: 
 Akin to traditional numerical methods where lower-order accurate methods can be extended to obtain higher-order accuracy by using multiple time steps, in this work we introduce a novel ensemble approach, termed Mixture of Neural Operators (MoNO), which synergizes a collection of neural operator modelsG (0) θ , . . . , G (h-1) θ, orchestrated by a gating network R θ . The MoNO framework is designed to dynamically weigh and integrate the predictions from each operator. Formally,MoNO θ (U n , . . . , U n-h+1 ) := h-1 i=0 R θ (U n , . . . , U n-h+1 ) i • G (i) θ (U n-i ),where R θ (z) i denotes the gating network's weighting for the i-th operator, and G (i) θ (z) denotes the corresponding neural operator's output.The training objective for MoNO optimizes a loss function L(θ) that minimizes the discrepancy between the aggregated prediction and the true residual:L(θ) = |MoNO θ (U n , . . . , U n-h+1 ) -(U n+1 -U n ) |.This design encourages a cooperative interaction among the neural operators, where adjustments in one operator's parameters induce compensatory adaptations across the ensemble, ensuring a cohesive evolution towards minimizing residual prediction errors.To realize the MoNO framework, we explore various neural operator architectures, including UNet and Fourier Neural Operator (FNO), allowing for a diverse representation of the solution mappings. The flexibility of MoNO extends to the configuration of each neural operator, enabling adjustments in parameter allocation, activation functions, and architectural nuances to tailor the ensemble's predictive capacity. Central to the efficacy of MoNO is the gating network, R θ , which determines the proportional contribution of each neural operator to the ensemble prediction. We use a softmax linear layer to ensure that the experts' weights are non-negative and sum to one, facilitating interpretability and maintaining balanced contributions.Furthermore, MoNO's framework is inherently amenable to distributed computing environments. Leveraging standard model parallelism techniques, the framework can distribute its components across multiple GPUs, routing inputs to their respective neural operators for processing and subsequently aggregating the outputs, thus capitalizing on computational resources efficiently. Step Size = 0.2 UNet MoNO(UNet)Step Size = 0.4Step Size = 0.8Step Size = 1.6  
experiments: 
 In order to test the performance of the proposed model we present experimental evaluation on the one-dimensional Kuramoto-Sivashinsky (KS) equation. The KS equation is a fourth-order nonlinear PDE derived to model diffusive-thermal instabilities in laminar flame fronts (Kuramoto, 1978; Sivashinsky, 1977) . Its one-dimensional variant can be expressed as:u t + u xx + u xxxx + uu x = 0.Here, the fourth-order derivative u xxxx and the nonlinear term uu x contribute to complex and chaotic behavior which present a challenge for traditional numerical solvers (Hyman & Nicolaenko, 1986; Kevrekidis et al., 1990; Smyrlis & Papageorgiou, 1991) , necessitating fine spatial and temporal discretizations (and therefore increased computational costs) for accurate and stable trajectories.We assessed the performance of MoNO using KS data across various step sizes and history lengths. The objective was to predict the residual: the difference between successive time steps, U n+1 -U n . As a baseline neural operator G θ we used the U-Net by Gupta & Brandstetter (2022) , which is frequently used for neural PDE solvers Lippe et al. (2023) . A softmax-MLP serves as MoNO's gating network. We conducted parallel evaluations using another prevalent model, the Fourier Neural Operator (FNO) with 8 Fourier layers and 32 Fourier modes (Li et al., 2020) . To ensure equitable comparisons, we adjusted the channels of each model so that the model has approximately 50 million trainable parameters. For training, we chose the Mean Squared Error (MSE) as a loss metric. To evaluate, we autoregressively apply the models on the initial conditions in the testing dataset, and report the high-correlation time: the duration for the Pearson correlation between the true trajectory and predicted trajectory to fall below thresholds of 0.8 and 0.9. At a step size of 0.2, the U-Net baseline obtained a high-correlation time of 75.0 (79.8) for a threshold of 0.9 (0.8). An increase in history length up to 8 steps resulted in reduced high-correlation times, manifesting as 59.4 (65.4) for a threshold of 0.9 (0.8); whereas, the MoNO framework demonstrated significant improvements, recording 14%, 28%, and 44% longer high-correlation times at history lengths of 2, 4, and 8, respectively. This trend persisted across all step sizes. It is noteworthy that at step size ∆t = 1.6, extending the history from 4 to 8 did not yield further improvements, suggesting that data from 4 steps prior (or 6.4 seconds) lacked relevance for subsequent predictions. Similar to U-Net, MoNO with FNO backbone consistently outperformed the baseline across all history lengths and step sizes. See Figure 2 , top (U-Net) and bottom (FNO) panels.We make slight modifications to their code (https://github.com/brandstetter-johannes/LPSDA) and set ∆t = 0.2. 
 
 Appendix:  


Decision: Both reviewers appreciate the novelty of the work and vote for acceptance. The authors propose to use further historical information within FNO to act as a temporal multi-step method from numerical analysis. It is nice to bring concepts from numerical analysis into machine learning and see them succeed. I also vote for acceptance but agree with Reviewer ABiq that the method should be compared to other methods that use the same amount of historical information rather than the prior step.

Reviews: review: title: title: A good paper demonstrating novel framework for learning NO for leveraging historical information.
main_review: main_review: This paper proposes a novel framework for learning neural operators that can leverage historical information from previous time steps to improve the accuracy and stability of long-term predictions for time-dependent partial differential equations (PDEs). The framework, called Mixture of Neural Operators (MoNO), consists of an ensemble of neural operators, each dedicated to processing the solution from a distinct previous time step, and a gating network that dynamically weighs and aggregates their outputs. The paper demonstrates the effectiveness of MoNO on the Kuramoto-Sivashinsky equation, a challenging nonlinear PDE that exhibits complex and chaotic behavior.

## Pros
1. The paper introduces a new and innovative approach for incorporating historical information into neural operators, inspired by traditional multistep methods in numerical analysis. The paper also provides a clear explanation of the motivation and intuition behind the proposed framework, as well as the challenges and limitations of existing methods.
2. Performance: The paper shows that MoNO outperforms standard single-step neural operator baselines, such as U-Net and Fourier Neural Operator, across various step sizes and history lengths, in terms of reducing one-step prediction errors and increasing high-correlation time.

## Cons
1. The paper only evaluates MoNO on one PDE, the Kuramoto-Sivashinsky equation, which limits the generalizability and applicability of the proposed framework. The paper could benefit from testing MoNO on other PDEs, especially those with higher dimensions, different boundary conditions, or different types of nonlinearities, to demonstrate its robustness and versatility.
2. The paper only compares MoNO with standard single-step neural operator baselines, which may not be the most relevant or competitive methods for incorporating historical information. The paper could benefit from comparing MoNO with other data-driven methods that use historical information, such as recurrent neural networks, attention mechanisms, or memory networks, to provide a more comprehensive and fair evaluation.
3. The paper could benefit from providing more insights and explanations of why and how MoNO works, such as the role of the gating network, the effect of the history length, the trade-off between accuracy and stability, or the potential sources of error and instability. The paper could also provide more quantitative and qualitative measures of the performance and the behavior of MoNO, such as error bounds, convergence rates, sensitivity analysis, or error propagation.
title: title: Solid evaluation of a decent method that explicitly considers multiple history frames in timeseries forecasting.
main_review: main_review: ### Summary
A Mixture of Neural Operators (MoNO) is introduced that trains multiple Neural Operators, each of them learning to predict a target frame $x_{t+1}$ on basis of a particular frame from the past. A gating network decides on the weighting of each Neural Operator's contribution to the final solution. This approach allows the explicit incorporation of multiple historic frames to base a prediction of the next frame.

### Strenghts
- Appealing idea that tackles a relevant problem in temporal dynamics forecasting.
- Clear problem statement and paper organization.
- Proper experimentation supports the claims.

### Weaknesses
- Training/inference time and memory trade-off between single vs multi NODE model would be informative to assess the applicability of the introduced method.
- The statement "This leads the neural operator to focus on modeling second-order deteriorating rollout performance" needs some verification (referencing prior work or experimental results).
- Can you show some empirical results to underline your statement "This design encourages a cooperative interaction among the neural operators, where adjustments in operator’s parameters induce compensatory adaptations across the ensemble, ensuring a cohesive towards minimizing residual prediction errors."? It would be great to see what actually happens in the background, i.e., how $R_{\theta}$ adjusts its weights, or whether a single Neural Operator basically dominates all contributions to a prediction.

### Questions
- How would you expect your method to compare to multi-step forecasting methods, where a single model is trained to predict multiple forecast horizons, such as described in **Dynamics Forecasting** paragraph in Section 4 of [[1]](https://arxiv.org/abs/2306.01984)?
- Have you tried to play with the inputs of the gating network $R_{\theta}$? I'm wondering, whether $R_{\theta}$ would benefit from receiving the respective outputs of the individual Neural Operators.

### Minor Comments
- Around the middle of page two, a reference seems to be missing ("See ?? for further details.").

### References
[1] https://arxiv.org/abs/2306.01984