DONIS: Importance Sampling for Training Physics-Informed DeepONet

1 Introduction

Solving partial differential equations (PDEs) is critical for describing and predicting natural phenomena such as fluid dynamics, heat conduction, and wave propagation. It also plays a key role in optimizing engineering designs, simulating complex systems, and advancing mathematical theories. However, for most PDE problems, obtaining analytical solutions is either infeasible or exceedingly difficult. Numerical methods, such as finite difference and finite element methods, offer efficient and approximate solutions for practical problems, providing effective approaches to addressing complex systems. Nevertheless, traditional numerical methods often rely on fine mesh discretization to enhance accuracy, resulting in exponentially growing computational costs, especially for high-dimensional problems.

In recent years, numerous methods leveraging deep learning have emerged to solve partial differential equations (PDEs) [Huang et al., 2022], aiming to circumvent some of the limitations of traditional numerical approaches by exploiting the fitting capabilities of neural networks. Among these, Physics-Informed Neural Networks (PINNs) [Raissi et al., 2019] integrate initial conditions, boundary conditions, and equation residuals into the loss function, thereby incorporating physical prior knowledge to supervise the neural network in approximating PDE solutions.

∗Corresponding author.

1 Code is available at https://github.com/ruihuang-1/donis.

Another line of research focuses on training neural networks to approximate operators, referred to as neural operators, enabling a single model to solve a class of PDEs. Examples include DeepONet [Lu et al., 2021a], Fourier Neural Operator (FNO) [Li et al., 2020] and further developments [Jin et al., 2022; He et al., 2024; Kumar et al., 2025; Wen et al., 2022; Li et al., 2023; Bonev et al., 2023]. The operator-fitting approach inherently resembles parallel fitting of a class of functions, posing greater challenges in terms of labeled data requirements compared to PINNs. To address this issue, subsequent studies have integrated operator learning with physical constraints to mitigate the reliance on costly labeled data, as demonstrated in the framework of physicsinformed DeepONets (PI-DeepONets) [Wang et al., 2021] and later studies [Goswami et al., 2022; Howard et al., 2023; Jiao et al., 2024; Li et al., 2024].

Building on the integration of operator learning with physical constraints, further advancements have explored scenarios where no labeled data is available, enabling fully selfsupervised operator learning. Existing research on fully selfsupervised PI-DeepONet can be broadly categorized into two main approaches. The first focuses on decomposing the spatiotemporal domain of the solution, reducing the complexity of the operator by employing multi-step solving techniques [Xu et al., 2023; Wang and Perdikaris, 2023]. The second approach addresses the challenge of multiple tasks in loss function. By incorporating initial and boundary condition constraints directly into the output layer or through postprocessing, approaches in [Brecht et al., 2023] avoids the imbalance among different terms in the loss function. While the aforementioned methods mitigate certain training challenges associated with PI-DeepONet, their reliance on specific PDE scenarios significantly restricts their generality and broader applicability.

Meanwhile, importance sampling has been adopted to accelerate the training of PINNs, owing to its ease of implementation and independence from the specific form of the PDEs. The work by [Nabian et al., 2021] incorporates importance sampling into the collocation point selection process of PINNs, and further enhances efficiency through a piecewise constant approximation. Subsequent work of the DMIS framework [Yang et al., 2023] introduces a dynamic triangulation strategy to optimize the accuracy of importance estimation, achieving an effective balance between speed and precision.

Motivated by these findings, we propose a more practical and versatile approach to accelerate the training of PI-DeepONet: PI-DeepONet with Importance Sampling (DONIS). By harnessing the architectural flexibility of DeepONet, DONIS introduces a two-step importance sampling framework that sequentially applies importance sampling to the function and collocation point inputs of DeepONet. This strategy prioritizes mini-batch samples with greater influence on the learning objective, enabling faster convergence and enhanced predictive accuracy.

Our main contributions can be summarized as follows:

