Synergistic Learning with Multi-Task DeepONet for Efficient PDE Problem Solving

The goal of operator learning is to learn a mapping between two infinite-dimensional spaces on a bounded open set $\Omega\subset\mathbb{R}^{D}$, given a finite number of input-output pairs. Let $\mathcal{U}$ and $\mathcal{S}$ be Banach spaces of vector-valued functions defined as:

where $\mathcal{U}$ and $\mathcal{S}$ denote the set of input functions and the corresponding output functions, respectively. The operator learning task is defined as $\mathcal{G}:\mathcal{U}\to\mathcal{S}$. The objective is to approximate the nonlinear operator, $\mathcal{G}$, via the following parametric mapping:

where $\mathbf{\Theta}$ is a finite-dimensional parameter space. In the standard setting, the optimal parameters $\bm{\theta}^{*}$ are learned by training the neural operator with a set of labeled observations $\mathcal{D}=\left\{(u^{(i)},s^{(i)})\right\}_{i=1}^{N}$, which contains $N$ pairs of input and output functions. When a physical system is described by PDEs, it involves multiple functions, such as the PDE solution, the forcing term, the initial condition, and the boundary conditions. We are typically interested in predicting one of these functions, which is the output of the solution operator (defined on the space $\mathcal{S}$), based on the varied forms of the other functions, i.e., the input functions in the space $\mathcal{U}$.

The deep operator network (DeepONet) is inspired by the universal approximation theorem for operators [17]. The architecture of DeepONet comprises two deep neural networks: the branch network and the trunk network. The branch network encodes the input functions $\mathcal{U}$ at fixed sensor points $\{x_{1},x_{2},\dots,x_{m}\}$, while the trunk network encodes the information related to the spatio-temporal coordinates $\zeta=\{x_{i},y_{i},t_{i}\}$ where the solution operator is evaluated. The trunk network takes these spatial and temporal coordinates to compute the loss function. The solution operator for an input realization $u_{1}$ can be expressed as:

where $\{b_{1},b_{2},\ldots,b_{p}\}$ are the output embeddings of the branch network and $\{tr_{1},tr_{2},\ldots,tr_{p}\}$ are the output embeddings of the trunk network. In Eq. (4), $\bm{\theta}=\left(\mathbf{W},\mathbf{b}\right)$ represents the trainable parameters of the network including weights, $\mathbf{W}$, and biases, $\mathbf{b}$. The optimized parameters $\bm{\theta}^{*}$, are obtained by minimizing a standard loss function ($\mathcal{L}_{1}$ or $\mathcal{L}_{2}$) using a standard optimization algorithm.

The DeepONet model provides a flexible framework that allows the branch and trunk networks to be configured with different architectures. For equispaced discretization of the input function, a convolutional neural network (CNN) can be utilized for the branch network architecture, whereas a multilayer perceptron (MLP) is often employed for a sparse representation of the input function. An MLP is commonly used for the trunk network to handle the low-dimensional evaluation points, $\zeta$. Since its inception, standard DeepONet has been applied to address complex, high-dimensional systems [18, 19, 20, 21, 22, 23, 24]. Recent extensions for DeepONet have explored multi-fidelity learning [25, 26, 27], integration of multiple-input continuous operators [28, 29], hybrid transferable numerical solvers [30, 31], resolution independent learning [32], transfer learning [33], physics-informed learning to satisfy the underlying PDE [34, 2, 35], and learning in latent spaces [36].

The primary modification in the MT-DeepONet framework occurs in the branch network of the standard DeepONet architecture. Each problem is addressed uniquely. For example, in the case of multiple source terms in the Fisher equation (see Figure 1), the source terms are represented as a polynomial to create a unique representation for each equation:

The coefficients of the dependent variable $u$ in this polynomial expression are used as inputs to the branch network, along with the random initial conditions that define the problem. A schematic of the framework is shown in Figure 2 for understanding. Further details of the problem are discussed in Section 3.1.

To accommodate varying geometries concurrently in a single training process, we use a binary mask (array of $0$’s and $1$’s). This mask is constructed by fitting the geometry within a unit square plate for $2$D problems and a box for $3$D problems. The masking function assigns a value of $1$ to points within the boundary of the desired geometry and $0$ to points outside the boundary but within the plate or box. Figure 3 illustrates the masking function with a triangular geometry within a square plate. The solution operator is defined as the product of $\mathcal{G}_{\bm{\theta}}$ (as described in Equation 4) and the binary masking function, ensuring the solution is confined within the geometry’s bounds. Algorithm 1 details the steps for training the MT-DeepONet with a binary mask to address problems across multiple geometries. This approach also tests our hypothesis that learning multiple geometries improves the source model’s ability to transfer knowledge to target models with different geometries (Section 3.2). However, given the complexity of such knowledge transfer in the diverse range of PDEs representing different physical systems, we demonstrate the effectiveness of the masking framework for learning the solution operator across unseen 3D geometries in a steady-state heat transfer problem (Section 3.3).

