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

TL;DR

A multi-task deep operator network (MT-DeepONet) to learn solutions across various functional forms of source terms in a PDE and multiple geometries in a single concurrent training session and offers a novel approach to solving PDE problems in engineering and science under a unified umbrella based on synergistic learning that reduces the overall training cost for neural operators.

Abstract

Multi-task learning (MTL) is an inductive transfer mechanism designed to leverage useful information from multiple tasks to improve generalization performance compared to single-task learning. It has been extensively explored in traditional machine learning to address issues such as data sparsity and overfitting in neural networks. In this work, we apply MTL to problems in science and engineering governed by partial differential equations (PDEs). However, implementing MTL in this context is complex, as it requires task-specific modifications to accommodate various scenarios representing different physical processes. To this end, we present a multi-task deep operator network (MT-DeepONet) to learn solutions across various functional forms of source terms in a PDE and multiple geometries in a single concurrent training session. We introduce modifications in the branch network of the vanilla DeepONet to account for various functional forms of a parameterized coefficient in a PDE. Additionally, we handle parameterized geometries by introducing a binary mask in the branch network and incorporating it into the loss term to improve convergence and generalization to new geometry tasks. Our approach is demonstrated on three benchmark problems: (1) learning different functional forms of the source term in the Fisher equation; (2) learning multiple geometries in a 2D Darcy Flow problem and showcasing better transfer learning capabilities to new geometries; and (3) learning 3D parameterized geometries for a heat transfer problem and demonstrate the ability to predict on new but similar geometries. Our MT-DeepONet framework offers a novel approach to solving PDE problems in engineering and science under a unified umbrella based on synergistic learning that reduces the overall training cost for neural operators.

Figures (14)

• Figure 1: A schematic representation of the operator learning benchmarks and MTL scenarios considered in this study.
• Figure 2: Schematic of the MT-DeepONet designed to learn a family of parametric PDEs (Fisher equations) defined by different forcing functions, $F(u)$. The parametric representation of this equation family, along with random initial condition fields, is input to the branch network. Spatio-temporal points are input to the trunk network. The multi-task operator network learns to predict the solution field, $u$ across this parameterized family of equations and random initial solutions concurrently.
• Figure 3: Schematic showing the binary masking function for a triangular geometry. A uniform grid $100 \times 100$ is sampled in $\Omega \in [0,1] \times [0,1]$ and used for generating the basis function in the trunk network. The binary mask is constructed by delineating the boundaries of the domain. Grid points within the domain boundary are denoted by $1$, while those outside are denoted by $0$. The binary mask is applied to the solution from the operator network, enforcing the solution outside the geometry to be $0$, thereby aiding with network convergence.
• Figure 4: Comparison between the reference solution and prediction obtained using the proposed multi-task operator network for representative test cases. The plot shows the operator network predictions against reference solutions for different initial conditions and forcing functions $F(u)$. The results demonstrate good overall accuracy across different initial conditions and forcing functions.
• Figure 5: Different geometries considered as tasks for MT-DeepONet for the Darcy flow problem. The source MT-DeepONet is trained using various combinations of the geometries S$1$ - S$7$ while the target domain geometries considered are T$1$ - T$3$. The boundary indicated in red denotes the Dirichlet boundaries where the solution $h(x) = 0$. The first objective of MT-DeepONet is to learn the hydraulic pressure heads across a combination of source geometries given a parametric family of spatially varying conductivity fields. The second objective is to transfer the knowledge of source MT-DeepONet to different geometries in the target domain using the transfer learning approach proposed in goswami2022transfer_learning to reduce overall computation time.