• An easily implementable importance sampling method for input functions in PI-DeepONet.
• A gradient-based importance estimation and sampling strategy for collocation points in PI-DeepONet.
• Extensive experiments across multiple widely used benchmark physical scenarios, demonstrating that our method significantly improves the convergence speed and prediction accuracy of PI-DeepONet.

2 Background

2.1 Deep Operator Network

Deep Operator Network (DeepONet) provides a neural network framework for learning operators, which are mappings between function spaces. Unlike traditional neural networks that map vectors to vectors, DeepONet learns the functional relationship between an input function f(z) and an output function u(x), as defined by an operator G:

DeepONet achieves operator learning by decomposing the task into two sub-networks: the branch network and the trunk network. These networks work in tandem to approximate the target function.

• The branch network encodes the input function f(z) by sampling it at fixed sensor locations {z1, z2, . . . , zm}. These samples are passed through a neural network BranchNet, producing latent outputs b = [b1, b2, . . . , bp].
• The trunk network encodes the target evaluation location x via a neural network TrunkNet, producing latent outputs t = [t1, t2, . . . , tp].

The final output of DeepONet is then computed as a linear combination of the branch and trunk outputs:

where p is the number of latent modes. This formulation allows DeepONet to capture the functional relationship encoded by the operator G.

The supervised training of DeepONet minimizes the loss function:

where {fi, G(fi)} are labeled input-output pairs.

2.2 Physics-Informed Deep Operator Network

Physics-Informed Deep Operator Network (PI-DeepONet) extends the DeepONet framework by embedding the governing physical constraints directly into the learning process. This is achieved by incorporating the governing partial differential equations (PDEs), boundary conditions, and initial conditions as constraints in the loss function.

Consider a PDE of the form:

where f is the input function, u is the solution, u0(x) is the initial condition, F represents the PDE operator, and B represents the boundary conditions.

Our work focuses on scenarios constrained purely by physical laws, which are incorporated into the following loss function:

where:

• {(yir , tir )}Nr i=1 : Collocation points in the domain Ω × (0, T ),
• {(xib , tib )}Nb i=1 : Boundary points on ∂Ω × (0, T ),
• {xi0 }N0 i=1 : Points for the initial condition.

2.3 Importance Sampling

Importance sampling is a statistical technique for estimating properties of a target distribution p(x) by sampling from a more convenient proposal distribution q(x). It reweights samples from q(x) to approximate expectations under p(x).

The expectation of a function f(x) under the target distribution is:

Instead of sampling directly from p(x), importance sampling draws samples {xi}Ni=1 from q(x) and approximates the expectation as:

where α(x) = p(x) q(x) are the importance weights. A good q(x) should closely approximate p(x) in regions where f(x)p(x) contributes the most, reducing the variance of w(x).

In deep learning, importance sampling is utilized to enhance the efficiency of training by prioritizing samples that have a greater impact on the learning objective. This approach accelerates convergence and improves overall training efficiency. Specifically , the convergence rate of stochastic gradient methods is optimized when the sampling distribution aligns with the 2-norm of the gradient of the loss function [Needell et al., 2014; Zhao and Zhang, 2015].

3 Our Method

In this section, we provide a detailed overview of the proposed DONIS framework, which consists of two key components: sampling of functions and sampling of collocation points. The overall framework is illustrated in Figure 1.

3.1 Method Overview

In practical DeepONet training, a single step typically involves mini-batch partitioning over both the function space and the collocation point space. Consequently, the sampling space is represented as the Cartesian product of the function space and the collocation point space, with a total size of |Nf | · |Nc|:

where fi ∈ Mf represents a mini-batch of functions sampled from Nf . The points xi ∈ Mc, xbi ∈ M b c , and x0i ∈ M 0 c correspond to mini-batches of collocation points for the PDE residual, BC, and IC respectively. The terms lr, lb, and l0 are the respective loss functions for each.

