# LatentPINNs: Generative physics-informed neural networks via a latent
representation learning

Mohammad H. Taufik
Department of Physical Science and Engineering
King Abdullah University of Science and Technology
Thuwal, 23955, Saudi Arabia
<EMAIL_ADDRESS>
&Tariq Alkhalifah
Department of Physical Science and Engineering
King Abdullah University of Science and Technology
Thuwal 23955, Saudi Arabia
<EMAIL_ADDRESS>

###### Abstract

Physics-informed neural networks (PINNs) are promising to replace conventional
partial differential equation (PDE) solvers by offering more accurate and
flexible PDE solutions. However, they are hampered by the relatively slow
convergence and the need to perform additional, potentially expensive,
training for different PDE parameters. To solve this limitation, we introduce
latentPINN, a framework that utilizes latent representations of the PDE
parameters as additional (to the coordinates) inputs into PINNs and allows for
training over the distribution of these parameters. Motivated by the recent
progress on generative models, we promote the use of latent diffusion models
to learn compressed latent representations of the PDE parameters distribution
and act as input parameters to NN functional solutions. We use a two-stage
training scheme in which the first stage, we learn the latent representations
for the distribution of PDE parameters. In the second stage, we train a
physics-informed neural network over inputs given by randomly drawn samples
from the coordinate space within the solution domain and samples from the
learned latent representation of the PDE parameters. We test the approach on a
class of level set equations given by the nonlinear Eikonal equation. We
specifically share results corresponding to three different sets of Eikonal
parameters (velocity models). The proposed method performs well on new phase
velocity models without the need for any additional training.

##  1 Introduction

Partial differential equations (PDEs) provide mathematical formulations that
describe many physical systems. They have been the cornerstone of scientific
and engineering applications for decades. In practice, these equations are
often discretized and used to solve for the state field or the unknown
physical parameters (e.g., velocity, pressure, heat, etc.). With complex
boundaries and large computational domains, conventional PDE solvers developed
over the years (e.g., finite-difference and finite element methods) suffer
from general inefficiency and require complex discretization procedures.
Neural network-based PDE solvers have shown considerable promise in addressing
these limitations, as well as providing more efficient solutions. In this
family of approaches, neural networks can be utilized to either learn the
governing PDE operator (neural operators) Li et al. (2020a); Lu et al. (2019);
Li et al. (2020b); Kovachki et al. (2021) or be regularized (by the governing
PDE) to learn a specific PDE solution (neural instance solvers) Raissi et al.
(2019); Sirignano and Spiliopoulos (2018). Neural operators often require the
solution labels (obtained, for example, from conventional numerical solvers),
in which the training process is then performed in a supervised manner. Thus,
at inference time, neural operators can accommodate different realizations of
the PDEs. This comes, however, with heavy computational costs for large-scale
problems, as in this case, acquiring the labels themselves might require
significant computing power. Several attempts have been made to relax the
data-driven nature of neural operators by incorporating physical laws into the
neural operators Li et al. (2021); these additional constraints, however,
still render the solution mesh dependent and generally inefficient to train.
One notable example of neural instance solvers, termed physics-informed neural
networks (PINNs), on the other hand, is driven by the governing PDE at hand
(independent of numerical PDE solvers). At inference time, it provides the
functional representation for a specific PDE parameter. These parametric deep
neural networks (DNNs) address the limitations of numerical methods by
facilitating flexible (mesh-independent) solutions that are generally more
accurate than their numerical counterparts, thanks to the functional nature of
the derivative evaluations.

In this quest to fully replace conventional PDE solvers, the current PINNs are
facing two main challenges. Firstly, randomly initialized PINNs often require
thousands of iterations (epochs) to converge, massively increasing the cost of
attaining the PDE solution. To this end, several attempts have been made to
address the training dynamics as well as improve its efficiency. Fang (2021)
promoted the use of a local fitting approach to approximate the governing PDE
operator and improve PINNs’ efficiency. Wang et al. (2022) proposed the use of
causality constraint that results in better convergence of PINNs from some
initial condition. Huang and Alkhalifah (2022) promoted the use of higher-
dimensional embedding space of the coordinate input to improve PINNs’
convergence. Huang and Alkhalifah (2023a); Rzepecki and Doran (2022) used
alternative embedding layers and, thus, demonstrated significant computational
speed-ups. Moreover, a slow convergence can also be attributed to the
challenging training dynamics. One major challenge comes from the need to
appropriately balance the various loss terms during the PINNs training.
Schiassi et al. (2021); Huang and Alkhalifah (2023b); Taufik et al. (2023a)
promote the use of hard-constrained data loss function to alleviate the burden
of this balancing procedure. Secondly, PINNs are trained for specific PDE
parameters; a different parameter set requires additional PINN training. To
accommodate for different PDE parameters efficiently, several authors suggest
the use of transfer learning Goswami et al. (2020); bin Waheed et al. (2021)
and meta learning Penwarden et al. (2021); Liu et al. (2022); Huang et al.
(2022). For transfer learning to reduce the cost of training effectively, the
difference between the new PDE parameters and the original ones in which the
PINN was trained needs to be small. Such a restrictive condition might end up
requiring almost identical training costs to that from randomly initialized
PINNs. Meta learning, on the other hand, comes with the objective of finding
the best starting point to accelerate the new PDE parameter PINNs training.
Both of these approaches, unlike neural operators, still require additional
training for new PDE parameters.

In this paper, we introduce the use of latent representation learning to
PINNs. Informed with this additional latent representation, the trained PINNs
can be regarded as generative models that produce solutions to PDEs as a
function of latent variables (representing the PDE parameters). We investigate
the possibility of overcoming additional training for new PDE parameters using
a pre-training and fine-tuning scheme. In the experiments, we test our
framework on a nonlinear PDE with various datasets and share the challenges we
faced.

The main contributions of this paper can be summarized as follows:

  * •

We introduce the use of latent representation learning to physics-informed
neural networks.

  * •

We promote the use of autoencoder-based models to learn the latent
representation of PDE parameters that can be combined with state-of-the-art
diffusion models for functional data representation.

  * •