MTL serves as an inductive transfer mechanism designed to enhance generalization performance compared to single-task learning. It achieves this by leveraging valuable information from multiple learning tasks and utilizing domain-specific insights embedded within training samples across related tasks. In our study, we demonstrate the effectiveness of MTL in learning solutions across diverse sets of PDEs, initial conditions, and geometries simultaneously.

The first example considers the Fisher equation proposed by Ronald Fisher in $1937$, which provides a mathematical framework for analyzing population dynamics and chemical wave propagation with diffusion [37]. The original reaction-diffusion equation is defined as:

where $u$ represents population density that varies spatially and temporally, $D$ and $r$ are scalar parameters denoting the diffusion coefficient and the intrinsic growth rate, respectively. In dimensionless form, this equation is written as:

Kolmogorov, Petrovsky, and Piskunov introduced a more general form, the Fisher-KPP equation [38]:

where the population density, $u$, varies along two spatial dimensions $(x,y)$, and the reaction term $F(u)$ must satisfy the following criteria:

Assuming that the density $u$ is invariant along the $y$-axis, the Fisher-KPP equation in dimensionless form is re-written as:

In the conventional operator learning task, the focus is typically on analyzing the change in density profile $u(x,t)$ over a one-dimensional spatiotemporal domain $x\in[0,1]$ and $t\in[0,1]$ for parameterized initial condition $u(x,t=0)$ drawn from a distribution. Due to diffusion, regions with higher density expand over time towards areas with lower density, based on the initial distribution. For our MT-DeepONet, we aim to evaluate the Fisher-KPP model to learn the density variation $u(x,t)$ over time for two tasks: ($i$) varying initial conditions of density $u(x,t=0)$, and (ii) three reaction functions $F(u)$ in a single training cycle, considering different functional of $F(u)$ that are separately parameterized by two scalar coefficients $a$ and $b$. We generate multiple initial conditions as a Gaussian random field and multiple forcing functions $F(u)$ using the Fisher-KPP general form with three reaction terms from the literature, as listed in Table 1.

To incorporate the different functional forms of $F(u,a,b)$ as inputs to the network, we express the functions as shown in Equation 5. For the Zeldovich form, we use the Taylor expansion and re-write it as:

The coefficients of $u$, $u^{2}$, and $u^{3}$ are represented as $\alpha$, $\beta$ and $\gamma$, respectively in Figure 2, with the constant term is denoted as $\delta$. To train the MT-DeepONet for the Fisher-KPP equations, the coefficients of the parameterized form of $F(u,a,b)$ are concatenated with the flattened initial conditions, $u(x,t=0)$ and used as inputs to the branch network. We use Adam optimizer with a progressively reducing learning rate and the network is trained for $100$,$000$ epochs using mean-squared error as the loss function. The relative $\mathcal{L}_{2}$ norm for the test samples is approximately $\sim 2.4\%$. Figure  4 presents four representative test case predictions using the MT-DeepONet. The results indicate that the multi-task operator network can capture large-scale features across the space-time domain with reasonable accuracy, under varying initial conditions and parameters of the Fisher-KPP equation. Additionally, the training time for individual DeepONet for each forcing term with varying initial conditions was $1250$ - $1300$ seconds on an NVIDIA A100 GPU. In contrast, the MT-DeepONet framework was trained in $\approx 1200$ seconds, showing the computational effectiveness of our approach.

In the second example, we consider the Darcy flow through a bounded domain and aim to train multiple $2$D geometries (bounded domains) simultaneously considering parameterized spatially varying conductivity fields. An earlier attempt, as noted in [42], focused on learning similar yet parameterized geometries in a concurrent training session. In contrast, our work demonstrates the learning of distinctly different geometries.

Darcy’s law describes fluid flow through a porous medium, relating pressure, velocity, and medium permeability. The pressure in the porous medium is expressed as [20]:

