Source-linked AI summary
Deep transfer operator learning for partial differential equations under conditional shift
Somdatta Goswami, Katiana Kontolati, Michael D. Shields, George Em Karniadakis
TL;DR
The paper addresses expensive labeled-data collection and overfitting when related PDE tasks have different conditional distributions. It introduces TL-DeepONet, which transfers a source operator and fine-tunes selected target layers with a hybrid regression and conditional-embedding loss. Across diverse PDE shifts, the approach often learns accurately and faster with scarce target data, while large source-target feature changes limit performance.
Problem
Related PDE tasks may have scarce labeled target data and differing conditional distributions, making scratch training costly and prone to overfitting.
Method
TL-DeepONet transfers a source DeepONet, freezes lower layers, and fine-tunes task-specific layers with regression and conditional embedding operator discrepancy losses.
Results
Across six benchmark problems, TL-DeepONet demonstrates effective transfer under shifts in geometry, dynamics, material properties, and nonlinearities, including accurate learning from scarce data and faster training.
Takeaways & Limitations
The framework supports fast, task-specific PDE learning and uncertainty quantification without rebuilding surrogate models from scratch.
Takeaways & Limitations
Accuracy deteriorates when source and target domains differ substantially, because fine-tuning cannot recover target features absent from frozen lower layers or learn nonlinearities through limited adaptable layers.
Abstract
from arXiv · showhide
Transfer learning (TL) enables the transfer of knowledge gained in learning to perform one task (source) to a related but different task (target), hence addressing the expense of data acquisition and labeling, potential computational power limitations, and dataset distribution mismatches. We propose a new TL framework for task-specific learning (functional regression in partial differential equations (PDEs)) under conditional shift based on the deep operator network (DeepONet). Task-specific operator learning is accomplished by fine-tuning task-specific layers of the target DeepONet using a hybrid loss function that allows for the matching of individual target samples while also preserving the global properties of the conditional distribution of target data. Inspired by the conditional embedding operator theory, we minimize the statistical distance between labeled target data and the surrogate prediction on unlabeled target data by embedding conditional distributions onto a reproducing kernel Hilbert space. We demonstrate the advantages of our approach for various TL scenarios involving nonlinear PDEs under diverse conditions due to shift in the geometric domain and model dynamics. Our TL framework enables fast and efficient learning of heterogeneous tasks despite significant differences between the source and target domains.
Main Text
The paper proposes TL-DeepONet for task-specific PDE operator learning under conditional shift, combining selective fine-tuning with a hybrid loss that matches target samples and conditional distributions. Across geometric, material, dynamic, and nonlinear shifts, the framework generally improves data efficiency and training speed, but accuracy deteriorates when source and target features differ substantially.
- Framework: The framework transfers a source DeepONet to a target task, freezes lower layers, and fine-tunes task-specific layers using a hybrid regression and conditional-distribution loss.The conditional embedding operator discrepancy is combined with regression loss to address both individual target samples and global conditional properties.
- Framework: Operator regression targets mappings between function spaces, making DeepONet suitable for PDE surrogates across changing domains and task conditions.The paper motivates operator learning as discretization-invariant and applies it to PDE solution operators such as conductivity-to-hydraulic-head mappings.
- Results: With 250 target samples, TL-DeepONet achieves relative L2 errors below 5% across the four Darcy transfer tasks, whereas scratch training with few samples overfits.With 2,000 samples, scratch-trained DeepONet reports 1.4%, 1.62%, 1.83%, and 1.87% mean relative L2 errors, but at higher computational cost.
- Results: Transfer learning provides accurate scarce-data elasticity predictions and approximately 20× training speedup for TL5 at 1,900 target samples.For TL5, transfer learning achieves accuracy similar to target-domain training while reducing training time substantially.
- Limitations: Performance deteriorates when source and target internal boundaries or material properties differ greatly, because frozen lower layers cannot capture target features.In TL6, relative L2 errors are 7.62 ± 0.04 for horizontal displacement and 6.1 ± 0.10 for vertical displacement with 1,900 training data.
Methods
The method transfers a source DeepONet to a data-scarce target task under conditional shift, fine-tuning selected layers with a hybrid loss. The loss combines sample-level regression with RKHS-based conditional-distribution matching using labeled and unlabeled target data.
- Problem setup: The target model uses few labeled samples and additional unlabeled target inputs to estimate and constrain its conditional predictions.The method compares labeled target outputs with surrogate predictions on unlabeled target data.
- Problem setup: The framework transfers a source DeepONet to a target domain with identical input marginals but different conditional QoI distributions.It addresses conditional shift where P(X_S)=P(X_T) while P(Y_S|X_S) differs from P(Y_T|X_T).
- Conditional embedding: Conditional distributions are embedded in reproducing kernel Hilbert spaces, and CEOD measures their discrepancy through conditional embedding operators.The conditional embedding operator maps a fixed input value to an RKHS embedding; empirical estimation uses covariance operators, a Gram matrix, and regularization.
- Transfer and training: The source DeepONet is trained with a standard regression loss, and its learned parameters are transferred to initialize the target model.Source training can use relative L2 error or MSE and is optimized with gradient descent, particularly Adam.
- Transfer and training: Task-specific target layers are fine-tuned while other parameters remain frozen to preserve transferable features and limit training cost.The procedure fine-tunes fully connected branch layers and the final trunk layer for target-task expressivity.
- Hybrid loss: The hybrid target loss combines regression matching with CEOD, while trainable coefficients adjust the relative importance of the two components during optimization.The target parameters minimize the hybrid loss, whereas λ2 is maximized to progressively penalize conditional-probability discrepancies.
Author Contributions Statement
The author contribution statement assigns broad research, software, analysis, and writing responsibilities across the listed authors, with senior authors additionally contributing supervision, resources, funding, or project administration.
- S.G. and K.K. are credited with conceptualization, data curation, formal analysis, investigation, methodology, software, validation, visualization, and writing.
- M.D.S. contributed to conceptualization, funding acquisition, investigation, project administration, resources, supervision, and writing.
- G.E.K. contributed to conceptualization, funding acquisition, project administration, resources, supervision, and writing.
- The authors declare no competing interest.
S1 Theoretical details
The theoretical setup distinguishes finite-dimensional function regression from operator regression for PDE solution mappings and formulates transfer from a well-sampled source domain to a data-scarce target domain.
- Problem formulation: The problem considers a nonlinear, high-fidelity PDE simulator and seeks to approximate its mapping from random inputs to quantities of interest.
- Transfer learning setup: The source surrogate is trained with sufficient labeled data, while the target domain has different conditions and only a small labeled dataset.Target conditions may include geometry, boundary conditions, or model parameters.
- Transfer learning setup: Conditional shift assumes identical input marginals but different conditional output distributions between source and target domains.
- Transfer learning setup: Training target surrogates with scarce labeled data can cause overfitting, motivating information transfer between source and target tasks.
- Operator learning: Operator regression maps functions between infinite-dimensional spaces, unlike function regression, which maps finite-dimensional discretized representations.
- DeepONet: DeepONet encodes input functions with a branch network and evaluation coordinates with a trunk network to approximate the solution operator.The architecture can incorporate PDE solutions, forcing terms, initial conditions, and other functions.
- DeepONet: DeepONet parameters are conventionally optimized using residual and initial-condition losses, while branch and trunk architectures can vary with input representation.
S2 Network architecture details
The target DeepONet uses the same architecture as the source model across the listed transfer-learning problems, with architectures documented in Table S1.
- Architecture specification: The target-domain DeepONet architecture is identical to the architecture used for source-domain training.
- Architecture specification: Table S1 lists DeepONet architectures for transfer-learning problems TL1–TL8.
S3 Data generation
The experiments generate PDE datasets across varied geometries, stochastic inputs, dynamics, and in- versus out-of-distribution parameter settings for transfer learning evaluation.
- Darcy flow and elasticity: Darcy flow and elasticity datasets use unstructured meshes spanning multiple source and target geometries.The geometries include triangular, square, cutout, notch, and irregular-polygon domains.
- Transfer scenarios: Source and target datasets vary through geometry, boundary conditions, material properties, and model-dynamics parameters across the transfer scenarios.These variations define the task-specific shifts used to evaluate transfer learning.
- Brusselator diffusion-reaction system: The Brusselator maps Gaussian-random-field initial concentrations to evolved concentrations, with source and target parameter sets defining in- and out-of-distribution data.The initial concentration is generated using Gaussian random fields and truncated Karhunen–Loève expansions.
S4 Additional experiments
Additional experiments test transfer across forcing functions, geometries, material properties, out-of-distribution dynamics, and nonlinear Burgers’ responses, exposing both successful regimes and limitations.
- Darcy flow: For Darcy TL9, target-domain training achieves 1.7 ± 0.33% relative L2 error, compared with 4.90 ± 0.17% for TL-DeepONet.The scenario transfers from a square source domain to a notched triangular target domain with substantially different external boundaries.
- Brusselator diffusion-reaction system: TL-DeepONet performs satisfactorily on both Brusselator out-of-distribution datasets even with fewer than 250 available data.Results are reported for OOD1 and OOD2 using repeated runs with different random seeds.
- Elasticity model: Elasticity experiments show that large geometry and material-property differences can deteriorate TL-DeepONet accuracy, whereas other boundary and material shifts remain accurate with few samples.TL10 exposes insufficient feature transfer, while TL11 and TL12 show stronger performance under their respective shifts.
- Burgers’ equation: Representative Burgers’ results compare random boundary-condition inputs with final-time model responses and point-wise error plots.The figure reports reference responses and TL-DeepONet errors for three realizations.
- Burgers’ equation: Burgers’ TL13 transfers from ν1 = 0.2 to ν2 = 0.001/π, where the target response contains sharp nonlinearities absent from the smoother source response.TL-DeepONet achieves satisfactory performance only when sufficient labeled target data are available because fine-tuning fully connected layers must learn the nonlinearities.