We test the proposed algorithm on three datasets to solve a nonlinear PDE
without the need to perform additional training for new PDE parameters.

The organization of the paper is as follows. We first discuss the related
works in the next section. We then provide the theoretical background and the
proposed framework in Section 3, followed by numerical experiments in Section
4. Finally, we discuss the current challenges and future directions, as well
as conclude in the last two sections.

##  2 Related Work

###  2.1 Generalizing PINNs

There have been many attempts to pre-train PINNs so that they converge faster
for new PDE parameters. Transfer learning was suggested in which the PINNs are
initialized with the weights of the training for a previous set of PDE
parameters bin Waheed et al. (2021). In transfer learning, the change in the
PDE parameters needs to be small to improve convergence compared to starting
with random initialization. Alternatively, meta learning was adopted in PINN
training by finding initial weights that can work for a range of tasks or PDE
parameters Penwarden et al. (2021); Liu et al. (2022); Huang et al. (2022).
Meta learning is, however, costly considering the pretraining on a range of
parameters. In all of these methods, additional training is required to fine-
tune the PINN to provide accurate solutions for a new set of PDE parameters.
In other words, these methods do not necessarily reduce the requirement for an
extra PINNs training step for different PDE parameters.

###  2.2 Generative models with PINNs

Uncertainty estimation has been the primary goal of incorporating generative
models into physics-informed neural networks Yang and Perdikaris (2018); Zhu
et al. (2019); Yang et al. (2020); Linka et al. (2022); Oszkinat et al.
(2022). For such purposes, the generative models are utilized such that they
can stochasticity be included in the _deterministic_ vanilla PINNs. Huang et
al. (2022) introduces the use of an auto-decoder network to accelerate PINNs
training from a meta-learning perspective. Different PDE parameters can be
viewed as different tasks. Though latent representation is also included as an
additional input parameter, the trained PINN still requires additional
training for new PDE parameters. From the perspective of reduced order
modeling, Kim et al. (2022) promotes the use of a masked autoencoder to learn
the lower-dimensional manifold of the PDE parameters. They use the trained
decoder and incorporate existing conventional PDE solvers instead of neural
networks to solve the PDE. Finally, Zou and Karniadakis (2023) introduces the
use of normalizing flow to perform density estimation for further downstream
tasks. Extra training for new PDE parameters, however, is still required for
those tasks.

###  2.3 Incorporating latent vectors into PDE solvers

Almost any machine learning algorithm inherently tries to learn the hidden
_latent_ representation between the input data that can be mapped into an
output for a certain task. In a more restrictive manner, these latent
variables can be constructed (and later accessed) such that they represent the
data with a significant reduction in size. Ranade et al. (2021) incorporates
latent vectors with existing numerical solvers to generalize the PINNs
solutions. The idea of operating on a much lower-dimensional space for solving
PDEs has been around under the banner of model order reduction Lumley (1967).
Motivated by this idea, Wu et al. (2022) accelerates PINNs for solving large-
scale time-dependent problems. Similarly, Kontolati et al. (2023) incorporates
such an idea into a class of neural operators to achieve better accuracy.

To the best of our knowledge, we promote the first attempt to make PINNs learn
solutions for a certain PDE parameter distribution without the need for
additional training to accommodate new PDE parameters within the distribution.
In such an approach, we aim to incorporate a generative feature into the
trained PINNs model, such that it can generate various PDE solutions for
various PDE parameters.

##  3 Methodology

Consider a general parametric PDE in the form of

$u_{t}(\boldsymbol{x},t;\boldsymbol{\xi})+\mathcal{A}(u(\boldsymbol{x},t;\boldsymbol{\xi}))=\mathcal{S}(\boldsymbol{x},t;\boldsymbol{\xi}),\quad\mathcal{B}(u;\boldsymbol{\xi})=0,\quad
t\in[0,T],\quad\boldsymbol{x}\in\Omega,\quad\boldsymbol{\xi}\in\mathcal{P},$
(1)

where $u(\boldsymbol{x},t;\boldsymbol{\xi})$ denotes the desired field as a
function of the spatio-temporal variables $[\boldsymbol{x},t]$, in which
$u_{t}$ corresponds to its first-order temporal derivative. Here,
$\boldsymbol{\xi}$ denotes the PDE parameters, and $\Omega$ is a subset of a
$d$-dimensional domain of real numbers, $\mathbb{R}^{d}$. $\mathcal{A}$ and
$\mathcal{F}$ denote the non-linear operator and the source function,
respectively, applied on the computational domain $\Omega$, while
$\mathcal{B}$ defines the boundary operator applied on $\partial\Omega$ (the
surrounding surface). Conventional PINNs will then try to learn a mapping
between the input ($\boldsymbol{x},t$) to the PDE solution $u$ for that input
location. This mapping is facilitated by a neural network
$\mathcal{N}(\boldsymbol{x},t;\boldsymbol{\xi},\boldsymbol{\theta})$, where
$\boldsymbol{\theta}$ denotes the network’s parameters. The output of this
network will then be regarded as a surrogate to the PDE solution $\hat{u}$.
During the PINNs’ training, the PDE residuals and their boundary conditions
are incorporated into the objective function in the form of

$\displaystyle\mathcal{L}$
$\displaystyle=\mathcal{L}_{\text{PDE}}+\mathcal{L}_{\text{boundary}},$ (2)
$\displaystyle=||\hat{u}_{t}(\boldsymbol{x},t;\boldsymbol{\xi})+\mathcal{A}(\hat{u}(\boldsymbol{x},t;\boldsymbol{\xi}))-\mathcal{S}(\boldsymbol{x},t;\boldsymbol{\xi})||_{2}^{2}+||\mathcal{B}(\hat{u}(\boldsymbol{x},t;\boldsymbol{\xi}))||_{2}^{2}.$
(3)

To this end, once the training is performed for specific PDE parameters
$\boldsymbol{\xi}$, additional training is required (either by transfer
learning or within a meta learning framework) to obtain PDE solutions for
other PDE parameters.

###  3.1 Latent Representation Learning