where $K(\bm{x})$ denotes spatially varying hydraulic conductivity field, $h(\bm{x})$ is the hydraulic head, and $g(\bm{x})$ is the source term. For simplicity, we set $g(\bm{x})=1$. The objectives for multi-task operator learning to predict the hydraulic head, $h(\bm{x})$, include ($i$) spatially varying conductivity fields $K(\bm{x})$ drawn from a Gaussian random field, and ($ii$) operation on varying $2$D geometries. In this task, we aim to demonstrate the generalization ability of MT-DeepONet through transfer learning, where the target domains differ geometrically from the source domains. The source and target geometries considered in this example are shown in Figure 5. The source geometries are labeled with S, while the target geometries are labeled with T. The source MT-DeepONet involves training on source domains S concurrently with sufficient labeled data, which is later transferred to related target domains T where only a small amount of training data is available. To capture geometric variation in a single training session, the conductivity field $K(\bm{x})$ is combined with a binary mask $\bm{M_{\text{binary}}}$ and used as input to the branch network, which employs a CNN architecture. The trunk network receives inputs from a uniformly discretized square domain $\Omega\in[0,1]\times[0,1]$, subdivided into a $100\times 100$ uniform grid. For all geometries, ground truth data is obtained using the MATLAB PDE Toolbox on an irregular mesh and thereafter interpolating the solution onto this $100\times 100$ regular grid. The trunk network’s basis functions (output embeddings) are used to evaluate the solution field at all domain points. The solution field outside the domain boundary is enforced to be zero by multiplying the solution operator with the binary mask. This mask ensures the solution is zero outside the geometry, thus improving convergence. Refer to Algorithm 1 showing the implementation details of this workflow.

In this section, we investigate how multi-task training in the source model enhances generalization across different geometries in target models. To assess the effectiveness of transfer learning, we train four source models using various geometric combinations: S1, S1+S2, S1+S3, and S1+S2+S3. Each source model is trained on a total of 5,400 samples. For single-geometry models (e.g., S1), all samples are from the same geometry. For multi-geometry models, we evenly distribute samples across the geometries: 2,700 samples from each geometry for S1+S2 and S1+S3, and 1,800 samples from each geometry for S1+S2+S3. We train the MT-DeepONet using a piece-wise constant learning rate scheduler with rates [0.001, 0.0005, 0.0001] over 5,000 epochs, employing mini-batching with a batch size of $1,000$ and optimizing with mean squared error. This experimental setup allows us to evaluate how the diversity of geometries in the source model affects the transfer learning process and subsequent performance on target models. Figure 6 presents three representative cases of predicted solutions when source MT-DeepONet was trained concurrently with geometries S$1$ + S$2$ + S$3$. The binary mask is applied to the output of the operator network to enforce that solutions outside the geometric domain are zero. The results demonstrate that the source MT-DeepONet has accurately learned all three geometries. Table 2 presents the $\mathcal{L}_{2}$ relative errors for various source model configurations. We observe a trend of increasing error as the number of geometries in the training data grows. This pattern is consistent with our expectations, given the multi-task feature set that the operator network must learn during multi-geometry training. While this approach may lead to a slight reduction in accuracy for individual tasks, it offers a more versatile representation across multiple geometries. This trade-off between task-specific performance and cross-geometry generalization is a key aspect of our multi-task learning strategy.

We employ transfer learning to adapt our trained source multi-task operator network for new geometries T1 and T2. The target models are considered distinct tasks. The transfer learning process begins by initializing the network with trained parameters from the source model. The layers updated during fine-tuning include the first input CNN layer of the branch network, three MLP layers following the convolution modules in the branch, and the linear output layer of the trunk network. We assess the prediction accuracy of the target model using varying numbers of training samples from the target domain: 50, 100, 200, 500, and 800. This approach significantly reduces computational costs for learning pressure heads on new geometries. Table 2 presents a summary of errors obtained from the target models for different test cases. Our analysis reveals distinct error patterns for target geometries T1 and T2. For T1, the source model combination S1 + S2 yields a lower $\mathcal{L}_{2}$ error compared to the single-geometry S1 source model. This demonstrates that training on multiple geometries can enhance the network’s transfer learning capabilities. Conversely, source models with geometrical combinations S1 + S3 and S1 + S2 + S3 result in higher errors in the target model compared to single-task learning on geometry S1. These findings underscore a crucial insight: naively combining all source tasks does not universally improve prediction performance for the target task. This phenomenon, known as negative transfer, occurs when source tasks unrelated to the target tasks are included. In our case, the results indicate that source geometry S3 introduces a negative transfer effect. Negative interference is a known challenge in MTL and we plan to explore mitigation strategies in future research. Conversely, for geometry T2, the source model S1 yields similar error values to the combination S1 + S2, indicating no significant improvement with the multi-task source model. This result may be attributed to the similarity between geometry T2 and S1. Figures 7 and 8 compare predictions from the transfer learning model against reference solutions for different source models trained with $50$ samples used to fine-tune the target model. The source model was trained with $5400$ samples where each geometry had equal representation.

