FBPINNs: A New Domain Decomposition-Based Method for Physics-Informed Neural Networks

FBPINNs (Finite Basis Physics-Informed Neural Networks) is a new method for physics-informed neural networks proposed by Ben Moseley et al. from the University of Oxford. Addressing the spectral bias problem of traditional PINNs (difficulty in capturing high-frequency/multi-scale features), it significantly improves the solution performance on high-frequency and multi-scale problems through domain decomposition, subdomain normalization, and flexible training scheduling, achieving a 10-1000x speedup. The project source code is available on GitHub (https://github.com/benmoseley/FBPINNs), and the related paper was published in Advances in Computational Mathematics in July 2023.

Physics-Informed Neural Networks (PINNs) solve complex physical problems without mesh generation by incorporating the residual of physical equations into the loss function. However, they face a fundamental challenge—spectral bias: neural networks prioritize learning low-frequency patterns and have weak ability to capture high-frequency components or multi-scale features. When dealing with high-frequency oscillations or multi-scale phenomena, extremely deep networks and long training times are required, and convergence may even fail, limiting their applications in fields such as wave equations, turbulence simulation, and quantum mechanics.

The core idea of FBPINNs is to decompose complex global problems into simple local problems:

1. Domain Decomposition: Divide into overlapping subdomains, deploy an independent neural network in each subdomain, and the solution is the superposition of outputs from all subdomains;
2. Window Function Constraint: Multiply the output of each subdomain network by a smooth and differentiable window function to limit its influence range and ensure the continuity of the overall solution;
3. Subdomain Normalization: Perform independent local normalization on the input of each subdomain to reduce the effective frequency and alleviate spectral bias.

FBPINNs were migrated from PyTorch to the JAX framework, bringing significant improvements:

• Parallel Acceleration: Use jax.vmap to parallelize subdomain computations, achieving a 10-1000x speedup;
• Scale Expansion: Support collaborative training of thousands of subdomains;
• Training Scheduling: The subdomain scheduler can define active subdomains during training steps, suitable for time-dependent problems (e.g., time-stepping training for the Burgers equation).

Experimental results validate the advantages of FBPINNs:

1. 1D Harmonic Oscillator Problem: Accuracy and computational efficiency are far superior to traditional PINNs;
2. (2+1)D Multi-Scale Wave Equation: Successfully captures the interaction of different frequency components, while traditional PINNs fail. These results prove that domain decomposition and subdomain normalization can effectively improve the solution effect of high-frequency/multi-scale problems.

FBPINNs provide a modular API, with usage steps as follows:

1. Define the solution domain (Domain class);
2. Define the PDE problem (Problem class, including boundary conditions and data constraints);
3. Define the domain decomposition (Decomposition class);
4. Define the network structure (Network class);
5. Configure hyperparameters (Constants object);
6. Start training (FBPINNTrainer). The project also provides examples for custom problems, inverse problems, subdomain scheduling, etc.

The related results of FBPINNs were published in Advances in Computational Mathematics in July 2023, which elaborates on the mathematical foundation, algorithm design, and experimental validation, providing important theoretical contributions to the field of physics-informed machine learning. This method combines classical numerical methods with deep learning, overcomes the limitations of PINNs, and provides inspiration for scientific machine learning problems: proper problem decomposition and preprocessing can improve the solution ability of neural networks.

FBPINNs solve the core bottleneck of PINNs through domain decomposition strategy, improve computational efficiency and accuracy, and open up new possibilities for solving complex physical problems with neural networks. The JAX version enhances performance and extensibility, making it a powerful tool in the intersection of scientific computing, physical simulation, and machine learning, which is worth in-depth exploration and application.