The objective of most representation learning algorithms is to find a
compressed representation of a dataset without losing any of its key
information that distinguish the samples in the dataset. Given an input x, the
objective function used for representation leanring can be broadly formulated
as

$\displaystyle\mathcal{L}$
$\displaystyle=\mathcal{L}_{\text{reconstruction}}+\mathcal{L}_{\text{regularization}}.$
(4)

Regularization is often introduced to balance the trade-off between semantic
and perceptual compression Esser et al. (2021). The reconstruction term can
come in many flavors depending on the choice of the network architecture. In
the simplest case, autoencoder networks Rumelhart et al. (1985); Hinton and
Salakhutdinov (2006) utilize deep neural networks to both encode (the input
data) and decode (the latent vector). Though they perform well in
dimensionality reduction, the learned latent vectors do not represent the data
distribution well for generative tasks. To solve this limitation, variational
autoencoders Kingma and Welling (2013) resort to a Bayesian representation of
the latent vector. Instead of mapping between vectors, they map the input
vectors to their latent distributions. This results in a better learned
posterior (data) distribution. Generative adversarial networks (GANs)
Goodfellow et al. (2020) replace this encoder-decoder scheme with a
discriminator-generator pair that are trained in an adversarial manner. This
results in a more representative learned posterior distribution than the
variational autoencoder at the cost of more challenging training dynamics
(e.g., mode collapse pathology). The rise of score-based diffusion models Song
and Ermon (2019); Song et al. (2020, 2021) in outperforming GANs in
representing the data distribution inspired us also to utilize this model and
incorporate it into physics-informed neural networks. So, we utilize such
powerful generative models to introduce a generative feature into physics-
informed neural networks.

###  3.2 PINNs with latent space inputs

Motivated by the recent success of latent-based models Rombach et al. (2022),
we introduce a two-stage training scheme into the PINNs training. During the
first stage, a Kullback-Leibler divergence regularized autoencoder is utilized
to learn the compressed latent representation of the PDE parameters. The
resulting latent vectors (representing a Gaussian distribution mapping of the
PDE parameters) will then be incorporated as an additional input into the
PINNs training (second stage).

Figure 1: The proposed workflow of the latentPINN.

Figure 1 illustrates the overall workflow of the proposed method. As opposed
to only using the position coordinates to train the PINNs
{$\textbf{x},\textbf{x}_{s}$}, we also incorporate the latent vectors z as an
additional input parameter. Specifically, the multi-layer perceptron used as
the main PINN model outputs the PDE solution denoted by $\hat{\tau}$. The PINN
objective function, $\mathbb{J}$, is formulated by the PDE residual and
parameterized by the lateral derivatives $\partial_{h}\hat{\tau}$ computed
using backpropagation through the network. In this study, as we will see
later, we consider modeling a solution for a class of level set equations
given by the Eikonal equation. The objective then becomes predicting the PDE
solution (traveltime fields) given knowledge of the PDE parameter (medium
phase velocity fields $v$). We replace the true PDE parameters by the
reconstruction result from the autoencoder $\hat{v}$.

##  4 Numerical tests

In this section, we describe the setup of the problem, specifically the
partial differential equation we aim to solve and the neural network model
used to solve the problem, including its training parameters. Finally, we
share the results of testing our approach.

###  4.1 The PDE problem

We test the proposed method on a level-set-based PDE, often used as a tool for
numerical analysis of surfaces and shapes. Specifically, the level-set model
often performs numerical computations for surfaces on a fixed Cartesian grid
without the need to parameterize these objects Osher and Sethian (1988).
Considering a surface, like a wavefront, moving with speed $V$ in the
direction normal to the surface, the level set function $\phi$ satisfies the
following nonlinear first-order PDE:

$\frac{\partial\phi(\boldsymbol{x},T)}{\partial
T}=V(\boldsymbol{x})|\nabla\phi(\boldsymbol{x},T)|,\quad\boldsymbol{x}\in\mathbb{R}^{3},$
(5)

where $T$ is the traveltime, and the operator $|.|$ is the Euclidean norm.

In wave theory, the level set function $\phi$ represents the phase of a wave,
defining its geometrical shape in 2D or 3D as a function of time. In the high
frequency asymptotic approximation, $\phi(\boldsymbol{x},T)=\omega
T(\boldsymbol{x})$, and equation 5 reduces to its Eikonal form as follows

$\left(\nabla
T(\boldsymbol{x})\right)^{2}=\frac{1}{V^{2}(\boldsymbol{x})},\quad\boldsymbol{x}\in\mathbb{R}^{3},$
(6)

where the solution of the PDE is represented by the traveltime field, $T$.
Thus, the PDE is parameterized by the phase velocity of the medium $V$, which
can be inhomogeneous.

In the following cases, we consider the problem of solving the Eikonal
equation, which is equivalent to traveltime modeling for an inhomogeneous
medium phase velocity model representing the PDE parameters. We, specifically,
aim to solve the equation for a point source in a 2D Cartesian computational
domain ($\textbf{x}\in\\{x,z\\}\in\mathbb{R}^{2}$). To mitigate source-related
singularity, which is a known problem in solving eikonal equations
numerically, we first substitute the traveltime field $T$ with
$T_{0}\hat{\tau}$, known as the factorization approach. For details of its
implementation in PINNs, we refer the readers to Taufik et al. (2023b). Here,
we consider a domain of size 5 km ($x\in[0,5]$) horizontally and 1 km
($z\in[0,1]$) vertically. The velocity models are discretized in this domain
with 128 points along each axis. For all examples, the source is located in
the middle ($\textbf{x}_{s}\in\\{2.5,0.5\\}$ km).

###  4.2 The neural network setup for training

As described earlier, we incorporate a two-stage training workflow for
latentPINN before direct inference on new PDE parameters (velocity models). In
the first stage, we train an autoencoder to learn a latent representation of
the PDE parameters (velocity models), and in the second stage, we train a PINN
that includes an added latent variables input. For the three examples below,
we consider a latent dimension of 96 to be used on the bottleneck of the
fully-connected layer.

