Learning Dynamical Coupled Operator For High-dimensional Black-box Partial Differential Equations

1 Introduction

Solving black-box parametric partial differential equations (PDEs) remains a significant challenge in many scientific and engineering disciplines [Zachmanoglou and Thoe, 1986]. Traditional methods, such as the finite element method, are computationally expensive and poorly scalable for large-scale or real-time problems. This is particularly problematic in scenarios where large amounts of data need to be processed quickly, such as uncertainty quantification, optimization under uncertainty, or optimal experimental design.

The emergence of neural networks has provided a powerful alternative for approximating solutions to complex problems. Some studies directly apply neural networks [Ruthotto and Haber, 2020; Huang et al., 2022], while others use Neural ODEs [Chen et al., 2018; Verma et al., 2024] to learn mappings between infinite-dimensional function spaces. However, both approaches face distinct challenges: neural networks often struggle with issues such as fixed query locations, whereas Neural ODEs require numerical integration, which can lead to significant error accumulation over time.

Therefore, the Deep Operator Network (DeepONet) [Lu et al., 2019] has gained prominence, offering a more specialized framework for learning mappings between function spaces. DeepONets represent a mapping F(u) from a function u to an output function at a query location y as:

where ϕi(y) are functions mapping the query location y to a finite-dimensional vector, and vi(u) are functionals mapping the input u to real numbers in R. They form a linear combination to approximate the true operator. DONs confront three fundamental limitations. First, the mapping ϕ from the low-dimensional input y to high-dimensional latent representations (Fig. 1) induces inevitable information bottlenecks, compromising approximation learning and inducing training instability. Second, existing approaches exhibit deficiencies in encoding mechanisms for input functions and query locations, constraining their capacity for approximation. Third, while the original DeepONet architectures have universal approximation guarantees, many derivative formulations sacrifice this theoretical foundation without rigorous justification, undermining their reliability in scientific computing applications.

For instance, MIONet [Jin et al., 2022] enhances the expressivity of the operator v by incorporating multi-input function information. However, it still fails to address training stability issues effectively. The LOCA method [Kissas et al., 2022] enables efficient approximation of the operator ϕ by calculating kernels between input and query locations. It struggles with early training stability and has shown limited practical effectiveness. On the other hand, IPOT [Lee and Oh, 2024] introduces lightweight operators based on attention mechanisms, aiming to improve efficiency. However, these approaches still do not guarantee the desired approximation properties. Additionally, related neural operator architectures and their variants, such as those in [Li et al., 2020; Li et al., 2022a], further highlight the challenges in achieving both stability and approximation guarantees.

To address these challenges, we propose the Dynamical Coupled Operator (DCO), a novel operator in the DeepONet class. We introduce a new approximation formulation for ϕ that incorporates temporal dynamics to learn an inputquery coupled function. By leveraging function information, we tackle the issue of mapping low-dimensional inputs to high-dimensional information bottlenecks and improve robustness. Additionally, we introduce an adaptive spectral input function encoder based on empirical mode decomposition to decompose non-stationary functions, thereby enhancing function representation. Furthermore, we propose a hybrid location encoding mechanism to improve query location encoding.

Overall, our method is compared with other operator method in Tab 1 and contributions are as follows:

• Dynamically Coupled Representations: We integrate temporal dynamics to learn an input-query coupled function, address the information bottleneck and and improve robustness.
• Adaptive spectral Input Function Encoder and Hybrid Location Encoding: We integrate EMD to handle non-stationary functions and utilize hybrid location encoding techniques to enhance model performance.
• Theoretical Guarantees of Universality: We prove that the proposed framework satisfies the universal approximation property, meaning it can approximate any continuous operator with arbitrary accuracy.
• Numerical Experiments: We conduct extensive numerical experiments on real-world physical systems, demonstrating that DCO outperforms models in the DON class.

2 RelatedWork

2.1 Operator Learning

Neural operators, which approximate mappings between input and output function spaces, include DeepONet [Lu et al., 2019], which utilizes branch and trunk networks to handle inputs and query positions, demonstrating strong approximation capabilities. The Fourier Neural Operator (FNO) [Li et al., 2020] operates in the frequency domain to capture global dependencies but lacks the ability to generalize to arbitrary query locations. Extensions like MIONet [Jin et al., 2022] improve flexibility of multiple input functions. LOCA [Kissas et al., 2022] calculates kernels between input function locations and query locations, enabling efficient approximation of query location representations. Some attentionbased frameworks, including Galerkin Transformers [Cao, 2021] and OFormer [Li et al., 2022b], employ cross-attention mechanisms to enhance adaptability to diverse input conditions. These models reduce computational costs while maintaining flexibility. IPOT [Lee and Oh, 2024] offers flexibility in processing arbitrary discretization. However, DONs do not leverage dynamic information and fail to address information bottlenecks, while non-DONs approaches lack universal approximation guarantees. In contrast, we address these issues by proposing a universal approximation guarantee and a more flexible operator formulation.

