Operator Learning with Domain Decomposition for Geometry Generalization in PDE Solving
Authors: Jianing Huang
Kaixuan Zhang
Youjia Wu
Ze Cheng
Venue: ICLR 2026 Poster (2026)
Keywords: Neural Operator
Domain Decomposition
Geometric Generalization
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: This paper introduces a method that makes use of domain decomposition to solve the geometric generalization problem in neural operators.

Abstract: Neural operators have become increasingly popular in solving partial differential equations (PDEs) due to their superior capability to capture intricate mappings between function spaces over complex domains. However, the data-hungry nature of operator learning inevitably poses a bottleneck for their widespread applications. At the core of the challenge lies the absence of transferability of neural operators to new geometries. To tackle this issue, we propose operator learning with domain decomposition, a local-to-global framework to solve PDEs on arbitrary geometries. Under this framework, we devise an iterative scheme Schwarz Neural Inference (SNI). This scheme allows for partitioning of the problem domain into smaller subdomains, on which local problems can be solved with neural operators, and stitching local solutions to construct a global solution. Additionally, we provide a theoretical analysis of the convergence rate and error bound. We conduct extensive experiments on several representative linear and nonlinear PDEs with diverse boundary conditions and achieve remarkable geometry generalization compared to alternative methods.These analysis and experiments demonstrate the proposed framework's potential in addressing challenges related to geometry generalization and data efficiency. The code is publicly available at this repository: https://github.com/questionstorer/sni.

Published as a conference paper at ICLR 2026
OPERATOR LEARNING WITH DOMAIN DECOMPOSITION FOR GEOMETRY GENERALIZATION IN PDE SOLVING
Jianing Huang" Kaaxuan Zhang, Youjia Wu, Ze Cheng
Corporate Research, Bosch (China) Investment Co., Ltd. (j dianing. huang, kaixuan. zhang, youjia. wu, ze. ze. cheng) 8cn. www.boscc.o..o
ABSTRACT
Neural operators have become increasingly popular in solving partial differential equations (PDEs) due their superior capability capture intricate mappings between function spaces over complex domains. However, the data- data-hungry nature of operator learning inevitably poses a bottleneck for their widespread applications. At the core of the challenge lies the absence of transferability of neural operators to new geometries. To tackle this issue, we propose operator learning with domain decomposition, a local-to-giobal framework to solve PDEs arbitrary geometries. Under this framework, we devise an iterative scheme Schwarz Neural Inference (SNI). This scheme allows for partitioning of the problem domain into smaller subdomains, on which local problems can be solved with neural operators, and stitching local solutions construct a global solution. Additionally, we provide theoretical analysis of the convergence rate and error bound. We conduct extensive experiments on several representative linear and nonlinear PDEs with diverse boundary conditions and achieve remarkable geometry generalization compared to alternative methods. These analysis and experiments demonstrate the proposed framework's 's potential in addressing challenges related to geometry generalization and data efficiency. The code is publicly available at this repository: https 7/githu's com quest ietpsstiinstorr////omm
INTRODUCTION
Partial differential equation (PDE) solving of paramount importance in comprehending natural phenomena, optimizing engineering systems, and enabling multidisciplinary applications (Evans 2022). The computational cost associated with traditional PDE solvers (Liu and Quek 2013) et al.
2019) has prompted the exploration of learning- -ased methods as potential alternatives to overcome these limitations. Neural operators (Li et al., 2020b 2023 2024 Liu et al., 2023 Hao et al. 2023), as extension of traditional neural networks, aim learn mappings between the functional depen- dencies of PDEs and their corresponding solution spaces. They offer highly accurate approximations to classical numerical PDE solvers while significantly improving computational efficiency. Despite its success, operator learning, as a data-driven approach, encounters the inherent 'chicken-and-egg" problem, revealing an interdependence between operator learning and the availability of data. This dilemma arises from the challenge of simultancously addressing the inefficiency of classical solvers and acquiring an ample amount of data for neural operator training.
Existing works in alleviating the above challenges explore symmetries of PDEs. Lie point symmetry data augmentation (LPSDA) (Brandstetier al., 2022) generates potentially infinitely many new solutions of PDE from existing solution by exploiting symmetries of differential operator defining the PDE. Subsequent work (Mialon al. 2023) applies LPSDA for self-supervised learning. However, LPSDA only partially alleviate the problem in data efficiency and the problem of how to quickly generalize new geometry untouched. While existing neural operators have shown capabilities in handling diverse geometries through approaches such as geometry parametrization (Li et al. 2023) or coordinate representation (Hao et al. 2023), they lack the ability to generalize to entirely novel geometries that differ significantly from those present in the training data distribution. The inability
"Corresponding author; wawningghung@@cn.osccc.c.

---

in Published as a conference paper at ICLR 2026
to quickly adapt neural operators to unseen geometries without further generating new data hinders the applicability of neural operator learning to real-world problems in industry.
To tackle this challenge, natural idea is to break down domain into some basic shapes where neural operator can generalize well. Domain decomposition methods (DDMs) (Toselli and Widlund
2004 Mathew 2008) provide the suitable tool for this purpose. Related efforts such Mao al. (2024) have combined operator learning with DDMs on uniform grids to accelerate classical methods. In contrast, our work aims extend this paradigm arbitrary geometries through a local-to-global framework. This framework consists of three parts: (1) Training data generation: creation of random basic shapes and imposition of appropriate boundary conditions on these shapes. This generated data serves as the training set for the neural operator in our framework. (2) Local operator learning: neural operator training to learn solutions on basic shapes. Data augmentation based on symmetries of PDEs is utilized to enable the neural operator to capture the intricate details and variations within these shapes. (3) Schwarz neural inference (SNI): a three-step algorithm for inference. Firstly, the computational domain is partitioned into smaller subdomains. Then, the learned operator is applied within each subdomain to obtain the local solution. Finally, an iterative process of stitching and updating the global solution is performed using additive Schwarz methods.
Our Contributions. We summarize our contributions below:
We introduce local-to-ghe gloogl framework that integrates operator learning with domain decomposition methods as an attempt in tackling the geometry generalization challenge in operator learning.
We design novel data generation scheme that leverages random shape generation and symmetries of PDEs to train local neural operators for solving PDEs on basic shapes.
We propose an iterative inference algorithm, SNI, built upon a trained local neural operator obtain solutions arbitrary geometries. We theoretically analyze the convergence and the eгror bound of the algorithm for a wide range of elliptic PDEs. Through comprehensive experiments, we empirically validate the effectiveness of our framework on generalizing to new geometries for both linear and nonlinear PDEs.
2 PROBLEM FORMULATION AND PRELIMINARIES
In this section, we provide an introduction to the problem formulation and essential background on domain decomposition methods, which will be utilized throughout the entirety of the paper.
2.1
PROBLEM FORMULATION
Our primary focus is stationary problems of PDEs defined in the following form:
C(u)= = ʃ
u = u.D
on Tn
(I)
on IN du = g dn
where L is a partial differential operator and Tp UT = aΩ denotes Dirichlet and Neumann boundary, respectively. We assume all the domains Q bounded orientable manifolds embedded in some ambient Euclidean space 1 (Li et al. 2023). Later we will extend our method to handle dependennndnt equations.
We consider situations where geometry of domain Qut at inference time is decoupled from that of Quain in training time, i.c. Qunt does not have to fall in resemble training geometries and can be of arbitrary shapes. For implementation we will mainly focus on S2 C IR?.
2.2 DOMAIN DECOMPOSITION METHODS
Domain decomposition methods (DDMs) solve Eq. by decomposing domain into subdomains and iteratively solve coupled system of equations each subdomain. An overlapping decomposition of
2

---