The autoencoder used for the first stage of training is made up of an encoder
part that takes three-channel images of size 128x128 pixels and performs
downsampling across five convolutional blocks to yield images of size 4x4.
Specifically, each convolutional block consists of a 2D transposed convolution
layer of stride size 2 followed by a Gaussian error linear unit (GELU)
activation function and a 2D convolution of stride size 1. To construct the
convolutional block for the decoder part, we use identical layers as in the
encoder in reverse order, including the GELU activation function. To obtain
the latent vectors, a fully-connected layer with a Tanh activation is used at
the end of the encoder network. For the training, an Adam optimizer is used
with an initial learning rate of 1e-3 and a learning rate scheduler reducing
the learning rate by half every 20 stagnating (non-decreasing reconstruction
loss value) epochs. Finally, by incorporating a diffusion model into the
trained autoencoder, we showcase the ability of latentPINN to facilitate a
more representative PDE parameters sampling. This setup is generally used for
all three PDE parameter sets tested below, except for the example where we
make some changes to the autoencoder model.

In the second stage, we train a modified PINN that takes in, in addition to
the coordinates of the medium, the latent variables that represent the PDE
parameters (the velocity models). We also utilize the decoder to reconstruct
the velocity model used in the loss function governed by the PDE in equation 6
(or specifically the factored variation of it) to train the PINN. The reason
for using the decoder output instead of the original model, in the loss
function of PINN, though they are close, is because the decoder output is a
function of the latent vector through the decoder. We utilize a multi-layer
perceptron network with an input layer of size 99
($\textbf{x},\textbf{x}_{s}$,z, and 96 latent vector variables), an output
layer of size 1 ($\hat{\tau}$), 12 hidden layers of size 128 neurons each, and
exponential linear unit (ELU) activation function. An Adam optimizer is used
with an initial learning rate of 5e-4 and a scheduler to reduce the learning
rate by half every 200 stagnating (non-decreasing PDE loss values) epochs. We
randomly select 9830 collocation points as the training samples. We train the
network with 100 velocity models for 10k epochs with a batch size of 163
points. We randomly sample these velocities (by their latent vectors) such
that every iteration corresponds to solving for a specific velocity field.
Again, this setup for the latentPINN is generally used to solve for the
eikonal equation for all three tests below, except we make some changes for
the last test.

All experiments were performed on a single NVIDIA Quadro RTX 8000 GPU card
utilizing the PyTorch machine learning framework Paszke et al. (2019). We will
assess the accuracy of the learned traveltime functions by comparing their
predictions with the numerical solutions. We also test the accuracy of the
PINN solution by comparing the reconstructed velocity models (plugging in the
predicted PDE solutions $T$ into equation 6) to the output of the decoder
$\hat{V}$.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 2:  Reconstruction analysis on three different datasets; Gaussian
random fields (a and b), Mechanical MNIST examples (c and d), and realistic
Earth models (e and f). The left column includes the original samples drawn
from the testing set (a,c,e), and the right column includes the reconstructed
samples (b,d,f).

(a)

(b)

(c)

Figure 3:  Predicted PDE solutions (traveltimes, solid black lines), compared
with conventional solvers (dashed yellow lines) from the three datasets; a)
Gaussian random fields, b) Mechanical MNIST, and c) realistic Earth models.
The leftmost column of each plot corresponds to solutions from a seen training
PINN sample (velocity model), while the rest of the columns correspond to
samples from the unseen testing set.

(a)

(b)

(c)

(d)

(e)

(f)

Figure 4:  Accuracy analysis on three different datasets from the
reconstructed PDE parameters (using equation 6). The left column plots show
the original samples drawn (a,c,e), and the right column plots include the
reconstructed samples (b,d,f). The left most samples from each plot correspond
to velocity models from the training set; the rest of the columns include
results for unseen PDE parameters (from the testing set).

(a)

(b)

(c)

(d)

(e)

(f)

Figure 5:  Loss curves comparison between our approach and several training
mechanisms. The top row (a to c) compares the required additional epochs to
train for a new PDE parameter between the trained latentPINN (blue stars) and
conventional PINN training from randomly initialized weights (blue lines) and
performing transfer learning (green lines). The bottom row (d to f) details
the training cost to train our latentPINN model for 100 randomly chosen
velocity models (PDE parameters) from the training set used previously in the
autoencoder training.

###  4.3 Numerical Results

The goal here is to test the ability of the latent implementation to
generalize on test PDE parameter examples not used in the training of
latentPINN. To do so, we consider three sets of PDE parameters (phase velocity
model distributions): 1) A set drawn from Gaussian random fields, 2) a set
drawn from a Mechanical MNIST distribution, and finally, 3) a set related to
examples of layering inside the Earth.

####  4.3.1 Gaussian Random Fields

In the first scenario, we use Gaussian random fields (GRFs) to generate the
set of PDE parameters for learning the latent space and training the
latentPINN. The isotropic 2D spatial Gaussian random field is given by

$\boldsymbol{x}=\mathcal{F}^{-1}\left(C\cdot\mathrm{n}^{2}\cdot\sqrt{2}\cdot\tau^{\frac{1}{2}(2\alpha-2)}\cdot\left(\frac{4\pi^{2}(k_{x}^{2}+k_{x}^{\prime
2})+\tau^{2}}{2}\right)^{-\alpha/2}\right),\quad\boldsymbol{x}\in\mathbb{R}^{2},\quad
C\in\mathcal{N}(0,\text{I}),$ (7)

where an inverse Fourier transform ($\mathcal{F}^{-1}$) is performed on a
correlation function parameterized by the pixel dimension $n$, spatial
wavenumbers $k_{x}\in[0,\frac{\text{n}}{2}]$, and parameters controlling the
variance, $\tau$ and $\alpha$ in its exponent argument
($\sigma=\tau^{0.5*(2*\alpha-\text{n})}$ where $\sigma$ is the intended
variance). Specifically, we generated 50k image samples
($\sim\mathcal{N}(0,\text{I})$) of dimension 128x128 with $\tau=7$,
$\alpha=2.5$, and a periodic boundary condition. Out of the 50k three-channel
image samples, 39k were used for the training set, 1k for validation, and the
rest for testing. These near zero-mean images that range between -1 and 1 are
then converted into velocity fields by first adding a constant of 2 km/s
followed by a multiplication by a factor of 2. Thus, the resulting velocity
range is [2,6] km/s. Figure 2(a) shows samples out of the 50k generated
velocity fields.