2.2 Approximation Theory

The approximation properties of neural networks have been extensively studied, with a focus on network structures such as feedforward neural networks [Cybenko, 1989; Hornik et al., 1989; Leshno et al., 1992] and residual networks [He et al., 2016]. These traditional networks often use a fixed architecture, with weights adjusted to approximate target functions. Recent theoretical advancements have extended universal approximation results to infinite-dimensional function spaces and nonlinear operators. For instance, Chen and Chen [Chen and Chen, 1995] demonstrated that neural networks can approximate nonlinear operators, providing a foundational theoretical basis for operator learning. Building on this, architectures like DeepONet [Lu et al., 2019], LOCA [Kissas et al., 2022], holomorphic operators [Adcock et al., 2024], and neural operators [Kovachki et al., 2021] have emerged, focusing on learning mappings between infinite-dimensional Banach spaces. While these methods primarily rely on static state representations, they lack flexibility in incorporating dynamic information, which limits their robustness and scalability.

3 ProblemFormulation

In this section, we aim to formalize the operator learning problem, specifically the relationship between input and output functions as defined by the symbols in Table 2.

Given N pairs of dynamic series of input functions and output functions {uℓt-1 (x), uℓt (x), sℓ (y)}Nℓ=1 , where uℓ ∈ C(X,Rdu ) and sℓ ∈ C(C(X,Rdu ) × Y,Rds ), the data pairs are assumed to be generated by an unknown ground truth operator G, defined as:

which maps an input function uℓt-1 (x) to its corresponding output function sℓ (uℓt (x), y).

Our goal is to learn an operator F such that for each function pair, the learned operator should satisfy:

and for ℓ = 1, . . . , N, we want:

We discuss this from the perspective of a simplified version, this formulation can be applied to various scenarios. For instance, consider a case where u represents a temperature field over the Earth's surface, and X is the spatial domain. The output s could be the pressure field sampled at specific locations Y. The operator F would then predict the pressure fields given new temperature distributions.

4 Method

We construct our model in two steps, inspired by attention mechanisms. First, we map the input function ut-1 to a feature vector v(ut-1) ∈ Rn . For each output location y and input function ut, we enhance the model by coupling the current state ut and y using an attention kernel, as in formula 1, to obtain an approximating function ϕ : C(X,Rdu ) × Y → Rn×ds through ψ(y) and k(ut). Here, σ denotes the softmax function. The forward pass of the proposed model is written as follows:

In this section, we will describe how the functions v and ϕ are constructed through ψ and k.

4.1 Adaptive Spectral Input Function Encoder

In our approach, the feature mappings k(·) and v(·) are constructed through two main operations. The first operation is the function E, which maps an input function u to a finite - dimensional vector space Rd :

Traditional spectral methods, such as the Fourier Transform and wavelet scattering networks[Bruna and Mallat, 2013] using in LOCA, offer effective feature representations but struggle with non-linear and non-stationary functions. For instance, Fourier Transform is sensitive to small deformations, while wavelet scattering relies on fixed basis functions that may not adapt to input signals' intrinsic characteristics. To address these challenges, we use Empirical Mode Decomposition (EMD) for feature extraction.

EMD is a data - driven technique that decomposes signals into multiscale oscillatory components, making it well - suited for non - linear and non - stationary signals. EMD decomposes an input signal u into Intrinsic Mode Functions (IMFs), denoted as χi(m), where each IMF captures an oscillatory component of the signal at a particular scale. The decomposition can be expressed as:

where d is the number of IMFs and M is the sequence length. Each IMF χi(m) satisfies: The difference between zero crossings and extrema is at most 1.

The mean of the upper and lower envelopes of χi(m) is zero.

This decomposition allows for the extraction of features at multiple scales. We define the mapping E(·) using the first d IMFs, along with the input signal u:

where ⊕ denotes concatenation.

After obtaining the finite - dimensional representation of u, we apply functions q and f, drawn from universal function approximators (e.g., fully connected neural networks), to obtain the feature representations:

4.2 Hybrid Query Locations Encoding

To encode the query locations y, we design ψ(y) to capture both intrinsic location information and the relative distance to the input function locations. To achieve this, we adopt a hybrid encoding scheme that combines a normalized radial basis function (RBF) kernel with discrete Fourier encoding.

The normalized RBF kernel is defined as:

where κ(y, y′ ) quantifies the similarity between y and y′ ∈ Y. Using a set of reference points {y′i }ni=0 , we compute the RBF - based features as:

where W1 ∈ R(n+1)×d is a learnable weight matrix.

To complement the localized features provided by the RBF kernel, we apply discrete Fourier encoding to capture periodic features of y. The Fourier features are given by:

where {ci}Ni=1 are predefined Fourier frequencies. Finally, the combined features a and b are concatenated and passed through a universal function approximator g to compute the final representation:ψ(y) = g(a ⊕ b) ∈ RD ,

This hybrid encoding effectively captures both localized and periodic characteristics of y, enhancing the expressivity of y representation.

4.3 Input-Query Coupled Function

To construct the state-dependent representation ϕ, we couple the input function ut with the query points y. First, we encode the current state ut using the feature encoder k, and encode the query points y ∈ Y using ψ. Inspired by by attention mechanism[Bahdanau, 2014], we compute a tensor-matrix product between ψ(y) and k(ut), followed by applying the softmax operation σ : Rn → ∆n to the rows. This softmax operation bounds the values of ϕ, where ∆n represents the probability simplex:

This formulation allows the representation ϕ(ut, y) to capture both the current state dependencies through the attention mechanism and spatial correlations via the positional encoding of y. By combining these components, the model is able to learn a richer and more expressive mapping.

4.4 Loss Function

To account for the geometric properties of a sphere, we employ a latitude-weighted mean squared error (MSE) loss function[Nguyen et al., 2023]:

where the latitude weight L(y) = cos(h) 1H ∑Hy′ cos(y′ ) , and h represents the latitude of the grid point y. Ω denotes the set of all grid points, and N is the total number of samples.

5 Theoretical Guarantees of Universality

In this section we give conditions under which the DCO model is universal. There exist multiple definitions of universality present in the literature, for example see[Sriperumbudur et al., 2010]. To be clear, we formally extend the definition we use below.

Definition 1 (Universality). Let X ⊂ Rdx , Y ⊂ Rdy , and U ⊂ C(X,Rdu ) be compact sets. A class of operators A is said to be universal if for any continuous operator

and any ϵ > 0, there exists an operator F ∈ A such that:

To explore the universality properties of our model we note that evaluation of the model can be written as

We will show that our model with the bounded normalization, state-query coupling ϕ and Adaptive spectral input function encoder v are universal by adding these components back one at a time. First, the following theorem shows that bounded normalization, state-query coupling ϕ does not reduce the approximation power of this class of operators.

Theorem 1 (DCO Preserves Universality). Let U ⊂ C(X, Rdu ) be a compact set, and let

be a continuous operator, where X ⊂ Rdx and Y ⊂ Rdy are compact. Then, for every ϵ > 0, there exist n ∈ N, functionals vj : U → R for j ∈ [n], and functions ϕj : C(X , Rdu ) × Y → Rn×ds , such that:

Lastly, we present a result showing that a particular architecture choice for the input feature encoder v also preserves universality. We show that EMD can be used to construct a universal class of functionals on C(X, Rdu ).

Proposition 1 (EMD Encoding Preserves Universality). Let Ad ⊂ C(Rd ,Rn ) be a set of functions dense in C(Rd ,Rn ), and {ei}∞i=1 a set of basis functions such that for some compact set U ⊂ C(X,Rd ). Let Ed : U → Rd denote in 4.1. Then for any continuous mapping h : U → Rn and any ϵ > 0, there exist d and f ∈ Ad such that:

Proof. The proof is provided in Appendix B.

For example, if our compact space of input functions U is contained in C 1 (X,Rdu ), the architecture proposed is expressive enough to approximate any functional from U → R.

6 Numerical Experiments

Unlike previous operator experiments, which typically involve low - dimensional and relatively simple datasets, we conduct experiments on a real - world, high - dimensional black - box PDE dataset. In this section, we present a comprehensive set of experiments to evaluate the performance of the proposed method. These experiments aim to assess the model's effectiveness in solving high - dimensional black - box partial differential equations. Beyond ablation studies, we focus on addressing the following key research questions 1 :

• How does the model perform in predicting various physical quantities compared to baselines?
• How well does the model train to predict unseen query locations compared to state-of-the-art DONs?

6.1 Experimental Setup

Datasets. We use the preprocessed 3 - hour increment ERA5 dataset from WeatherBench [Rasp et al., 2020], which contains governing black - box PDEs. We consider a dimensionality of du = 54 for input function and ds = 5 for output function from the ERA5 dataset: ground temperature (t2m), atmospheric temperature (t), geopotential (z), and ground wind vector (u10, v10). We use data at three different scales for evaluation. More details can be found in Appendix C.