The steady-state heat transfer equation for an isotropic medium is defined as:

where $K$ is the thermal conductivity, $q$ is the internal heat source, and $\bm{u}$ is the spatially varying temperature field. For multi-task operator learning, we solve this equation on a parametric $3$D circular plate as outlined in [43]. The plate has an outer diameter of $4$ inches, a central protrusion of height $\frac{1}{4}$ inch, and an overall thickness of $1$ inch. The location and the number of holes are parameterized as follows: the holes are placed at distances $\mathbf{d}\in[0.9,1.6]$ inches from the center of the protrusion, with the number of holes $\mathbf{n}$ varying between $2$ and $9$. All holes are equispaced in all the considered geometries, while the location of one hole is fixed in all the geometries in the plate (see Figure 10(a)). The radial location of the holes, $\mathbf{d}$ is varied in steps of $0.1$ inch. Heat sources $\mathbf{q}$ are placed in each hole, with a convective heat flux on the protrusion and side face, and adiabatic conditions elsewhere. These natural convective boundary conditions are shown in Figure 10(b). The objective is to learn the temperature distribution across the $64$ geometries that result from combinations of $\mathbf{n}$ and $\mathbf{d}$ (details of data generation in Section S1.3). Figure 11 presents four representative geometries and the corresponding temperature fields obtained from MATLAB’s PDE Toolbox.

Our objective here is to learn the operator mapping between the geometry parameters defining the different plate configurations and the resulting temperature field due to the heat sources placed inside the holes located in the plate. Since the number and location of the holes change based on the parameters $n$ and $d$, the resulting temperature field varies for each plate configuration. For the operator learning purpose, we select a total of $24$ training samples that encompass the minimum, maximum, and median values of the parameter $d$ for different numbers of holes. For instance, for a plate with $n=3$, we choose three samples with holes located at $d=0.9$, $d=1.2$, and $d=1.6$ as training samples corresponding to $n=3$. The sample choice ensures that the extreme ends of the design space are well represented in the training dataset while avoiding oversampling. The remaining geometries are used as test cases for network inference. A list of the cases used for training and testing is provided in Table S1.

The parametric representation for the plate ($n$ and $d$) is used as input to our branch network, while the trunk network receives uniformly sampled points in the domain $\Omega=\{(x,y,z)\in\mathbb{R}^{3}:-2\leq x\leq 2,-2\leq y\leq 2,-0.25\leq z\leq 1\}$ as input. Unlike the previous example on the 2D Darcy problem where the binary mask was considered as input to the branch network, here only the geometry parametrization values are employed. The binary mask is used as additional information in the loss function. We utilize the same grid points for all plate configurations with the trunk network and later apply a binary mask to constrain the temperature field to be zero outside the geometry, similar to the $2$D Darcy problem. The masking matrix is a $3$D grid with values of $0$ (outside) and $1$ (inside) the domain boundary. By utilizing this binary mask, we establish the associativity between the grid points and geometry points, enabling the operator network to learn the temperature field across multiple geometries simultaneously. The training step employs an exponentially decaying learning rate, starting at $1\times 10^{-3}$ with a decay rate of $0.9$ every $1000$ iterations.

The binary mask, applied to the network output, ensures that the solution at points outside the geometry is zero, thereby improving convergence and accuracy. Figure 12 compares the convergence of the MT-DeepONet with and without the masking function in the loss function. The application of a mask to the solution output demonstrates better error convergence. Figure 13 shows the prediction of MT-DeepONet without the masking operation, where higher inaccuracies are observed across the domain due to insufficient geometry representation in the training dataset. Applying a binary mask imposes an extra constraint that assists the network in learning solutions across different geometries. This masking process supplies essential geometric details and improves the network’s capacity to generalize to new plate configurations (with the same $n$). Figure 14 presents representative examples of geometries comparing the reference temperature field with the network prediction. These predictions align well with the reference numerical solution. Higher error values are observed near the holes due to their being under-represented in the training data. We observed better results when more geometries are added to the training dataset but this leads to oversampling in the training dataset. The masking framework functions similarly to transfer learning between tasks, albeit without involving network re-training, and is critical for predicting solutions across different geometries in our experiments.

For this problem, the goal is to learn the operator mapping from the random initial density condition $u$ across a parameter range $a$ and $b$, as defined in Table 1, for its entire time evolution. This mapping is expressed as $\mathcal{G}:u(\bm{x},a,b,t=0)\mapsto u(\bm{x},t)$, where $t>0$ and $\bm{x}\times t\in[0\times 1]\times[0\times 1]$. The initial density is modeled as a Gaussian random field (GRF), defined by:

where $\mu(\bm{x})$ and $\text{Cov}(\bm{x},\bm{x}^{\prime})$ are the mean and covariance functions, respectively. We set $\mu(\bm{x})=5+0.1\times\sin(\pi x)$, while the covariance matrix is defined by the squared exponential kernel:

where $\ell=0.4$ is the correlation length, and $\sigma^{2}=2$ is the variance. We utilize Karhunen-Loéve expansion (KLE) to generate $1,000$ random initial conditions. The temporal points $t$ and spatial points $\bm{x}$ are discretized into a $20\times 64$ grid, resulting in a total of 1,280 collocation points. Parameters $a$ and $b$ are discretized in steps of $0.1$ and $0.15$, respectively, within the defined parameter space. Combining different parameter values for $a$ and $b$ with the random initial conditions, we generate a total of $5,000$ training samples. This dataset is split into training and testing sets using an $80\%-20\%$ split, with $N_{\text{train}}=4,000$ and $N_{\text{test}}=1,000$. For preparing data for the multi-task operator network, we concatenate the initial condition data at $64$ domain points with the parameter values for the three Fisher equations defined in Table 1, resulting in a $67$ dimensional input vector ($64$ initial condition points + $3$ equation parameters).

The multi-task objective in this problem is to learn the operator for the Darcy flow equation described in Equation 12 over different $2$D spatial domains under randomly generated initial states for conductivity field $K(\bm{x})$. The hydraulic conductivity field is modeled as a stochastic process, using truncated Karhunen-Loéve expansion. The conductivity field is generated in a reference square domain $\Omega\in[0,1]\times[0,1]$ with $100$ grid points using a Gaussian random field with length scale, $\ell=0.05$. A total of $6000$ solution fields are generated using the multiple initial values of conductivity fields for each geometry. A Dirichlet boundary condition of $h(\bm{x})=0\quad\forall\bm{x}\in\partial\Omega$ is applied on all the edges of each geometry, shown in Figure 5. The solution to Darcy’s equation for all samples is generated using MATLAB’s PDE Toolbox solver. The domain mesh is generated using the MATLAB PDE Toolbox mesh generation algorithm with a maximum mesh size of $0.03$ inches for all geometries. We use unstructured triangular elements for the mesh. The hydraulic pressure head solution obtained is linearly interpolated on a regular grid in $\Omega\in[0,1]\times[0,1]$ to utilize a common trunk network input. Points in the square grid $\Omega$ that lie outside the triangular geometry are extracted and the solution field at such points is manually set to $0$ throughout the dataset. We sub-divide the $6000$ samples into $N_{train}=5400$ and $N_{test}=600$ for each source model. For training the MT-DeepONet with multiple geometries, we use an equal number of samples from each geometry to obtain a combined total of $N_{train}=5400$, thus ensuring equal representation.

To generate data for the steady-state heat transfer problem with multiple plate geometries, we create a parametric CAD model of the plate where the parameters $\bm{n}$ and $\bm{d}$ determine the number of holes and their location relative to the central axis of the plate protrusion (see Figure 10), respectively. To generate unique geometries with a combination of $n$ and $d$, we fix the location of one hole across all the geometries (see Figure 10(a)). We solve the steady-state heat transfer Equation 14 using MATLAB’s PDE Toolbox for various design configurations that result from varying the design parameters $\bm{n}$ and $\bm{d}$ in the plate geometry, where $2\leq\bm{n}\leq 9$ and $0.9\leq\bm{d}\leq 1.6$. We increment $\bm{d}$ in steps of 0.1 resulting in a total of $64$ distinct design configurations. As discussed earlier, we separate the training and test cases based on the hole location, $\bm{d}$ as shown in Table S1. For each design, we use PDEToolbox for mesh generation with a maximum element size of $0.1$ inch and tetrahedral elements. A heat source $\bm{q}=1$ is placed inside each hole, simulating a cartridge heater to heat the plate. Convective boundary conditions are applied on the protruding face and side walls with an ambient temperature $u_{\infty}=6$ and a convective heat transfer coefficient $h_{c}=0.3$. The output solution, which is the temperature field $u$, is evaluated at all node points of the mesh. Figure 11 shows representative examples of the different geometries and their corresponding temperature fields obtained by solving Equation 14. The temperature solutions are initially obtained on unstructured grid points. For simplifying the inputs for the trunk network, we interpolate these solutions over a regular $3$D grid ($100\times 100\times 100$) around the plate. Points in the regular grid that fall outside the domain of the $3D$ plate are set to $0$ before training the network.