For this PDE parameter distribution, we train an autoencoder to learn the
corresponding latent representation for 30 epochs. Though slightly smoothed
out, we can see the well-reconstructed images from the test set in Figure
2(b), reflecting that it learned the compressed latent representation. Out of
the 39k training samples, 100 are used during the second stage, the main PINN
training. Shown in the left-most column of each plot in Figures 4 and 3 are
velocities (out of the 100 samples) used during the PINNs training, while the
rest correspond to the unseen (testing) velocities. Going from left to right,
for each row, we order the testing samples based on an increasing value of
point-wise distance in the latent space. With such an order, we expect to have
more dissimilar images with increasing point-wise distance. We can see in
Figures 4(b) and 3(a) that the trained PINNs generalize to new velocities
well. It is noteworthy to mention that the solutions for the velocities from
the test set (of the Gaussian Random field set) are without any additional
training, unlike for vanilla PINNs (Figure 5(a)).

####  4.3.2 Mechanical MNIST

Next, we consider a slightly more complex dataset in the form of the solutions
to the Cahn-Hilliard equation Cahn and Hilliard (1958). The fourth-order
parabolic PDE details the phase separation dynamics of a two-phase mixture. In
our case, we utilize such a dataset to understand the behavior of our approach
in the face of a sparser (in pixel values) data distribution. We consider the
use of a smoothed version of the Mechanical MNIST dataset Kobeissi and Lejeune
(2022). We consider training, validation, and testing sets of sizes 15784,
1000, and 9420, respectively. The three-channel 128x128 are converted into
velocity fields using the same addition and multiplication constants as in the
previous case ($=2(x+2)$ km/s, where $x$ is the Mechanical MNIST pixel value,
between -1 and 1). Figure 2(c) illustrates samples drawn from the training
set.

Using the same training configuration and autoencoder model as in the previous
GRF case, we obtain after training well-reconstructed image samples from the
testing set in Figure 2(d). We assess the learned latent vectors after the
main PINNs training. For this training, we utilize a similar training
configuration as in the previous case with a slight modification. Considering
the complexity of this set (more variance), we double the hidden layers (from
12 to 24). We train the network for 10k epochs of batch size 163 points. We
test the accuracy of the predicted PDE solutions on test samples given by the
four rightmost columns in Figures 4(c) (input) and 4(d)) (reconstruction). We
utilize the trained encoder to provide latent vectors for the training of the
PINN, and we use the decoder to provide the velocity model used in the PDE
loss function. There is no clear distinction in the efficiency between
randomly initialized and transfer learning training (Figure 5(b)). This might
be related to the inherent limitations of the Eikonal equation in dealing with
such complex PDE parameters set (velocity models). As we are dealing with
sharper boundaries between two phases in the Mechanical MNIST dataset, the
PINNs training becomes slightly more challenging than for the previous
smoother spatial Gaussian random fields. Nonetheless, our trained latentPINN
still provides new PDE solutions given new PDE parameters without additional
PINN training.

####  4.3.3 Realistic Earth Models

Finally, we consider a PDE parameter distribution for applications in
traveltime modeling of the Earth subsurface in the form of realistic synthetic
subsurface velocity models (Figure 2(e)). We compile patches of images from
several 2D velocity images Billette and Brandsberg-Dahl (2005); Martin et al.
(2006) and 3D velocity cubes Aminzadeh et al. (1997); Naranjo et al. (2011)
resulting in a 70,522 256x256 three-channel images. The dataset entails a wide
range of geological settings ranging from highly faulted models to anomalous
geological bodies (e.g., salts and low-velocity layers). We utilize a 90-5-5
percentage split to generate the training-validation-testing datasets. Figures
2(e) and 4(e) illustrate samples drawn from the datasets.

For the first stage of the latentPINN training, we modify the encoder-decoder
network to perform down- and upsampling of the higher-resolution images.
Specifically, each convolutional block now consists of a 2D transposed
convolution layer of stride size 4 followed by a GELU layer and a 2D
convolution of stride size 1. We utilize two of these blocks in the encoder
and decoder (reversing the order within the encoder convolution block)
network. The same optimizer, latent vector size, and learning rate scheduler
(as in the previous two cases) are used for the training. We train this model
for 45 epochs with a batch size of 256 samples. The well-reconstructed images
shown in Figure 2(f) demonstrate the well-represented data distribution from
the compressed latent vectors. The quality of these vectors can be further
attested from the PINNs training results shown in Figure 4(f). For this
training, we utilize almost identical training dynamics (as in the Mechanical
MNIST) with a slight change in the initial learning rate. We start the
training with a 3e-4 learning rate. Our trained latentPINN manages to
generalize to the new testing samples (the four rightmost columns in Figures
4(e) and 4(f)). The subtle changes (from left to right) can be captured by the
trained PINN without additional training. Figure 5(c) further details the
saving in training costs that would be otherwise needed in conventional
approaches. Here, we only have to contend with the overhead cost of training
the latentPINN.

###  4.4 Functional Data Representations

One notable feature of our proposed latentPINN framework is its readily
extensible to sample functions over the continuous domain (e.g., PDE
parameters), termed _functional data_ representation. In such a task, we aim
to perform infinite-dimensional sampling of the functional data accommodated
by a neural network function. Such representation is an important feature for
a wide range of PDE applications. For example, being able to perform this
sampling benefits application related but not limited to imaging the Earth’s
interior Yang et al. (2021), weather forecasting Pathak et al. (2022), and
fluid flow simulation Wen et al. (2022).

To perform this sampling, we extend the trained autoencoder with a diffusion
model. Specifically, we utilize the recent latent diffusion model in which the
diffusion (generative modeling) process is performed over the learned latent
space. We utilize a diffusion model with 1000 diffusion steps, a linear noise
scheduler with initial and final constants of 0.0015 and 0.0195, respectively,
attention resolution of {8,16}, an initial learning rate of 2e-6, a channel
multiplier of size {1,2,3,4}, a channel head of size 32, and a batch size of 5
image samples. The training process requires 27 epochs for the Gaussian
spatial random fields, 32 epochs for the Mechanical MNIST dataset, and 289
epochs for the realistic Earth models. Figures 6 showcase the unconditional
sample images from the trained latent diffusion model.

