💻 computer science

Domain decomposition of large neural network surrogate models

This paper proposes a domain decomposition framework for training large neural network surrogate models by partitioning the input space into local subdomains with interface constraints enforced via Lagrange or augmented Lagrange methods, demonstrating that the augmented Lagrange approach offers superior scalability and faster convergence for large-scale problems while maintaining high accuracy.

Original authors: Timm Gödde, Eisso H. Atzema, Bojana Rosić

Published 2026-03-30

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Timm Gödde, Eisso H. Atzema, Bojana Rosić

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a giant, super-smart robot how to predict how a piece of metal will bend and stretch when you squeeze it. This is a classic engineering problem. Usually, to get this right, you need to run thousands of complex physics simulations (like a digital wind tunnel for metal) to generate training data.

The problem is that training a single, massive "brain" (a Neural Network) to learn all of this at once is incredibly slow and expensive. It's like trying to teach a student to memorize the entire encyclopedia in one sitting. They might get the general idea, but they'll miss the tiny, tricky details in specific chapters, or they'll take forever to study.

The Big Idea: The "Team of Specialists" Approach

This paper proposes a clever solution: Don't build one giant brain; build a team of smaller, specialized brains.

Think of the metal object not as one big puzzle, but as a jigsaw puzzle. Instead of one person trying to solve the whole thing, you split the puzzle into smaller sections. You assign a small, fast, and simple neural network to solve just one section.

The Old Way: One giant, slow network trying to learn everything.
The New Way (Domain Decomposition): A team of small networks, each working on their own piece of the puzzle, all at the same time (in parallel).

The Problem: The "Seams"

Here's the catch: If you just let these small networks work alone, the edges where their pieces meet will look messy. One network might say, "At this edge, the metal bends 5mm," while its neighbor says, "No, it bends 6mm." The result would look like a jagged, broken piece of metal.

To fix this, the authors introduce Interface Constraints. Think of this as a strict rulebook for the team. The networks are told: "You can solve your own piece, but you must agree with your neighbor exactly where your pieces touch."

The Two Methods: The "Strict Teacher" vs. The "Flexible Coach"

The paper tests two different ways to enforce these rules:

The Strict Teacher (Lagrange Multipliers):
Imagine a teacher who stops the class every time a student makes a mistake at the boundary. The teacher says, "Stop! You can't move forward until you perfectly match your neighbor's answer." This ensures perfect accuracy, but it's very slow. The teacher has to constantly check, re-check, and adjust the rules, which takes a lot of time. It's precise, but it gets bogged down.
The Flexible Coach (Augmented Lagrange Multipliers):
This coach is smarter. Instead of demanding perfection immediately, they say, "Try to match your neighbor, but if you're slightly off, I'll give you a gentle penalty (a fine) to encourage you to do better next time." As the teams get closer to the right answer, the coach tightens the rules.
- Why it's better: This method is much faster. It allows the networks to learn quickly and converge on a good solution without getting stuck in endless "stop-and-check" loops. It's slightly less perfect at the very edges than the Strict Teacher, but it gets the job done 10 to 20 times faster.

The Results: Why This Matters

The authors tested this on a 2D metal block and a 3D cylinder. Here is what they found:

Speed: By splitting the problem, they could train the models in parallel. The "Flexible Coach" (Augmented Lagrange) was the clear winner, solving large problems in hours instead of days.
Accuracy: Even though the models were smaller, they were actually more accurate at the tricky, non-linear parts (the corners and edges) than the giant single model.
Handling Missing Data: In the real world, sensors often fail or data is missing. The team found that even if there were "gaps" in the data between the sections, the method could still figure out the solution, acting like a bridge between the known data points.

The Bottom Line

This paper is about efficiency through specialization. Instead of trying to force one massive computer model to do everything, we break the problem down, let smaller models do the heavy lifting in parallel, and use a smart "coach" to make sure they all agree on the boundaries.

It's the difference between trying to move a mountain with one person and a shovel, versus organizing a team of 100 people with shovels, giving them a clear map, and a foreman to make sure the piles of dirt line up perfectly. The result is a job done faster, with better precision, and less stress on the system.

1. Problem Statement

The paper addresses the computational challenges associated with training Neural Network (NN) surrogate models for high-dimensional engineering problems, particularly in design optimization and uncertainty quantification.

The Challenge: Standard global NNs often require a massive number of parameters to capture local nonlinearities accurately, leading to prohibitively long training times and over-parametrization.
The Limitation of Existing Methods: While Physics-Informed Neural Networks (PINNs) have utilized domain decomposition (DDM) to solve partial differential equations (PDEs), most existing DDM approaches are either overlapping (introducing computational redundancy) or rely heavily on penalty-based loss functions that struggle to balance constraints and data fidelity.
The Goal: To develop a non-overlapping, data-driven domain decomposition framework that partitions the input feature space into local subdomains. This approach aims to reduce the number of trainable parameters, enable parallel training, and maintain global accuracy by enforcing strict continuity constraints at the interfaces between subdomains.

2. Methodology

