[Paper Review] Bayesian neural networks for weak solution of PDEs with uncertainty quantification

Xiaoxuan Zhang, Krishna Garikipati|arXiv (Cornell University)|Jan 13, 2021

Model Reduction and Neural Networks64 references17 citations

TL;DR

This paper proposes a physics-constrained Bayesian neural network (BNN) framework that solves partial differential equations (PDEs) using the weak form discretization via a vectorized, convolutional operator-based implementation. By formulating the loss function from the discrete residual and employing encoder-decoder architectures with deterministic and Bayesian layers, the method enables high-throughput PDE solution with uncertainty quantification—capturing both epistemic and aleatoric uncertainty—while generalizing to unseen boundary conditions and geometries.

ABSTRACT

Solving partial differential equations (PDEs) is the canonical approach for understanding the behavior of physical systems. However, large scale solutions of PDEs using state of the art discretization techniques remains an expensive proposition. In this work, a new physics-constrained neural network (NN) approach is proposed to solve PDEs without labels, with a view to enabling high-throughput solutions in support of design and decision-making. Distinct from existing physics-informed NN approaches, where the strong form or weak form of PDEs are used to construct the loss function, we write the loss function of NNs based on the discretized residual of PDEs through an efficient, convolutional operator-based, and vectorized implementation. We explore an encoder-decoder NN structure for both deterministic and probabilistic models, with Bayesian NNs (BNNs) for the latter, which allow us to quantify both epistemic uncertainty from model parameters and aleatoric uncertainty from noise in the data. For BNNs, the discretized residual is used to construct the likelihood function. In our approach, both deterministic and probabilistic convolutional layers are used to learn the applied boundary conditions (BCs) and to detect the problem domain. As both Dirichlet and Neumann BCs are specified as inputs to NNs, a single NN can solve for similar physics, but with different BCs and on a number of problem domains. The trained surrogate PDE solvers can also make interpolating and extrapolating (to a certain extent) predictions for BCs that they were not exposed to during training. Such surrogate models are of particular importance for problems, where similar types of PDEs need to be repeatedly solved for many times with slight variations. We demonstrate the capability and performance of the proposed framework by applying it to steady-state diffusion, linear elasticity, and nonlinear elasticity.

Motivation & Objective

• To develop a high-throughput surrogate PDE solver that reduces computational cost compared to traditional FEM for repeated simulations.
• To enable uncertainty quantification in PDE solutions by integrating Bayesian neural networks (BNNs) with physics-informed learning.
• To formulate the loss function based on the discretized weak form residual using efficient convolutional operations, avoiding strong-form constraints.
• To allow generalization to new boundary conditions and problem domains without retraining, using a single neural network.
• To support design and decision-making in engineering and materials science by enabling fast, reliable, and uncertainty-aware PDE solutions.

Proposed method

• The method uses a discretized residual of the weak form of PDEs as the core loss function, derived from finite element assembly principles.
• A convolutional operator-based, vectorized implementation computes the residual across all elements using 2D convolutional kernels to extract local nodal values.
• The residual is computed per Gauss point using a B-operator matrix and weighted integration, then reduced to a global residual via summation.
• Dirichlet and Neumann boundary conditions are encoded as input tensors and enforced directly in the solution matrix via masking and reassignment.
• An encoder-decoder neural network architecture is used, with deterministic layers for deterministic solutions and Bayesian layers for uncertainty quantification.
• For BNNs, the likelihood function is constructed from the discretized residual, enabling posterior inference over network weights and uncertainty propagation.

Experimental results

Research questions

• RQ1 Can a physics-constrained neural network trained on the weak form of PDEs achieve high accuracy and generalization across diverse boundary conditions and geometries?
• RQ2 How effectively can Bayesian neural networks quantify both epistemic and aleatoric uncertainty in PDE solutions without requiring labeled data?
• RQ3 To what extent can a single trained neural network interpolate or extrapolate to new boundary conditions and problem domains not seen during training?
• RQ4 How does the vectorized, convolutional implementation of the residual compare in efficiency and accuracy to traditional FEM or strong-form PINNs?
• RQ5 Can the proposed framework be generalized to nonlinear PDEs such as those in nonlinear elasticity and diffusion?

Key findings

• The proposed method achieves accurate solutions for steady-state diffusion, linear elasticity, and nonlinear elasticity problems using only the weak form residual as a loss function.
• The model generalizes to new boundary conditions and problem domains not seen during training, demonstrating strong extrapolation capability.
• Bayesian neural network variants successfully quantify both epistemic uncertainty (from model weights) and aleatoric uncertainty (from data noise) in PDE solutions.
• The vectorized, convolutional implementation enables efficient computation of the residual across all elements without explicit loops, matching traditional FEM assembly in accuracy.
• The framework supports high-throughput PDE solving by training a single surrogate model that can be rapidly applied to new configurations.
• The method maintains high accuracy even for complex geometries and nonlinear material behaviors, as demonstrated on benchmark problems in elasticity and diffusion.