PINN Balls: Scaling Second-Order Methods for PINNs with Domain Decomposition and Adaptive Sampling
Authors: Andrea Bonfanti
Ismael Medina
Roman List
Björn Staeves
Roberto Santana
Marco Ellero
Venue: NeurIPS 2025 poster ()
Keywords: Physics-Informed Neural Networks
Second-order optimization
Domain Decomposition
Adversarial Adaptive Sampling
Primary area: machine_learning_for_sciences
TL;DR: A method to scale second-order training for PINNs, based on domain decomposition and adversarial adaptive sampling.

Abstract: Recent advances in Scientific Machine Learning have shown that second-order methods can enhance the training of Physics-Informed Neural Networks (PINNs), making them a suitable alternative to traditional numerical methods for Partial Differential Equations (PDEs). However, second-order methods induce large memory requirements, making them scale poorly with the model size. In this paper, we define a local Mixture of Experts (MoE) combining the parameter-efficiency of ensemble models and sparse coding to enable the use of second-order training.  Our model -- PINN Balls -- also features a fully learnable domain decomposition structure, achieved through the use of Adversarial Adaptive Sampling (AAS), which adapts the DD to the PDE and its domain. PINN Balls achieves better accuracy than the state-of-the-art in scientific machine learning, while maintaining invaluable scalability properties and drawing from a sound theoretical background.

PINN Balls: Scaling Second-Order Methods for PINNs with Domain Decomposition and Adaptive Sampling
Andrea Bonfanti
Ismael Medina
BMW AG, Digital Campus Munich
University of Göttingen, Campus Institute Data Science
Basque Center for Applied Mathematics
University of the Basque Country
abonfanti101@@kaasl.. ehu.eus eus
Roman List
Björn Staeves BMW AG, Digital Campus Munich
BMW AG, Digital Campus Munich
Roberto Santana
Marco Ellero
University of the Basque Country, Intelligent Systems Group roberto. san tana@ehu. eus
Basque Center for Applied Mathematics, CFD Group IKERBASQUE, Basque Foundation for Science Swansea University, Complex Fluids Research Group mellero@bcamath.c org
Abstract
Recent advances in Scientific Machine Learning have shown that second-order methods can enhance the training of Physies-Informed Neural Networks (PINNs), making them suitable alternative to traditional numerical methods for Partial Differential Equations (PDEs). However, second-order methods induce large memory requirements, making them scale poorly with the model size. In this paper, we define a local Mixture of Experts (MoE) combining the parameter-efficiene of ensemble models and sparse coding to enable the of second-order training. Our model PINN BALLS- also features a fully learnable domain decomposition structure, achieved through the use of Adversarial Adaptive Sampling (AAS), which adapts the DD to the PDE and its domain. PINN BALLS achieves better accuracy than the state-of-the-art in scientific machine learning, while maintaining invaluable scalability properties and drawing from sound theoretical background.
1 Introduction
Obtaining accurate solutions to partial differential equations (PDEs) is a ubiquitous problem in countless engineering applications, as it enables reliable simulations of real-world scenarios. While discretization-basse methods for solving PDEs offer convergence guarantees as well as a solid theo- retical underpinning, their implementation and runtime be resource intensive. Novel techniques developed within the domain of Scientific Machine Learning (SciML) strike a highly sought-after balance between accuracy and runtime, which has led rapid growth of the field in recent years 77.. A milestone SciML. architecture is the Physics-Informed Neural Network (PINN), which was theorized in the 90s [9] and gained the interest of the research community following the publication that coined its name 28 The salient aspect of PINNs is the inclusion of physical laws in the loss function, computed via automatic differentiation. While their training process is not simple, researchers proposed alternatives to fix pertinent issues of the PINN architecture, such as the complex loss landscape generated by the PDE residuals (36), unbalanced optimization of the components of the loss functions [17), and the spectral bias intrinsic to neural networks 229 The combination of
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