In this sense, a straightforward and ideal approach to importance sampling is to jointly compute the importance of all samples of size |Nf | · |Nc| , and then samples a mini-batch of size |Mf | · |Mc| accordingly. However, this not only incurs significant computational costs, but also interfere with vectorization [Lu et al., 2022]. Instead, we propose a two-step sampling approach that separates the sampling of functions and collocation points, striking a balance between effectiveness and computational cost. First, importance sampling is applied to the function set, selecting functions with higher importance to form the function batch for the current step. Next, batches of collocation points are sampled from regions

Algorithm 1 Importance Sampling of Functions

• 1: Input: Dataset of functions Nf = {fi}|Nf | i=1 and seed points Sf = {an}|Sf | n=1 .
• 2: Parameter: model weights θ.
• 3: Output: Batch of functions Mf = {f ′i }|Mf | i=1
• 4: for i = 1 to |Nf | do
• 5: Calculate {lr (fi, an; θ) |an ∈ Sf}
• 6: Compute qf i according to Eq (11)
• 7: end for
• 8: Mf sampled according to qf i
• 9: returnMf

with higher importance within the domains of the selected functions.

Apart from the two-step scheme described above, we further reduce the computational overhead in two aspects. Firstly, we replace gradient-based importance with the loss value to avoid gradient computation. Both theoretical analyses [Katharopoulos and Fleuret, 2017] and practical applications [Nabian et al., 2021; Yang et al., 2023] have shown that using the loss value as a proxy for the 2-norm of the loss gradient is a reasonable approximation. Secondly, because boundary and initial conditions usually act as penalty terms, the importance sampling methods discussed here focus solely on collocation points within the spatiotemporal domain, corresponding to the PDE loss.

3.2 Importance Sampling of Functions

First, we consider sampling on the function set (DONIS-F), where an input function and its corresponding set of collocation points are treated as a single sample. As mentioned earlier, we measure the importance of a sample based on its loss value. Therefore, to measure the importance of a function sample, we need to obtain the overall loss value of the spatiotemporal domain Ω corresponding to this function. Here, instead of computing the loss at every point of Nc, we adopt an approximate method to bypass the costly process, as demonstrated in Algorithm 1.

Here, we define the sampling probability qf i for a function sample fi ∈ Nf as proportional to the loss (i.e., PDE residuals) lr over a set of seed points Sf = {an}|Sf | n=1 :

where seed points Sf is a sparse uniform grid across Ω.

By this means, we estimates the overall loss value for each function in the function set, thereby ensuring that more important functions are selected with higher probability for training.

3.3 Importance Sampling of Collocation Points

After function sampling, the sampling space for the minibatch is significantly reduced. However, the sampling of

Algorithm 2 Importance Sampling of Collocation Points

• 1: Input: Batch of functions Mf = {f ′i }|Mf | i=1 , sets of seed points S = {Sc,i}|Mf | i=1 where Sc,i = {bi,l}|Mf |,|Sc| i=1,l=1 .
• 2: Parameter: Collocation points Nc = {c}|Nc| k=1 , model weights θ.
• 3: Output: Batch of collocation points Mc = {Mc,i}|Mf | i=1 , updated sets of seed points S ′ .
• 4: Mc ← {}, S ′ ← {}
• 5: for i = 1 to |Mf | do
• 6: Calculate {lr (fi, bi,l; θ) |bi,l ∈ Sc,i}
• 7: Interpolate to approximate {ˆlr (fi, ck; θ) |ck ∈ Nc}
• 8: Compute qc i,j according to Eq (12)
• 9: Mc ←Mc,i sampled according to qc i,j
• 10: Compute gi,j according to Eq (13)
• 11: S ′ ←S′c,i sampled according to gi,j
• 12: end for
• 13: returnMc, S ′

functions primarily reflects the overall error. To better capture localized errors, we extend this regime to collocation points (DONIS-C).

Specifically, we need to estimate the importance of each element in the collocation point set. Contrary to the input functions in operator learning, the collocation points are distributed across a continuous spatiotemporal domain, which facilitates the use of interpolation methods to accelerate this process. Evidently, the effectiveness of interpolation heavily relies on the distribution of the source points. We refer to the source points as the seed points Sc, similar to Sf introduced in the previous section.

