← Back to Scimba Jax Module · scimba_jax

Physics-Informed Neural Networks (PINNs)

A Physics-Informed Neural Network (PINN) uses neural networks as non-local basis functions to approximate PDE solutions. To do that, it minimizes the residual of the PDE at a batch of sampled points. In scimba_jax, that optimization loop uses a Projector, which is exactly the same object a collocation basis or any other approximation space uses, but with a different loss function.

1. PINN initialization and optimization

The example below shows a 2D Poisson problem as a physical model, posed on a square domain. The ApproximationSpace wraps the network (an MLP here) to be later used in a PINN or another approximation method. Finally, the Projector ties model, space and sampler into one object whose .project(...) runs the optimization loop and returns the optimized space. 

from scimba_jax.domains.meshless_domains.domains_2d import Square2D
from scimba_jax.nonlinear_approximation.approximation_spaces.approximation_spaces import (
    ApproximationSpace,
)
from scimba_jax.nonlinear_approximation.integration.monte_carlo import (
    DomainSampler,
    TensorizedSampler,
)
from scimba_jax.nonlinear_approximation.networks.mlp import MLP
from scimba_jax.nonlinear_approximation.numerical_solvers.projectors import Projector
from scimba_jax.physical_models.elliptic_pde.laplacians import LaplacianDirichletND
domain = Square2D([(-1.0, 1.0), (-1.0, 1.0)], is_main_domain=True)
sampler = TensorizedSampler([DomainSampler(domain)], bc=False)
nn = MLP(in_size=2, out_size=1, hidden_sizes=[16, 16], key=key)
# post_processing strongly enforces the BC: u_h = 0 on the boundary,
# for any network weights.
space = ApproximationSpace(
    {"x": 2},
    [(nn, "scalar", None)],
    model_type="x",
    post_processing=lambda approx, xy: approx * (1 - xy[0] ** 2) * (1 - xy[1] ** 2),
)
model = LaplacianDirichletND(domain, lambda x: source(x), bc="strong", model_type="x")
# A PINN is optimized by a Projector, like every approximation space in scimba_jax.
pinn = Projector(model, space, sampler)  # default optimizer: see below
key, pinn = pinn.project(key, space, n_epochs=200, n_colloc=1000)
# ENG, 200 epochs: loss 95.2 -> 5.5e-09.

2. Brief overview of the math

Optimizing a PINN reduces to the problem of minimizing a sum of squared residuals, one term per piece of the problem. Pieces can be the interior PDEs, the boundary conditions, data, etc. Hence, the loss function reads

\[\mathcal{L}(\theta) = \sum_{\text{label}} \frac{1}{N_{\text{label}}} \sum_{i=1}^{N_{\text{label}}} \left\lVert R_{\text{label}}(u_\theta, x_i) \right\rVert^2.\]

Therefore, we have as many residual terms as sub-domains and boundary pieces the model declares. For instance, a Dirichlet boundary and a Neumann boundary each get their own label and their own term, weighted and summed into one scalar loss. The network itself can be either scalar-valued for a single unknown (temperature, pressure) or vector-valued when the problem couples several fields at once. The ApproximationSpace's "scalar"/"vec" tag says which. For instance, a 2D Navier-Stokes velocity-pressure system is a length-3 vector output [u_x, u_y, p].

3. Optimization algorithms: first-order vs. quasi-second-order

To optimize the Projector, we rely on a single keyword (optimizer), which is swappable without touching the model, the space or the sampler. A first-order optimizer only ever uses the gradient:

\(\theta^{(k+1)} \gets \theta^{(k)} - \eta\, \nabla \mathcal{L}(\theta^{(k)})\).

Adam, AdamW and Muon are examples of such optimization algorithms. Conversely, a quasi-second-order optimizer also uses some model of curvature C to rescale that gradient:

\(\theta^{(k+1)} \gets \theta^{(k)} - \eta\, C^{-1} \nabla \mathcal{L}(\theta^{(k)}).\)

For instance, L-BFGS, SS-BFGS and SS-Broyden build C from successive gradients (in a quasi-Newton manner). ENG, its matrix-free twin MFENG, ANaGRAM, RATATAM and RATATAM-LM build it instead from the PDE's own natural-gradient/Gauss-Newton metric. It means that the optimizer uses a curvature that comes from the physics, rather than from an optimization history. The default is ENG (Energy Natural Gradient), owing to its stability and convergence speed. 

# Values of the `optimizer=` argument of the `Projector`:
# "ENG"          Energy Natural Gradient (default), metric
# induced by the PDE energy
# "MFENG"        Matrix-free ENG, without assembling the Gram
# matrix
# "Adam"         Usual deep learning optimizer, when the ENG
# Gram matrix is too costly
# "AdamW"        Adam with decoupled weight decay (Loshchilov
# & Hutter)
#   "Muon"         Sets every singular value of the update to 1
#   "L-BFGS"       Quasi-Newton, second order without a Hessian
# "SS-BFGS"      Self-scaled BFGS variant, tuned for PINN
# losses
#   "SS-Broyden"   Same family, with Broyden's update
#   "ANaGRAM"      Natural gradient method for large networks
#   "RATATAM"      Levenberg-Marquardt-type PINN optimizer
# Changing the optimizer does not change the model, the space
# or the sampler.
pinn = Projector(model, space, sampler, optimizer="Adam")
key, pinn = pinn.project(
    key, space, n_epochs=1000, n_colloc=1000
)
# a dict, one entry per residual plus "total"
print(pinn.best_loss["total"])
# Adam, 1000 epochs: loss 95.2 -> 0.32, against 5.5e-09 for
# ENG in 200 epochs.
# The best optimizer depends on the problem.

A PINN and a mesh-based solver both approximate a PDE, but they have different trade-offs. Indeed, for PINNs, there is no mesh to build, but the linear systems inherent to mesh-based solvers are replaced by nonconvex optimization problems. In low dimensions (up to 4), if you already have a suitable mesh, consider using a Mesh-based solver instead.