Metrics. We evaluate our benchmarks using two commonly used metrics: latitude - weighted RMSE and latitude

1 https://github.com/ResearchGroupHdZhang/DCO

• weighted MSE. More metrics details can be found in Appendix D.

Baselines. To evaluate the effectiveness of our proposed method, we compare it with several baseline methods:

• DeepONet[Lu et al., 2019]: A foundational neural operator model that approximates functions using a branch - trunk network architecture, offering a rigorous mathematical framework based on approximation theory.

where g is the function representing query location y, and v is the function representing the input function u.

• LOCA[Kissas et al., 2022]: An extension of DeepONet that integrates the Wavelet Scattering Network and attention mechanisms to enhance the coupling between the input and query coordinates

where κ : Y × Y → R is the kernel, σ is the softmax function, v is the input feature encoder, and g is the proposed score function.

• IPOT[Lee and Oh, 2024]: A neural operator framework designed for optimal transport problems, which employs a differentiable iterative solver to approximate transport maps.

Implementation Details. To comprehensively evaluate our method, we conducted experiments using different lead times (3, 6, 9, 12 hours) in both the black-box PDE solving and query location generation tasks. All the training details can be found in Appendix E.

6.2 Experiments on Solving Black-box PDEs

To answer first question, We evaluate our model's performance in predicting multiple forward time steps of black-box PDEs. As shown in Table 3, our model significantly outperforms DONs and IPOT in terms of MSE and RMSE across all time steps and physical quantities in small scale PDEs. For instance, at 3h, our model achieves an average MSE of 0.0619, compared to 0.0750 for DeepONet and 0.0643 for IPOT. Similarly, at 6h, our model's MSE is 0.0883, while DONs and IPOT report 0.1037 and 0.0905, respectively. In terms of RMSE, at 3h, our model achieves a value of 23.797, which is a notable improvement over the 29.3809 from DeepONet and 26.307 from IPOT. At 6h, the RMSE of our model (36.96) is also significantly lower than that of DeepONet (44.6334) and IPOT (40.973). These results demonstrate the superior predictive accuracy and generalization ability of our model.

6.3 Experiments on Query Locations Generation

To answer the first question, we evaluate our model's ability to generalize to unseen query locations, specifically predicting from small-scale to medium-scale. As shown in Table 4, our model consistently outperforms baselines in both MSE and RMSE for unseen queries. At 3h, our model achieves an average MSE of 0.0619 and RMSE of 23.797, significantly lower than DeepONet's 0.0750 and 29.3809, and IPOT's 0.0643 and 26.307. Similarly, at 6h, the average MSE and RMSE are 0.0883 and 36.96, compared to DeepONet's 0.1037 and 44.6334, and IPOT's 0.0905 and 40.973. These results demonstrate the superior generalization capability of our model for unseen query points. As shown in Fig. 2, we also conduct zero-shot unseen query generation experiments to further highlight the superior performance of our model, specifically from small-scale predictions to large-scale ones. More details can be found in Appendix F.

6.4 Ablation Study and Visualization

As shown in Table 5, we conducted an ablation study to evaluate the impact of different components on model performance. Removing the Wavelet scattering (w/o B) leads to a decrease in performance across all time intervals, with higher MSE and RMSE values. Similarly, removing the hybrid query location encoding (w/o A) results in a significant drop in performance, especially in the RMSE metric. Our full model, which includes both components A and B, achieves the best performance, demonstrating the importance of both components in optimizing model accuracy and robustness. As shown in Fig 1, and the complete comparison of training stability can be found in Appendix G. In Fig 3, we campare to 4 physical quantities. it can be seen that our method also outperforms in terms of intricate scale, as indicated by the black box.

7 Conclusion

In this work, we introduced the Dynamical Coupled Operator (DCO), a novel framework for learning mappings between continuous function spaces. DCO overcomes key limitations of traditional neural operators by incorporating temporal dynamics and coupling input representations with query locations. This design reduces information loss, enhances the model's ability to capture complex spatio-temporal dependencies, and improves prediction accuracy in dynamic environments. We established a theoretical foundation for DCO, proving its universal approximation capabilities. Extensive benchmarking against existing models demonstrated DCO's superior performance in prediction accuracy, robustness, and generalization across black-box partial differential equations (PDEs). Additionally, the integration of Empirical Mode Decomposition enables DCO to outperform wavelet scattering methods. In future work, we hope to develop a unified and universality guarantee operator architecture based on this paper.

Acknowledgments

The research is supported by the Guangdong Provincial Key Lab of Integrated Communication, Sensing and Computation for Ubiquitous Internet of Things (No.2023B1212010007). Lin Ma and Haodi Zhang are the corresponding authors.