Here, we present a dynamic interpolation algorithm to accelerate importance sampling for collocation points demonstrated in Algorithm 2, where the distribution of Sc is determined by the spatial gradient of the the loss value (i.e., PDE residuals lr). Specifically, in each training step, we calculates each loss value of points in Sc. Then, interpolation is employed to estimate the the loss value over the spatiotemporal domain Ω, from which the mini-batch of collocation points is sampled. We define the sampling probability qc i,j for a collocation point cj ∈ Nc corresponding to selected function

f ′i ∈ |Mf | as proportional to the estimated PDE residuals ˆlr:

where ˆlr denotes the approximate of lr by interpolation over Sc.

Lastly, we compute the spatial gradient of the loss value, and update Sc based on the magnitude of the gradient:

where∇s denotes the Sobel operator, and gi,j is the probability of the j-th point of Nc being selected as a seed point for the i-th function.

Compared to a fixed uniform grid, seed points updated based on the spatial gradient allows the interpolation results to better align with the true data distribution, thereby enhancing the accuracy of the importance sampling algorithm. In practice, this calculation is performed via the Sobel operator, which gives an approximation of the spatial gradient, making the algorithm more efficient with minimal hyperparameters.

3.4 Importance Reweighting

As shown in Eq (9), reweighting is necessary to ensure that samples from the proposal distribution accurately represent the target distribution. Hence, we modify the PDE loss in Eq (10) when utilizing importance sampling for functions and/or collocation points:

where αi is the importance weight of the i-th function, and βi,j is the importance weight of the j-th point for the i-th function:

4 Experiments

In this section, we validate the proposed methods through a series of experiments and evaluate their performance using the specified metrics.The implementation of these methods is built upon the DeepXDE library [Lu et al., 2021b], with the complete code available in the supplementary materials for reference. Table 1 presents the hyperparameters used in our experiments.

4.1 Experimental Setting

We consider three partial differential equation scenarios: the Allen-Cahn equation, viscous Burgers' equation, and a nonlinear diffusion-reaction equation. Within each PDE problems, we compare the results of vanilla PI-DeepONet with three variants: PI-DeepONet with importance sampling on collocation points (DONIS-C), importance sampling on functions (DONIS-F), and importance sampling on both collocation points and functions (DONIS).

Metrics. We first train the model without any labeled data within the spatiotemporal domain Ω and then evaluate its performance using results obtained from traditional numerical solvers. The evaluation is conducted based on the Mean Square Error (MSE) and the L2 relative error as metrics.

Hyperparameters. General hyperparameters are shown in Tabel 1.

4.2 Allen-Cahn Equation

We first consider the one-dimensional Allen-Cahn equation in the form of:

with periodic boundary conditions. ϵ, which controls interface thickness, is set to 0.01.

We aim to train an operator mapping the initial condition u0(x) to the solution, i.e., G : u0(x) ↦→ u(x, t), t ∈ [0, 1], where the initial conditions are generated based on the Exponential Sine Squared kernel:

where l = 2, T = 1.

We set the size of seed points as 50 for Sf and 200 for Sc for this problem.

Figure 2 illustrates that DONIS accelerates loss convergence and achieves lower prediction error over the course of the training procedure. Figure 3 and the metrics in Table 2 demonstrate a consistent reduction in prediction error when applying DONIS-C, DONIS-F, or both in combination. Overall, DONIS reduces MSE by 18% and brings L2 relative error down to below 10%. By leveraging importance sampling, DONIS prioritizes functions and collocation points that exert a greater influence on the learning objective, thereby facilitating faster convergence and improved predictive accuracy.

4.3 Burgers' Equation

Next, we consider the one-dimensional Burgers equation:

with zero Dirichlet boundary conditions. The viscosity ν is set to 0.01. Similarly, the goal is to train an operator mapping the initial condition u0(x) to the solution, i.e., G : u0(x) ↦→ u(x, t), t ∈ [0, 1], where the initial conditions are generated based on the Radial Basis Function (RBF) kernel:

where l = 0.5.

We set the size of seed points as 100 for Sf and 500 for Sc for this problem.

As shown in Figure 4, for the Burgers equation, DONIS achieves a significant improvement in prediction accuracy, even though its acceleration of convergence is less pronounced. Table 3 reports a 39% reduction in MSE achieved by DONIS. Notably, DONIS-F demonstrates superior performance compared to DONIS-C in this experiment, which can be attributed to the Burgers equation's higher sensitivity to variations in input functions (i.e., initial conditions).This sensitivity arises from its nonlinear convection term, particularly under conditions of weak diffusion.

4.4 Nonlinear Diffusion-Reaction Equation

Lastly, we consider a nonlinear diffusion-reaction equation in the form of:

subject to zero Dirichlet boundary conditions and an initial condition given by u(x, 0) = 0. Here, we train an operator mapping the source term f(x) to the solution, i.e., G : f(x) ↦→ u(x, t), t ∈ [0, 1], where f(x) are generated based on the RBF kernel shown in Eq (20).

We set the size of seed points as 50 for Sf and 200 for Sc for this problem.

Figure 6 and 7, along with Table 4, illustrate the experimental results for the nonlinear diffusion-reaction equation. By fixing the boundary and initial conditions, the convergence difficulty of the corresponding loss terms was significantly reduced, which greatly simplified the convergence process. Notably, despite the changes in problem formulation and the reduction in training complexity, DONIS achieved faster convergence and lower prediction error. Specifically, DONIS-F, DONIS-C, and DONIS all demonstrated further improvements with DONIS demonstrating a 29% reduction in MSE as shown in Table 4.

4.5 Hyperparameter Analysis

In this section, we explore the effect of the seed points of DONIS-F and DONIS-C to further optimize the performance of DONIS.

Table 5 and 6 reveal that the accuracy of DONIS generally improves as the sizes of |Sc| or |Sf | increase. This is because a larger number of seed points enhances the accuracy of importance estimation, thus improving the overall performance of the method. However, once the number reaches a certain threshold, the improvement becomes less significant, and the computational cost increases considerably. Hence, specific sizes of Sc and Sf |Sc| were chosen in our experiments to strike a balance between accuracy and computational cost.

4.6 Ablation Study

In the following section, we present ablation experiments on the dynamic interpolation algorithm used in DONIS-C. The algorithm updates the seed points Sc by leveraging the spatial gradient of the loss value, with the goal of enhancing the accuracy of DONIS-C by ensuring that the interpolation results better reflect the true data distribution. Figure 8 provides a visualization of the updating process of Sc.

The results are summarized in Table 7, where methods that do not employ the dynamic interpolation algorithm are labeled with "fixed", indicating that the seed points Sc remain unchanged after initialization. As evidenced by the results, incorporating dynamic interpolation improves the performance of DONIS-C, both when used independently and in conjunction with DONIS-F.

5 Conclusions

In this paper, we propose a novel approach to accelerate the training of physics-informed DeepONet: PI-DeepONet with Importance Sampling (DONIS). Leveraging the architectural flexibility of DeepONet, DONIS introduces a two-step importance sampling framework that sequentially applies importance sampling to the function inputs and collocation points of DeepONet. In function sampling, the importance of function samples is approximated using the loss function evaluated over a uniform grid. For collocation points, their importance is estimated using a dynamic interpolation algorithm. Evaluated on three widely tested PDEs and across different operator mapping configurations, DONIS consistently enhances the convergence speed of self-supervised operator training while reducing prediction errors.These results demonstrate the practical versatility of DONIS, making it a robust and efficient framework applicable to a wide range of operator learning tasks.

Acknowledgements

This work is supported by the National Natural Science Foundation of China under Grant 62306199; the National Key Research and Development Program of China (2022YFC3801304).