The authors propose a non-overlapping domain decomposition strategy where the global domain $G$ is split into $M$ disjoint subdomains $G_i$ . Each subdomain is modeled by a local, simpler NN ( $\hat{u}_{\theta_i}$ ). To ensure the global solution remains continuous and differentiable ( $C^0$ and $C^1$ continuity), interface constraints are enforced between adjacent subdomains.

Two specific algorithms are developed to handle these constraints:

A. Lagrange Multiplier Algorithm (LMA)

Formulation: The constrained optimization problem is reformulated using standard Lagrange multipliers ( $\lambda$ ). The local loss function includes a term $\lambda^T Q_i(\theta_i)$ , where $Q_i$ represents the interface constraints (function value and normal gradient continuity).
Linearization: Since NNs are highly non-convex, the constraints are linearized around reference points using Taylor expansions.
Optimization: The method employs a Dual Ascent strategy:
1. Primal Update: Local NN parameters ( $\theta_i$ ) are updated using the L-BFGS optimizer to minimize the linearized Lagrangian.
2. Dual Update: Lagrange multipliers ( $\lambda_i$ ) are updated via gradient ascent to enforce constraints.
3. Interface Training: Interface NNs are trained to match the predictions of adjacent local domains.
Complexity: This approach requires iterative linearization and feasibility checks, which can be computationally expensive and sensitive to the nonlinearity of the problem.

B. Augmented Lagrange Multiplier Algorithm (ALMA)

Formulation: This method introduces an augmented Lagrangian term, adding a quadratic penalty $\rho \|Q_i(\theta_i)\|^2$ to the loss function.
Advantage: The penalty term regularizes the constraints, allowing the algorithm to converge without the strict need for iterative linearization of the constraints at every step.
Optimization: Similar to LMA, it uses a primal-dual loop, but the primal update minimizes the augmented loss directly using L-BFGS. The penalty parameter $\rho$ acts as a step size for the dual update.
Scalability: ALMA is designed to be more robust and computationally efficient for large-scale problems compared to LMA.

3. Key Contributions

General Data-Driven Framework: Unlike previous works focused on PINNs (solving PDEs), this method is designed for general data-driven surrogate modeling, allowing for parallel computation of local nonlinearities.
$C^1$ Continuity Enforcement: The method explicitly enforces continuity of both the function values and their normal gradients across subdomain interfaces, ensuring smooth physical approximations.
Comparative Analysis of Solvers: The paper provides a rigorous comparison between standard Lagrange multipliers (requiring linearization) and Augmented Lagrange multipliers (using penalty terms) in the context of NN training.
Handling Sparse Data: The study investigates the robustness of the method when data is missing near interfaces (simulating real-world sensor limitations), showing that while gaps introduce stochasticity, the method remains stable.

4. Results

The methods were validated using Finite Element Method (FEM) data from two scenarios:

A. 2D Linear-Elastic Compression Problem

Setup: A 2D cylinder compressed vertically, split into varying numbers of subdomains (2 to 10).
Performance:
- ALMA demonstrated superior scalability. It achieved minimal training times (5.8 hours for 2 domains) and showed nearly linear growth in computational cost as domains increased.
- LMA suffered from scalability issues, with training times reaching up to 140 hours and failing to converge for multi-directional splits within reasonable timeframes. This was attributed to the difficulty of finding feasible linearizations for highly nonlinear regions.
- Accuracy: Both methods significantly reduced discontinuities at interfaces compared to unconstrained approximations. ALMA achieved a maximum relative error of ~4% with 3 vertical subdomains, offering the best trade-off between runtime and accuracy.
Gap Analysis: Introducing gaps (missing data) between domains increased uncertainty, causing the solution to vary stochastically, but the method remained stable.

B. 3D Cylinder with Uncertainty Propagation

Setup: A 3D cylinder with stochastic material properties (bulk and shear moduli) used for optimization. The domain was split into two halves.
Results:
- The surrogate model successfully approximated the mean and standard deviation of the displacement field.
- Mean Error: Negligible (<2.5% for displacement components).
- Variance Error: Slightly higher (~10% for standard deviation) due to sensitivity to second-order moments, but acceptable for optimization purposes.
- Interface Consistency: Kernel density estimates confirmed high continuity across the interface, with predictions from both subdomains closely matching the reference FEM data.

5. Significance and Conclusion

Efficiency: The proposed Augmented Lagrange Multiplier Algorithm (ALMA) is identified as the preferred approach for large-scale NN surrogate models. It outperforms standard global training and the standard Lagrange method by significantly reducing training time while maintaining high accuracy.
Scalability: ALMA's ability to handle parallel training of local subdomains makes it suitable for high-dimensional problems where global NNs fail due to the "curse of dimensionality" and excessive parameter counts.
Practical Application: The method enables the use of NNs in complex design optimization and uncertainty quantification tasks where computational resources are limited, and data may be sparse.
Future Work: The authors suggest applying ALMA to even larger-scale problems and exploring modeling interface gaps as random fields to further reduce computational costs associated with high-dimensional stochastic variables.

In summary, the paper demonstrates that domain decomposition with augmented Lagrange multipliers is a promising, scalable strategy for training efficient, high-fidelity neural network surrogates for complex engineering systems.