(a)

(b)

(c)

Figure 6:  Samples from the latent diffusion model for the functional data
representation task. Shown from left to right are the resulting analysis on
different datasets; a) Gaussian random fields, b) Mechanical MNIST, and c)
realistic Earth model.

##  5 Discussion

Incorporating neural networks in solving PDEs is an emerging field in the
scientific machine learning community. So far, two broad approaches have
emerged, differing in whether the NN facilitates the learning process to
replace either the PDE operator (neural operators) or a solution to a
particular PDE (PINN). The latter offers more flexibility and can be done
completely independently of conventional solvers. Thus, it offers features
that motivate us to use it to replace the less flexible, mesh-oriented,
conventional solvers. However, to realize this replacement, we need to improve
the training efficiency of PINNs to at least rival the numerical ones. This is
needed because, unlike neural operators, the training of PINNs is part of the
inference, as PINNs are often trained for a fixed set of PDE parameters.
Further solutions for new PDE parameters, unlike neural operators, require
additional training. In other words, PINNs are trained _deterministically_ for
a specific parametric PDE. To bridge this gap, orthogonal to previous
approaches in speeding up the training process, latent representation learning
is introduced by appending the conventional coordinates input with learned
latent vectors representing certain PDE parameters. With this additional
information, we encourage the neural network to use its interpolation
capabilities (some of which shown in Taufik et al. (2022)) in its lower-
dimensional manifold representation of the solutions to generate the
corresponding solution. In other words, when such input latent vectors are
sufficiently learned to compress and express the distribution of the PDE
parameters, the trained PINNs will possess _generative_ features. By doing so,
though we have to sacrifice an increase in the training cost compared to the
usual PINN training, such an increase does not scale one-to-one with the
number of representative PDE parameters needed to train our latentPINN. For
example, as depicted in Figure 5, the training cost for a hundred PDE
parameters increases the training cost of the vanilla PINN (blue lines in
Figure 5) by 13 times at best (the GRF case) and 33 times (the mechanical
MNIST) at worst. This increase is, however, justifiable, as once the
latentPINN is trained, a new PDE solution can be predicted instantly for a
sample PDE parameter from the trained distribution ($\sim$ 40k velocity
fields) without additional NN training.

Our latentPINN implementation relies on the assumption of a well-represented
PDE parameters distribution by its compressed latent representations. The main
challenge in achieving such representations has been to harmonize the training
samples reconstruction and generalization process to unseen samples. One
significant limitation might arise when the new unseen PDE parameter lies far
away from the training distribution. Our approach, like most of the neural
operator models, will then either require re-training (from the first stage)
or we resort to transfer learning to calculate the new PDE solution. In such
scenarios, the calculated PDE solutions can be viewed as _ansatz_ for further
PINNs training. With the recent development of latent diffusion models used in
this study, we partially solve these trade-offs by promoting the use of an
autoencoder to learn the compressed representation and diffusion models to
sufficiently represent the posterior distribution of the PDE parameters.
Specifically, a Kullback-Leibler divergence regularization is introduced for
better generalization. Different means of regularization can also be used to
achieve generalization during the latent compression Esser et al. (2021),
which is worth investigating in future work.

As a byproduct of such a training scheme, not only does our framework
facilitate PINN to learn a distribution of PDE parameters, its trained
autoencoder can be further utilized for downstream tasks. One such task can be
in the form of functional data representation. In this task, instead of
facilitating surrogate maps for the PDE solution, the learned NN facilitates a
functional representation for the infinite-dimensional vector spaces dubbed
_functional data_. Using the trained autoencoder as its backbone, a latent
diffusion model can be deployed to perform generative modeling on such
functional data. Thus, we obtain much more realistic drawn samples compared to
the previous approach Seidman et al. (2023). Finally, in this study, we only
consider unconditional sampling with such a model; further progress can
include classes of the PDE parameters (i.e., geologic conditions for the Earth
models or pattern classes for the mechanical MNIST dataset) to perform
conditional sampling.

##  6 Conclusion

We proposed using learned compressed latent representations in physics-
informed neural network training. More specifically, we utilize a latent
diffusion model to learn compressed representations of the PDE parameters.
These compressed representations are used as input samples, along with samples
representing the coordinates of the solution space, to train a PINN to provide
functional solutions of PDE for a range of PDE parameters. Though the proposed
framework should comply with a more general class of parametric PDEs, we
tested our algorithm on a first-order nonlinear PDE on three different
datasets. With the additional latent vector input, the trained physics-
informed neural networks are able to learn to generate PDE solutions as a
function of PDE parameters. Thus, this approach retains the flexibility and
accuracy features of the functional representation of PINN solutions while
gaining the generalization abilities of a neural operator.

##  7 Acknowledgement

The authors thank King Abdullah University of Science and Technology (KAUST)
for supporting this research and the Seismic Wave Analysis group for the
supportive and encouraging environment. This work utilized the resources of
the Supercomputing Laboratory at KAUST in Thuwal, Saudi Arabia.