those with second-order optimizers has proven to be an invaluable approach to surpass the limitations of PINNS 4 24. Second-order methods are crucial step to enable highly accurate solutions, which can be a strict requirement in some applications. However, these methods scale poorly with an increasing number of parameters.
Scaling to large number of parameters is a crucial component for PINN architectures as the majority of problems simulated in industrial scenarios typically feature highly complex solutions and arbitrarily large domains. The majority of approaches in the literature rely on domain decomposition (DD), which is a common approach adopted by traditional PDE solvers. DD divides the task of predicting the PDE solution into different subdomains, which are routed to different submodels. In the machine learning community, architectures that combine multiple models are regarded as ensemble models. A widely known example of an ensemble model is the Mixture of Expert (MoE) [22]. which approximates a function as linear combination of submodels (or experts) driven by - possibbly learnable - weighting factor, which is referred to as "gating mechanism". This allows the MoE to scale well to larger problems and often to parallelize the training of the submodels 32] strongly limiting the computational overhead when several models are combined [3]. While a natural choice for scalability in the field of PDEs, the use of DD involves several challenges. A major issue appears when equations contain non-local quantities eg. the pressure in incompressible Navier-Stokes equations - that cannot be easily inferred from local information. This approach is therefore typically complemented by the addition of a loss component the interface of neighboring models. However, these interface conditions can strongly hinder the training of the model and force the predicted solution to oscillate between local minima. Furthermore, there are several ways of defining interface conditions for ensemble models; the choice is typically priori 21
and can not be defined a
An additional complexity introduced by DD is intrinsic to the decomposition itself. Since world scenarios typically include multi-scale and finely detailed domains, it is imperative to be able to adapt the decomposition to the underlying domain and PDE. Indeed, problems in the field of fluid dynamics are typically inhomogeneous and they require adaptive meshes for both stability and scalability. Therefore, the DD can not always be decided a priori. It has to be learned alongside the PDE solution, and possibly be adapted both to the geometry of the problem and to the nature of the PDE that is being solved. However, most of the literature works combining DD and PINNs do not instantiate such a domain decomposition. In most cases, the DD is not trainable but rather fixed beforehand. To address this issue we adopt the Adversarial Adaptive Sampling framework (AAS) proposed in [3), which provides solid theoretical framework for dynamic sampling that supports the training process and guides the DD.
Related Work The synergy between PINNs and DD has already been explored for example by Extended-PINNSs (XPINN) [16], which relies on non-overlapping subdomains and interface conditions to match the prediction of neighboring submodels. Another well-known example is the Finite- Finite-Basis PINN (FBPINN) 233, which adopts predefined overlapping sub-domains and leverages the sparsity of the model's output to compute parameter updates without relying on interface conditions. This model has been extended to Extreme Learning Machines [2] with the goal of achieving highly accurate solutions. A few additional works also propose to include a trainable domain decomposition, such as the Augmented PINN [14, which represents an extension of the XPINN a soft pre-trained domain decomposition. Similarly, the MoE PINN (3), which takes advantage of the ensemble structure and the agreement of the submodels as a natural approach for uncertainty quantification. We refer to [B3 for a non-exhaustive overview of existing methods.
Motivation and Goal In our research, we propose to combine the benefits of DD, second-order training methods and Adversarial Adaptive Sampling (AAS). Indeed, the paradigm of ensemble models for DD synergizes with the use of second-order methods: while the resource efficiency of second-order methods scales poorly with the number of parameters, ensemble models are parameter- efficient and their structure allows for parallel evaluation of the model and its gradients (23 Moreover, we propose to base the DD on overlapping radial basis functions to emphasize locality, hence restricting the prediction of each submodel to ball within the PDE domain. This approach ensures that the Jacobian matrix required for computing the update step of the second-order method is sparse, relieving the memory-intensive nature of second-order optimizers as Natural Gradient Descents (24 and wwrenberg-Marquarr 4. Last but not least, we leverage the theoretical backbone of AAS 3
to train the DD alongside the ensemble model, which requires priori knowledge of the features
2

---

Table 1: Difference in features between PINN BALLS, the APINN 14 and the FBPINN 23
APINN FBPINN
PINN BALLS
Sparse No interface loss
Trainable DD
Adaptive Sampling
Second-order training
of the PDE solution or the computational domain. This allows the DD to direct its attention to PDE features that are harder to compute, enhancing the accuracy and robustness of the method.
The paper is structured as follows: Section introduces the notation for PINNs and DD. Section introduces AAS, showcasing theoretical advancements in 3.1 The complete PINN BALLS model is presented in Section 4 and its scalability with respect to the number of parameters of the models is discussed in 4.1 The full training pipeline is then highlighted in Section [42] and its performance evaluated in Section 5
Limitation are discussed in Section 5.3 and the paper is concluded with Our contributions are gathered in Table To summarize:
Section 6
To the best of our knowledge, we are the first to combine second-order training methods with domain decomposition on PINNs, including scaling to large parameters,
We adopt AAS introduce a highly adaptable framework to train the domain decomposition;
We advance the mathematical foundation of AAS, ensuring robustness of the model.
Our numerical results show outstanding performance benchmark PDEs, demonstrating consistent behavior of the parameter efficiency and global accuracy with respect to the size of the ensemble model. This also holds for cases such the Navier-Stoker equations, which involves particular challenges for DD due to the aforementioned non-local features.
Physics-Informec Neural Networks and Domain Decomposition
The PINN architecture leverages the power of neural networks as functional approximators to address classical PDE solution instances bounded domain Q R4,, Without loss of generality, consider the problem of finding a function u Q2 >> IR dous which satisfies the following equations:
Ru(z) x f(x).
x ZE
u u(x) g(z), x an.
where the underlying PDE is fully defined by the differential operator R. and boundary and initial conditions collected in known function g. PINNs aim to approximate the solution of the aforementioned PDE through a neural network ug + Rdont with layers, defined as:
u(z; 0) WL-1 (02-1-0(----- - oWWz oz + bb) +....)+ b
(2)
where W. E R and b; E Ah denote respectively the weights matrix and bias vector of the i-th hidden layer with dimension h, For the sake of compactness, represents the collection of all the trainable parameters of the network, i.e. (W; by !-!!! The activation function o IR R is smooth foordinate-wist function, such as the hyperbolic tangent, or the sine function which are common choices for PINNs. The "Physies-Informed" nature of the neural network ug lies in the choice of the loss function chosen for minimization, which is typically given by
[ : J/ Ru(x; 0) f fx)) dx
|u(x; 0) - g(x)) dw)())
(3)
Jan
In the above formulation, dz and du represent infinitesimal volume and surface element respectively. The integrals in (3) are typically approximated as:
No
1 Nan Ru([;; f(x))1 Naa
1
((O)
Ju(āx:
- g
No.
where No represents the number of collocation points [ and Nag the number of points k
No in Ω to minimize the PDE residuals Nao used to fit the initial and/or boundary conditions at 0.

---

2.1 Domain Decomposition
The concept of domain decomposition is ubiquitous the field of numerical solution of PDEs and has been seamlessly used in order to tackle the scalability of PINNs. The rationale is to divide the }= j=1 such that و S2 Ω. Each subdomain can be associated with local solution or, in the case of PINNs, with a local neural network (x) == (sr 0z) parameterized by a subset of denoted by 0z. The aforementioned splitting allows parallelizable training routines, which are beneficial in terms of scalability [3] with respect to the number of parameters and domain size, and in the effectiveness and accuracy of the model [10]. Therefore, the overall model ue is obtained combination of the local models:
PDE domain Ω into M possibly overlapping
subdomains {
uo(x) Σ (x(x))((T)) S.L. - (x) Vj =j=1.-m M, Vo Q.
(5)
Here, {(x) = (Z) (r)( may be interpreted as vector of weighting factors indicating which submodel has more impact in the prediction ug at x, and is commonly referred as the gating function. The latter property of À is typically enforced to impose a proper partition of unity.
The most popular works domain decomposition in PINNS [16 19 also rely on the introduction of interface conditions to align the prediction of submodels between subdomains. Choosing the correct interface condition a nontrivial, comain-dependee task 211 Typical interface conditions require to evaluate the difference between each pair of overlapping bases Uj and U,, While this step can be computed efficiently through parallelization, interface conditions require an additional loss component. As the nuanced interplay between loss components is known to introduce unfavorable complexity in the training of PINNs (20), we advocate for a global optimization step, which can still be computed efficiently due to the locality of each submodel 23
3 Adversarial Adaptive Sampling
Adversarial Adaptive Sampling was introduced 33 technique enhance the training of
PINNs. Its core idea is to sample the collocation points from probability density p E P(2)) to be learned simultaneously with the PINN, following the blueprint of generative adversarial neural networks. This corresponds to the minimax problem:
with J(0. pp) = / [Ru( 0) - f(x BG(p)),
min max J(0,Po), 8 P(EP(N)
(6)
where Ps is parameterized version of p and G is a regularizing term that prevents the collapse of ps Dirac delta. The intuition is that Pa should emphasize regions where the residuals are difficult to minimize, while ensuring full coverage of the PDE domain. The founding paper [333 shows existence of solutions for the minimax problem (6) with chosen to constrain the Lipschitz constant of pạ- Nevertheless, [ employs in the numerical examples the more practical regularization Sa (2)(x))ddr, for which a convergence proof is not provided. We bridge this gap in Section 3.1
3.1 Convergence Result
To lay down sound theoretical framework, we extend the theoretical results in 33 to practical regularization functionals for the probability distribution p. In particular we complete the rigor- ous treatment of the regularization term Sg |p(x)) proposed by [33), as well as an entropic regularization term motivated by the field of optimal transport [8]
Theorem 3.1. Assume one of the following settings:
1. G the Dirichlet energy Gp(p) = Jq (p) and info |(Ruu f) or
2. G is the negative entropy G=(p) == Sn p(x) log p(x), and infe [[[Rup f)²llc(o)) = 0.
Then the optimal value of the minimax problem (6) is In particular, for any sequence (0, }m converging to the infimum assumed above, there exists sequence of probability measures {pm (with Pa maximizing J J(Bn..) such that
lim J(An. Pn) 0. n-000

---

Moreover, the probability density functions {Px h. converge to the uniform distribution on Q
Idea of the proof. The uniform density ş 1/|52] constitutes lower bound for the nested maximiza- tion step in ( This fact implies certain coercivity of the maximization problem, which can be
used to show existence of maximizer p": combining this aspect with a minimizing sequence for the residuals completes the proof. The detailed proofs can be consulted in Sections A. and A.2
The entropic regularization GE proposed in Theorem 3.1 and the gradient-based loss GD of [ share its aim to smooth out p. In addition, the entropic regularization GE features an infinite slope at the origin that inhibits the probability density from reaching zero anywhere within the domain. Hence, positivity of Pa is baked into GE. while it needs to be additionally imposed for GD.
Remark 3.2. As shown by 18 26 both the gradient flow of the Dirichlet energy with respect to the Lo-oopoloyy and the gradient flow of the entropy with respect to the L Wasserstein distance generate the same dynamics: the heat equation. This suggests compelling interpretation: gradient ascent step in ( may be seen as the application of a heat kernel that smooths out with respect to the landscape of the residuals. For more context and implications of this intuition see Section A.3]
Remark 3.3. The novelty of Theorem 3.1 is twofold: on the one hand we extend the convergence results of 33 to regularization functions that are more practical to implement, since the Lipschitz show the
constraint considered in 3 is hard implement in practice. On the other hand
convergence of to a uniform distribution along any converging subsequence, which in particular implies that the sequence of maximizing distributions {paIn tends to cover the full domain as the residual of Un converges to zero. This result further validates the robustness of our method, ensuring an adequate generalization to arbitrarily complex PDE domains.
4 Sparse Ensemble of Physics-Informed Neural Networks
A crucial step in the proposed model is the usage of Radial Basis Functions (RBFs) to identify the domain decomposition. In particular, opt for quartic basis functions as in Equation g instead of the more traditional Gaussian RBF - to emphasize locality of the submodels.
(x) MM (1(2)
( 12
if
x-- <1 8j
(x(x)=
where
())
(8)
elsewhere
The centers and radii of each ball 6, are indicated respectively as ej and 8j, and identifies the trainable parameters of the domain decomposition. We adopt the notation $2, to indicate the compact support of the basis function Φj. i.e. the set {a E f2 s.t. Φ (x) > 0}. The final gating function X, is obtained by normalizing Φj with respect to the total sum of RBFs. This ensures proper partition of unity while maintaining the sets 22 compact and small. Ensuring that EM di (x) > 0 for all x can be done by choosing sufficiently large (s, - at initialization.
Connecting Domain Decomposition and Adversarial Adaptive Sampling In order to align the training of the parameters of each Φ, with the training of the PINN parameters, we devise connection between the domain decomposition and the sampling distribution ps. In particular, we opt approximate p& through the combination of the basis Φj as:
1
(²(x)=
(x)x(T).
M
where wi chosen so that wa is probability distribution. This parametrization of pa enforces the domain decomposition to focus on areas where the PDE residuals are higher, which typically occurs in regions of the domain where the PDE solution presents steep gradients and/or high frequency components. Morcover, it implicitly defines suitable training proceedure for the DD, given by the gradient ascent step of the minimax training routine.
Definition 4.1 (PINN BALLS). We define PINN BALLLS tuple of the form (uo, pT), where uo is the ensemble model defined in 5) with domain decomposition structure M given by (8) and Ps is the probability distribution of the form 9 The loss function is given by discretizing the AAS
loss 6 set of points {x!] sampled from Pst
5

---

Nn
Nan
No log po (z,)) (10)
((0 6) = No >
Ru( 0) f(xx))
Non l
Σ
Σ
Na
Remark 4.2. Ensemble models of the form 5 are universal approximators. This holds trivially as each submodel is universal approximator itself.
4.1 Scalable second-order training of the parameters
To train our model we adopt the Levenberg Marquardt update rule, which is second-order quasi- Newton method that represents a stabilized version of the Gauss-Newtor method [25).
g++1 8 (D,I." D + n (oT.
(11)
Here D. denotes the Jacobian of with respect to the model parameters, and the superscripts and represent respectively the training iteration and the Moore- Poorr-PPenooee pseudo-inverse. The scalar value 1j determined heuristically during training, balancing the interplay between gradient descent step and second-order update. The latter is necessary to efficiently reach suitable minimum, whilst the former is useful practice at early iterations.
The most prominent drawback of this training routine is given by the storage and inversion of the matrix D&L" " D&L in GPU memory, which is highly impractical as its size scales quadratically with the number of parameters of the model. A practical solution to this is to ensure that the aforementioned matrix can be stored in a sparse format, indicating only a linear growth with respect to the number of parameters. This can be done by leveraging the locality of the submodels of PINN BALLS, since the effect of each U, is restricted to the compact support 22-.
Indeed, the components of D&L are defined as the partial derivatives of with respect to the parameters O. Outside of the compact support Ω2j. hence where \ and all its derivatives are identically 0, the jacobian components of U, vanish. This implies that in the aforementioned region it is not necessary to backpropagate obtain the partial derivatives necessary to compute the PDE residuals.
Oo M J=
XJU;
Σ 1 = 0 Σド
DaL =
O₂
M |
(12)
On top of sparsity, the PINN BALLS model also enables parallel implementations for the computation of DoL. Sparse and parallel implementations combined allows to adopt the exact formulation of the LM update step by computing D. D. in fast and memory efficient fashion. The inversion remain the ultimate challenging aspect, which can be done through sparse solvers. Another valuable alternative is to consider matrix-free multiplication methods conjugate gradient, whose implementation can benefit from the sparse nature of our Jacobian. This class of methods is however out of the scope of the present paper, since they are often sensitive poor conditioning of the underlying matrix, which is often the case when training PINNs. Finding a viable preconditioning strategy remains thus an interesting direction of future study.
4.2 Min-max Training Routine
The training of the PINN BALLS model follows the alternating descent-ascent approach proposed for AAS in [33]. For the parameters 0 of each submodel we adopt the Levenberg-Marquarii (LM) update step defined in (T).. When the problem is relatively simple - e.g. linear PDEs one can also resort to more efficient approach of alternating gradient descent step on 0 and second-order fine tuning of the parameters of the last layer of the model, which is closely connected to the approach of Extreme Learning Machines [15] and close to that employed in 2
The training step for the domain decomposition á consists in a gradient ascent step of 10 proposed in 33 is convenient to rewrite (T)p using importance sampling:
As
No F(Pe) = No
Ru(x
[(x!\) po(a.)
No
(a())) log))(z;))
(13)
Pz(x7)
No
p5(07)
Here P& stands for the current value of Ps. which taken to be constant and therefore does not enter the computational graph. [13 is evaluated on the same sample points as the 0-update, and the

---

boundary points are ignored, since they do not depend on Φ. The choice of the learning rate for the parameters of Pe is important for stabilizing the overall training routine. particular, we notice that the LM training shows better performance when Ps does not strongly vary over training iterations. However, its fluctuations during training are important to ensure proper convergence of the model. Hence, we recommend an exponential decay for its learning rate to emphasize variability in early iterations and favor convergence of the model in late training stages. The complete training pipeline is summarized in Figure]
Domain Instance Ω
Training Routine
Parallei Step
Initialize collocation points
Sample collocation points (xx %
P
(,000
Build D DopL
Initialize 12 ابر and pa (K-means)
Update Ps
xx00)
Build Doy
Φ
+ V V
Build sparse DyL
Trained PINN BALLS
Maximization Step
Mrd 2 )
elsee Stopping Criterion met
g++1
(DyL'' D'L 7l) V.L
Minimiiatiion Stop
Figure I: High-level diagram of the minimax training routine devised for the PINN BALLS model.
Numerical Evaluation
We begin by evaluating the behavior of the Levenberg-Marquard (LM) update step on a simple supervised problem - approximating the function f(x, y) = sin 4r sin For various number of parameters, we instantiate a standard dense neural network and PINN BALLS architectures with 10, 20, and 50 PINN Balls. The disposition of the centers is determined by K-means over the collocation points, and s, selected be the minimum possible cover all training samples, which is identified by the standard deviation of each cluster. For each model, we perform a single LM update and collect statisties on runtime, Jacobian sparsity and memory consumption; results are summarized in Figure 2 In the case of PINN BAL.S,, we exploit the sparsity of the Jacobian by constructing it Compressed Sparse Row (CSR) format using the scipy.sparse package, solving the resulting system either in the scipy package 35 via direct dense sparse solvers depending on the structure. The statistics shown do not refer to the accuracy the model nor include the computation of the matrix L DaL, respectively discussed in Sections 5.1 and 5.3
Efficiency of the matrix inversion Although dense solvers generally achieve faster wall-clock times for moderate problem sizes, models employing large number of PINN Balls achieve competitive runtimes. Notably, highly sparse models (50 PINN Balls) present comparable runtimes even in moderate problem size, which is mainly due to numerical heuristic of sparse matrices: visible performance improvements are typically noticeable when the nonzero elements of the matrices are 10% or less of the total.
Memory Efficiency From a memory perspective, using large number of PINN Balls provides significant gains, which is to be expected due to the sparse nature of the Jacobians. However, the dense architecture more memory efficient than the 5-PINN Balls model, whose Jacobian presents 30% of nonzero elements. Indeed, in the sparse formulation the nonzero values stored alongside their indices, which results in roughly 3 times the number of nonzero values in the matrix. Notably, for the simple supervised problem considered, the dense solver reaches out-of-memory error at roughly 8000 parameters., while the sparse ones can handle much larger models.
To assess the robustness of these findings, we replicated the experiment across a range of target functions of varying complexity. Remarkably, the trends observed in Figure persist almost indepen- dently of the choice of f. While in theory the condition number of D L. could impact the difficulty of solving the LM update, we observe that the damping effect inherent to the LM formulation sufficiently regularizes the system, mitigating potential instabilities due to ill-conditioning. Note that PINNs
7

---

DOYGE NM 20 PYUN 52 PMUN Dalls
Dorio NC
Do SC
Ball
Vomory (or updite 1452
FWS bulla PRL Dal's
25 (146
Conas NV
50 71N Da
Ends
lod:
PMN kol:
2023
6305
6306
0030
1800
3300
4030
3030
1003
10090
200 4003
800 8030
33009
2305
1200
6030
8009
1030
Number
Mumber
Figure 2: Mean value with shaded variance, for different ensemble sizes, various statistics of the LM update step on PINN BALLS. From left to right, the runtime in seconds, the percentage of non-zero elements of D&L, and the CPU memory consumed to invert and store the matrix DøL D&L
trained with second-order methods show very strong performance even with few parameters when compared to larger models trained with first-order methods LBFGS [ Notably, the sparsity structure of the Jacobian may altered during training by the DD, due to changes in the basis radii. However, when preserving high sparsity is critical, the basis radius can be capped or even fixed after initialization.
5.1 Training on Benchmark PDEs
We now showcase the results obtained by the PINN BALLS model on three well-known test cases for PINNs: The Helmholtz. Equation, Burgers' Equation and time-dependent, incompressible Navier Stokes Equations in 2D). We show the relative 1.2 error achieved during training the number of submodels increases, keeping the same architecture for each individual submodel. These results collected in Figure and commented below. For the sake of brevity, gather details of the PDEs solved and additional results in Appendix B We do not include comparison with first order optimization. We refer to for a thorough comparison between first and second order optimization in general. Combined with the DD settings, tendentially first-order optimization is slower, especially in a scenario like ours which includes an adversarial minimax formulation. All models are implemented in JAX [5] using double precision and trained on single NVIDIA A10 GPU.
Helmholtz Equation The leftmost plot of Figure 3 showcases the training loss achieved on the Helmholtz equation. Increasing the ensemble size also enhances the fluctuations in the loss function. This behavior expected, as increasing the number of basis makes the maximization step in (13) more accurate. However, the very small relative L error achieved shows that the PINN BALLS model is capable of obtaining extremely accurate solutions. Remarkably, the L error with respect to the ground truth does not present high frequency components, which is common artifact in PINN solutions. Moreover, the I,2 error the order of 10-10, which notably low for a PDE solver in general. The final distribution of the basis oy does not present any noteworthy pattern.
The center plot of Figure showcases the 1 Error achieved during training on Burgers' equation. It is possible to notice that strong oscillation in accuracy occurs during early iterations, which is often due the higher complexity that can be captured in the maximization step by Ps. This behavior proves to be beneficial for the model, as it allow to reach lower L errors, to the order of 10 for a model with 20 PINN Balls.
Burgers' Equation
Navier-Stokes Equations Navier Stokes equations represent truly challenging scenario for PINNs. as the traditional implementation fails to converge on the classical benchmark of a fluid flow past cylinder in 2D 6 The behavior of PINN BALLS is shown in the rightmost plot of Figure 3 Once more, increasing the number of PINN Balls increases the accuracy of the model. This result non-trivee noivtrivial for DD-based approach, due the strong nonlinearity of Navier Stokes' equations and the presence of the pressure, which is a nonlocal quantity. This favorable behavior is due to the global nature of the LM update, which has the scaling advantages of DD, while keeping the invaluable benefits of second-order method, which is crucial for solving nonlinear PDEs with PINNs 4. Notably, in our GPU implementation, global LM update results an memoory error at approximately 2000 parameters, which are not sufficient capture the deeeddennnynncy of the whole
8

---

solution. Indeed, capturing phenomena as the vortex shedding in the cylinder wake requires a model with consistent number of parameters.
30
20
PRIIIHT
13
PRODUCT
6.5
10:g
7000
1742 1922
2000
750
1966
179
300
Figure 3: Median Loss value and minimax range over 10 independent trainings of PINN BALLS with an increasing number of submodels on Helmholtz (left), Burgers (center) and Navier Stokes' (right) equations. Each submodel has 3 hidden layers and 10 neurons per layer.
5.1.1 Behavior of the Domain Decomposition
Hereinafter, we highlight the behavior of the DD under the chosen initialization and optimization pipeline. Under reasonable conditions, model performance is not highly sensitive to the initial number or placement of basis functions. However, careful initialization is essential for preserving the sparsity structure of the Jacobian DgL. Excessively many or poorly placed basis functions can result in significant computational overhead when constructing the matrix Da D&L. Utilizing first-order optimizer for the global update generally yields slow convergence for the majority of the PDE considered. On the other hand, in case of lavier-Stokes equation, convergence is not achieved as the predicted pressure oscillates during training and oftentimes diverges.
Furthermore, our initialization procedure yields probability density function Pe that is approximately uniform. In relatively simple or linear scenarios, such as the Helmholtz equation, we observe that the domain decomposition remains mostly unchanged during training. which aligns with theoretical convergence expectations. However, in nonlinear settings, such as the Navier-Stoker equations, the AAS paradigm introduces meaningful deviations from the uniform distribution. These deviations serve to allocate more submodels in regions with higher solution complexity (c.g., steep gradients) and can be observed in the results Appendix B While theoretical convergence to uniform distribution assumes that the residuals vanish entirely the case in practice. Consequently, the final distribution does not always converge to uniformity, but this deviation often enhances the model's 's accuracy. We view this discrepancy between theory and practice as an important and interesting topic that merits deeper exploration future work.
reasonable but idealized assumption this is rarely
5.2 Comparison with existing approaches
We compare the performance of PINN BALLS with the XPINN and the APINN through the results reported in the respective publications. We train PINN BALLS of variable ensemble size that maintain a comparable number of parameters roughly 5000 each model Comparing with existing architectures, the results in Table clearly highlight the superiority of model by large margin in the Helmholtz equation. Notably, this is achieved with a lower number of parameters with respect to the two models highlighted. Regarding Burgers equation, we largely outperform the FBPINN model as they showcase an error on the order of 10 The results obtained by PINN BALLS is comparable that of the APINN architecture. However, it important to mention that the models used for comparison directly instantiate DD that divides the domain in O, which is where the shock locates the solution of Burgers' equation, and represents a somewhat optimal division of the domain as it captures discontinuity while maintaining continuous submodel predictions- Our model, on the other hand, learns the decomposition autonomously. Opposed to the results observed in Figure 3 larger ensembles do not necessarily outperform smaller ones for a fixed parameter count. In practice, it is important that each submodel mantains a reasonable number of parameters to capture the behavior of the target PDE solution in the related subdomain.

---

Table 2: Comparison of the relative L Error achieved by PINN BALLS, XPINN and APINN models on Helmholtz and Burgers Equation, for comparable number of total parameters.
Model
Rel. 2 (Burger)
Rel. L (Helm.)
XPINN
1.3e = 7 3
1.2e-3 1.7e 4
APINN
3.66
1.2e-3 4.70-4
PINN BALLS (5)
2.
- 1.8c-
PINN BALLS (20)
3
1.3e- 9 1.4e-9
5.3 Limitations
While PINN BALLS offer memory-efficient alternative for second-order training of PINNs, they also present some limitations. Efficiently computing of the Jacobian matrix remains challenging. While sparsity accelerates matrix operations at larger scales, the overhead of assemblying and managing the sparse Jacobian becomes non-negligible at moderate problem size. Therefore, it is crucial restrict the application of PINN BALLS to cases where the usage of large number of model parameters is a strict requirement. Furthermore, to the best of our knowledge, existing antoomati-differentiallt packages do not support backpropagation through sparse matrices. Indeed, the applicability to large industrial use cases assumes the availability of efficient sparse solvers, which although mature, are still generally less optimized for GPU than dense linear algebra routines.
Another limitation lies in the optimization dynamics. Although the sparse structure of the Jacobian improves memory efficiency, it can lead to slow convergence if the individual experts do not have enough parameters or are poorly initialized. In our experiment, we noticed that the training dynamics of the DD strongly depends on the initialization. While this may be attributed to the imbalanced training - first order for the DD and second-order for the PINN BALLS- -, properly adapting the DD to highly complex domains can be crucial in various applications since each PINN Ball is responsible for a narrow portion of the input space. We noticed that the quality of the LM optimizer largely bypasses these issues in the DD training. However, further analysis is required to match the quality of the training step of the DD to the effectiveness of the LM step for the PINN BALLS' submodels.
6 Conclusion
This paper introduced PINN BALLS, an efficient, scalable, second-order MoE for solving challenging PDEs in general domains. We achieve favorable scaling properties in the number of parameters using DD to construct an ensemble of local models. These models are effectively trained using second-order methods, while the DD partition is updated employing AAS, a framework for which we establish novel existence and stability theoretical results. PIŇN BALLS achieves state-of-the-as performance on challenging PINN benchmarks, showing improved model accuracy the number of submodels increases, and effectively addressing the memory requirements of the second-order methods without compromising their efficiency. We show that combining several submodels is effective strategy to achieve high accuracies when facing model size constraints, either due to hardware or memory limitations. Our results hint trade-off between the number of submodels and their individual capacity, motivating future research into the dynamical scaling of the number of experts.
Incorporating established PINN enhancements, such as Random Fourier Features 22), Temporal Causality 36), and Curriculum Training 200 could further boost the expressiveness and efficiency of the model, particularly given their strong synergy with second-order training methods 4 Another critical extension involves exploring anisotropic domain decomposition strategies and alternative energy penalization schemes, both of which are crucial for high-fidelity fluid dynamics simulations (30. Finally, performing the AAS ascent update with respect to the Wasserstein distance may result improving the convergence of the DD partition scheme in complex
in smoother updates [8] 12
PDE scenarios.
Impact Statement The model represented enables second-order training of PINN-styled architec- tures. Moreover, the theoretical backbone ensures a uniform coverage of the PDE domain through fully learnable domain decomposition. This model represents a step towards enabling PINNs as valu- able alternative to traditional PDE solution methods, and their inclusion into industrial applications.
10

---

Acknowledgements
We acknowledge fundings from BMW through the ProMotion programme; funding from SORR_KK__22__ project of the Basque Science Foundation, from the Spanish Ministry of Science, Innovation and Universities (Projects PID2023- 1100023-14999555 19119999-BBBI 100 and 11902022-134422B---2 Basque Government (Grant Nos. KK-2023/00012, KK- IK-2202/00030, and IT (T1504-22). Partial funding by BERC 2022-2025 program and the Spanish State Research Agency through BCAM Severo Ochoa excellence accreditation TEXX00210001 22-S/ICIN/AEE/1,, 13039/501 10001 1033, and through the project PID2024-1589940B 1589940B-CCC funded by MICIU/AEI/IO, 13019//////1011109//10000101 10001 1033 and cofunded
and the
by the European Union is also acknowledged. IM was supported by the Emmy Noether program of
the DFG, project number 403056140.
References
[1] L. Ambrosio, N. Gigli, and G. Savare. Gradient Flows: In Metric Spaces and in the Space of Probability Measures. Lectures in Mathematics. ETH Zürich. Birkhäuser Basel, 2005.
[2] S. Anderson, V. Dolean, B. Moseley, and J. Pestana. Elm-fbpinn: efficient finite-basis physics- informed neural networks. arXiv preprint arXiv:2409.01949, 2024.
[3] R. Bischof and M. A. Kraus. neural networks. In Proceedings of 33. Forum Bauinformatik. Forum Bauinformatik, 2022.
-oauur--o--xppett--eeeebbee meta-learnin for innsice-informed
[4] A. Bonfanti, G. Bruno, and C. Cipriani. The challenges of the nonlinear regime for physics- informed neural networks. Advances in Neural Information Processing Systems (NeurIPS), 2024.
[5] J. Bradbury, R. Frostig, P. Hawkins, Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milng al. JAX: composable transformations of Python+ NumPy programs. 2018.
[6] P.-Y. Chuang and L. A. Barba. Experience report of Physics-informed neural networks in fluid simulations: pitfalls and frustration. arXiv preprint arXiv:2205. 2022.
[7] S. Cuomo, V. S. Di Cola, Giampaolo, G. Rozza, M. Raissi, and F. Piccialli. Scientific machine learning through pwysics-informe informed neural networks: Where we and what's next. Journal of Scientific Computing, 92(3):88, 2022.
[8] M. Cuturi. Sinkhorn distances: Lightspeed computation of optimal transportation distances. In Advances in Neural Information Processing Systems 26 (NIPS 2013), pages 2292-2300, 2013.
[9] M. W. M. G. Dissanayake and N. Phan-Thien. Neural network-based approximations for solving partial differential equations. Communications in Numerical Methods in Engineering, 10(3): 1955)195-201, 1994.
[10] V. Dolean, A. Heinlein, S. Mishra, and B. Moseley. Finite basis physics-informed neural networks as a schwarz domain decomposition method. arXiv preprint N.Xiv:22/110555: 2022.
[11] L. C. Evans. Partial Differential Equations. American Mathematical Society, Providence, R.I., 2010.
[12] Feydy, Glaunès, B. Charlier, and M. Bronstein. Fast geometric learning with symbolic matrices. In H. Larochelle, M. Ranzato, R. Hadsell, M. Balcan, and H. Lin, editors, Advances Neural Information Processing Systems, volume 33, pages 14448 1188-1446 14462. Curran Associates, Inc., 2020.
[13) A. Heinlein, A. Klawonn, M. Lanser, and J. Weber. Combining machine learning and domain decomposition methods for the solution of partial differential equations-4 review. GAMM- Mitteilungen, 44(1) 2202222000,0 2021.
[14] Z. Hu, A. D. Jagtap, G. E. Karniadakis, and K. Kawaguchi. Augmented physies-informed neural networks (apinns): A gating network-based soft domain decomposition methodology. arXiv preprint arXiv:2211.0880 2022.

---

[15] G.-B. Huang, Q.-Y. Zhu, and C.-K. Siew. Extreme learning machine: theory and applications. Neurocomputing, (0(1-3):489-533 501, 2006.
[16) A. Jagtap and G. Karniadakis. Extended Physics-Informed Neural Networks (XPINNs): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations. Communications in Computational Physics, 28:2002-204) 11 2020.
[17] X. Jin, S. Cai, H. Li, and G. E. Karniadakis. NSFnets (Navier-Stokee Stokes flow nets): Physics- informed neural networks for the incompressible Navier-Stokes equations. Journal of Computa- tional Physics, 426: 109951, Feb 2021.
[18] R. Jordan, D. Kinderlchrer, and F. Otto. The variational formulation of the Tokker-Planch equation. SIAM Journal on Mathematical Analysis, 29(1):1-17, 1998.
[19] E. Kharazmi, Z. Zhang, and G. F. Karniadakis. hp-VPINNs: Variational Physics-informed neural networks with domain decomposition. Computer Methods in Applied Mechanics and Engineering, 374: 13547, 2021.
[20] A. Krishnapriyan, A. Gholami, S. Zhe, R. Kirby, and M. W. Mahoney. Characterizing possible failure modes in physics-informed neural networks. Advances in neural information processing systems, 34:26548-26566 2021.
[21] S. Li, M. Penwarden, R. M. Kirby, and S. Zhe. Meta learning of interface conditions for multi-domain physics-infomin informed neural networks. arXiv preprint arXiv:2210.1266 2022.
[22] S. Masoudnia and R. Ebrahimpour. Mixture of experts: a literature survey. The Artificial Intelligence Review, 42(2):275, 2014.
[23] B. Moseley, A. Markham, and T. Neysen Finite Basis Physics-Informed (FBPINNs): a scalable domain decomposition approach for solving differential equations. arXiv preprint arXiv:2107.07871, 2021.
Informed Neural Networks
[24] Müller and M. Zeinhofer. Achieving high accuracy with pinns via energy natural gradient descent. In International Conference on Machine Learning, pages 25471-25488 PMLR, 2023.
[25] Nocedal and J. Wright. Numerical Optimization. Springer, 1999.
[26] F. Otto. The geometry of dissipative evolution equations: The porous medium equation. Communications in Partial Differential Equations, 26(1-2):101-174, 2001.
[27] M. Raissi. Deep hidden physics models: Deep learning of nonlinear partial differential equations. The Journal of Machine Learning Research, 19119932-9555, 2018.
[28] M. Raissi, P. Perdikaris, and G. E. Kamiadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378:686-707, 2019.
[29] O. Sallam and M. Fürth. On the use of Fourier Seatures-Physieee Informed Neural Networks (FF-PINN) for forward and inverse fluid mechanies problems. Proceedings of the Institution of Mechanical Engineers, Part M: Journal of Engineering for the Maritime Environment, 237(4) 866, 2023.
[30] I. F. Sbalzarini, A. Hayer, A. Helenius, and P. Koumoutsakos. Simulations of (an) isotropic diffusion on curved biological surfaces. Biophysical journal, 90(3):878-88)) 885, 2006.
[3]] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le. G. Hinton, and J. Dean. Outra- geously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv: 1701.0638, 2017.
[32] K. Shukla, A. D. Jagtap, and G. F. Karniadakis. Parallel physics-informed neural networks via domain decomposition. Journal of Computational Physics, 447: 106683, 2021.
12

---

[33] K. Tang, J. Zhai, X. Wan, and C. Yang. Adversarial adaptive sampling: Unify pinn and optimal transport for the approximation of pdes. In The Twelfth International Conference on Learning Representations, 2024.
[34] C. Villani. Topics in Optimal Transportation, volume 58 of Graduate Studies in Mathematics. American Mathematical Society, Providence, 2003.
[35] Virtanen, R. Gommers, T. E. Oliphant, M. Haberland, T. Reddy, D. Cournapeau, E. Burovski,
P. Peterson, W. Weckesser, J. Bright, S. J. van der Walt, M. Brett, J. Wilson, K. J. Millman, N. Mayorov, A. R. J. Nelson, E. Jones, R. Kern, E. Larson, C. J. Carey, Polat, Y. Feng, E. W. Moore, J. VanderPlas, D. Laxalde, Perktold, R. Cimrman, I. Henriksen, A. Quintero, C. R. Harris, A. M. Archibald, A. H. Ribeiro, F. Pedregosa, P. van Mulbregt, and SciPy 1.0 Contributors. SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python. Nature Methods, 17:261-272, 2020.
[36] S. Wang. S. Sankaran, and P. Perdikaris. Respecting causality for training physics-informed neural networks. Computer Methods Applied Mechanics and Engineering, 111:11681,, 2024.
13

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Yes, the main claims in the abstract and introduction accurately reflect the paper scope and contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
The paper reserves specific section to discuss limitation of the model proposed.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g. if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
14

---

Answer: [Yes]
Justification: All the theoretical results are accompanied with solid proofs which are included in the appendix, for the sake of brevity, and sketched in the manuscript, in order provide an intuition to the reader. The assumptions made for each proof are also fully included (at times in the appendix).
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof' provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: The technical details used for all the experimental results are fully disclosed throughout the paper. Moreover, the paper does often rely on existing algorithms and methods which are properly referenced. Furthermore, the majority of the methods referenced are also available in common Python packages.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
15

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: Making the code public is currently in discussion with the partner institutions. Due to legal reasons, it might not possible to have released as open source. However, make the code available per request to the corresponding author.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips cc/
publicc/uiides/Codeie Seaedddse/CodeSSmmmisiiiiiiii
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ //nips.co httppss.coppbllc/guuueeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Details of the training and testing proceedures of all the experimental results obtained in the paper are provided in the paper and thoroughly discussed in the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: All the tests which include variability (such as initialization of the networks) are obtained for several runs, and are showcased alongside the variability obtained during training.
Guidelines:
16

---

The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error that would yield results that are out range (e.g. negative error rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: The experimental set up used to obtain the numerical results provided in the paper is fully described.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
The research conducted in the paper conforms with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
17

---

Justification:
The paper includes an impact statement
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g.. disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out, For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
The methods and models discussed in the paper do not pose such risk.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual- should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question:
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [NA]
Justification:
The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
18

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the datasets.paperssiii paperswi thcode. www.com//aaaa datasets package should be provided. For popular datasets,
has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not introduce any new specific asset.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
19

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: The core method development in this research does not involve LLMS as any important, original, non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy (https://neurips /att::///earrrpsssc//oooerrreee/ cc///oofeeeeeecs//22/////.. for what should or should not be described.
20

---

A Extension of the Adversarial Adaptive Scheme
For simplicity in this Section we will assume without loss of generality that 12 C Rdin is a compact set with volume [ and a regular boundary. We also define r(u(x)) = Ru(x)) f(z), and will ignore the boundary conditions for simplicity.
A.1 Gradient penalization
We consider the problem of this paper in abstract form: let U be space of functions Ω - R and let be the space of probability densities with weak gradient L², i.c. V = P(22) n n (Q), where we identify absolutely continuous probability measure with its density. Consider the minimax problem
min max (u, p), with J(u.p) := r(((())) / -3 weU PEV
Vv(x))
(14)
Remark A.1. Since we want consider the case where the optimal solution of (4 can be approximated by functions in U, it is natural to assume that the infimum of the residuals is zero. Moreover, so that the integration against the T function p in (1) is meaningful, we will assume that this infimum is meant in the 2 sense. This assumption is weaker than that in the original paper [33] which assumed to be surjection from U to C (S2); in the latter case a sequence {un]n can be easily constructed whose residuals converge to 0 even uniformly.
Theorem A.2. Assume that inf (1)2()))l2 = O. Then the optimal value of the minimax problem (6) is O. In particular, for any sequence {un, In minimizing [|r (u) there exists corresponding sequence of probability density functions {Pr In. each maximizing I (un:)) :), such that
lim I Pn). 0. 70000
(15)
Moreover, the probability density functions (puin converge strongly in H' (S)) the uniform distri- bution on Ω.
Remark A.3. Theorem A.2 improves upon previous results in at least two aspects. First, shows convergence solutions of the minimax problem using the H' regularization, which is much more widespread and practical than the Lipschitz regularization considered in the theory of [33]. Moreover, show that the sequence of marginally optimal probability distributions {pr converge to the uniform distribution. This is a general fact, not depending on the specific minimizing sequence (u, Jn. and rules out concentration of pa on regions with higher residual the total error converges to zero.
Proof of Theorem A.2 To begin with, note that for a given u. E U, a natural lower bound for max, I [u, p) is given by choosing the uniform density p(x) 1, this is:
max J(u,p) J(u / (u,,())))) 3 J Vi(x) = 1 ((n((t))de (16) PEV
This shows that 14 must be non-negative. Let now (u,), be minimizing sequence for | 2 Fix momentarily a given n; we will first show that there exists an optimizer for the maximization problem maxpev J (u. p). Let p E V be a candidate with a higher or equal score than the uniform density 5; then, by (16) we must have:
2 (un(())) < [ (((n(()))(()) - / \p(x))
8 8/1 [p(a)) S
nn (z),(()) (p(x) dz S 11p 1 (
s [[IVllla) (u,))lz
where C is the constant of Poincare's inequality I1 (since p a probability measure, and Ω has unit volume, the average value of p is I; also note that Ω has a regular boundary by assumption). Dividing by the norm of Vp then yields:
Vpll2 < (l,))(u,))l
(17)
21

---

and using the Poincaré's 's inequality once again on the left-hand side results in an estimate for the deviation p with respect to the uniform density:
dlp - 1112 - VVlle & (un)l(,))l..
(18)
We have shown an H' bound for any probability density that has a higher score than the uniform density ş by I ). Consequently, the superlevel sets of J",, compact for the H ( weak convergence, given maximizing sequence {ph he there exists (weak) cluster point Pa. which must maximizer for J (un:)) -) due to the upper-semicontinityy
(un::))
sup I [un, p) lim sup I (Uns Ph) S I (un: S sup [Un p). PEV
(19)
k
PCV
Finally, since Ph satisfies 18 holds:
t llpn -1112 S /VPol22 fllr2
(20)
Since (u,) converges to zero in 1,2,, we obtain both strong 1,2 convergence of pa to the uniform distribution p = 1/|S2] and strong T,² convergence the gradient Vp, to zero as n - co. Combining the convergence behavior of (u,)), and {pr Jn yields (15)
A.2 Entropy penalization
Although intuitive, the gradient penalization term in (1) assumes that the probability density has gradient, which may exclude some interesting cases. Besides, there is no mechanism preventing p to reach zero, which could imply that the residuals in some regions may not be accounted for at all. Inspired by the field of optimal transport, and in particular Wasserstein gradient flows (see Section A.3 for a discussion), we propose an alternative penalization term that solves these issues. We consider:
min max J(u,)), wEU PEP())
with J Ju, = r(u(()) - p(x) logp(x)dz. (21)
An interesting observation is that in this case the optimizer of the nested maximization step can be given in closed form, given by Lemma This is elementary result; however gather the proof in Section A.4 for completeness.
Lemma A.4. Assume r(u(z))² E C()). Then the optimizer of the problem
min 3 / log - r(u(()))p(z)ddz PEP(N)
(22)
given by
p" (x(x) Cexp (r(((()))))))
(23)
with C = (Sn exp (r() a normalization constant ensuring pd has unit mass.
Note that an assumption on the regularity of r²(((x)) is necessary in order obtain existence of optimizers, though tighter assumption may in principle be used (e.g. upper lower-semicontinuity and boundedness).
As in the previous Section, we will need assume certain regularity of the residuals. Since in this case the regularization term only requires to be probability distribution (which is a rather weak assumption), we need a stronger setting for (u). In this case, the appropriate assumption that inf,, (w))()))) = 0, this is, along at least some sequence, the residuals are continuous and converge uniformly to zero. Again, this assumption weaker than the assumptions in (33, since the latter implies the existence of some u" such that 7² (u") = 0.
22

---

Theorem A.5. Assume that infuev (1)(lc))) = 0. Then the optimal value of the minimax problem (1 is O. In particular, for any sequence {ua)n minimizing (1)((n)))))) there exists corresponding sequence of probability density functions {Pa Iw each maximizing I (Uny 1) such that
limn I (Un: Pn) =
(24)
T 70000
Moreover, the probability density functions (poIn are everywhere positive and converge uniformly to the uniform distribution P
Proof. First note that, as in the proof of Theorem A.2 natural lower bound for max, J (u, p) is given by choosing the uniform density p(x) 1. Therefore the optimum of (21) must be non- negative.
Now consider [4,, Ju minimizing sequence for (1))())) Under the assumptions of the theorem, by Lemma ea. un associated with the optimal probability distribution given by (23). that we denote Pn (with corresponding normalization constant Cn 1/ Ja exp(r?(u, (x) ())) Then the minimax problem becomes:
(ln, / (r( (2)) - Blog Pm (1))dd
PEP(O)
(n()) (x)) - 8 log C₂ - (un = -B On.
(25)
where we used that pn has unit mass. Finally let us bound 25 noting that S exp(r (nn))))BB)))))))))))))))))))){)) (xp())r(ux))lle
1/C..
/ exp(1 (up (x))/ dx S /
1 1=
ldr <
= exp((lr((un))
jn
which after taking the logarithm becomes:
|| ||=
0 < log Cn <
(26)
13
Plugging this bound into (25 and using that (u,>(un)) value (2T) is zero. The associated Cn converges to I, since r² (u,) converges uniformly to zero, and therefore P.. converges likewise uniformly to the uniform density.
0 concludes that the minimax optimal
A.3 Relation to Wasserstein gradient flows
The theory gradient flows deals with the generalization of curves steepest descent the form
a Fx)
(27)
to general metric spaces, where neither the left-hand side nor the right-hand side may be well defined. This framework has proven be extremely fruitful in the field of PDEs, since the founding work of
Jordan, Kinderlehrer and Otto [8] 26) showed that many PDEs of interest can be described as curves of steepest descent with respect to certain functionals the space of probability measures endowed with the Wasserstein distance (see e.g.
for an overview).
To allow for general formulation in this abstract framework, it is customary to employ the formalism of minimizing movements. This is, for a metric space (M, d). a functional F M + R U {+oo), an initial condition In E M and time step
considers the discrete trajectory generated by:
xf E arg min read 2T
d(x, xt 21 + F(x),
A 1
(28)
When (M, d) is diite-dimensional Euclidean space, 28 can be shown to reduce to an implicit Euler scheme approximating (27). But the benefit of this general framework stems from its application to functional or measure spaces, well its ability to handle non-smooth functionals. Then, under suitable assumptions on F, M and d, the discrete trajectories can shown to converge to limit trajectory, which can then be characterized with tools from the theory of partial differential equations.
23

---

In relation to the matter of this paper, we have considered two regularization terms: the Dirichlet energy (p( Sa Vpl and the entropy (p(() Sa p(x:) log p(z)dr. The gradient flow of the Dirichlet energy with respect to the (2t) topology yields the heat equation [18]. The heat equation tends to homogenize a given initial condition, which is desired property in the minimax scheme we studied. Interestingly, the heat equation can be also obtained as the gradient flow of the entropy Gg. in this case on the space of probability measures ((2)) endowed with the - Wasserstein distance [1].. Subsequently, we can expect both functionals to perform similarly a gradient-descom framework, as long gradients and Lassessertenin for Gg). It remains somewhat surprising that in our numerical experiments the entropy seems to be better hehaved even using the importance sampling update rule (33), which has an T,2 flavour. We consider this further proof of the robustness of the method, and plan to investigate purely Wasserstein updates for the probability mesaure p in future work.
computed with respect to the right metric (this is, L for GD
A.4 Additional proofs
Proof of Lemma A.4 For simplicity let us define E(p) : 3 Sa p(ar) logp(x)) - -
minimizer p" exists by sheer lower-semicontiniity of & and compactness of P(S2) with respect to the weak topology. Moreover, any minimizer must have positive density; otherwise the objective value would be infinite, since the uniform density a feasible candidate with finite objective. This means that by slight abuse of notation we can use p* also to refer to its density.
Let 0 be an arbitrary function in L (2, satistying Sa Oy" O, and consider the perturbation of p" given by Pe : (1 s0)p'. It follows that for small e, Pe remains in P(2). Let compute how E (p")) changes under a perturbation:
Epp)) - E(p") = 3
Pe Pe(r) - (2)) de - (=)()))0()))) (c))d,
g
&
which, using the monotonicity of the function log(s), by dominated convergence yields in the limit s 0:
0 S / ((z)pp [B)@p) (u(x)))
(29)
Since 29 must hold also by replacing 0 by 8, the inequality turns into equality. Moreover, since 0p" is an arbitrary zero-mean L function, the expression between brackets must constant. All in all, parametrizing this constant offset by B log C, with C 0, yields:
Blog "es""(ar (x) (u(()) = 3 log C = p" = C exp())(u(z))/3)
(30)
Thus identify C' normalization constant enforcing that p" is probability distribution, i.c.
c- (=((xxx("(uu)))))
The boundedness of (u) guarantees that C' is finite, and that p" is bounded away from zero, while focusing the attention on the regions with larger residual. Whenever the residual is uniform, p* will become uniform as well.
B
Details of the PDEs
B.1 Helmholtz Equation
The Helmholtz equation is diffusive-reactive second-order equation similar to the wave equation. Although linear, the Helmholtz equation represents a challenging PDE instance for PINNs due to the high-frequency components included in its solution. In our paper, we choose to solve the bidimensional formulation, which is highlighted Equation 33 the domain l = =0,1]²
Au kokyu f. u(x, v) = 0, x, y E a
x, y Ω
(31)
with s(x,y) y) -ki sin (k.yar sin (kyy)) For our test case we choose k, and ky 1, analogous to that presented in [[4] for comparability. Each submodel consists of as little as 2 hidden
24

---

layers with 10 neurons each. Training is performed with N, 10* collocation points for training the PDE residuals, initially sampled with latin hypercube sampling, and N,, = 3 10 points for training the boundary and initial condition in 0R. We train PINN BALLS with 2 hidden layers and 10 neurons per layer for 2000 iterations, each iterations consists on single LM step on and 20 Adam iterations on Ps. Figure 4 showcases the result of the best performing model on Helmholtz equation, alongside the learnt DD.
Predicted Solution
True Solution
L - Error
0.24
0.16
0.16
36
0.6
0.08
0.08
3.0
0.00
0.00
24
0.4
-0.06
-0.00
1.8
0.2
-0.16
-0.16
12
-0.24
-0.24
0.6
8.0
0.2
0.4
0.6
0.8
1.0
-0.32
0.0
0.2
0.4
0.6
0.8
1.0
-0.32
0.0
0.2
0.4
0.6
0.8
1.0
Figure 4: Prediction with learnt center and radii of the PINN Balls (left); correct solution (center); and 1 Error achieved by our best performing model on Helmholtz equation.
B.2 Burgers' Equation
Burgers' equation ID version of Navier-Stokes equations without pressure. Its solution high times present discontinuity, which makes it challenging for spectrally biased architectures. The specific instance chosen our numerics for Burgers' equation the same [27) In particular, we refer to the exact same data provided by the authors. In particular, given (I, T) E Ω = [-1, x [0, 1]. we solve for u R the following equation:
0,u + ud,u voLu 0,
(z, E # [-11),
u(x, 0) = - -in((
(32)
u m(-1,7) 7(1,7)= =0 T [0.
with the diffusivity v being equal publicly by the authors of (27). Training is performed with N, = -0 collocation points for training the PDE residuals, initially sampled with latin hypercube sampling, and N₂ training the boundary and initial condition aQ We train PINN BALLS with hidden layers and 10 neurons per layer for 2000 iterations, each iterations consists single LM step and 20 Adam iterations on Ps. Figure showcases the result of the best performing model on Burgers' equation, alongside the learnt DD.
for this specific instance. The correct solution is provided
103 points for
B.3 Navier-Stokes Equation
The last and most crucial PDE instance considered is given by the incompressible Navier-Stokes equations. particular, solve the fluid flow past a 2D cilinder presented [17]. We consider the same instance proposed in the aforementioned paper, given by (ar, V. E Ω 2.5. [- 2.5, 2.5 x [0, 16) with the goal of finding Ω - T defined as ([( y, [u(x,, y, 1), o(x, 1), p(x, 0) In this formulation and represents respectively the horizontal and vertical components of the fluid
25

---

Predicted Solution
Solution
| ule. x)
1e-2.
1.00
0.50
0.25
0.00
0.25
0.25
-0.25
-0.50
-0.50
-0.75
0.75
0.75
0.0
0.2
0.4
0.6
0.8
1.0
-1.00
0.0
0.2
0.4
0.6
0.8
1.0
-1.00
0.0
0.2
0.4
0.6
0.8
1.0
Figure 5: Prediction with learnt center and radii of the PINN Balls (left); correct solution (center); and L Error achieved by our best performing model on Burgers' equation.
velocity; and the pressure. Navier-Stokes equations are expressed in vectorized form as follows:
0,u ud v0,
(0.u + 033)) + 0.p 0, (z, t) S2
Re
1
0,0 ud,u + v0yr
(0,u + 030) + 0,p
0, (z, Ω
Re
0,u + 0,1 0,
(I, EΩ,
(33)
u(z, y, 0) guo (x, y),
(z,y) E |-2,5, 7.5] 22, 2.5].
v(z v, 0) = Goo (a,, v),
(x, (-255,77 x 2 2.5].
u(2.5, y, t) 1,
(w,t) [-2.5,25) x [0, 16),
v(2.5, U. = O.
(w.t) E (-2.5,2.5) 2.5] x 16].
where Re is the Reynolds' number, which is an adimensional quantity defined by the problem and is to 100 for our case. The initial conditions (9ua: 9v,) can be found in the repository of [27], as well as an high fidelity solution which is used ground truth. At = -2.5 the fluid velocity at the inlet is imposed. Further conditions are given by the presence a cylinder centered in (ar, y) (0,0) 0) with radius 0.25. Furthermore, where the no-slip condition can be chosen (u e or the correct solution can be given as boundary condition. Since the simulation provided in 27 refers to a free-flow stream, we use the correct solution the boundaries.
additional condition appears at the borders, namely where y +2.5.
To train our PINNs, we use N, = 5 10 collocation points for training the PDE residuals, sampled with latin hypercube sampling. and N₂ 2- 10 points for training the boundary and initial condition in 0R. Morever, at every iteration. we minimize the loss on random batches of the training data, respectively 10 points for the residuals and 10a for boundary and initial condition. Figure showcases the result of the best performing model on Navier-Stokes' equations, alongside the learnt DD. Notice that we not normalize the pressure predicted by the model, but do compute the absolute error by comparing the normalized pressures.
26

---

5375
633
62
02
E031
E003
00
0.
E24
00
-2.6
-0.
66
533
64
611
0.00
-1.0
Figure 6: Prediction with learnt center and radii of the PINN Balls (left); correct solution (center); and Absolute Error achieved by our best performing model Navier Stokes' equations. The top row refers to the horizontal velocity u, the central row to the vertical velocity v, and the bottom row to the pressure p.
27


Reviews:
[{"review_id": "bdE4TezNyK", "reviewer": "Reviewer_gnjE", "summary": "The paper introduces PINN BALLS, a sparse Mixture-of-Experts architecture for Physics-Informed Neural Networks (PINNs).\nKey ideas are:\n- trainable radial-basis “balls” that decompose the domain and create a partition of unity.\n- second-order Levenberg-Marquardt updates applied locally so Jacobians stay sparse and fit GPU memory.\n- Adversarial Adaptive Sampling (AAS) used as the gating distribution, with new convergence theorems that replace the original Lipschitz constraint by Dirichlet, and entropy-based regularizers.\n- extensive empirical study on Helmholtz, Burgers, and 2-D Navier–Stokes benchmarks plus ablations on memory/runtime.\nA comparison table highlights how the approach differs from XPINN, FBPINN and APINN (sparsity, no interface loss, adaptive sampling, second-order training)", "questions": "- Have you attempted to scale PINN BALLS to 3D problems or complex multi-physics PDEs? Would sparse Jacobian construction and solver performance remain tractable in such high-dimensional or coupled systems?\n- Could you elaborate on why you chose quartic radial basis functions for the gating balls? Did you explore other compactly supported kernels or make the RBF shapes learnable? How sensitive is the model to this design choice?\n- How sensitive is model performance to the number or initial placement of balls? \n- Would your method require significant architectural or training changes to handle time-dependent PDEs.\n- Can you provide more quantitative analysis of the time and memory trade-offs of sparse Jacobian construction and sparse LM solving compared to dense LM or first-order optimizers?", "limitations": "- Reproducibility is limited by the lack of a code release at submission time, though the authors state their intent to share it. Open access would allow others to test and verify scalability claims.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- While mixture-of-experts in PINNs is not a new idea, and prior works already explored sparse gating and soft domain decomposition, the idea proposed here is to combine the benefits of DD, second-order training methods and Adversarial Adaptive Sampling (AAS)\n- A radial basis “ball” formulation is derived for soft domain decomposition that leads to both partition-of-unity gating and Jacobian sparsity. \n- The paper provides theoretical convergence results for Adversarial Adaptive Sampling (AAS), relaxing prior assumptions and generalizing earlier proofs using Dirichlet and entropy regularizers.\n- In addition to empirical evaluation, Authors provide ablation studies across ball count, sampling, and sparsity.\n- Transparent discussion of limitations such as sparse Jacobian overhead and current AD support issues, which raises my respect and trust to their science work quality.\n\nWeaknesses:\n- No ablation study on the effect of different components the authors proposed in their study. For example, comparing the following variations can shed light on which components are truly necessary and and where most of the gains come from: Full PINN BALLS, vs without AAS, vs Without LM update step, etc.\n- Lack of comparison with non-PINN state-of-the-art models: Since the ultimate goal is to train an accurate surrogate for PDE solutions, it is important to evaluate how PINN BALLS performs against other state-of-the-art models in terms of predictive accuracy, regardless of the underlying optimization algorithm or architecture type. I suggest that the authors include a comparison with data-driven models (such as FNO or DeepONet) as part of their experimental evaluation.\n- No public code at the time of submission exists and their answer to NeurIPS checklist question must be \"No\". This can limit the reproducibility. Even an anonymized repo for the results reported in this paper would help.\n- The notation can be dense and intimidating for readers unfamiliar with second-order methods and sparse linear algebra. A high-level algorithm box or diagram of the training pipeline would improve clarity.", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the reviewer for the positive feedback and for the constructive comments that help us improve the quality of our manuscript. In particular, we acknowledge the lack of an ablation study analyzing the contributions of the algorithmic components, specifically the AAS and the LM optimization step. From our experience, the LM step plays a critical role in achieving the reported results. In the context of domain decomposition, first-order optimization methods tend to show a significantly slower convergence as the number of basis functions increases. This is especially problematic in nonlinear regimes such as those arising in the Navier-Stokes equations, where convergence is often not achieved. Conversely, the AAS component can improve model accuracy by one to two orders of magnitude in such nonlinear settings, although its impact is minimal in linear cases. We plan to explicitly incorporate this discussion, along with a high-level diagram of the algorithm - as suggested - , in a revised version of the paper.\n\nBelow, we provide specific responses to the reviewers’ questions:\n* The central focus of our work is to demonstrate scalability with respect to the number of trainable parameters — arguably one of the primary bottlenecks in applying second-order optimization to overparameterized physics-informed models. While we do not directly address time-dependent 3D problems or complex multiphysics scenarios in this study, the proposed methodology is designed keeping such extension in mind. In fact, high-dimensional and multiphysics problems often demand the precision and convergence speed of second-order methods to be tractable. Even in the time-dependent 2D setting of the  Navier-Stokes equations considered here, we encounter severe memory limitations: the computation of the dense Jacobian $D_\\theta L$ and its inversion leads to an out-of-memory error at approximately 2000 parameters on our GPU implementation, due to the need to backpropagate through all collocation points and equations of the PDE. Our domain decomposition strategy directly addresses this scalability issue by enabling a block-wise construction of the matrix $D_\\theta L^T D_\\theta L$. Notably, the dimensionality of $D_\\theta L^T D_\\theta L$ depends only on the number of model parameters, which means that at this point in the algorithm memory growth is decoupled from domain/PDE complexity. This allows our framework to remain tractable in principle even as the problem dimensionality or physical complexity increases — a direction we view as a natural extension of this work.\n* The selection of RBFs in our framework is relatively flexible. Our aim was to ensure compact support and avoid steep derivatives. We experimented with Gaussian kernels (non-compact support) and second-order finite-element bases (with discontinuous first derivatives) and observed no significant impact on model performance. This robustness can be attributed to the efficacy of second-order optimization in the training process.\n* Under reasonable conditions, the model’s performance is not highly sensitive to the initial number or placement of basis functions. Nevertheless, a well-chosen initialization—such as the strategy proposed in our paper—is important to maintain the sparsity of the Jacobian $D_\\theta L$. Excessive or poorly placed basis functions can lead to substantial computational overhead during the construction of $D_\\theta L^T D_\\theta L$, as discussed in the Limitations section.\n* This is an excellent point. In principle, domain decomposition in PINN Balls is not restricted to spatial variables. In the presented Navier-Stokes experiment, we address a time-dependent case, demonstrating that the approach is compatible with temporal domains. One could implement DD either in space alone or across both space and time, depending on the application. We believe that future work could significantly benefit from more sophisticated time-marching strategies, and we intend to explore this in subsequent research.\n* The Jacobian is constructed by iterating over each submodel, requiring backpropagation to compute partial derivatives for each. This cost is roughly proportional to computing a standard physics-informed loss, multiplied by the number of submodels. For a standard implementation that loops over the model, this means that the computational time of each step would grow linearly with the number of submodels. Therefore, adopting an efficient and parallel implementation is essential to keep this step tractable and the overhead minimal. Regarding solver performance, Figure 1 illustrates the advantage of a sparse LM solver compared to its dense counterpart. In principle, a dense solver is more effective, however, the dense implementation fails before reaching 8000 parameters, in contrast to the sparse model which deals with more than 10'000 with ease. When contrasted with first-order methods, the LM step achieves several orders of magnitude improvement in loss with a fraction of the iterations (few thousands vs. several hundreds of thousands). Moreover, in the DD setting—especially for complex problems like the Navier-Stokes equations first-order methods fail to converge altogether. We will include this discussion, in a revised version of the paper.\n\nLastly, we acknowledge the reviewer’s note regarding code availability. The release of the code is currently under discussion with our primary partner institution. However, we are happy to share it upon request through the corresponding author."}, {"review_id": "QVdQTzhJf6", "reviewer": "Reviewer_DRVa", "summary": "PINN Balls is a model that combine second-order training with domain decomposition on PINNs. An overlapping RBF-based domain decomposition (\"balls”) is used to create a sparse Jacobian. Sparse Levenberg–Marquardt updates are adopted for second-order optimization. Adversarial Adaptive Sampling (AAS) is used to learn the decomposition along with the PINN training. Experimental results are presented on 2D Helmholtz, Burgers, and Navier–Stokes equations.", "questions": "1. Consider adding experiments on larger-scale PDEs, e.g., 3D or fine-resolution 2D, where dense LM cannot run due to GPU memory limits, but sparse LM succeeds and outperforms other models. \n\n2. How sensitive is the method to the initialization of RBF centers and radii?\n\n3. PINN Balls relies on AAS to train the DD. What are the benefits of AAS to residual-based methods? An ablation study of PINN Balls with residual-based sampling and AAS (with Dirichlet/entropic regularization) reporting final errors and convergence speed would be a good addition in the appendix.", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n-- Scaling higher-order methods is an open problem. This paper considers this problem along with learning the adaptive domain decomposition.\n\n-- Results on Helmholtz especially is impressive. \n\n-- Extension of AAS theory, proving convergence to sampling under Dirichlet or entropic regularization.\n\nWeaknesses:\n\n-- Although sparse Jacobians allow for larger models, all benchmarks are on small 2D PDEs solvable by dense Levenberg–Marquardt. There is no demonstration of a problem that dense second-order methods OOM but PINN BALLS completes. \nMoreover, assembling the CSR matrix has non-negligible overheads at moderate to large scale and could outweigh the gains from sparsity. GPU sparse kernels are also not as efficient as dense kernels.\n\n-- Fluctuating losses with increasing balls point to unstable AAS maximization.\n\n-- Benchmarks are three textbook PDEs on 2D domains. No demonstration on multiphysics and heterogeneous boundary conditions where PINNs often struggle.\n\n-- Comparisons to XPINN/APINN use reported numbers rather than reproduced results under identical settings.\n\nThis paper presents good small-scale results, but falls short on demonstrating practical scalability which is the motivation for sparsifying second-order PINNs.", "was_revised": true, "final_justification": "I appreciate the authors response but will stick to my original score. For the canonical problem of flow past a cylinder, authors note that PINN (dense LM) runs on only a relatively small number of parameters (unclear how small?) and in this setting, the average velocity field is reasonably approximated but the unsteady dynamic of vortex shedding is not captured. However, there are no visualizations to validate these claims. In general, the scaling of PINN balls needs stronger validation.", "rebuttal": "We would like to express our sincere appreciation to the reviewer for highlighting several important areas for improvement in our work. We have carefully considered the points raised, and we address them in detail below.\n1.  Domain decomposition enables tractability in large-scale or fine-resolution 3D problems through the parallel, block-wise construction of the Jacobian matrix. Importantly, the main computational bottleneck in the Levenberg-Marquardt (LM) optimization step — the inversion of the matrix $D_\\theta L^T D_\\theta L$ — depends solely on the number of trainable parameters in the model, which is the core focus of our paper. Notably, the time-dependent 2D Navier-Stokes example already illustrates a  scenario where applying a dense LM solver leads to out-of-memory errors. For quantitative context, in our GPU implementation, the LM update fails for architectures with more than approximately 2,000 parameters due to limitations in computing the dense Jacobian. While prior work has shown that second-order methods enable accurate training of PINNs even with a low number of parameters, the strategy we propose can extend this framework to a more consistent number of parameters, which is at times necessary to predict certain PDE solutions. For instance, in our experience, 2000 parameters are insufficient to model the vortex shedding behind a cylinder that we present. We intend to emphasize this point more clearly in a revised version of the manuscript.\n2.  Under typical conditions, model performance is not highly sensitive to the initial number or placement of basis functions. However, careful initialization—as described in our paper—is essential for preserving the sparsity structure of the Jacobian $D_\\theta L$. Excessively many or poorly placed basis functions can result in significant computational overhead when constructing the matrix $D_\\theta L^T D_\\theta L$, a limitation we have discussed in the paper but will clarify further.\n3.  To the best of our understanding, the AAS mechanism functions as a residual-based method. It adaptively concentrates model capacity on regions with large residuals, thus facilitating a more effective domain decomposition. In addition to its conceptual alignment with DD, AAS also offers desirable convergence behavior. Empirically, we find that AAS improves model accuracy by one to two orders of magnitude in certain nonlinear regimes, such as near the cylinder in the Navier-Stokes setup. We will make this mechanism and its empirical benefits more explicit in the revised version of the paper.\n\nWhile we did not directly reproduce the results of XPINN and APINN, we aimed to ensure a meaningful and fair comparison by adapting our model to align closely with their training setups — for instance, by matching the number of collocation points, model parameters, and activation functions wherever applicable.\nIn addition, the reviewer raises a valid concern regarding the oscillatory behavior of the loss function when PINN Balls are instantiated with multiple submodels. We acknowledge that these fluctuations often stem from instability in the AAS maximization. However, in our experiments, such behavior typically occurs only in very low-loss regimes, where the optimization is approaching convergence. Employing a stronger exponential decay in the AAS update tends to mitigate these fluctuations effectively. We will ensure these points are better articulated in the manuscript."}, {"review_id": "33zTRRLFwl", "reviewer": "Reviewer_MXFf", "summary": "This paper introduces an ensemble model for PINNs, where each submodel learns a compactly supported Radial Basis Function (RBF). By ensuring the supports of RBFs cover the entire domain, the ensemble model autonomously learn domain decomposition. Employing the Levenberg-Marquardt update rule and Adversarial Adaptive Sampling (AAS) with a specified probability density function, it achieves notable performance. Furthermore, the paper includes theoretical extension of [33] to a more general class of regularizers.", "questions": "1. If there is an empirical result on complex domain such as dolphin-shaped interface, it would strengthen the generalizability of the proposed method.\n2. The theorem said that probability density function $p_n$ converges to a uniform distribution. If this is observed in real experiments, could we expect the similar results by using uniform distribution rather than using AAS?\n3. Considering the ensemble nature of the proposed framework, could the submodels be added in an adaptive manner? For example, after training PINNBalls-10, would it be possible to introduce an additional set of 10 submodels whose RBF centers are initialized in regions exhibiting high residual error? If the support of these new submodels is restricted to those high-error regions, this strategy might enable further improvement in accuracy without degrading performance in areas that are already well-learned.", "limitations": "yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. The proposed method autonomously learns domain decomposition without requiring any a priori knowledge and eliminates the need to explicitly manage interfaces between subdomains.\n2. The model outperforms existing baselines across several benchmarks.\n3. The theoretical results on Adversarial Adaptive Sampling (AAS) are extended to more general settings, increasing their practical relevance.\n\nWeaknesses\n1. While accuracy improves with an increasing number of submodels, this comes at the cost of learning stability. Such instability may hinder deployment in real-world applications.\n2. Although the concept of learnable domain decomposition is appealing, the paper lacks experimental validation on complex domains where such decomposition would be most advantageous.\n3. Several important implementation details are provided only in the appendix and are not referenced in the main text. For example, the network architecture used in experiments is described in the appendix, but not mentioned or cited in the main body of the paper.", "was_revised": true, "final_justification": "Although my questions are not fully satisfied, they are not so significant to decrease ratings: Original paper include applications on non-trivial domain (simpler than dolphine, but sufficiently non-trivial) and applying adaptive strategy is potential benefit of the proposed method rather than necessary experiment to prove its performance. Consistency between theoretical and experimental result on trained distribution should be declared and rebuttal sufficiently has addressed it. I still suggest acceptance and maintain my rating.", "rebuttal": "We sincerely thank the reviewer for the positive feedback and for the insightful comments that help improve the clarity and rigor of our manuscript. Regarding the concern raised about training stability, we observe that fluctuations in the loss function primarily occur in low-loss regimes, typically when the training is nearing convergence. In practice, applying an exponential decay to the learning rate of AAS effectively mitigates this instability and stabilizes the training process. We will make this aspect more explicit in the final version of the paper and, in line with the reviewer’s suggestion, include additional implementation details in the main body of the manuscript.\n\nWe now address the specific points raised:\n1. While we did not explicitly consider cases involving highly irregular interfaces such as dolphin-shaped domains, the reviewer can refer to our Navier-Stokes experiments for a qualitative understanding of how the learned basis behaves in more complex scenarios. In such settings, the domain decomposition (DD) mechanism typically allocates higher resolution to regions where the PDE solution exhibits sharp gradients or localized complexity. This behavior is naturally guided by the residual-based adaptation mechanism embedded in the training process.\n2.  We appreciate the reviewer’s insightful observation regarding the distribution of submodels. Indeed, our initialization procedure yields a probability density function that is approximately uniform. In relatively simple or linear scenarios—such as the Helmholtz equation—we observe that the domain decomposition remains mostly unchanged during training, which aligns with theoretical convergence expectations. However, in nonlinear settings, such as the Navier-Stokes equations, the AAS paradigm introduces meaningful deviations from the uniform distribution. These deviations serve to allocate more submodels in regions with higher solution complexity (e.g., steep gradients). While theoretical convergence to a uniform distribution assumes that the residuals vanish entirely — a reasonable but idealized assumption — this is rarely the case in practice. Consequently, the final distribution does not always converge to uniformity, but this deviation often enhances the model's accuracy. We view this tension between theory and practice as an important and interesting topic that merits deeper exploration in future work.\n3. We find the reviewer’s proposed approach—dynamically adding submodels when convergence plateaus—both elegant and highly promising. This strategy could help avoid overparameterization by preventing the initialization of an unnecessarily large number of basis functions, and instead adaptively finding an optimal balance between submodel count and resolution. Furthermore, leveraging transfer learning for the initialization of newly introduced submodels could prevent disruption of ongoing training. We believe this direction holds great potential and plan to pursue it in future research.\n\nWe are also grateful to the reviewer for pointing out typographical errors in the manuscript. These have been promptly corrected in the revised version."}, {"review_id": "7cjCCMFAPY", "reviewer": "Reviewer_MZup", "summary": "The submission suggests algorithmic improvements to physics-informed neural networks (PINNs): it proposes to combine domain decomposition, adaptive sampling, and second-order optimisation.\nSpecifically, the approach is to replace a single neural network with an ensemble, combining the outputs of each ensemble member using a locally-supported weighting function (akin to a partition of unity).\nThe collocation points for the PINN are sampled from the same weighing function, and the resulting model is trained using the Levenberg--Marquardt optimiser. \nThe submission also includes an analysis of the existence of solutions to the minimax problem underlying the adaptive sampling algorithm, which extends a previous result. Numerical results show that the resulting algorithm is competitive with prior approaches.", "questions": "Here is a collection of some more minor questions. They don't affect my score, but they might be helpful feedback for revising the paper. (I'd still be interested in their answers, though I understand that rebuttal time and space may be limited.)\n\n1. Line 187: \"Ensuring \\phi > 0 can be done by choosing a sufficiently large $s$ at initialisation\" -> is this always true?\n2. Equation 10: Does this loss assume the negative entropy regularisation? \n3. Equation 10 (again): How does the domain decomposition interact with the boundary?\n4. Equation 11: I thought Equation 10 (\"J\") was the preferred loss, not Equation 4 (\"L\"). What have I missed? Same in Line 228.\n5. Line 241: The exponential decay in learning rate feels quite aggressive. How important is that? And what kind of exponential decay is used?\n6. Line 245: What's the domain of the problem? The target function $f(x) = sin(4x)sin(y)$ has 4 \"bumps\" if x and y are between 0 and $\\pi$. Depending on the domain, should there be a \"correct ensemble size\" of four, or is the number of ensembles less straightforward to imagine?\n7. Line 247: What are $K$-means applied to?\n8. Line 266: Are those results in the paper? I can't find them, but I would like to see them.\n9. Line 268: To what extent is the condition number of $DL$ relevant for Figure 1? Doesn't the experiment measure a single step of Levenberg-Marquardt implemented with a direct solver? In my understanding, this should be independent of any condition numbers. Have I missed something?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\nThe submission is a nice paper, and I enjoyed reading it. It does not just contain a well-explained combination of ideas; merging domain decomposition and adaptive sampling appears to be effective, and the resulting algorithm is sufficiently simple to motivate reimplementation in future PINN studies. Furthermore, the theory is novel, convincing, and a valuable addition to the literature on adaptive sampling for PINNs.\n\n## Weaknesses\nWhile I recommend accepting the paper, I am currently assigning a borderline score due to the weaknesses explained below. I am open to revising my evaluation if my concerns turn out to be easy to resolve.\n### Second-order optimisation\n\nThe second-order optimisation parts of the paper are not as convincing as the other components. The main optimisation contribution of the paper is replacing [33]'s stochastic gradient descent with Levenberg--Marquardt. However, the paper does not demonstrate that for the PINN-Balls, second-order methods outperform first-order methods. Instead, the motivation for choosing second-order methods is that they are promising for PINNs, as indicated by recent studies [4, 24]. \n- If using second-order optimisers for PINN-Balls is considered a contribution of the paper, there needs to be a corresponding benchmark or analysis. \n- If referencing [4, 24] is considered a sufficient argument for using second-order methods, then I think that using second-order methods for PINN-balls is not a contribution of this paper.\n\nSection 4.1 is dedicated to discussing the sparsity of the Gauss--Newton matrix $D L^\\top DL$ with the corresponding savings in memory and runtime. This perspective overlooks matrix-free linear solvers, such as conjugate gradients, which do not need sparsity as long as matrix-vector products are efficient. For $G=  DL^\\top DL$, these matrix-vector products can be implemented with automatic differentiation. I believe this omission is a weakness in Section 4.1; a discussion or benchmark of alternative linear algebra approaches would be appropriate.\n\n### Experimental clarity\n\nThe numerical experiments leave several open questions, and I believe this hinders the appreciation of the results to the extent they should be appreciated. For example, I would like the following information added to the context of Figures 1 and 2 and Table 2:\n- Figure 1:\n    - The sparsity of the $DL^\\top DL$ matrix should depend on the distribution of points, so it should vary quite a lot during optimisation. Why is this experiment, which benchmarks only a single Levenberg--Marquardt step, representative of the sparsity of the system matrix?\n    - All but the left-centre plot would be much easier to read on a logarithmic scale. As is, I find trends difficult to identify.\n    - The far-left figure suggests that the sparsity does not yield any runtime improvements. Is this because the problem is too small? What happens to larger matrices?\n    - All those numbers are on CPU, do I understand correctly? Do the authors have any intuition about what they look like on GPU? In my experience, CPU and GPU runtimes differ considerably when running linear algebra routines.\n    - Do I read the figure correctly that the more PINN balls, the more efficient the linear algebra? How many balls are \"too many\" in terms of approximation quality?\n- Figure 2:\n    - What does $L^2$ error mean here? Is this a comparison to a reference solution or the $L^2$ norm of the residual? If it is the former, where does the reference come from?\n    - What precision is used? L2 errors of $10^{-10}$ sound like double precision, but in my experience, this is difficult on GPUs.\n    - It appears that the more PINN balls are used, the lower the loss. What happens to more than 20 Pinn balls?\n- Table 2: The results seem to disagree with Figure 2 in the sense that in Figure 2, 20 balls achieve lower error than five balls, but in Table 2, the roles are reversed. Where does this discrepancy come from?\n\nI list some additional, more minor questions under \"Questions\" below. (For reference, I don't expect all to be answered in the rebuttal. It would be nice, but I understand that time and space for the rebuttal are limited.)\n\nDespite these issues, the paper's strengths outweigh the weaknesses, and I recommend acceptance. However, resolving the issues I discuss above (or explaining how I misunderstood something) would lead me to improve my score.", "was_revised": true, "final_justification": "My original assessement gave a score of 4, mainly listing experimental clarity and gaps in the numerical-linear-algebra discussion as weaknesses. The rebuttal alleviated some of these concerns, but a significant part of the numerical-linear-algebra-related criticism persists, which is why the evaluation remains a 4. Details are in the reply to the rebuttal.", "rebuttal": "We sincerely thank the reviewer for the positive feedback, appreciation of our work, and the thorough and constructive comments, which we believe will significantly enhance the quality of the manuscript. First, we address the reviewer’s major concerns in detail.\n\n**Second Order Optimization:**\n* As the reviewer rightly noted, the primary contribution of our work is not the use of second-order optimization per se, but rather the development of a framework that enables such optimization to scale beyond current memory limitations. While prior work has shown that second-order methods enables accurate training of PINNs even in low parameter regimes, we are the first to propose a strategy which can extend this framework to a consistent number of parameters, which are at times necessary to predict certain PDE solutions. For instance, in the case of the time dependent Navier-Stokes equations, the computation of the Jacobian $D_\\theta L$ leads to an out-of-memory error at approximately 2000 parameters on our GPU implementation. This is due to the need to backpropagate through each collocation point and each PDE equation. However, in our experience, 2000 parameters are insufficient to model complex behaviors such as vortex shedding behind a cylinder. The combination of a compactly supported RBF-based domain decomposition — which induces sparsity in the Jacobian — and a sparse solver for the LM step — which keeps the optimization computationally feasible — forms the core of our contribution. Together, these components enable the use of second-order training for physics-informed models in settings where high model capacity is not just beneficial, but necessary. An additional reason why we reference [4, 24] is to highlight the general dominance in performance of second-order optimization with respect to first order optimizers. Nevertheless, we have tested first-order optimization and found that it is consistently outperformed by second-order approaches in all cases, as in the cited papers. For the Navier-Stokes equations in particular, first-order optimization fails to converge and exhibits significant oscillatory behavior - especially in the pressure field - even after 500'000 iterations. These findings will be emphasized more clearly in the revised version of the manuscript.\n\nWe appreciate the reviewer’s suggestion regarding matrix-free solvers such as conjugate gradient. In principle, these methods are attractive due to their memory efficiency. However, the construction of the Jacobian $D_\\theta L$ can trigger itself out-of-memory errors with as few as a few thousand parameters. However, matrix-free methods could benefit from the sparse structure introduced by our method, as it enables efficient sparse matrix-vector products that are central to such solvers. We will extend Section 4.1 to incorporate this valuable point raised by the reviewer and more clearly position our method in relation to matrix-free approaches.\n\n**Experimental Clarity** \n* **Figure 1**\n* As noted by the reviewer, the domain decomposition may alter the sparsity structure of the Jacobian over the course of training due to changes in the basis radii. If preserving high sparsity is critical, the basis radius can be capped or even fixed after initialization. The figure in its current form provides insight into the sparsity pattern at initialization, which can be an indicator of the sparsity during training, when required.\n* We chose to use a linear scale in Figure 1 to better highlight trends in the high-parameter regime, which is the primary focus of the paper. While a logarithmic scale offers finer resolution for lower parameter counts, we believe the linear scale more clearly conveys the relevant trends for our purposes.\n* It is true that sparse solvers are generally slower than their dense counterparts, particularly due to limitations in available open-source implementations. However, as also noted in the Limitations section, the advantage of sparse solvers lies in their scalability and memory efficiency. Dense solvers often run out of memory before solving large systems—whereas the sparse solver we used successfully handled systems with over 10,000 parameters. We will clarify in the revised manuscript that the dense solver in our experimental setup fails before 8000 parameters, whereas the sparse solver remains tractable.\n* The experiments were conducted on CPU to ensure fair comparison, given the lack of GPU-native, open-source sparse solvers. While GPU acceleration would benefit both dense and sparse solvers, the performance advantage may differ based on implementation efficiency and the size of the problems considered (typically GPU advantage will be more noteworthy for larger problems, both for dense and sparse solvers).\n* Indeed, sparse solvers perform better when the matrices involved are more sparse. However, using too many submodels (PINN Balls) for a fixed total number of parameters reduces the expressivity of each submodel. Quantitatively, we observed a drop in performance when submodels are only able to capture low complexity (e.g., one or two layers with at most eight neurons per layer).\n\n* **Figure 2**\n* Two of the subfigures report the Loss value (so the norm of the residuals) since the $L^2$ error against the exact solution is computationally expensive to evaluate at every iteration — especially for the Navier-Stokes equation. However, we provide the error against the exact solution for Burgers' equation, since the norm of the residuals is at times misleading in that case. Reference solutions for Burgers’ and Navier-Stokes equations are computed trhough high-fidelty simulations and are available in several public GitHub repositories, referenced in the Appendix.\n* The optimization is indeed performed using double precision. This choice is particularly important for second-order methods, where precision is crucial. This aspect further enhances the challenge in using a large number of parameters.\n* Increasing the number of PINN Balls would further increase the accuracy of the model. Considering a limit scenario, where an infinite number of PINN Balls is used, the number of parameters of the model would be infinite too. In this scenario, the universal approximation theorem holds, and it should be possible to fit a PDE solution with arbitrary accuracy. However, this is always subject to numerical limitations. Regarding the specific PDEs included in our paper, for the case of Helmholtz equation, we do not expect the loss to decrease for more than 20 PINN Balls, since the error is already approaching machine precision. However, for the cases of Navier Stokes' and Burgers' equations, it is likely that using a larger number of PINN Balls will further increase the accuracy of the overall model.  Finally, note that increasing the number of PINN Balls entails an increase in the numerical complexity of the training step, so one needs to consider the trade-off between accuracy and computational cost, which in our case is reasonably satisfactory at 20 PINN Balls. \n\n**Table 2** The observed discrepancy between PINN Balls 5 and PINN Balls 20 is primarily due to architectural choices. To ensure fair comparison with XPINN and APINN, we constrained both models to have approximately the same number of parameters. In this configuration, the submodels of PINN Balls 20 appeared to be too small for this scenario - roughly as those mentioned in the answer above -.\n\nWe hope that our responses clarified the major doubts of the reviewer. We now address the minor concerns in a more compact way:\n1) Yes. In practice, when using K-means clustering, the exact value of sigma that ensures $\\phi >0$ for every collocation point is given by the standard deviation of each cluster.\n2) Yes, the loss function in Equation 10 is based on the negative entropy.\n3) The loss term for the boundary condition is evaluated on points sampled uniformly over the boundary, and as such it does not interact directly with the domain decomposition in the minimax problem. The mechanism by which boundary errors affect the DD is that high errors at the boundary propagate to the domain, which trigger changes in the DD.  As we show this results in an efficient training; alternatively, one can also impose the output of PINN Balls to match the exact boundary conditions, as done in several existing works on PINNs.\n4) The use of $L$ instead of $J$ is indeed a typo. We will correct this in the revised version.\n5) The exponential decay is important for maximizing the performance of the LM optimization in late training stages. The core concept behind the exponential decay is to conclude the optimization with few LM iterations free from the AAS maximization step. We use the exponential decay implemented in optax, with end value $10^{-10}$.\n6) The problem is indeed solved over the domain $[0, \\pi]$. We agree with the reviewer that an optimal ensemble size likely exists, although it is difficult to predict a priori for general PDEs. Identifying the optimal number of submodels — perhaps via progressively adding submodels during training — is an important topic for future investigation.\n7) K-means is applied to the collocation points in the PDE domain.\n8) We did not include those results as they appeared nearly identical to the current ones.\n9) In our intuition, a sparse ill-conditioned matrix could yield a dense inverse matrix, which would hinder the effectiveness of the sparse solver. Notably, the Jacobian of a PINN is often ill-conditioned, especially for nonlinear PDEs. However, we have not noticed this issue numerically, attributing this to the regularization of the LM step. However, this could also suggest that the condition number may not play a decisive role in this context, as per the reviewer's intuition.\n\nWe thank the reviewer once again for the thoughtful feedback and constructive suggestions. Their input has contributed meaningfully to improving the clarity, rigor, and completeness of our work."}]

Author rebuttal:

Meta-review: (1) Summary of this work: This paper introduces PINN BALLS as a sparse MoE architecture for Physics-Informed Neural Networks (PINNs), which can effectively unite the domain decomposition, second-order optimization and adversarial adaptive sampling. The authors provide extensive experiments on different PDEs to verify the effectiveness of PINN BALLS. 

(2) Strengths and weaknesses: One of the main strengths of this work is its seamless and sound unification of several techniques. Especially, the idea of scaling higher-order methods through learning the adaptive domain decomposition is novel, which is also acknowledged by several reviewers. An in-depth discussion of theoretical property is also included, along with a transparent description of limitations, showing favorable scientific rigor. However, one main concern of this paper is the term "scaling". Although the authors have provided experiments about "scalability" in model parameters, I think the "scaling" term is more about generalizability, namely, one model to different tasks, which is really hard to achieve in PINNs. Besides, the evaluation of scaling significance is also mentioned by Reviewer DRVa.

(3) Summary of rebuttal: Although Reviewer DRVa expresses his/her slight concerns about this paper's evaluation part, other reviewers think it should not be a reason for rejection, which I also agree with. Also, during the rebuttal, the authors have provided detailed information about the effect of each algorithmic component and clarification of experiments (They need to add this into the final version of this paper).  

**Final decision:** After a comprehensive discussion with reviewers and carefully reading the rebuttal, I think the proposed method can be inspiring for future applications of PINN in handling complex problems, and it can tackle the bottleneck in the adoption of second-order methods for training PINN. Thus, I would recommend acceptance, but not a higher score, due to the limitation in evaluation. And I strongly suggest the authors add more baseline methods and a more comprehensive evaluation of the scaling property in the final paper.
Decision: Accept (poster)