Published as a conference paper at ICLR 2026
Ω is collection of open subregions (1 9 S for k = l...., K such that U S4 Ω.
We denote V and [V be finite element space associated with domain f2 and {2t } 1= We
can define restriction operators {Rs V + Va restricting functions on Ω to {Ax jkm and extension operators {R] Vx VI extending functions on {_ to S by zero.
In the subsequent discussion, we revisit the idea of additive Schwarz method (ASM) in DDMs for overlapping decomposition. The additive schwarz-Richardse iteration (Mathew 2008 has the following form:
n+1 ut
[ Ru R2 Rge""]
where <<< R
is a hyperparameter contrrll the convergence rate, and wx ! is the solution of the following equation:
[(w) 1 = 0
in S2
wk n++ uD on 82 nT
Ow n+
(3)
Ön g
on nrN
w n+I = u"
on 02g 0o
We denote the local operator S₁ (2", 2.D g) w n+1 Note that the first two boundary conditions in Eq. is is the boundary condition on the global boundary part of an, and is not updated during iteration. The last boundary condition is along the artificial boundary created by decomposition and the value is updated through iteration. Hence { can considered as a single-input operator when the global boundary condition and decomposition determined. This iterative process be shown to converge for FEMs under mild assumption on properties of equation and decomposition. Please refer to Appendix A for more details.
3 OPERATOR LEARNING WITH DOMAIN DECOMPOSITION
In order to solve PDE on arbitrary geometry with neural operator, a natural idea is to decompose domain into a prescribed family of building blocks (basic shapes) since it is not feasible to explicitly consider arbitrary shapes during training stage. For that purpose, propose to train neural operator to solve local problems on basic shapes and stitch local solutions together to get a global solution. An illustration of the proposed framework is presented in Figure A detailed implementation will be discussed in the following subsections.
3.1 TRAINING DATA GENERATION
Data generation serves the purpose of operator learning, which fundamentally aims to approximate the local solution operator G P x H. + U. Here, P denotes the space of basic shapes, H. represents boundary conditions and other input functions, 24 represents the solution space. Next we will delve into a comprehensive examination of how P and H are determined separately.
Choice of basic shapes. The selection of basic shapes cannot be arbitrary due the requirement of ensuring the neural operator's capability in solving local problems across wide range of shapes. Moreover, three necessary criteria should be set forth for basic shape generation: (1) sampling feasibility: it should be tractable to sample from and solve boundary value problems on basic shapes in P. (2) complete coverage: basic shapes in P should be flexible to cover any shape of domain.
For implementation, we focus on Ω C IR.. We propose to use the space P. (n) of simple poly- gons with at most n vertices (i.e. planar polygon without self-intersection and holes) uniformly bounded by compact region in IR2 Simple polygons are Lipschitz domains with straightforward sampling method (Auery and Heldz 2019) and flexible enough to constitute any discretized planar domain (Preparata and Shamos 2012). We note, however, that this is not the only choice of these basic shapes. We could equally use convex polygons, star-shaped polygons, etc. as long as the two aforementioned criteria satisfied.
3

---

Published as a conference paper at ICLR 2026
(6) Training
Roarday
Nount opensor
Aggrartad purple
(b) Schwarz Neural Inference (SNI)
Z
Aa
fix
Daratire Schone
update local artificial boundary condition gtt
Wasling
Tp
EGR contiiiinn addeendde
telamenn durtel
5.6. crotluping pubgapia
2.P. miter engepph
Figure I: An illustration of Operator Learning with Domain Decomposition Framework. (a) During training stage, the goal is to ensure that the neural operator can effectively model the local solution operator various building blocks of shapes. These building blocks selected and generated based on specific criteria, allowing for a more efficient and targeted learning process. Proper boundary conditions then imposed to generate local solutions which operator: (b) During inference, for an arbitrary given domain, an automated decomposition algorithm employed decompose the domain into subdomains. By leveraging the trained local operator and Schwarz Neural Inference (SNI), global solution can be obtained by stitching local solutions on subdomains.
training data for neural
Imposing boundary conditions. The imposition of boundary conditions presents two complications: (1) Types of boundary conditions. Neumann boundary conditions in Eq. [ will inevitably result in mixed boundary conditions in local subdomains. To generate solutions with mixed boundary conditions, we randomly divide the boundary of a basic shape into two connected components, representing the Dirichlet and Neumann boundaries. respectively. During inference, we have to carefully set hyperparameters for decomposition to make sure boundary of subdomains have at most two connected components for Dirichlet and Neumann boundaries. (2) Functional range of boundary conditions. In general, the inference process for subdomains will encounter arbitrary ranges boundary conditions. However, it is practically infeasible to train the neural operator handle unbounded boundary values. Instead, generate random functions with values normalized within bounded range for both boundary conditions and other input functions such as coefficient fields and source terms. We will handle this complication with symmetries of PDEs during inference.
3.2 LOCAL OPERATOR LEARNING
now train a neural operator gl to approximate the mapping G. Our focus is not on design of neural operators, but on ensuring that the neural operator can solve local problems accurately.
Choice neural operator architecture. Our framework orthogonal to the choice of neural operator architecture long the architecture accommodate flexible input/output formats and possesses sufficient expressive power to solve local problem with randomly varying domain and input functions. For implementation, transformer-base neural operator. We note that, however, training neural operator on highly varying geometries presents challenges both design of architectures and training schemes.
adopt GNOT (Hao al. 2023) which is a highly flexible
Data augmentation. To enhance the generalization capabilities of the neural operator, Lie point symmetry data augmentation (LPSDA) (Brandstetter al. 2022) can be naturally applied to local
4

---

Published as a conference paper at ICLR 2026
Algorithm 1 Schwarz Neural Inference
Input:
Domain St; Global boundary data B (Dirichle/Neumann) Other input functions H (e.g.. coefficient fields, source terms, PDE parameters); Number of subdomains K; Depth of extension d; Local operator G Step size т; Convergence criterion C;
Output: Global Solution u;
1: Apply METIS and extension to get overlapping decomposition {Ωx J obtain restriction operators {R j and extension operators {R
2: Initialize the global solution u";
3: while convergence criterion C' not satisfied do
4: Form local inputs {(R,, Bb Hx)] from Ω B, H and the last-step solution w":
5: Obtain preprocessing transforms (Tx and postprocessing transforms {Ty 1
6: Local inference on each subdomain: 10 up+t Ti. (gt (Ig( , H.))));
7: Global update (additive Schwarz form): u=+1 w" + Ek= R2 ( Ryu");;
8: n n
9:
end while
10:
return u":
solutions during training. Examples of such transformations are rotation and scaling. is crucial appropriately extend these transformations boundary conditions and other input functions, taking into account the symmetries inherent in the PDEs. Please refer to Appendix for detailed discussion.
3.3 SCHWARZ NEURAL INFERENCE
Inspired by additive Schwarz method, we introduce a similar iterative algorithm called Schwarz Neural Inference (SNI), which is outlined in Algorithm In the subsequent discussion, we will explore several important considerations.
At each iteration SNI assembles, for every subdomain 22 a local boundary input B by combining the prescribed global boundary data on 802g n as and the interface Dirichlet data from the previous iterate on 022 n Q. For PDEs with additional input functions parameters (c.g., coefficient fields and source terms), we denote by H₃ their restriction to 2g- In Algorithm i denotes the post-processed local solution on 12g-
Decomposition into overlapping subdomains. In general, there no canonical way to decompose an arbitrary domain into prescribed shapes. Here adopt common practice in the DDM litera- ture [Mathew] 2008). We assume there exists a pre-defined triangulation Tr, (S2) of the domain Ω2 and graph can be constructed represent the connectivity of this triangulation. A graph partition algorithm such as METIS (Karypis and Kumar
1997) is then employed to partition this graph into K non-overlapping connected subgraphs with index sets , Ix. To achieve an overlapping decomposition, each subgraph is then extended iteratively by including neighboring vertices for iterations. This process generates index sets Z ................................ IIx that, together with the original mesh, form an overlapping decomposition denoted as (Ω4 1 k=1- intuitive illustration of this process depicted
in Figure
For implementation, partition number K and extension depth are hyperparameters that should be carefully set to ensure that the resulting subdomains resemble shapes in P.
Normalization. During inference on an arbitrary decomposed subdomain, the range of geometry and boundary conditions may differ from that of the generated training data. We thus leverage the symmetry properties of PDEs to handle this mismatch. More specifically, we can apply problem- dependent transformations Tj (c.g.. spatial translation/scaling and value shift/calling) to map local input (2x, BT H2) that lies outside the training range back into the training range. When the PDE symmetry acts additional inputs (e.g., coefficient fields, source terms, parameters), the transformation is extended to H₃ accordingly. After neural operator inference, we map the predicted local solution back via an inverse transformation T... We implement these as preprocessing and postprocessing steps in the inference pipeline (see Appendix E.
5

---

Published as a conference paper at ICLR 2026
Time Complexity. Suppose the single inference time of local operator and the number of iterations are denoted as and N. Let v, e, denote the number of vertices, edges and subdomains respectively. Our Algorithm I consists of two main parts: mesh partition using the METIS algorithm, the time complexity of which is approximately (v + + K log K) (Karypis and Kumar 1997); and an iterative scheme with time complexity roughly O(bKN). Therefore, the overall time complexity of our algorithm can be approximated as O + e K log K + 6K N In practice, may not be independent of and More vertices can sometimes lead to more partitions required depending on the property of mesh. While providing an exact time complexity analysis for FEM can be challenging due to the complexity and variability of different problem setups, is worth noting that FEM is generally considered to be computationally demanding.
3.4 THEORETICAL RESULTS
Here provide theoretical analysis of our proposed algorithm by stating the following result:
Theorem 1. Assume the operator C in Eq. is self-adjoint and coercive elliptic (Mathew 2008 Let u" be the iterate of the classical additive Schwarz-Richardse iterates of SNI in Algorithm ! initialized with the same ud = =..
iteration in Eq. 2 and let it be the
Let S& denote the exact local solution operator on Suk and define the learned (pre/post-processed)
local solver S₂ - Ty o Gt Tg. Assume:
(Classical contractivity) the classical iteration mapping is contractive in a norm || with rate 0 p 1;
(Uniform local error) (|$x - Sx S cfor all k;;
(Overlap factor) each degree of freedom is covered by most I subdomains.
Then we have:
Convergence: SNI converges to fixed point u
Error bound: for all n, ||""" un c'. and
|| - u² S c', where may take
d' = Tt
The theorem suggests that if learned local operator maintains a uniform error bound, the algorithm converges and exhibits minimal approximation error. See Appendix C a proof. This result relies on the assumption on operator C. In general, such convergence validate the effectiveness of our framework for nonlinear differential equation through experiment.
guaranteed and empirically
EXPERIMENTS
In this section, we perform comprehensive experiments to showcase the effectiveness of our method on various challenging datasets.
4.1 EXPERIMENTAL. SETUP
Datasets. To demonstrate the scalability and superiority of method, construct several datasets on multiple PDEs. We also extend our framework to a time-dependent problem, heat conduction. To aggregate training sets, we generate random simple polygons bounded by the unit square [-0.5, 0.5) C 2 Boundary/initial conditions and coefficient functions are piecewise linear functions determined by random values within [0, 1). For each of the following problems, we test on datasets based on three different domains A, B and C shown in Figure 2 Details of these datasets are given Appendix H.1
With neural operators implementing línear transformers, e.g., GNOT applied in this work, b O($).

---

Published as a conference paper at ICLR 2026
Equation
Domain
GNOT(%)
SNI(%)
A
22+2
2.2++0.6
Dapiace2d-Dirichlet
B
22+2
2.1+0.4
C
28:13
2.1
A
10.7
6
Laplace2d-Mixed
B
10.74.0.8
7+1
C
38+6
6+1
A
16+1
8+2
Darcy2d
B
63-E3
8+2
C
167-88
5.440.66
A
11.5+00.6
5.3+0.2
Heat2d
B
304:10
1142
C
20
5.8.1..3
A
22+2
2.0+-0.4
Nonlinearl.aplace
B
26+22
2.2+0.4
C
28:12
2.2
Table 1: Main results. The la relative errors along with standard deviation over different random boundary/initial conditions on three domains are reported.
Laplace2d- Laplace equation in 2d with pure Dirichlet boundary condition on various shapes.
-aieeee2d-Mixed: Laplace equation in 2d with mixed Dirichlet and Neumann boundary condition on various shapes.
Darcy2d: Darcy flow 2d with coefficient field a(x), source term f(x) and pure Dirichlet boundary condition various shapes.
Heat2d: Time-dependent heat equation in 2d with a coefficient a for thermal diffusivity, initial condition and time-varying pure Dirichlet boundary condition various shapes.
NonlinearLaplave2d:: A nonlinear Laplace equation in 2d with pure Dirichlet boundary condition various shapes.
Baseline. Our baseline is direct inference of the trained neural operator on domains shifted and scaled [-0.5, 0. with boundary/initial conditions and coefficient functions adjusted accordingly.
Evaluation Protocol. The evaluation metric we utilize is the mean /2 relative error. See Appendix H.2 for details.
4.2 MAIN RESULTS AND ANALYSIS
The main results for all datasets are shown in Table I More details and hyperparameters are summarized Appendix H.2 due to limited space. Based on these results, we have the following observations.
0.5
Figure 2: Illustration of experiment domain A, B, C from left right respectively.
7

---

Published as a conference paper at ICLR 2026
Stationary Problems. First, find that method performs significantly better on all stationary problems compared with baseline. On all domains, we reduce prediction error by 34,8 34.8%-96.8% The excellent performance shows the effectiveness of our framework in dealing with arbitrary geometries unseen during training. In partieular, our framework usually leads by a larger margin on more complicated domain, due the fact that simple polygons used in the training data fail adequately resemble the complex testing domains. Solutions on multiply connected domains usually exhibit characteristics that are not present on simple domains.
Second, we find that the the performance of our method is consistent across various geometries during inference. On all types of PDEs in our datasets, the difference in prediction error over various geometries is within 3.25%, showing the ability to solve PDE with consistent accuracy on various geometries with single trained neural operator. This also provides evidence for our theoretical result in Theorem I where we show that the SNI ensure the convergence an approximation of the ground-truth solution with error bound determined by the generalization error of the neural operator.
Third, we find that complexity of the PDE together with types of boundary condition affect the generalizability of the neural operator in solving local problems and thus also the accuracy of our method. For simple problem such Laplace2d-Dirichhes our method achieve 59.8% lower error compared to other problems. For Laplace2d-Min Mixee,,, neural operator struggles to capture subtlety in presence of both Dirichlet and Neumann boundaries. The complexity of Darcy2d lies in the need to capture changes in coefficient and source term in addition to geometry and boundary condition. We argue that having strong neural operator that can generalize well all basic shapes and boundary conditions is necessary for our framework to work with reasonable accuracy.
Timeedependent Problems. There natural way extend our framework time-dependent problems (Li and Cail 2015) where a space-time decomposition is constructed by taking the product of spatial decomposition and a temporal decomposition. We train a neural operator that can predict heat conduction on multiple time steps and the same SNI is applied during inference on this 3d problem. Our framework works well on this problem and reduce prediction error by 54. 1 14.... This demonstrates the potential of our framework to handle dimeddependent problems. We refer to Appendix D for detailed implementation.
Data Efficiency. The exploration results on data efficiency of SNI are shown in Figure 3 implying the following observations: (1) At all abundances data, the l2 relative errors SNI are significantly lower than those of GNOT direct inferences; (2) Errors of SNI are comparable to even lower than validation errors at large data volumes. (3) SNI requires much smaller datasets to achieve comparable results to GNOT direct inference. Overall, these results demonstrate that SNI has substantial advantages in terms of data efficiency. Our proposed framework possesses remarkable ability extract more insights from limited data and scale more effectively as data volumes increase. More supplementary results are provided in Appendix H.4
4.3 ABLATION EXPERIMENTS
Hyperparameter Exploration. The number of partitions (K), the depth of extension (d) and step size (r) are the key main hyperparameters that can affect the performance of SNI. Based the
relative entor
0.2
0.33
0.2
relattvve
relative
8022
0.1
0.1
0.1
GNOT
0.3
0.5
1.0
1.5
2.0 x10
0.0
0.5
1.0
1.5
2.5
0
0.5
1.0
1.5
2.0
Training set size
Training set size
x10*
Training set size
K10*
(a) Lap2d-D domain A
(b) Lap2d-D domain B
(c) Lap2d-D domain C
Figure 3: Comparison between the lo relative errors from SNI (blue), GNOT direct inference (orange) on Laplace2d- Diriichet upon three domains (A, B and C) with different numbers of training samples. The results of SNI and GNOT direct inference are presented based on 100 inferences with different boundary conditions. The best validation emrors during training (red) are also provided reference.
8

---

Published as a conference paper at ICLR 2026
20 parts
parts
depth
depth
relagive error
tay 0.02
0.6
parts parts
0.6
depth
depth
tau 0.04
50 parts
902
025
0.0
1000 2000
3000 4000 5000
0.0
1000 2000 3000 4000 5000
0.00
1000 2000 2000 4000 3000 Nubmer of iterations
Nubmer of Iteratins
Nubmer lterations
(a) = 2, r 0.04, varying K
(b) K 20, = 0.04, varying
(c) K = 20, d = 2, varying r
Figure 4: Comparison between convergence rate of SNI Laplace2d-Dirichee domain A.
results presented in Figure factors analyzed have no significant impact on the accuracy of our algorithm, but they do influence the convergence rate. Specifically, increasing the number of partitions leads to a smaller l2 relative error but slower convergence. Once the partition number surpasses 20, the algorithm's 's final performances become comparable. Regarding the depth of extension, it does not affect the performance on the tested domains. The convergence curves for depth of extension 1, 2, 4, and are nearly identical. When it comes T, larger value results in faster convergence. However, it is important to note that there exists maximum limit 1/ K beyond which cannot be set.
Data Augmentation Exploration. To explore the effects of data augmentation, compare the performances of models trained with different degrees of data augmentation for Laplace2d-Dirichhe Dirrccchlet demonstrated in Table For models trained without data augmentation. the variation of performances on different domains is large, ranging from 2.8% to 4.4%. Specifically, it reports 4.4 + 1.6% l2 relative error on domain while this error can be reduced to 1.9 + 0.4% with a rotation+(O.8, 1] scaling augmentation. While rotation can generally be beneficial, the effectiveness of scaling can sometimes be limited or even detrimental. Hence, important to apply data augmentation with caution and consider its suitability for different types of PDEs.
Choice of neural operator architecture. To explore the choice neural operator architecture in our framework, we train Geo-FNO (Li et al. 2023) Laplace2d-Dirichlet Dirichlet and apply SNI for inference on domains A, B and C to get le relative error of 3%. 13 = 1% and 13 3%. This result is comparable to that achieved by SNI with GNOT and demonstrates that works with various choices neural operator architecture. However, an error gap does exists between SNI with GNOT and GNo-FNO due to variations in their generalizability. This is also reflected in their respective best validation errors, as detailed Appendix H.4 Supplementary results on Darcy2d are also provided there.
proposed framework
5 RELATED WORK
Operator Learning. The idea of operator learning is first introduced in Lu et al. (2019). This work proposes a notable architecture called DeepONet, which employs a branch network for processing input functions and a trunk network for handling query points. Adopting the munk-branch architecture and utilizing the attention mechanism, Hao et al. (2023) develops GNOT to handle irregular mesh, multiple input functions, and different input data types. The high accuracy and versatility makes GNOT the benchmark in our work. In the other direction, Fourier neural operator (FNO) (Li et al. 2020b) leverages the Fast Fourier Transform (FFT) to learn operators in the spectral domain, and
Validation %))
Domain A(%)
Domain B(%)
Domain C(%)
No Data Aug
3.79
4+2
3.0++0.6
3+1
Rotation Only
2.50
2.240.66
2.1:1.44
2.1.10.9
[0.2, 1]
5.31
4+11
3.4.10.5
3.4
Rotation + Scale
[0.5,
3.62
2.740..5
3.7.1..6
3.2-10.6
[0.8. I]
2.86
1.8:10.4
3.3=0.7
2.8+0.8
Table 2: Comparison between models trained with different data augmentations for Laplace2d- Dirichlet.

---

Published as a conference paper at ICLR 2026
achieves a favorable trade-off between cost and accuracy. Variants of FNO are proposed to reduce computational cost (FFNO in Tran et al. (2021)). handle irregular mesh (Geo-FNO in et al. (2023) ),
and improve expressivity (UFNO in Wen et al. (2022).
Methods to Deal with Complex Geometry. Several approaches have been proposed to tackle the challenge of complex geometry and save the process efforts in operator learning. One
decoder framework called CORAL (Serrano et al.) 2023) is able to encode a complex geometry into lower dimensional representation to save the computational efforts and solve different types of problems, In (Wu et al., 2024), one mechanism called physics attention is proposed to aggregate complex input geometry and functions into several physics-aware tokens to reduce the number of tokens to deal with. AROMA (Serrano et al., (2024) introduces a diffusion refiner in latent space to solve temporal problems with complex geometries.
Domain Decomposition Methods Applied in Deep Learning. In general. the integration of deep learning and DDMs can be categorized into two groups (Heinlein et al.] 202T]] Klawonn et al.) 2024). The first category involves using deep learning techniques to improve the convergence properties or computational efficiency DDMs. For instance, Mao et al. (2024) proposes to combine operator learning with DDMs uniform grids in order to accelerate traditional DDMs. Several methods (Heinlein et al. 2020 2019) have also been proposed to reduce the computational cost in adaptive FETI-DP solvers by incorporating deep neural networks while ensuring the robustness and convergence behavior. The second category is centered around the substitution of subdomain solvers in DDMs with neural networks. There have been multiple endeavors to employ PINNs (XPINNs (Jattap and Karniadakis 2020), parallel inference with cPINNs and XPINNs (Shukla et al. 2021)) Deep Ritz methods alternatives to subdomain solvers or discretization techniques in traditional DDMs (Li al. 2020a 2019; Jiao el al. (02TT). These approaches leverage the universal
approximation capabilities of neural networks represent solutions of PDEs, subject to specific assumptions regarding the activation function and other factors.
Data Augmentation Techniques in Operator Learning. Different types of data augmentations are proposed to improve the generalization capabilities in operator learning. A Lie point symmetry framework is introduced in [Brandstetter et al. (2022). which quantitatively derives comprehensive set of data transformations, to reduce the sample complexity. Motivated by this approach, Mialon et al. (2023) learn ceneral-purposs representations of PDEs from heterogeneous data by implementing joint embedding methods for self- self-supervised learning. An alternative research approach (Fanaskov et al.) 2023) introduces a computationally efficient augmentation strategy that relies on general covariance and straightforward random coordinate transformations. In general, applying data augmentation techniques for PDE operator learning can be challenging due to the unique nature of PDE theory.
6 CONCLUSION AND FUTURE WORKS
We presented a local-to-global framework based on DDMs to address the geometry generalization and data efficiency issue in operator learning. Our framework includes novel data generation scheme and an iterative inference algorithm SNI. Additionally, we provided theoretical analysis of the convergence and error bound of the algorithm. We conducted extensive experiments demonstrate the effectiveness of our framework and validate our theoretical result. For future works, the rich literature of DDMs when combined with operator learning provides many potential directions to handle higher- dimensoai problems, non-overlapping decomposition and more challenging types of equations.
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
Detailed descriptions the experimental setup, task definitions, and evaluation metrics are provided in section and Appendix H Source code is attached in the submission.
REFERENCES
Auery, T. and Heldz, M. (2019). Rpg software package for the generation of random polygons.
Baratta, I. A., Dean, P., Dokken, J. S., Habera, M., Hale, J. S., Richardson, C. N., Rognes, M. E., Scroggs, M. W., Sime, N., and Wells, G. N. (2023). DOLFINx: the next generation FEniCS problem solving environment. preprint.
Brandstetter, J., Welling, M., and Worrall, D. E. (2022). Lie point symmetry data augmentation for neural pde solvers. In International Conference on Machine Learning, pages 2241-2256. PMLR.
Evans, L. C. (2022). Partial differential equations, volume 19. American Mathematical Society.
Fanaskov, V.. Yu, T., Rudikov. A., and Oseledets. I. (2023). General covariance data augmentation for neural pde solvers. arXiv preprint arXiv:2301.12730.
Gander, M. al. (2008). Schwarz methods over the course of time. Electron. Trans. Numer. Anal, 31(5):228-25 255.
Geuzaine, C. and Remacle, J.-F. (2008). Gmsh: a three-dimensional finite element mesh generator with built-in pre-and post-processing facilities.
Hao, Z.,, Wang, Z., Su, H., Ying, C., Dong, Y., Liu, S., Cheng, Z.., Song, J., and Zhu, J. (2023). Gnot: A general neural operator transformer for operator learning, In International Conference Machine Learning, pages 12556-12569, PMLR.
Heinlein, A., Klawonn, A., Lanser, M., and Weber, J. (2019). Machine learning in adaptive domain decomposition methods predicting the geometrie location of constraints. SIAM Journal Scientific Computing, 41 41(6))33877-3311
Heinlein, A., Klawonn, A., Lanser, M., and Weber, J. (2020). Machine Learning in Adaptive FETI-DP-A Comparison of Smart and Random Training Data. Springer.
Heinlein, A., Klawonn, A., Lanser, M., and Weber, J. (2021). Combining machine learning and domain decomposition methods for the solution of partial differential equations-a review. GAMM- Mitteilungen, 44(1) (4(1):=202100001.
Jagtap, A. D. and Karniadakis, E. (2020). Extended physics-informed neural networks (xpinns): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations. Communications in Computational Physies, 28(5):2002 2041.
Jiao, A., He, H., Ranade, R., Pathak, and Lu, L. (2021). One-shot learning for solution operators of partial differential equations. arXiv preprint arXiv:2104.05512.
Karypis, G. and Kumar, V. (1997). Metis: A software package for partitioning unstructured graphs, partitioning meshes, and computing fill-reducing orderings of sparse matrices.
Klawonn, A., Lanser, M., and Weber, (2024). Machine learning and domain decomposition methods-a survey, Computational Science and Engineering, 1(1):2
Kovachki, N., Li, Z., Liu, B., Azizzadeneshelli K., Bhattacharya, K., Stuart, A., and Anandkumar, A. (2023). Neural operator: Learning maps between function spaces with applications to pdes. Journal of Machine Learning Research, 24(89):1-97.
Langtangen, H. P. and Logg, A. (2017). Solving PDEs in python: the FEniCS tutorial I. Springer Nature.
Li, K., Tang, K., Wu, T., and Liao, Q. (2019). D3m: A deep domain decomposition method for partial differential equations. IEEE Access, 8:5283-5294.

---

Published as a conference paper at ICLR 2026
Li, S. and Cai, X.-C. (2015). Convergence analysis of two-level space-time additive schwarz method for parabolic equations. SIAM Journal on Numerical Analysis, 53(6):2727-2751.
Li, W., Xiang, X., and Xu, Y. (2020a). Deep domain decomposition method: Elliptic problems. In Mathematical and Scientific Machine Learning, pages 269-286. PMLR.
Li, Z., Huang, D. Z.,, Liu, B., and Anandkumar, A. (2023). Fourier neural operator with learned deformations for pdes on general geometries. Journal of Machine Learning Research, 24(388):1- 26.
Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., Bhattacharya, K., Stuart, A., and Anandkumar,
A. (2020b). Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895.
Li, Z., Kovachki, N., Choy, C., Li, B., Kossaifi, J., Otta, S., Nabian, M. A., Stadler, M., Hundt, C., Azizzadenesheli, K., et (2024). Geometry-infornet neural operator for large-scale 3d pdes. Advances in Neural Information Processing Systems, 36.
Liu, G.-R. and Quck, S. S. (2013). The finite element method: a practical course. Butterworth- Heinemann.
Liu, S., Hao, Z., Ying, C., Su, H., Cheng, Z., and Zhu, (2023). Nuno: A general framework for learning parametric pdes with non-uniform data. In International Conference on Machine Learning, pages 21658-21671, PMLR.
Lu, L., Jin, P., and Karniadakis, G. E. (2019). Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators. arXiv preprint arXiv: arXiv:1910.03: 1910.03193.
Mao, C.,, Lupoiu, R., Dai, T., Chen, M., and Fan, J. A. (2024). Towards general neural surrogate solvers with specialized neural accelerators. In Proceedings of the 41st International Conference Machine Learning, (CML'24. JMLR.org.
Mathew. P. (2008). Domain decomposition methods for the numerical solution of partial differential equations. Springer.
Mialon, G., Garrido, Q.. Lawrence, H., Rehman, D., LeCun, Y., and Kiani, B. (2023). Self-supervised learning with lie symmetries for partial differential equations. Advances in Neural Information Processing Systems, 36:28973-29004.
Pfaff, T., Fortunato, M., Sanchez- Gonzalez, A., and Battaglia, P. (2020). Learning mesh-based simulation with graph networks. In International conference on learning representations.
Preparata, F. P. and Shamos, M. I. (2012). Computational geometry: an introduction. Springer Science & Business Media.
Richardson, L. (1911). Ix. the approximate arithmetical solution by finite differences of physical problems involving differential equations, with an application to the stresses in masonry dam. Philosophical Transactions of the Royal Society of London. Series Containing Papers of Mathematical or Physical Character, 21 2100459-470:3307-3337
Scroggs, M. W., Baratta, I. A., Richardson, C. N., and Wells, G. N. (2022a). Basix: runtime finite element basis evaluation library. Journal of Open Source Software, 7(73):3982.
Scroggs, M. W., Dokken, J. S., Richardson, C. N., and Wells, G. N. (2022b). Construction of arbitrary order finite element degree-of-freedon maps on polygonal and polyhedral cell meshes. ACM Transactions on Mathematical Software, 18:))18::111822
Serrano, I.., Le Boudec, I.., Kassaï Koupaï, A., Wang, T. X., Yin, Y., Vittaut, J.-N., and Gallinari, P. (2023). Operator learning with neural fields: Tackling pdes on general geometries. Advances in Neural Information Processing Systems, 36:70581-70611
Serrano, L., Wang. T. X., Le Naour, E., Vittaut, J.-N., and Gallinari, P. (2024). Aroma: Preserving spatial structure for latent pde modeling with local neural fields. Advances in Neural Information Processing Systems, 1352-1
12

---

Published as a conference paper at ICLR 2026
Shukla, K., Jagtap, A. D., and Karniadakis, G. E. (2021). Parallel physics-informed neural networks via domain decomposition. Journal of Computational Physics, 407:1
Toselli, A. and Widlund, O. (2004), Domain decomposition methoos--llorrthmss and theory, vol- 34, Springer Science & Business Media.
Tran, A., Mathews, A., Xie, L.., and Ong, C. (2021). Factorized fourier neural operators. arXiv preprint arXiv:2111.
Wen, G., Li. Z., Azizzadenesheli, K., Anandkumar, A., and Benson, S. M. (2022). U-fno an enhanced fourier neural operator-based deep-learning model for multiphase flow. Advances in Water Resources, 163:104180.
Wu, H., Luo, H., Wang, H., Wang, J., and Long, M. (2024). Transolver: A fast transformer solver for pdes on general geometries. arXiv preprint arXiv:2402.02366.
13

---

Published as a conference paper at ICLR 2026
A BACKGROUND ON DOMAIN DECOMPOSITION
Domain decomposition is a widely used technique in computational science and engineering that enables the efficient solution of large-scale scale problems by dividing the computational domain into smaller subdomains. This approach is particularly beneficial when dealing with complex problems that cannot be solved using a single computational resource. The main idea behind domain decomposition is to break down large computational domain into smaller, more manageable subdomains. These subdomains can be arranged in a variety of ways, such as overlapping or -verlaaaa,,,, depending on the specific problem and the desired computational approach.
In this work, we decompose our domain into subdomains and adopt the hybrid formulation of Eq. following Mathew (2008) Section 1.1). A decomposition of ? is collection of open subregions {5 22 S S for k I..... K such that U S2 12. This decomposition is referred to as non-overiapping if in addition, S4 n n = Ø for any i # j. Alternatively, an overlapping decomposition one satisfying U Qg Ω. Typically, a non-overlapping decomposition is one where subdomains do not intersect with each other in the interior while an overlapping decomposition constructed in practice has overlapping neighboring subdomains.
Given a decomposition of Ω2 a hybrid formulation of Eq. D a coupled system of local PDEs on subdomains Q4 equivalent to Eq. [ satisfying two requirements. First, the restriction ux(²) of the solution u(xz) of Eq. [ to each domain 24 must solve the local PDE, thus ensures that the hybrid formulation is consistent with the original problem in Eq. [ Second, the hybrid formulation must be well posed as coupled system of PDEs in the sense of Evans 2022), i.e. its solution must exist, be unique and depend continuously on given input function and boundary/initial conditions. Intuitively, hybrid formulation consists of a local problem posed on each subdomain and matching conditions that couples the local problems.
In this work we focus on the earliest and most elementary formulation termed Schwarz hybrid formulation (Mathew] 2008 Section 1.2) based on overlapping decomposition and is applicable to wide class of self-adjoint and coercive elliptic equations. Given overlapping decomposition, 022 can be decomposed into two disjoint parts. One (possibly empty) part I& = 0S2, n an is located in the boundary of Ω and the global boundary condition should be imposed. The other part B₁ = 02g N is nonempty artificial boundary from the overlapping decomposition and a Dirichlet boundary condition from the coupling of local problems is imposed.
We refer to Mathew (2008) for a strict definition. As an illustrative example, assume we have an overlapping decomposition with K 2 and consider as the original problem Laplace equation with mixed Dirichlet and Neumann boundary conditions. The following coupled system of two local PDEs is Schwarz hybrid formulation of the original problem and solving the original equation is equivalent to solving this coupled system.
Au1 = 0
Au2 = 0
Ω2 on 02g 0
on 02, Uu 01 U2 10,
U22 ui 180,
U1 = up
on a2, nID
and
122 = uD
on 0Z2 n
du
du2 On
g on 02, n'N
g on 02g n''
On
Based on the Schwarz hybrid formulation, there are various iterative schemes with different paral- lelism and convergence rate. In the subsequent discussion, focus is primarily introducing the additive Schwarz methods (ASM). The ASM is a highly parallel algorithm (Mathew) (008)) in solving the coupled system from Schwarz hybrid formulation. We briefly introduce ASM with finite element methods and refer to Gander et al. (2008 and Mathew (2008) for details.
Assume that under weak formulation of Eq. ann finite element space V. Eq. T has the form Au where A is the stiffness matrix. Given an overlapping decomposition {2,} compatible with the
finite element space on Ω we have V = Ek= VA sum of local finite element subspaces Vi2 on Pg and we can define local stiffness matrices A& Va VR. restriction operators (R& )k_ restricting V to V& and extension operators [RT]] ] extending Va to V by zeros extension. We then define operators Pe V - V by Pt- R A R A. Additive Schwarz operator is then defined as the
14

---

Published as a conference paper at ICLR 2026
sum Pad Pa - Pk. This operator can be show be self-adjoint and coereive and have the following equivalence.
K Σ
Au = f Pa
R
(4)
k=1
We note that the right hand side of Eq. is preconditioned version of the left hand side. The Richardson iteration for this preconditioned problem has the following form.
K
u n++ u² +
R A
R& -Au)
(5)
In the composite operator A Rg. the operator Rg first restrict a function to S2 A solve the local problem and R extend the local solution to f. This iterative process can be shown to converge by estimating bound on condition number of Pas under mild assumptions on equation and decomposition.
B REVISIT ON OPERATOR LEARNING
The goal of operator learning is learn mapping 9 A U between two infinitely dimensional spaces (Kovachki et al.] 2023). When applied to PDEs, / is the solution space of a PDE and A is the space functions that determine unique solution of PDE. Examples of A are coefficient functions or boundary/initial conditions that defines the PDE and parameters that determine the geometry of domain.
In our study, we decompose any domain into subdomains each of which lives in a distinguished class of basic shapes P. We assume all shapes in have Lipschitz boundary and are uniformly bounded, i.e., they are all bounded by a ball D) C R We are interested in solving boundary value problems in Eq. [ any domain f2 E with any appropriate boundary condition. We thus separate geometry and boundary conditions from other inputs and represent the input function space of the operator as P x (D) x where (D) is the Sobolev space (D). The space x H² represents the geometry of the domain together with boundary/initial conditions, H represents any other input functions such as coefficient function field or source term in the PDE. The neural operator thus approximates the following mapping. Note that in the case of time dependent problem, the space H*(D) D) represents the space of initial condition together with time varying boundary condition and the solution space L/ represents time series up to time span. The solution operator G thus has the following form.
P:P x H H
(6)
For learning the operator, we assume P, H² (D) and H. are probability spaces and thus we can sample observations from A. practice, we randomly sample geometry from P and random boundary conditions are imposed, then solution is generated from numerical solver to get solutions. It is important to highlight that, unlike the usual setting for neural operators, there is significant variation in the shape of input domains.
C PROOF THEOREM
Theorem. Assume the operator C in Eq. 7is self-adjoint and coercive elliptic (Mathew 2008 Let w" and nh denote the iterates of the classical additive Schwarz-Richandson iteration (Eq. and SNI (Algorithm 7) respectively, with the same initialization ui .
Let Sy denote the exact local solution operator on O2 and define the learned (pre/post- local solver S, 2 Тī o G o Ty. Assume the classical iteration mapping is contractive in a norm || with rate 0 < < 1. and the uniform local error bound [|5a(-) - Sx(-)] s holds for all k. Let t be the maximum number of subdomains covering any degree of freedom. Then we have:
Convergence: SNI converges fixed point ut.
15

---

Published as a conference paper at ICLR 2026
Error bound: for all n, ||@" uIll d' and || u²Il I c. where may take C.
Proof. (1) Write the classical and learned iterations in the same additive Schwarz form:
K
w n+
u²" +
R (Sx("") - R&u"),
(7)
ildeu u1" + (S((()))) R₂
where Sz 2 TT G o Th.
(2) Let - n" u², Subtracting the two updates in Eq. yields
K
-+1=
(F(i") (++==(F(ū))
F(u")) +
-
where F(u) + K k= R (S) Rx is the classical iteration mapping. By the contractivity assumption of the classical iteration, (T*) F(u")ll & p||e"||. Moreover, using the uniform local approximation error bound and the overlap factor (the maximum number of subdomains covering any degree freedom), we have
R (Sx(i)) Sa()))
rtc.
1=
Therefore,
|e S p|lc"1| Tlc.
lterating the inequality gives ||e"|| < L=p"
T-p TLC. Taking d'
stc I-p yields ||ū"" uIll S for all n.
(3) Since ut converges by assumption, it is Cauchy. For any m > n, we have
l ū"ll < |u" wlll + a n + 1 n" < || u" + 2c'.
Hence i is also Cauchy and thus converges to some ä". Passing to the limit in || un! S d
yields ||u u² S c', completing the proof.
we apply matrix form of neural operator, namely, the neural operator aims to approximate ===1 and assume, then we can have the following result:
Corollary Consider the exact operator AE' and inexact neural operator Az =.,k= K. Let n and nn represent the solutions updated by A and Å respectively at the n-th step, where the updating rule is given by Eq. 5 with T and both sharing the same initialization, Suppose that
ll - ||< c, for k = I,... K, and p(T MA) < 1, where M === =E R. Ar Re. then
have:
Convergence: the algorithu converges to a fixed point;
Error bound: there exists a constant (c) such that un un ||< MA-
Condition number: x(MA) S minnt(( - 1), 1 + max; H
where I, K, Hy and d denote the maximal number of overlapping subdomains, the number of subdomains, the diameter of k-th subdomain, and the number of extensions, respectively.
Note that the condition p(I MA) < AA) employ the Richardson iteration trick (Richardson, 19IT) in order to ensure the convergence of the proposed algorithm (Al (Algorithm]]
is generally challenging to satisfy. To address this issue,
16

---

Published as a conference paper at ICLR 2026
D TIME-DEPENDENT PROBLEMS.
We consider the time-dependent PDE with the following form:
Ut Lu f u(2, t) (2, u(x, 0) uz(
[0,
on 0n x [0,
(8)
on 52 {0}
where C is again self-adjoint and coercive elliptic operator. The additive Schwarz method can be naturally extended to a space-lime additive Schwarz method (Li and Cail 2015) by consider- ing decomposition of the space-time domain 92 [o, [] by taking the product of overlapping decomposition of ? and [0, Ij respectively. The space-tima domain decomposition has the form
Q x | 5T, 1 + 6r] where 8T the temporal depth and represent overlap in time domain. Once such a decomposition is constructed, the same additive Schwarz method can be applied to the space-time decomposition
get a global solution on the space-time domain, allowing parallel iteration in both space and time domain. Local problems for the above decomposition are again of the form in Eq. B
In our implementation on heat equation, we discretize the time domain with fixed time step 4. fix rollout length of and train a neural operator to map initial and boundary conditions to time series for the k steps at = 0, tg..... (k 1) More precisely, the neural operator is trained to map up and uo to time series of the form ux: 0), u(x, ts)," ... u(x, (k 1)t,).
E SYMMETRIES OF PDES
The symmetry group of a general partial differential operator C refers to a set of transformations that map a solution to another solution, forming mathematical group. Lie point symmetry is subgroup of the symmetry group that has a Lie group structure and acts on functions pointwise as transformations on coordinates and function values (Brandstetter et al. 2022). In this work, we will in addition be concerned with not just single operator L, but a family of operators depending on various coefficient fields (e.g., Darcy flow) and various boundary/initial conditions. Symmetries have to properly extended to these input functions so that solution with an input function is transformed to another solution with a different input function.
Leveraging these symmetries allows for the generation of an infinite number of new solutions based on given solution. The idea of utilizing these symmetries as a data augmentation technique for operator learning was initially introduced in |Brandstetter et al. (2022). However, we apply these data augmentation solutions on basic shapes in training local operator and this usage of symmetries echos point mentioned in Brandstetter et al. (2022) Section 3.2) where the authors point out that these data augmentation can be applied on local patches of solutions instead of the solution on the entire domain.
There is another direct usage of symmetries in our framework. Instead of incorporating symmetries as a form of data augmentation in training time, one can directly apply transformations to input and output of neural operator during inference time. We implement these transformations as preprocessing and postprocessing steps in the inference pipeline. We summarize the symmetries of each PDE applied in our implementation in Table Normalizations applied as preprocessing and postprocessing for each of the equations are summarized in Table
F
TIME COMPLEXITY
F.1 EMPIRICAL TIME COMPLEXITY
We provide empirical results on runtime for our main results. We discuss how better initialization can accelerate the whole iterative process in the next point.
We first provide empirical runtime on single sample for each stationary problem and each domain reported in our main result. We use the following metrics:
Time convergence (TTC).
17

---

Published as a conference paper at ICLR 2026
Equation
Lap2d-D
Lap2d-M
Darcy2d
Heat2d
NonLap2d
Spatial
(01, (x) f1, 12)
Shift
Spatial
(211, x22 + (x) cos 1/2 sin 0, 21 sin + 22 cos
Rotation
(11. 2222 7 (s2 S22)
Spatial
U. 7
uD uD
0 + s²u
u + u U0 +
Scaling
UD - SuD
a(I) + a(r)
uD + UD s²a
9 0 9
[(z f(x)
Value
- +
Shift
Up + up + L
Value
Scaling
uD + sup
g - 5g
Ug
sug
Table 3: Symmetries of various PDEs applied in our implementation.
Time 15/10/5% relative error following the practice in Mao et al.
2024
Equations
Domains
TTC(s)
TS15%(s)
TT10%(()
TT5%(s)
Soplace2d-Dirichhie
100
70
Diiiice2d-Diricclle
B
269
107
137
194
Laplace2d-Dirichlet
C
28
8
11
17
Laplace2d-Mixed
A
162
103
137
Laplace2d-Mixxn
B
714
511
620
Laplace2d-Mixx
C
68
43
52
Darcy2d
A
84
Darcy2d
B
247
144
176
Darcy2d
C
26
12
15
Table Empirical runtime for different equations and domains
Factors that affect the runtime are:
1. Type of equations. We observe that Laplace2d-Mixed takes longer on all domains. We also observe that the existence of Neumann boundary condition leads larger range of function values for the solution of Laplace2d-Mixx This leads to more iterations steps required to reach convergence.
2. Number of subdomains K and step size In the above table, domain B takes longer for all equations because has 40 subdomains compared to 20 for and C. A large number of subdomains leads more time consumption for iteration. We illustrated how choice of affects the number of iterations to convergence in section 4.3 our paper.
3. Local operator architecture. While GNOT gets better results in accuracy, drawback of transformer-bas methods is that they are usually slower than FNO (Hao et al. 2023).
4. Initialization. This is discussed in the next point.
We note that our implementation is not optimized to fully parallelize the iterative process; for example, the normalization process is not parallelized in implementation.
Numerical solvers are very fast generating solutions for the domains we tested on and we do not expect our approach to be faster than these highly optimized numerical solvers these (still) simple domains. As a reference, generating a solution for Laplace2d- 2iiricchhe on domain A using classical FEM solution takes 6.15 10 seconds and performing GNOT inference the same domain takes 1.26 10 seconds. We can see that even classical numerical solver is faster than GNOT inference. However, DDMs conventional approach implemented in commercial software
18

---

Published as a conference paper at ICLR 2026
designed to solve PDEs on large-scale and complicated domains. We replace the local FEM solver in DDMs by a data-driven neural operator and thus expect our approach to show superiority when the problem domain is large and complicated.
F.2 ACCELERATION THROUGH BETTER INITIALIZATION
We discuss how to accelerate the iterative process by starting with a better initialization. In the original implementation, we always start with zero solution in the interior of the domain. To accelerate the process, we initialize with solutions from GNOT direct inference and find that it considerably saves our time. We report the time consumption on Saslace2d-Dirichhee using the same metrics as the previous table. The only difference is in initialization.
Equations
Domains
TTC(s)
TT15%(s)
TT10%(s)
TT5%(s)
haplace2d-Dirichiee
A
28
11
17
Daolace2d-Dirichaee
B
70
Laplace2d-Dirichle
C
20
2
7
Table 5: Runtime results with improved initialization for Laplace2d-Dinn -iirichlet equations
However, coming up with better initialization is not trivial and can be an interesting future work.
G DISCUSSIONS
Message passing in DDMs. In our framework, we solve a coupled system of local problems by an iterative algorithm SNI. Through iteratively solving local problems based on boundary values from the last iteration and thus from neighboring subdomains, SNI is essentially performing message passing between subdomains. This message passing operation may be implemented in other forms, e.g., through graph neural network.
Higher-dimennional PDEs. Our framework can be extended to higher-dimensional cases as long as basic shapes and corresponding solutions can be properly generated. For 3-d problems, one potential selection of basic shapes is the class of polytopes.
The fundamental "local-to-global" principle of our method remains unchanged: decompose complex 3D domain into smaller, canonical 3D subdomains, solve locally with a neural operator, and iteratively stitch the solutions. However, implementation faces additional challenges in 3D. For example, generating building blocks in 3D. One way is to use polyhedrons as building blocks in 3D. Another choice is to generate building blocks from "growing" tetrahedra in 3D. Both be very interesting ways to explore. In summary, extending our method to 3D is a challenging but highly promising research trajectory. It is not simple "plug-and-play" extension but requires careful research the intersection of computational geometry, operator learning, and high-performance computing.
Other formulations of DDMs. Schwarz hybrid formulation discussed in this work is of the most elementary formulation in DDMs. There are many other more ad advanced DDMs (Mathew] 2008). Sieklov-Poincard framework based on non-overlapping decomposition and transmission condition as coupling condition for local problems. Langrange multiplier framework leads to the well-known FETl method and is also based on nor-overlgp decomposition.
Other types of PDEs. The additive Schwarz method in classical DDMs works for self self-adjoint and coercive elliptic equations. Non-self-adjoin elliptic equations, parabolic equations, saddle- saddle-point problems and non-linear equations requires separate treatment. Addressing these cases presents challenges in both training the local operator and designing the iterative algorithm.
Future Works. Based on the above discussion, there are many potential directions for future works. First, would be interesting to implement this framework using a message-passing framework instead of an iterative algorithm accelerate the convergence. Second, extending our framework to address digher-dimensional problems important, particularly since industrial problems often involve 3-d simulation. Third, more advanced DDMs such as 2008) may also be explored. Lastly, other types of PDEs such as saddle point problems and non-linear equations such Havier-Stoker equation is out of the scope of current work, and present unique
BDDC and FETI (Mathew
19

---

Published as a conference paper at ICLR 2026
challenge. Tackling these challenges requires not only expertise operator learning, but also deep understanding of PDEs themselves. We speculate that it would be fruitful to combine rich literatures of DDMs with operator learning.
H EXPERIMENTS
H.1 DATASETS
Here we introduce more details of our datasets in both training and testing stage. For training data, we generate random simple polygons with mesh using Gmsh (Geuzaine and Remacle 2008). We prepare separate dataset for validation during training. For testing data, we generate the three domains depicted in Figure togethher with mesh using the Gmsh UI. We argue that the complexity of geometric domains is fundamentally determined by their underlying topological and geometrical properties. Based on this intuition, we considered three domains of increasing complexity for evaluation: (I) Domain A: This domain is simply connected, representing the simplest class of geometries; (2) Domain B: This domain has two holes and is multiply connected, indicating higher level of complexity compared to the simply connected Domain A; (3) Domain C: This domain has one hole with comnrs,, further increasing the geometrical complexity compared to the previous two domains. Through systematical evaluation across this spectrum of domains, from the simple geometry to more intricate multiply connected domains with holes and corners, believe the results provide comprehensive understanding of our framework's 's capabilities.
vertices within [-0.5, 0.5]2 and create uniform
Once the geometries and meshes created, we specify boundary/initial conditions for various equations and domains and generate solutions using FEniCSx (Baratta et al. 2023 Scroggs et al. 2022a b). popular open-source platform for solving PDEs with the finite element method (FEM). We adopt Lagrange element of order 1 (linear element) as our finite element space in generating boundary/initial conditions and solutions. Next we give details on how these boundary/initial conditions and solutions are generated for each type of PDEs. We also summarize these details in Table[andd
Aoplacedd-Dirichle Laplace equation in 2d with pure Dirichlet boundary condition. The governing equation
Au in Ω
up 80
For both training and testing data, we specify piecewise linear Dirichlet boundary condition with randomly generated values within [0, 1] on boundary nodes.
Laplace2d-MMxed. Laplace equation in 2d with mixed Dirichlet and Neumann boundary condition on NQ Ip U U.N. The governing equation is
Tp Tx
Figure 5: Illustration of mixed Dirichlet and Neumann boundaries for domain A, B, Laplace2d- Mixed.
20

---

Published as a conference paper at ICLR 2026
Au = in Ω
uD Tp
(10)
Ou On
on Tn.
For training data, 20% of the data have pure Dirichlet condition and the generation process is the same as in Saplace2d-Dirichlle The rest 80% of the training data have mixed boundary condition with non-empty connected Neumann boundary and Neumann boundary is randomly specified to be less than half of the entire boundary. Then a random number is sampled from (0.55 to specify functional range for Dirichlet and Neumann boundaries as described next. Among data with non-empty Neumann boundary, 50% have up E [0,r] and g E [0, 1] and the other 50% have up [0, 1J and g E r0,
For testing data, Dirichlet and Neumann boundary is specified for each of the domain A,B and C as shown in Figure Boundary conditions up and g are both piecewise linear with randomly generated values within [0, 1].
Darey2d. Darcy flow in 2d with coefficient field a(z), source term f(x) and pure Dirichlet boundary condition. The governing equation is
V(a(I)Vu)= = f in Ω
(II)
= up an
For training data, Dirichlet boundary condition is specified with a random range [0.3, 1] and boundary values are generated as up E [0. The coefficient function a(x) and source term (*) are specified as piecewise linear functions with randomly generated values within [0, 1] on nodes.
For testing data, Dirichlet boundary condition, coefficient function and source term are piecewise linear functions with randomly generated values within [0, 1].
Heat2d. Time-dependent equation of heat conduction in 2d with coefficient a denoting the thermal diffusivity, vary-narying boundary condition and initial condition. The governing equation is
tu a&u
in Ω x [0,
Øt
(12)
u(x, t) = up( t)
on ON x [0,
u(r, 0) uo (1)
on Ω (0)-
For training data, we discretize the time domain with fixed time step = 0.01, generate piecewise linear initials and time-varying boundary conditions with values randomly generated within [0, 1). a is random number within [0.8, 1]. We adopt the backward Euler method (Langtangen and Logg 2017) and generate a time series of 10 time steps. During training we separate these 10 time steps into 2 time series of times steps and training the neural operator to predict time steps.
For testing data, we fix a = 1. Initial condition is piecewise linear with values randomly generated within [0, 1]. Boundary condition is specified be constant over time and varied randomly within [0, across boundary nodes. We also adopt the backward Euler method and generate a time series of 50 time steps.
Nonlinearl MonlinearLaplaceed.. A nonlinear Laplace equation in 2d with pure Dirichlet boundary condition following an example in Langtangen and Logg (017). The governing equation
2ppacee2.
V ((u2 + 1)Vu) = 0
(13)
uD on 02.
For both training and testing data, we specify piecewise linear Dirichlet boundary condition with randomly generated values within (0, 1] on boundary nodes.
21

---

Published as a conference paper at ICLR 2026
H.2 EVALUATION PROTOCOL AND HYPERPARAMETERS.
Evaluation Protocol. The evaluation metric we utilize is the mean 12 relative enror. Let ug uf R" represent the ground truth solution and predicted solution for the i-th sample, respectively. Considering dataset of size D, the mean 2a relative error is computed as follows:
D D
04 l1
E
(14)
U1 12
Hyperparameters. All experimental hyperparameters used in the paper are listed in Table 6 For data generation, the number of vertices of simple polygons are uniformly chosen between 3 to 12. And x b in configurations denotes the generation of b shapes, each having distinct boundary/initial conditions. For investigating data efficiency issue, we only vary the number of various shapes while keeping the number of random input functions per shape constant. For boundary condition imposition, we summarize the details in Table 7 and 8
Computing Resource. We run our experiments
Tesla V100 GPU.
Lap2d-I>
Lap2d-M
Darcy2d
Hear2d
Nonlinearl.ap2d
Polygon
[3,12]
Training
10x2000
20x0000
10x4000
50x 1600
10x 2000
Configuration
Data
Validation
Generation
10x250
20x200
10x250
50x240
10x 2000
Configuration
Testing
100
100
100
10
100
Configuration
GNOT
expert and layers of width 128
Optimization
Adam
Learning rate
cycle learning rate strategy with 0.001
Operator
Epoch
500
1000
500
200
500
Learning
Rot+
Rot.+
Data Aug.
Rot.
No
Rot.
Sca. [0.8.1]
Sca. [0.8,1]
Time steps
5
Partition K
20
20 16
20
Depth d
2
Temp. Depth dT
Step size
0.04
0.002125
0.04
Spa. Shift+Scale
Spa. Shift+
Scift+
Pre/Post-pro.
Spa. Shift
Val. Shift+Scale
Scale
Val. Shift+Scale
Partition K
40
40x16
40
Depth d
2
Inference (SNI)
Temp. Depth &y
Step size
0.024
0.0014625
0.024
Spa. Shift+Scale
Spa. Shift+ Spa. Shift+
Pre/Post-pro.
Spa.Shift
Val. Shift+
Scale
Val. Shift+Scale
Partition K
20
20x 16
20
Depth
Temp. Depth dr
Step size
0.04
0.002125
0.04
Spa. Shift+Scale
(Spa. Shift+ Shift+Scale
Pre/Post-pro.
Spa. Shift
Val. Shift+Scale
Scale Val. Shift+Scale
Table 6: Key hyperparameters of main experiments. Configuration under Data Generation is specified as (number of random input functions per shape) x (number various shapes). Partition K for Heat2d is specified as (number of spatial partition) (number of temporal partition).
22

---

Published as a conference paper at ICLR 2026
PDE
Description
Lap2d-D
range of boundary condition: 0 1]
20% pure Dirichlet condition: range U 1]
40% mixed boundary condition with Tp/00 U 005,,
Lap2d-M
range of Dirichlet:
U(0, r] where ~~ 000.5
U 11, range of Neumann:U/0,, 1]
40% mixed boundary condition with /n/00 ~ U 00,,1
range of Dirichlet: U(0, 1), range of Neumann: 0 where U(0.5,
range of boundary conditions: U(0,r) where r U,0.3, 1].
Darcy2d
range of a(x) and f: [0,
Heat2d
range of initial/boundary condition: U|0, U
NonlinearLap2d
range of boundary condition: 0 1
Table Details of boundary/initial condition and input function generation in training data.
PDE
Description
Lap2d-D
range of boundary condition: U
I'D and I'N as in Figure 5
Lap2d-M
range of Dirichlet: U (1, IJ. range Neumann: U[0,1]
range of boundary conditions: 0,0
Darcy2d
range of a(x) and f(x): U[
Heat2d
range of boundary/initial condition: U[0,1 1] range of boundary condition: U00,
boundary condition do not vary with time
Nonlinearl.app ap22
Table 8: Details of boundary/initial condition and input function generation in testing data.
H.3 VISUALIZATION OF BASIC SHAPES AND DOMAIN DECOMPOSITION.
Visualization of basic shapes. We provide examples of generated basic shapes for training in Figure
Figure Examples of basic shapes.
Visualization of decomposed domains: A, B, and C We provide the visualization of decomposed domains A, B, C in 7 Please be noted that this is just rough visualization in that we do not correctly plot the overlapping part of subdomains. So this is only for an intuitive understanding the decomposed domain.
H.4 OTHER SUPPLEMENTARY RESULTS.
Data Efficiency. The results for data efficiency Laplace2d-Mixed and Darcy2d are shown in Figure The average performance of SNI is better than GNOT on all of three domains, while the margins between the two methods on domains A and B are not statistically significant due to the high variance in the la relative errors of SNI on these two domains. GNOT struggles the generalization to domain C, while SNI can still handle it with good performance.
23

---

Published as a conference paper at ICLR 2026
Figure 7: Visualization decomposed domains A, B and C.
relativv error
0.3
0.0
SNI GNOT
Vol
0.2
relative
0.2
relative
0.2
0.1
0.0
Training set size
x10
Training set size
Training set size
(a) Lap2d-M domain A
(b) Lap2d-M 22p2ddMMo domain B
(c) Lap2d-M - domain C
12 reletive arror
2.0
Voll
relative
2.5
relative
3.0
GNOT
0.2
1.5
0.5
0.0g
0.0
0.0
Training set size
x10
Training set size
x10
Training set size
x10
(d) Darcy2d domain A
(e) Darcy2d domain B
(1) Darcy2d domain C
Figure 8: Comparison between the l2 relative errors from SNI (blue), GNOT direct inference (orange) and validation (red) on Saplace2d-Mixet and Darcy2d upon three domains (A, and C) with different numbers of training samples.
Irregular and More Complicated Domain: Dolphin-shape Domain. We provide the result of Laplace2d-Dirichle on a dolphin-shapp domain in Table 9 which has more complex boundary. The mesh and decomposed domain is illustrated in Figure 9 We can see that this irregular and more complicated domain, the SNI with GNOT still gets reasonably good result. The error is relatively higher than that of simpler domains in Table [ This error gap be caused by the gap between training and testing shape distribution.
Figure 9: Visualization dolphin-shape domain.
24

---

Published as a conference paper at ICLR 2026
Equation
Domain
SNI with GNOT (%)
Dolphin
4.5 0.9
Laplace2d-Diin Diricchee
Disk
2.1
Table 9: Laplace2d-Dirichlett on dolphin-shape and disk domains.
Equation
Domain
valaNoT
SNI with GNOT
valGeo.FNo
SNI with Geo-FNO
A
2.2.:0.6
9#3
Laplace2d-Dirichee
B
2.5
2.1
5.6
14:11
C
2.1
12:22
A
9422
11::2
Darey2d
B
3.8
8+:22
6.4
15.7++0.9
C
5,4
20122
Table 10: SNI with GNOT or Geo-FNO as different choices of local operator for Laplace2d- and Darcy2d. Validation enrors are provided for reference.
Simple Domain: Disk. We provide the result of Laplace2d-Dirichle on simple disk domain in Table 9
This result is comparable to these in Table T
Comparison with Graph-based Neural Networks (GNN). We provide the result of direct inference Laplace2d training data and
with MeshGraphNets (Pfaff et al. 2020). We train the GNN evaluate it on the domains A. We get relative le error of 11.5%. We find that GNN does provide better generalization across different domains compared GNOT and can be potentially used to accelerate our iterative algorithm in our future work.
Choice of neural operator architecture. We provide results of SNI with Geo-FNO Li et al. 2023 on Laplace2d-Diricee Dirichhet and Darcy2d in Table I0
Solution and Error Visualization. We provide visualization for stationary problems in Section
We visualize ground-truth solution from testing data, absolute error from SNI and GNOT direct inference in Figure 10 II and 12 We also provide the error visualization in Figure 14 along with the
decomposed domain to understand where the error is located.
25

---

Published as a conference paper at ICLR 2026
1 BROADER IMPACTS
First, the proposed framework holds the potential to serve as an alternative to conventional PDE solving tools. Through its ability to address challenges related to ceometry-generalizaiee and data efficiency, the framework offers advantages that can significantly improve the efficiency of PDE solving. This improvement can have positive impact on various industries, including engineering, physics, and finance, where PDEs are extensively employed for modeling and simulation purposes.
Second, the proposed three-level hierarchy for PDE generalization provides researchers with valuable directions for future exploration in neural operator research. This hierarchical structure offers framework systematically address the challenges associated with generalizing neural operators to new geometries and PDEs. By considering these three levels, researchers can focus developing techniques and methodologies that improve the adaptability, flexibility, and scalability neural oper- ators. Furthermore, current operator learning methods in the neural operator field are predominantly driven by data and do not adequately consider the underlying PDE information. In our research, we introduce domain decomposition into the neural operator domain to tackle the issue of geometric generalization, incorporating traditional PDE approaches. This research direction presents significant potential for further investigation.
ceb
cac
-30
45
Figure 10: Visualization of test dataset of Laplace2d- Dirrcchhee domain A, B and C. The three columns from left to right display the ground-truth solution, absolute error from SNI and GNOT direct inference.
26

---

Published as a conference paper at ICLR 2026
430
230
Figure 11; Visualization of test dataset of Laplace2d-Mixxs on domain A, and C. The three columns from left to right display the ground-truth solution, absolute error from SNI and GNOT direct inference.
Figure 12: Visualization of test dataset of Darcy2d on domain A, B and C. The three columns from left to right display the ground-truth solution, absolute error from SNI and GNOT direct inference.
27

---

Published as a conference paper at ICLR 2026
Figure 13: Visualization of test dataset of NonlinearPoisson24 on domain A, B and C. The three columns from left to right display the ground-truth solution, absolute error from SNI and GNOT direct inference.
Decompooit
622
Figure 14: Visualization of error distribution together with decomposed domain for Laplace2d- Dirichlet domain B.
28


Reviews:
[{"review_id": "j0uccL26go", "reviewer": "Reviewer_HhMK", "summary": "This paper proposes operator learning with domain decomposition, a novel framework designed to improve the geometry generalization and data efficiency of neural operators for solving PDEs. By integrating domain decomposition methods with neural operator learning, the authors train local neural operators on randomly generated basic shapes and then use an iterative inference algorithm called Schwarz Neural Inference to stitch local solutions into a global one for arbitrary geometries. The paper provides theoretical guarantees on convergence and error bounds and demonstrates the method’s strong performance across a variety of linear and nonlinear PDEs.", "questions": "See the weakness", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- Develop a local-to-global framework with strong theoretical grounding that seamlessly integrates domain decomposition and neural operator learning, offering flexibility to support multiple neural operator architectures for solving PDEs on complex geometries.\n\n- Demonstrates robust geometry generalization, effectively overcoming one of the major limitations of existing neural operator methods.\n\n- Presents comprehensive experiments encompassing both linear and nonlinear, as well as steady and transient, PDEs.", "weaknesses": "- The paper does not justify that domains A, B, and C adequately represent the full diversity of 2D geometries. Moreover, all experiments are limited to 2D cases, leaving the scalability and applicability to 3D geometries unverified. At least one experiment on a 3D domain would have strengthened the validation of the proposed framework.\n\n- Since the computational cost heavily depends on the number of iterations required by the iterative inference process, efficiency can become a concern. Moreover, the number of subdomains Kmay depend on the number of vertices and edges, especially when field variations within the domain are significant. However, the authors treat Kas independent of the number of vertices and edges when estimating the time complexity, which can be misleading. This assumption suggests a linear relationship with the number of vertices, while in practice, the computational complexity may scale multiplicatively with both the number of vertices and the number of iterations.\n\n- Although the proposed framework represents the computational domain as a graph, the paper does not include any comparison or discussion with graph-based neural networks or their variants (e.g., [1, 2, 3]), which could provide relevant baselines or complementary perspectives.\n\n[1] Neural Operator: Graph Kernel Network for Partial Differential Equation, 2020.\n\n[2] Learning Mesh-Based Simulation with Graph Networks, ICLR, 2021.\n\n[3] HAMLET: Graph Transformer Neural Operator for Partial Differential Equations, ICML, 2024.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "o1ED7l0Jqr", "reviewer": "Reviewer_1tUX", "summary": "This paper addresses generalization to novel geometries in neural operator-based PDE solving. The approach builds on domain decomposition theory, specifically proposing a neural version of Schwarz inference with theoretical guarantees. The authors train models (GNOT) on simple geometries, then apply them at inference to complex domains that can be decomposed into subdomains. The iterative method that the authors propose consistently outperforms the vanilla direct inference.", "questions": "- How do you embed the boundary conditions in the GNOT architecture? Does it differ from the way you do it with GeoFNO? Could this explain the performance gap? Is your method really model agnostic or do you assume specific capabilities (e.g. boundary conditions encoding) in the neural operator?\n- Could you provide visualisations for some training domain examples? \n- Could this method be extended to 3d domains?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- As seen in Table 1, the method obtains excellent results on elliptic problems and on the temporal Heat equation.\n- Figure 4 is particularly enlightening and shows that the error consistently decreases as the number of iterations grows, which highlights the stability of the method.\n- The choice of shapes at test time is interesting and appears non trivial, even if they are in 2d.", "weaknesses": "- The authors don't provide in the main paper illustrations for the simpler domains that they choose during training.\n- The method assumes that a partitioning is already available at test time, which suggests that the method relies on some preprocessing. This needs to come with further explanations and would help a reader less familiar with domain decomposition method how it works in practice.\n- Figure 4 shows that in some cases the method requires up to 5000 iterations to converge. How much inference time overhead does this represent compared to the direct inference.\n- The role of the data augmentation is not very clear: Table 2 shows that this has a little effect on the generalization performance and I would therefore insist less on this aspect than on the analysis of the core of the method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "h6fsPcJcI0", "reviewer": "Reviewer_bCzG", "summary": "This paper introduced an operator learning strategy based on domain decomposition to handle complex geometries. The paper proposes to use data augmentation to improve the learning of the PDE on simple geometries. The solving method is based on domain-decomposition and locally uses the neural operator. Finally the method reconstruct dynamics on complex geometries by merging sub-solutions, using a SNI algorithm. The paper proposes a theoretical study of the algorithm as well as some experiments.", "questions": "### Questions\n-\tCould you compare the proposed method to more recent baselines ? consider public/standard datasets of the litterature? Do you plan on releasing the dataset for future work? \n-\tLine 093-095, you mentionned that the method is restricted to parabolic equations. Could you detail about this ? What are limitations in terms of PDE equations/BCs (also reference to lines 198-199)? \n-\tCould you provide some insight on how to chose the hyper parameters mentioned in sections 3.1 (ie n?, K, d?) It is specified that hyper parameters have to be carefully chosen, why ? Doesn’t the method converge in these cases ? It does not seems to be very sensitive to hp in terms of final performance in the ablation on figure 4.\n-\tCould you detail about the results on Heat2d – B ? While the method seems coherent and robust in all scenarios, on this geometry however, results drops x2. \n-\tHow does the method behaves on simple geometries ? \n-  could you elaborate on the condition required for the algorithm to converge? Line 160 refers to Appendix A, but Appendix A finally states that convergence depends on the PDE and decomposition. \n- Have you studied where are located the errors? Is there a higher (or lower) error around the boundary of the local shapes? \n- How does the method scales in term of memory consumption with respect to the number of points ? Would it be tractable for 3D PDEs ? Industrial complex geometries ? \n-\tHave you compared inference time with respect to baselines ? I saw the method’s inference time in the appendices, but I could not find baselines’. \n- Regarding the theoretical section: I think I am not qualified to review such theoretical results, but I have question regarding the applicability of the proof: When are the assumptions met? For example, could you cite some PDE respecting the hypothesis on L? \n\n### Minor comments :\n-\tLegend of figure3 : it it hard to understand what validation refers to, I think it is detailed in the last sentence of description? just adding a (validation) at the end would remove any questioning. \n-\tFigure 4 : I think that log scales for errors would be more appropriate for these figures. As is, it is hard to visualize the differences between all ablation (at least the 2 firsts).", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The proposed paper is well-writen and carrefully introduces a lot of concepts.\n- A theoretical analysis is proposed as well as experimental evaluations.", "weaknesses": "- Scope of applicability: My concerns about applicability relies in the fact that the problem formulation states that PDE are tested in parabolic equations (with time). Is it a limitation of the method? Why is the method limited to such PDEs? \n- Missing references: a lot a references dealing with complex geometries and generalizations are missing, see eg [1-5]\n- No experiments on public datasets: I think evaluating the proposed method on well-known dataset helps the reader positioning the method with respect to existing baselines. As an example, some dataset considering irregular/complex geometries are in [6-7].\n- While the paper is very rich, several key components of the method are detailed in appendices, making the reading as is hard to understand (eg algorithm). \n- Inference time and scaling : I have some concerns regarding these 2 aspects of the method, see questions.\n\n### References\n[1] Operator Learning with Neural Fields: Tackling PDEs on General Geometries, Louis Serrano, Lise Le Boudec, Armand Kassaï Koupaï, Thomas X Wang, Yuan Yin, Jean-Noël Vittaut, Patrick Gallinari, 2023\n\n[2] Transolver: A Fast Transformer Solver for PDEs on General Geometries, Haixu Wu, Huakun Luo, Haowen Wang, Jianmin Wang, Mingsheng Long, 2024\n\n[3] A.D.Jagtap, G.E.Karniadakis, Extended Physics-Informed Neural Networks (XPINNs): A Generalized Space-Time Domain Decomposition Based Deep Learning Framework for Nonlinear Partial Differential Equations, Commun. Comput. Phys., Vol.28, No.5, 2002-2041, 2020. \n\n[4] K. Shukla, A.D. Jagtap, G.E. Karniadakis, Parallel Physics-Informed Neural Networks via Domain Decomposition, Journal of Computational Physics 447, 110683, (2021).\n\n[5] AROMA: Preserving Spatial Structure for Latent PDE Modeling with Local Neural Fields, Louis Serrano, Thomas X Wang, Etienne Le Naour, Jean-Noël Vittaut, Patrick Gallinari, 2024\n\n[6] Learning Mesh-Based Simulation with Graph Networks, Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-Gonzalez, Peter W. Battaglia, 2020\n\n[7] A comprehensive and fair comparison of two neural operators (with practical extensions) based on FAIR data, Lu Lu, Xuhui Meng, Shengze Cai, Zhiping Mao, Somdatta Goswami, Zhongqiang Zhang, George Em Karniadakis, 2021", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5dUvXwcg2v", "reviewer": "Reviewer_kNbL", "summary": "The authors propose a new framework for training neural operators to generalize to unseen geometries at inference with domain decomposition. Instead of generating a variety of complex geometries and train a neural operator such that it can generalize to an unseen complex geometry at test time, the authors instead propose to train a neural operator on a diverse number of simple geometries that can be seen as a decomposition of a complex geometry. At inference, they introduce the SNI algorithm, that decompose a complex unseen geometry into a set of simpler shapes, that can be solved iteratively.", "questions": "Could you provide runtime comparisons? It would be helpful to provide actual runtime comparisons: (a) pure neural operator direct inference, (b) SNI (with partition & iteration), (c) classical PDE solver.\n\nThe method has been tested under specific boundary conditions and in a specific setting. Could you elaborate a bit on how general this method for PDE modeling? How could it be applied to more challenging physical systems?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I really like the tackled problem and find it very interesting and well motivated. Data driven approaches are known to be robust to the geometries observed in the training distribution, but they naturally tend to fail when tested on geometries that shift from the training distribution. Decomposing a complex geometry into a set of very simple shapes (that thus are more aligned with the training distribution) is an interesting approach to generalize any neural operator to arbitrary geometry.\n\nExperiments successfully demonstrate that the framework proposed help to boost performance of neural operators on unseen geometries via geometry decomposition. The framework is general and ca be applied to any neural operator that accepts arbitrary mesh points.\n\nAblations are provided to see how sensible is the method with respect to the number of sub-domains K, data augmentation, etc.", "weaknesses": "While the motivation is clear, one of the potential benefit of neural operators is to overcome the computational cost of traditional methods. Decomposing the domain into sub-domains naturally increase the computational cost of data-driven approaches. Ablations provided in appendix seem to show that the computational time needed is high for very accurate predictions. It would be nice to compare this with the time needed for solving the PDE with a numerical method, and put it as limitations based on the results.\n\nThe experiments have been tested on relatively simple and a limited number of geometries. It would be nice to test the method on more challenging geometries and see how well it performs with respect to the number of sub-domains. If the geometry is much more complex, we can imagine that it would need much more sub-domains to recover the simple geometries seen during training. \n\nI think the paper would clearly benefit from a related work that mentions existing methods for geometry generalization. The method should be compared to existing works when possible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)