## References

  * Li et al. (2020a) Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Fourier neural operator for parametric partial differential equations.  _arXiv preprint arXiv:2010.08895_ , 2020a.
  * Lu et al. (2019) Lu Lu, Pengzhan Jin, and George Em Karniadakis.  Deeponet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators.  _arXiv preprint arXiv:1910.03193_ , 2019.
  * Li et al. (2020b) Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Graph kernel network for partial differential equations.  _arXiv preprint arXiv:2003.03485_ , 2020b.
  * Kovachki et al. (2021) Nikola Kovachki, Zongyi Li, Burigede Liu, Kamyar Azizzadenesheli, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar.  Neural operator: Learning maps between function spaces.  _arXiv preprint arXiv:2108.08481_ , 2021.
  * Raissi et al. (2019) Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  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\.
  * Sirignano and Spiliopoulos (2018) Justin Sirignano and Konstantinos Spiliopoulos.  Dgm: A deep learning algorithm for solving partial differential equations.  _Journal of computational physics_ , 375:1339–1364, 2018\.
  * Li et al. (2021) Zongyi Li, Hongkai Zheng, Nikola Kovachki, David Jin, Haoxuan Chen, Burigede Liu, Kamyar Azizzadenesheli, and Anima Anandkumar.  Physics-informed neural operator for learning partial differential equations.  _arXiv preprint arXiv:2111.03794_ , 2021.
  * Fang (2021) Zhiwei Fang.  A high-efficient hybrid physics-informed neural networks based on convolutional neural network.  _IEEE Transactions on Neural Networks and Learning Systems_ , 33(10):5514–5526, 2021.
  * Wang et al. (2022) Sifan Wang, Shyam Sankaran, and Paris Perdikaris.  Respecting causality is all you need for training physics-informed neural networks.  _arXiv preprint arXiv:2203.07404_ , 2022.
  * Huang and Alkhalifah (2022) Xinquan Huang and Tariq Alkhalifah.  Pinnup: Robust neural network wavefield solutions using frequency upscaling and neuron splitting.  _Journal of Geophysical Research: Solid Earth_ , 127(6):e2021JB023703, 2022.
  * Huang and Alkhalifah (2023a) Xinquan Huang and Tariq Alkhalifah.  Efficient physics-informed neural networks using hash encoding.  _arXiv preprint arXiv:2302.13397_ , 2023a.
  * Rzepecki and Doran (2022) Jaroslaw Rzepecki and Chris Doran.  Fast neural network based solving of partial differential equations.  _arXiv preprint arXiv:2205.08978_ , 2022.
  * Schiassi et al. (2021) Enrico Schiassi, Roberto Furfaro, Carl Leake, Mario De Florio, Hunter Johnston, and Daniele Mortari.  Extreme theory of functional connections: A fast physics-informed neural network method for solving ordinary and partial differential equations.  _Neurocomputing_ , 457:334–356, 2021.  ISSN 0925-2312.  doi: https://doi.org/10.1016/j.neucom.2021.06.015.  URL https://www.sciencedirect.com/science/article/pii/S0925231221009140.
  * Huang and Alkhalifah (2023b) Xinquan Huang and Tariq Alkhalifah.  Microseismic source imaging using physics-informed neural networks with hard constraints.  _arXiv preprint arXiv:2304.04315_ , 2023b.
  * Taufik et al. (2023a) Mohammad Taufik, Tariq Alkhalifah, and Umair Waheed.  A stable neural network-based eikonal tomography using hard-constrained measurements.  _Authorea Preprints_ , 2023a.
  * Goswami et al. (2020) Somdatta Goswami, Cosmin Anitescu, Souvik Chakraborty, and Timon Rabczuk.  Transfer learning enhanced physics informed neural network for phase-field modeling of fracture.  _Theoretical and Applied Fracture Mechanics_ , 106:102447, 2020.
  * bin Waheed et al. (2021) Umair bin Waheed, Ehsan Haghighat, Tariq Alkhalifah, Chao Song, and Qi Hao.  Pinneik: Eikonal solution using physics-informed neural networks.  _Computers & Geosciences_, 155:104833, 2021.
  * Penwarden et al. (2021) Michael Penwarden, Shandian Zhe, Akil Narayan, and Robert M Kirby.  Physics-informed neural networks (pinns) for parameterized pdes: a metalearning approach.  _arXiv preprint arXiv:2110.13361_ , 2021.
  * Liu et al. (2022) Xu Liu, Xiaoya Zhang, Wei Peng, Weien Zhou, and Wen Yao.  A novel meta-learning initialization method for physics-informed neural networks.  _Neural Computing and Applications_ , 34(17):14511–14534, 2022.
  * Huang et al. (2022) Xiang Huang, Zhanhong Ye, Hongsheng Liu, Shi Ji, Zidong Wang, Kang Yang, Yang Li, Min Wang, Haotian Chu, Fan Yu, et al.  Meta-auto-decoder for solving parametric partial differential equations.  _Advances in Neural Information Processing Systems_ , 35:23426–23438, 2022.
  * Yang and Perdikaris (2018) Yibo Yang and Paris Perdikaris.  Physics-informed deep generative models.  _arXiv preprint arXiv:1812.03511_ , 2018.
  * Zhu et al. (2019) Yinhao Zhu, Nicholas Zabaras, Phaedon-Stelios Koutsourelakis, and Paris Perdikaris.  Physics-constrained deep learning for high-dimensional surrogate modeling and uncertainty quantification without labeled data.  _Journal of Computational Physics_ , 394:56–81, 2019.
  * Yang et al. (2020) Liu Yang, Xuhui Meng, and George Em Karniadakis.  B-pinns: Bayesian physics-informed neural networks for forward and inverse pde problems with noisy data.  _J. Comput. Phys._ , 425:109913, 2020.
  * Linka et al. (2022) Kevin Linka, Amelie Schäfer, Xuhui Meng, Zongren Zou, George Em Karniadakis, and Ellen Kuhl.  Bayesian physics-informed neural networks for real-world nonlinear dynamical systems.  _ArXiv_ , abs/2205.08304, 2022.
  * Oszkinat et al. (2022) Clemens Oszkinat, Susan E Luczak, and IG Rosen.  Uncertainty quantification in estimating blood alcohol concentration from transdermal alcohol level with physics-informed neural networks.  _IEEE Transactions on Neural Networks and Learning Systems_ , 2022\.
  * Kim et al. (2022) Youngkyu Kim, Youngsoo Choi, David Widemann, and Tarek Zohdi.  A fast and accurate physics-informed neural network reduced order model with shallow masked autoencoder.  _Journal of Computational Physics_ , 451:110841, 2022.
  * Zou and Karniadakis (2023) Zongren Zou and George Em Karniadakis.  L-hydra: Multi-head physics-informed neural networks.  _arXiv preprint arXiv:2301.02152_ , 2023.
  * Ranade et al. (2021) Rishikesh Ranade, Chris Hill, Haiyang He, Amir Maleki, and Jay Pathak.  A latent space solver for pde generalization.  _arXiv preprint arXiv:2104.02452_ , 2021.
  * Lumley (1967) John Leask Lumley.  The structure of inhomogeneous turbulent flows.  _Atmospheric turbulence and radio wave propagation_ , pages 166–178, 1967.
  * Wu et al. (2022) Tailin Wu, Takashi Maruyama, and Jure Leskovec.  Learning to accelerate partial differential equations via latent global evolution.  _arXiv preprint arXiv:2206.07681_ , 2022.
  * Kontolati et al. (2023) Katiana Kontolati, Somdatta Goswami, George Em Karniadakis, and Michael D Shields.  Learning in latent spaces improves the predictive accuracy of deep neural operators.  _arXiv preprint arXiv:2304.07599_ , 2023.
  * Esser et al. (2021) Patrick Esser, Robin Rombach, and Bjorn Ommer.  Taming transformers for high-resolution image synthesis.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 12873–12883, 2021.
  * Rumelhart et al. (1985) David E Rumelhart, Geoffrey E Hinton, and Ronald J Williams.  Learning internal representations by error propagation.  Technical report, California Univ San Diego La Jolla Inst for Cognitive Science, 1985.
  * Hinton and Salakhutdinov (2006) Geoffrey E Hinton and Ruslan R Salakhutdinov.  Reducing the dimensionality of data with neural networks.  _science_ , 313(5786):504–507, 2006.
  * Kingma and Welling (2013) Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  _arXiv preprint arXiv:1312.6114_ , 2013.
  * Goodfellow et al. (2020) Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial networks.  _Communications of the ACM_ , 63(11):139–144, 2020.
  * Song and Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in neural information processing systems_ , 32, 2019.
  * Song et al. (2020) Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  _arXiv preprint arXiv:2011.13456_ , 2020.
  * Song et al. (2021) Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon.  Maximum likelihood training of score-based diffusion models.  _Advances in Neural Information Processing Systems_ , 34:1415–1428, 2021.
  * Rombach et al. (2022) Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer.  High-resolution image synthesis with latent diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 10684–10695, 2022.
  * Osher and Sethian (1988) Stanley Osher and James A Sethian.  Fronts propagating with curvature-dependent speed: Algorithms based on hamilton-jacobi formulations.  _Journal of Computational Physics_ , 79(1):12–49, 1988.  ISSN 0021-9991.  doi: https://doi.org/10.1016/0021-9991(88)90002-2.  URL https://www.sciencedirect.com/science/article/pii/0021999188900022.
  * Taufik et al. (2023b) Mohammad H Taufik, Umair bin Waheed, and Tariq A Alkhalifah.  A neural network based global traveltime function (globenn).  _Scientific Reports_ , 13(1):7179, 2023b.
  * Paszke et al. (2019) Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al.  Pytorch: An imperative style, high-performance deep learning library.  _Advances in neural information processing systems_ , 32, 2019.
  * Cahn and Hilliard (1958) John W Cahn and John E Hilliard.  Free energy of a nonuniform system. i. interfacial free energy.  _The Journal of chemical physics_ , 28(2):258–267, 1958.
  * Kobeissi and Lejeune (2022) Hiba Kobeissi and Emma Lejeune.  Mechanical mnist - cahn-hilliard, 2022.  URL https://open.bu.edu/handle/2144/43971.
  * Billette and Brandsberg-Dahl (2005) FJ Billette and Sverre Brandsberg-Dahl.  The 2004 bp velocity benchmark.  In _67th EAGE Conference & Exhibition_, pages cp–1. EAGE Publications BV, 2005.
  * Martin et al. (2006) Gary S Martin, Robert Wiley, and Kurt J Marfurt.  Marmousi2: An elastic upgrade for marmousi.  _The leading edge_ , 25(2):156–166, 2006.
  * Aminzadeh et al. (1997) F. Aminzadeh, J. Brac, and T. Kunz.  _3-D Salt and Overthrust models_.  SEG/EAGE 3-D Modeling Series No.1, 1997.
  * Naranjo et al. (2011) John C Naranjo, Eric J Ekstrand, John Etgen, Kenneth L Hargrove, James G Brewton, Oswaldo Garcia, Georgiy Astvatsaturov, Damian Hite, Bill Howieson, Bing Bai, et al.  Survey design and implementation of a multiple wide-azimuth towed streamer seismic acquisition strategy at the tiber discovery; deep water gulf of mexico, usa.  In _2011 SEG Annual Meeting_. OnePetro, 2011.
  * Yang et al. (2021) Yan Yang, Angela F Gao, Jorge C Castellanos, Zachary E Ross, Kamyar Azizzadenesheli, and Robert W Clayton.  Seismic wave propagation and inversion with neural operators.  _The Seismic Record_ , 1(3):126–134, 2021.
  * Pathak et al. (2022) Jaideep Pathak, Shashank Subramanian, Peter Harrington, Sanjeev Raja, Ashesh Chattopadhyay, Morteza Mardani, Thorsten Kurth, David Hall, Zongyi Li, Kamyar Azizzadenesheli, et al.  Fourcastnet: A global data-driven high-resolution weather model using adaptive fourier neural operators.  _arXiv preprint arXiv:2202.11214_ , 2022.
  * Wen et al. (2022) Gege Wen, Zongyi Li, Kamyar Azizzadenesheli, Anima Anandkumar, and Sally M Benson.  U-fno—an enhanced fourier neural operator-based deep-learning model for multiphase flow.  _Advances in Water Resources_ , 163:104180, 2022.
  * Taufik et al. (2022) Mohammad Hasyim Taufik, Umair bin Waheed, and Tariq A Alkhalifah.  Upwind, no more: Flexible traveltime solutions using physics-informed neural networks.  _IEEE Transactions on Geoscience and Remote Sensing_ , 60:1–12, 2022.
  * Seidman et al. (2023) Jacob H Seidman, Georgios Kissas, George J Pappas, and Paris Perdikaris.  Variational autoencoding neural operators.  _arXiv preprint arXiv:2302.10351_ , 2023.