Scientific Machine Learning:
Regularization, MAP Estimation, and Loss Weighting
A loss function with more than one term is a statement about which errors matter and by how much. This notebook starts from regularization and its reading as maximum-a-posteriori estimation, then turns to the multi-term losses of a physics-informed neural network, where the terms compete for the same parameters and the weights between them decide whether the boundary conditions are satisfied at all. Fixed weighting, adaptive weighting, and the hard constraint that removes a term from the competition are compared on one Poisson problem.
Learning Objectives¶
After studying this notebook, you should be able to
- explain why a data-only objective is often insufficient, and what an added term contributes;
- read a regularized loss as a measure of fit plus a measure of the magnitude of the coefficients, and explain why an overfitted model has large coefficients in the first place;
- derive the regularized loss from Bayes' theorem as a maximum a posteriori estimate, and show that ridge regression is MAP estimation under a Gaussian prior;
- describe the PDE residual as a prior over functions, and state precisely in what sense it is and is not a regularizer;
- recognize the failure of a naive PINN, a total loss that falls convincingly while neither condition is satisfied, on the one-dimensional Poisson problem;
- explain why a fixed weight can recover the solution or converge to a wrong one, and why even a right value does not stay right;
- explain why the optimizer responds to the gradient of each term rather than to its value, and derive the gradient-based adaptive weighting rule from that fact;
- explain why the measured weights should be smoothed by a moving average, and what the smoothing does and does not change;
- build a boundary condition into the network output with a masking function, and say what the construction removes from the objective and what it costs.
In Physics-informed Neural Networks (PINN), the governing equation, the boundary conditions and the initial condition were all treated alike. Each was written as a residual, squared, averaged over a set of points, and added to the loss:
$$ \mathcal{L} = \mathcal{L}_{\text{PDE}} + \mathcal{L}_{\text{BC}} + \mathcal{L}_{\text{IC}} + (\mathcal{L}_{\text{DATA}}) . $$
Nothing in this construction forces the network to satisfy any of the conditions. The optimizer is only asked to make the sum of the violations small, and a condition enforced in this way, by a penalty that is minimized rather than an equation that is imposed, is called a soft constraint. Every term of a PINN loss is a soft constraint.
Adding a penalty term to a loss is a well-established idea in machine learning, where it is known as regularization. It comes with a theory: what a penalty term expresses, why it improves the solution, and how its weight should be chosen. The purpose of this notebook is to read the PINN loss through that theory. The physics terms will turn out to be regularizers in the precise sense of the word, and the theory of regularization will tell us what the weights of the loss mean and how to set them, which the earlier notebooks left open.
To see the mechanism in its simplest form, we begin with an ordinary regression problem that contains no physics at all, and return to the PINN in Section 4.
In supervised learning the parameters of a model are chosen by minimizing a measure of disagreement with observations,
$$ \theta^\star=\arg\min_\theta\;\mathcal{L}_{\text{data}}(\theta), \qquad \mathcal{L}_{\text{data}}=\frac{1}{N}\sum_{i=1}^{N}\left(y_i-\hat{y}_i(\theta)\right)^2 . $$
This is the whole of the objective in a first course, and for a large dataset drawn from a well-behaved distribution it is often enough. Engineering data are rarely like that. Measurements are few, because each one costs an experiment. They are noisy, because instruments are. They cover the operating region that happened to be tested rather than the one of interest. And the model typically has more parameters than there are measurements to constrain them.
Choosing among them requires information from somewhere other than the data, and the standard way to supply it is to add a term to the objective,
$$ \mathcal{L}=\mathcal{L}_{\text{data}}+\lambda\,\mathcal{L}_{\text{reg}} . $$
This is regularization, and in its usual presentation it arrives as a trick: a penalty on large parameters that happens to improve test error, with $\lambda$ tuned by cross-validation. That description is operationally correct and conceptually empty. It says nothing about what the second term means, where it should come from, or how $\lambda$ should be chosen when there are four such terms rather than one, which is exactly the situation in a physics-informed neural network.
A better account is available, and the rest of this section is built on it.
2.1. A Measure of Fit and a Measure of Magnitude¶
Regularization exists to control overfitting: a model that reproduces the data it was fitted to almost exactly, yet predicts badly everywhere else. It arises whenever a model has enough freedom to follow the noise as well as the signal, which in practice means whenever the number of adjustable parameters approaches the number of observations.
The mechanism is easiest to see in a small regression problem. Let $\{(x_i, y_i)\}_{i=1}^{N}$ be $N = 10$ noisy observations of a smooth function, and approximate that function by a linear combination of radial basis functions (RBFs),
$$ \hat y(x) \;=\; \sum_{j=1}^{d}\theta_j\,\phi_j(x) \;=\; \Phi(x)\,\theta , \qquad \phi_j(x) \;=\; \exp\left(-\frac{(x - u_j)^2}{2\sigma^2}\right) , $$
where $u_j$ is the centre of the $j$-th basis function and $\sigma$ its width. A radial basis function depends on its input only through the distance $\lvert x - u_j \rvert$ from that centre, which is what the name records; the Gaussian form above is the usual choice.
Here $d = 10$ centres are placed one at each observation, with width $\sigma = 1$ against a centre spacing of $1$. The model is linear in its parameters, so fitting it is the ordinary least-squares problem
$$ \min_\theta \; \lVert \Phi\theta - y \rVert_2^2 , \qquad \Phi_{ij} = \phi_j(x_i) , $$
and no neural network is involved. Because $d = N$, the design matrix $\Phi$ is square and, for these centres and this width, well conditioned, so the fit reproduces every observation exactly. Whatever goes wrong next is therefore a property of the model and not of the arithmetic.
The computation below fits the ten radial basis functions by least squares and shows the result twice: the fitted curve against the data, and the fitted coefficients as a stem plot. Look at the second panel before reading on.
import numpy as np
from matplotlib import pyplot as plt
%matplotlib inline
x = np.linspace(-4.5, 4.5, 10).reshape(-1, 1)
y = np.array([0.9819, 0.7973, 1.9737, 0.1838, 1.3180,
-0.8361, -0.6591, -2.4701, -2.8122, -6.2512]).reshape(-1, 1)
d = 10
u = np.linspace(-4.5, 4.5, d)
sigma = 1
xp = np.linspace(-5, 5, 501).reshape(-1, 1)
Phi = np.exp(-(x - u)**2/(2*sigma**2))
Phip = np.exp(-(xp - u)**2/(2*sigma**2))
theta = np.linalg.solve(Phi.T@Phi, Phi.T@y)
print('coefficients from %.1f to %.1f' % (theta.min(), theta.max()))
print('data from %.1f to %.1f' % (y.min(), y.max()))
print('RSS = %.2e' % np.linalg.norm(y - Phi@theta)**2)
plt.figure(figsize = (10, 4))
plt.subplot(1, 2, 1)
plt.plot(x, y, 'bo')
plt.plot(xp, Phip@theta, 'k')
plt.axis([-5, 5, -12, 6])
plt.grid(True)
plt.xlabel('x', fontsize = 12)
plt.ylabel('y', fontsize = 12)
plt.title('least squares, no penalty', fontsize = 12)
plt.subplot(1, 2, 2)
plt.stem(np.arange(1, d+1), theta, linefmt = 'k', markerfmt = 'ko', basefmt = 'k')
plt.ylim([-20, 20])
plt.grid(True)
plt.xlabel(r'coefficient index', fontsize = 12)
plt.ylabel(r'$\theta_i$', fontsize = 12)
plt.title('magnitude of the coefficients', fontsize = 12)
plt.show()
The fitted curve passes through all ten observations, with a residual sum of squares of zero to machine precision, and oscillates between them. The coefficients account for that oscillation. They run from about $-14.9$ to $+15.6$ and alternate in sign from one centre to the next, while the observations themselves lie between $-6.3$ and $+2.0$. The model therefore produces outputs of order one as a sum of contributions of order fifteen that very nearly cancel.
This near-cancellation is the mechanism of overfitting, not an incidental symptom of it. A curve built from smooth basis functions can be forced through scattered points only if large positive and large negative contributions cancel almost everywhere and fail to cancel exactly at the observations. Such cancellation is delicate: a small change in the data requires a large change in the coefficients, which is what it means for an estimator to have high variance. The model has fitted the noise, and the price of doing so is the magnitude of its coefficients.
This observation suggests its own remedy. Overfitting is accompanied by parameters of large magnitude, so the magnitude of the parameters should itself carry a cost. The objective then balances two quantities rather than one,
$$ \begin{align*} \text{total cost} \;=\;&\underbrace{\text{measure of fit}}_{\text{RSS}(\theta)} \;+\;\lambda\cdot\underbrace{\text{measure of magnitude of coefficients}}_{\lVert\theta\rVert_2^2} \\ \\ \implies\;&\min_\theta\;\lVert\Phi\theta-y\rVert_2^2+\lambda\lVert\theta\rVert_2^2 , \end{align*} $$
where the coefficient $\lambda\ge0$ fixes the rate of exchange between the two terms. For $\lambda=0$ the second term is absent and the estimate reduces to ordinary least squares; in the limit $\lambda\to\infty$ the first term is ignored and every coefficient is driven to zero. The useful regime lies between these two extremes. The resulting estimator is known as ridge regression.
The same problem is now solved with the penalty added, at $\lambda = 0.1$.
lamb = 0.1
theta_r = np.linalg.solve(Phi.T@Phi + lamb*np.eye(d), Phi.T@y)
print('largest coefficient %.1f -> %.1f' % (np.abs(theta).max(), np.abs(theta_r).max()))
print('RSS %.2e -> %.2f' % (np.linalg.norm(y - Phi@theta)**2,
np.linalg.norm(y - Phi@theta_r)**2))
print('number of exactly zero coefficients: %d' % np.sum(theta_r == 0))
plt.figure(figsize = (10, 4))
plt.subplot(1, 2, 1)
plt.plot(x, y, 'bo')
plt.plot(xp, Phip@theta, 'k--')
plt.plot(xp, Phip@theta_r, 'k')
plt.legend(['data', 'no penalty', 'ridge, $\\lambda = 0.1$'], fontsize = 11)
plt.axis([-5, 5, -12, 6])
plt.grid(True)
plt.xlabel('x', fontsize = 12)
plt.ylabel('y', fontsize = 12)
plt.title('ridge regression', fontsize = 12)
plt.subplot(1, 2, 2)
plt.stem(np.arange(1, d+1), theta_r, linefmt = 'k', markerfmt = 'ko', basefmt = 'k')
plt.ylim([-20, 20])
plt.grid(True)
plt.xlabel(r'coefficient index', fontsize = 12)
plt.ylabel(r'$\theta_i$', fontsize = 12)
plt.title('magnitude of the coefficients', fontsize = 12)
plt.show()
With the penalty in place the fitted curve is smooth and passes near the observations rather than through them. The largest coefficient has fallen from about $15.6$ to about $5.0$, and the alternation of signs from one centre to the next has largely disappeared. The residual sum of squares has risen from zero to about $4.60$, so the fit to the training data is strictly worse than before. This is the intended effect and not a failure: a small amount of bias has been accepted in exchange for a large reduction in variance, and the smooth curve is the better description of whatever process generated the observations.
The coefficient required to bring this about is small. At $\lambda = 0.1$ the penalty term contributes $2.81$ against a residual sum of squares of $4.60$, so the two are of comparable size at the solution.
One feature of the result deserves attention: not a single coefficient is exactly zero. The $L_2$ penalty shrinks every coefficient toward the origin and, except by coincidence, sends none of them all the way there. It therefore governs how large the model may be, not which basis functions it may use.
Sweeping $\lambda$ over five decades and recording all ten coefficients at each value gives the plot below, with $\lambda$ on a logarithmic axis. It is called a regularization path.
lambs = np.logspace(-4, 1, 200)
theta_record = []
for k in range(len(lambs)):
th = np.linalg.solve(Phi.T@Phi + lambs[k]*np.eye(d), Phi.T@y)
theta_record.append(th.ravel())
theta_record = np.array(theta_record)
plt.figure(figsize = (6, 4))
plt.semilogx(lambs, theta_record, 'k', linewidth = 1)
plt.grid(True)
plt.xlabel(r'$\lambda$', fontsize = 12)
plt.ylabel(r'weight $\theta$', fontsize = 12)
plt.title('ridge coefficients as a function of regularization', fontsize = 12)
plt.show()
Every trace begins at its unregularized value on the left, where the coefficients are spread from about $-14.9$ to $+15.6$, and every one of them contracts toward zero as $\lambda$ increases.
Section 2 left two questions unanswered. The penalty $\lVert\theta\rVert_2^2$ was chosen for convenience, and nothing was said about why that particular measure of magnitude should be the right one; the coefficient $\lambda$ was set by hand, and nothing in the fit itself determines it. Both are answered by describing the same regression problem in the language of probability, which is the subject of this chapter.
In traditional supervised learning, data is assumed to be fixed, and the model is optimized to fit it as closely as possible. In contrast, the probabilistic viewpoint assumes that the generative model comes first. This model, though unknown, gives rise to the observed data via a stochastic process.
From this perspective, the dataset is seen as a realization of a random process governed by an underlying distribution. Our goal is to infer this distribution - or aspects of it, such as its parameters - using the observed data. This leads to the broader framework of generative modeling, where learning becomes a process of probabilistic inference.
We begin with a generative perspective: assume there exists a true, underlying linear relationship between $x$ and $y$, but our observations are corrupted by noise.
Given observed data:
$$ D = \{(x_1, y_1), (x_2, y_2), \cdots, (x_m, y_m)\} $$
We assume each output $y_i$ is generated by a linear function of the input $x_i$ with additive Gaussian noise:
$$ y = \omega^T x + \varepsilon, \quad \varepsilon \sim \mathcal{N}(0, \sigma^2) $$
This leads to the conditional likelihood:
$$ P(y \mid x \,;\, \omega) = \mathcal{N}(\omega^T x, \sigma^2) $$
That is, for a given input $x$, the output $y$ is normally distributed around the mean $\omega^T x$, with variance $\sigma^2$ reflecting observation noise.
In the frequentist approach, the parameters $\omega$ and $\sigma^2$ are treated as fixed but unknown. Our goal is to estimate them using Maximum Likelihood Estimation (MLE).
Given the dataset:
$$ D = \{(x_1, y_1), (x_2, y_2), \cdots, (x_m, y_m)\} $$
and the generative model:
$$ y_i = \omega^T x_i + \varepsilon_i, \quad \varepsilon_i \sim \mathcal{N}(0, \sigma^2) $$
We assume that each $y_i$ is independently sampled:
$$ P(y_i \mid x_i; \omega, \sigma^2) = \frac{1}{\sqrt{2\pi \sigma^2}} \exp\left( -\frac{1}{2\sigma^2} (y_i - \omega^T x_i)^2 \right) $$
Likelihood and Log-Likelihood
The likelihood of the parameters given the dataset is:
$$ \begin{align*} \mathcal{L}(\omega, \sigma^2) = P(D \mid \theta) = P(D \,;\, \theta) &= P\left(y_1,y_2,\cdots,y_m \mid x_1,x_2,\cdots,x_m; \; \underbrace{\omega, \sigma}_{\theta}\right) \\ &= \prod_{i=1}^m P(y_i \mid x_i; \omega, \sigma^2) \end{align*} $$
Taking the logarithm (log-likelihood):
$$ \begin{align*} \ell (\omega,\sigma^2) = \log \mathcal{L}(\omega,\sigma^2) &= \log \prod\limits^m_{i=1}P\left(y_i \mid x_i\,; \omega,\sigma^2 \right)\\ &= \sum\limits^m_{i=1}\log P \left(y_i \mid x_i\,; \omega, \sigma^2 \right)\\ &= \sum\limits^m_{i=1}\log \frac{1}{\sqrt{2\pi\sigma^2}}\exp{\left(-\frac{\left(y_i - \omega^Tx_i \right)^2}{2\sigma^2}\right)}\\ &= \sum\limits^m_{i=1}\left\{ -\frac{1}{2}\log \left(2\pi\sigma^2 \right) - \frac{\left(y_i - \omega^Tx_i \right)^2}{2\sigma^2}\right\} \end{align*} $$
Maximum Likelihood Estimation (MLE)
To estimate parameters, we maximize the log-likelihood:
$$ \hat{\theta}_{MLE} = {\arg\max_{\omega, \sigma^2}} \;\; \ell(\omega, \sigma^2) $$
(1) Focusing first on maximizing with respect to $\omega$:
$$ \begin{align*} \hat{\omega}_{MLE} &= \arg\max_{\omega} \;- \frac{1}{2\sigma^2}\sum\limits^m_{i=1} \left(y_i-\omega^Tx_i \right)^2\\ &= \arg\min_{\omega} \frac{1}{2\sigma^2}\sum\limits^m_{i=1} \left(y_i-\omega^Tx_i \right)^2\\ &= \arg\min_{\omega} \sum\limits^m_{i=1} \left(y_i-\omega^Tx_i \right)^2 \end{align*} $$
This is exactly the least-squares objective from classical linear regression. (Amazing !)
Maximum likelihood estimation under a Gaussian noise model leads to the least-squares solution.
(2) Focusing then on $\sigma$:
$$ \begin{align*} \ell (\omega,\sigma^2) &= -\frac{m}{2}\log{2\pi}-m\log{\sigma}-\frac{1}{2\sigma^2}\sum\limits_{i=1}^m\left(y_i-\omega^T x_i\right)^2 \\\\ \frac{\partial \ell}{\partial \sigma} & =-\frac{m}{\sigma}+\frac{1}{\sigma^3}\sum\limits_{i=1}^m\left(y_i-\omega^T x_i\right)^2=0 \quad \\\\ & \implies \; \hat \sigma^2_{MLE}=\frac{1}{m}\sum\limits_{i=1}^m\left(y_i-\omega^T x_i\right)^2 \end{align*} $$
Big Lessons
- MLE under Gaussian noise leads directly to least-squares regression
- Probabilistic modeling justifies the use of least-squares from first principles
Lab:
The following example demonstrates that if a linear model is correctly specified and well estimated, then the residual errors (i.e., the differences between predicted and observed values) should approximately follow a Gaussian distribution:
$$ \varepsilon = y - \omega^T x \sim \mathcal{N}(0, \sigma^2) $$
m = 200
a = 1
x = 3 + 2*np.random.uniform(0, 1, [m, 1])
noise = 0.1*np.random.randn(m, 1)
y = a*x + noise
plt.figure(figsize = (6, 6))
plt.plot(x, y, 'b.')
plt.axis('equal')
plt.grid(alpha = 0.3)
plt.show()
# y = theta[1]*x + theta[0]
A = np.hstack([np.ones([m, 1]), x])
theta = np.linalg.solve(A.T@A, A.T@y).ravel()
# to plot the fitted line
xp = np.linspace(np.min(x), np.max(x))
yp = theta[1]*xp + theta[0]
plt.figure(figsize = (6, 6))
plt.plot(x, y, 'b.')
plt.plot(xp, yp, 'r', linewidth = 3)
plt.axis('equal')
plt.grid(alpha = 0.3)
plt.show()
yhat0 = theta[1]*x + theta[0]
err0 = yhat0 - y
yhat1 = 1.2*x - 1
err1 = yhat1 - y
yhat2 = 1.3*x - 1
err2 = yhat2 - y
plt.figure(figsize = (10, 6))
for k, (yhat, err) in enumerate([(yhat0, err0), (yhat1, err1), (yhat2, err2)]):
plt.subplot(2, 3, k + 1)
plt.plot(x, y, 'b.', x, yhat, 'r')
plt.axis([2.9, 5.1, 2.9, 5.1])
plt.grid(alpha = 0.3)
plt.subplot(2, 3, k + 4)
plt.hist(err, 31)
plt.axvline(0, color = 'k')
plt.xlabel(r'$\epsilon$', fontsize = 15)
plt.yticks([])
plt.axis([-1, 1, 0, 15])
plt.grid(alpha = 0.3)
plt.show()
plt.figure(figsize = (10, 6))
for k, (yhat, err) in enumerate([(yhat0, err0), (yhat1, err1), (yhat2, err2)]):
plt.subplot(2, 3, k + 1)
plt.plot(x, y, 'b.', x, yhat, 'r')
plt.axis([2.9, 5.1, 2.9, 5.1])
plt.grid(alpha = 0.3)
plt.subplot(2, 3, k + 4)
plt.plot(err[1:], err[:-1], '.')
plt.axis([-0.7, 0.7, -0.7, 0.7])
plt.gca().set_aspect('equal')
plt.grid(alpha = 0.3)
plt.xlabel(r'$\epsilon_i$', fontsize = 15)
plt.ylabel(r'$\epsilon_{i-1}$', fontsize = 15)
plt.show()
In the frequentist approach, we estimate the parameters $\omega$ by maximizing the likelihood (MLE), which implicitly assumes no prior information or a uniform prior over $\omega$. In contrast, the Bayesian approach treats $\omega$ as a random variable and incorporates prior beliefs about its distribution.
Bayesian linear regression not only provides a point estimate of $\omega$, but also quantifies uncertainty through its posterior distribution. This enables a principled approach to regularization and confidence assessment.
3.3.1. A Prior on $\omega$¶
Suppose we assume a Gaussian prior over the weight vector $\omega = \begin{bmatrix} \omega_0 \\ \omega_1 \end{bmatrix}$:
$$ P(\omega) = \mathcal{N}(0, \Sigma) = \mathcal{N}\left(0, \lambda^{-1} I\right) $$
This expresses our belief that $\omega$ is centered at zero, with uncertainty controlled by $\lambda$. A large $\lambda$ implies strong confidence that weights should be small (i.e., a stronger regularization effect).
The prior probability density is:
$$ P(\omega) = \frac{1}{(2\pi)^{D/2}} \exp\left( -\frac{\lambda}{2} \omega^T \omega \right) $$
This means that before seeing any data, we believe that the weights are centered around zero, with uncertainty (variance) controlled by the parameter $\lambda$.
Understanding this assumption is important: it formalizes the idea that we prefer simpler models (with smaller weights), unless the data provides strong evidence otherwise. This is the Bayesian interpretation of regularization.
3.3.2. Posterior Distribution¶
By Bayes' Rule, the posterior distribution over $\omega$ given data $D$ is:
$$P(\omega \mid D) = \frac{P(D \mid \omega) P(\omega)}{P(D)}$$
Taking logarithms, we get the log-posterior:
$$\log P(\omega \mid D) = \log\frac{P(D \mid \omega)P(\omega)}{P(D)} = \log P(D \mid \omega) + \log P(\omega) - \underbrace{\log P(D)}_{\text{constant}}$$
Since $P(D)$ does not depend on $\omega$, we can treat it as a constant in optimization.
3.3.3. Maximum-a-Posteriori (MAP)¶
To find the most likely value of $\omega$ given the data (i.e., the mode of the posterior), we use MAP estimation:
$$ \begin{align*} \hat{\omega}_{MAP} &= \arg\max_{\omega} \; \log P(\omega \mid D)\\\\ &= \arg\max_{\omega} \; \left\{\log P(D \mid \omega) + \log P(\omega)\right\}\\\\ &= \arg\max_{\omega} \; \left\{ \sum\limits^m_{i=1}\left\{ -\frac{1}{2}\log\left(2\pi\sigma^2\right) - \frac{\left(y_i-\omega^Tx_i\right)^2}{2\sigma^2} \right\} -\frac{D}{2}\log (2\pi) - \frac{\lambda}{2}\omega^T\omega \right\}\\\\ &= \arg\min_{\omega} \; \frac{1}{2\sigma^2}\sum\limits^m_{i=1}\left(y_i - \omega^Tx_i\right)^2 + \frac{\lambda}{2}\omega^T\omega \quad \\&\text{ (ignoring constants and changing max to min)} \end{align*} $$
Regularized Least Squares
If we assume $\sigma^2 = 1$, the MAP objective becomes:
$$\hat{\omega}_{MAP} = \arg\min_{\omega} \left\{\sum\limits^m_{i=1}\left(y_i - \omega^Tx_i\right)^2 + \lambda\omega^T\omega \right\}$$
This is equivalent to ridge regression, also known as $L_2$-regularized linear regression. (Amazing !).
Big Lesson
- MLE assumes no prior and yields the least-squares solution
- MAP incorporates a Gaussian prior and yields $L_2$ regularization
- Bayesian inference provides a natural framework to balance data fit and model complexity
MAP estimation illustrates how Bayesian reasoning not only extends classical approaches but also provides a principled foundation for regularization and improved generalization.
4.1. What Engineers Know Before Seeing Data¶
The prior in Section 3 was weak and generic: parameters are probably not enormous. In an engineering problem, far stronger prior knowledge is usually available, and it has nothing to do with the parameters of any model. Mass is conserved. Momentum is conserved. Energy is conserved. Stress and strain are related by a constitutive law. The field satisfies a partial differential equation with prescribed boundary and initial conditions.
None of this is contained in a table of measurements, and all of it is more trustworthy than the measurements. The question is how to write it into an objective function.
Consider a governing equation written with all terms on one side,
$$ \frac{\partial u}{\partial t}+\mathcal{N}[u]=0, $$
where $\mathcal{N}$ collects the spatial operators. Let a neural network with parameters $\theta$ stand in for the solution, $u(x,t)\approx u_\theta(x,t)$. Because the network is differentiable and its derivatives are available exactly by automatic differentiation, the equation can be evaluated at any point of the domain. What comes back is the residual
$$ r_\theta(x,t)=\frac{\partial u_\theta}{\partial t}+\mathcal{N}[u_\theta], $$
which is zero wherever the network happens to satisfy the physics and nonzero elsewhere. Its mean square over a set of sampled collocation points is a number that measures how badly the physics is violated,
$$ \mathcal{L}_{\text{physics}}=\frac{1}{N_f}\sum_{i=1}^{N_f}\left\lvert r_\theta(x_i,t_i)\right\rvert^2 , $$
and adding it to the objective expresses a preference for networks that obey the equation:
$$ \mathcal{L}=\mathcal{L}_{\text{data}}+\lambda_f\,\mathcal{L}_{\text{physics}} . $$
This is the basic physics-informed neural network, and read in the light of Section 3 it is a MAP estimate whose prior is a physical law rather than a statement about parameter magnitudes.
The question is worth answering carefully, because both possible answers mislead if stated flatly.
In the broad sense, yes. It is an additional term that restricts which solutions the optimizer is encouraged to find, it encodes knowledge not present in the data, and it enters the objective in the same place a regularizer does. Without it the network fits the observations; with it the network fits the observations while respecting the governing equation.
In the narrow sense used in statistical learning, no. A regularizer such as weight decay is a generic, problem-independent preference for small parameters, whose purpose is to control the effective capacity of the model and reduce variance. The physics term does none of that. It is problem-specific, it says nothing about capacity, and it can just as easily increase the difficulty of the optimization as reduce it. Its purpose is not to prevent overfitting but to impose a known constraint.
The accurate description is the middle one: the physics loss is a soft constraint, imposed by penalty, carrying domain knowledge into the objective.
The hierarchy $\text{regularization}\rightarrow\text{prior knowledge}\rightarrow\text{physical prior}\rightarrow\text{PINN}$ is a good guide to how the idea developed, provided the last arrow is understood as a change in the nature of the information and not merely in its name.
5.1. The Problem¶
The problem used throughout is the one-dimensional Poisson equation
$$ -u''(x) = (\alpha\pi)^2\sin(\alpha\pi x), \qquad x\in(0,1), $$
with homogeneous Dirichlet boundary conditions
$$ u(0) = u(1) = 0 , $$
and exact solution
$$ u(x) = \sin(\alpha\pi x) . $$
The parameter $\alpha$ sets the frequency of the solution and, through the factor $(\alpha\pi)^2$ in the forcing term, the scale of the PDE residual. The main experiments use $\alpha = 7$: four oscillations on the interval, and a forcing term of size $(7\pi)^2 \approx 484$.
Setup
NumPy, Matplotlib and PyTorch. All runs in this notebook work on a CPU; each of the long runs takes a few minutes.
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device:', device)
Model, points and loss terms
The network takes $x$ and returns $u_\theta(x)$ directly: one input, three hidden layers of $64$ units with $\tanh$, one output. The collocation points are the $N_f = 100$ interior points of a uniform grid, so that no collocation point coincides with a boundary point, and the forcing term is evaluated at them once. Two details matter for every comparison in this notebook. new_pinn() resets the random seed before building the network, so that every model starts from the same weights; and derivative wraps the one call to automatic differentiation that is needed, returning $\partial u/\partial x$ at the collocation points, so that calling it twice gives $u_{xx}$. The optimizer is Adam with learning rate $10^{-4}$, and every long run takes $30000$ iterations.
alpha = 7.0
n_iter = 30000
lr = 1e-4
n_pde = 100
x_pde = torch.linspace(0, 1, n_pde + 2, device = device)[1:-1].reshape(-1, 1)
x_bc = torch.tensor([[0.0], [1.0]], device = device)
f_pde = (alpha*np.pi)**2*torch.sin(alpha*np.pi*x_pde)
x_pde.requires_grad_(True)
x_plot = torch.linspace(0, 1, 400, device = device).reshape(-1, 1)
x_np = x_plot.cpu().numpy().ravel()
u_true = np.sin(alpha*np.pi*x_np)
iteration = np.arange(1, n_iter + 1)
def new_pinn():
torch.manual_seed(3) # every model starts from the same weights
return nn.Sequential(nn.Linear(1, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1)).to(device)
def derivative(u, x):
return torch.autograd.grad(u, x, torch.ones_like(u), create_graph = True)[0]
The boundary values stay at zero whatever $\alpha$ is; what $\alpha$ changes is the interior, which oscillates faster and needs a larger second derivative to match the forcing. The two conditions therefore live on very different scales, and the gap widens with $\alpha$. The cell below draws the exact solution for three values of $\alpha$.
x_line = np.linspace(0, 1, 1000)
plt.figure(figsize = (8, 4))
for a in [1, 5, 10]:
plt.plot(x_line, np.sin(a*np.pi*x_line), linewidth = 2, label = r'$\alpha = %d$' % a)
plt.plot([0, 1], [0, 0], 'ko')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.legend()
plt.show()
The two conditions become two mean squared residuals, the equation over the collocation points and the boundary condition over the two end points,
$$ \mathcal{L}_{\mathrm{PDE}} = \frac{1}{N_f}\sum_{i=1}^{N_f}\bigl(-u_{\theta,xx}(x_i) - (\alpha\pi)^2\sin(\alpha\pi x_i)\bigr)^2 , \qquad \mathcal{L}_{\mathrm{BC}} = \frac{u_\theta(0)^2 + u_\theta(1)^2}{2} , $$
and the most natural thing to do is to add them,
$$ \mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{PDE}} + \mathcal{L}_{\mathrm{BC}} . $$
This is the naive PINN, and it is what most first implementations do. Every iteration of the loop below does the same five things: clear the old gradients, evaluate the network and its second derivative at the collocation points, form the two losses and their sum, backpropagate, and take a step. The two losses are recorded at every iteration and plotted under the prediction.
net_naive = new_pinn()
optimizer = torch.optim.Adam(net_naive.parameters(), lr = lr)
loss_pde_hist = np.zeros(n_iter)
loss_bc_hist = np.zeros(n_iter)
for k in range(n_iter):
optimizer.zero_grad()
u = net_naive(x_pde)
u_x = derivative(u, x_pde)
u_xx = derivative(u_x, x_pde)
loss_pde = torch.mean((-u_xx - f_pde)**2)
loss_bc = torch.mean(net_naive(x_bc)**2)
loss = loss_pde + loss_bc
loss.backward()
optimizer.step()
loss_pde_hist[k] = loss_pde.item()
loss_bc_hist[k] = loss_bc.item()
print('Naive PINN | PDE %.3e | BC %.3e' % (loss_pde.item(), loss_bc.item()))
with torch.no_grad():
u_naive = net_naive(x_plot).cpu().numpy().ravel()
plt.figure(figsize = (7, 7))
plt.subplot(2, 1, 1)
plt.plot(x_np, u_true, 'k--', linewidth = 2, label = 'Ground Truth', alpha = 0.5)
plt.plot(x_np, u_naive, linewidth = 2, label = 'Prediction')
plt.title('Naive PINN')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.ylim(-1.2, 1.2)
plt.legend()
plt.subplot(2, 1, 2)
plt.semilogy(iteration, loss_pde_hist, label = r'$\mathcal{L}_{\mathrm{PDE}}$')
plt.semilogy(iteration, loss_bc_hist, label = r'$\mathcal{L}_{\mathrm{BC}}$')
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.legend()
plt.tight_layout()
plt.show()
The result is not a slightly wrong solution; it is a different function altogether. The prediction is a steep curve that leaves the plotting window, the PDE residual is $3.4\times10^{4}$, and the boundary loss is $58$, which means the two end values are of order $\pm 10$. After $30000$ iterations neither condition is satisfied, and yet nothing in the run would have raised an alarm: the total loss started above $10^{5}$ and fell for the whole run. The lower panel shows the two terms separately. The PDE loss is three orders of magnitude larger than the boundary loss throughout, and it is the only one that descends; the boundary loss starts near $10^{-2}$, because an untrained network returns values near zero at the ends, and is driven up to about $60$, with a few brief dips, as the optimizer trades the boundary values away for a smaller residual. A loss that is falling says that the optimizer is doing its job. It does not say that the job was well posed. The simplest remedy is to weight the two terms differently, which is the subject of the next chapter; what the optimizer was actually doing is examined in Section 7.1.
Fixed weighting sets the coefficient once and never touches it again. If the two gradient norms are mismatched by a known factor, a constant can bring the two contributions to the same length, and the method corrects the imbalance when the factor is known. The two runs below use the same network, starting weights, points, optimizer and $30000$ iterations as the naive run of Section 5.2, and differ from it in one constant.
Balanced at iteration $k$. The unweighted gradients (grey) are badly mismatched. The weights $\omega_{\mathrm{PDE}} = 0.5$ and $\omega_{\mathrm{BC}} = 3$ shorten one and lengthen the other until the two scaled vectors (black) are comparable, and the step is a compromise between the two conditions.
With a large enough boundary weight,
$$ \mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{PDE}} + 10^{4}\,\mathcal{L}_{\mathrm{BC}} , $$
the two objectives are brought to a comparable footing. The value is of the order of the ratio of the two gradient norms over the course of training, which the adaptive run of Section 7 will measure.
weight_bc = 1e4
net_fixed_good = new_pinn()
optimizer = torch.optim.Adam(net_fixed_good.parameters(), lr = lr)
for k in range(n_iter):
optimizer.zero_grad()
u = net_fixed_good(x_pde)
u_x = derivative(u, x_pde)
u_xx = derivative(u_x, x_pde)
loss_pde = torch.mean((-u_xx - f_pde)**2)
loss_bc = torch.mean(net_fixed_good(x_bc)**2)
loss = loss_pde + weight_bc*loss_bc
loss.backward()
optimizer.step()
print('Fixed, weight_bc = 1e4 | PDE %.3e | BC %.3e' % (loss_pde.item(), loss_bc.item()))
with torch.no_grad():
u_fixed_good = net_fixed_good(x_plot).cpu().numpy().ravel()
plt.figure(figsize = (7, 4))
plt.plot(x_np, u_true, 'k--', linewidth = 3, label = 'Ground Truth', alpha = 0.5)
plt.plot(x_np, u_fixed_good, linewidth = 1.5, label = 'Prediction')
plt.title(r'Fixed Weighting, $\omega_{\mathrm{BC}} = 10^4$')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.ylim(-1.2, 1.2)
plt.legend()
plt.show()
The solution is recovered. The prediction lies on the exact solution to the width of the line, with a maximum error of $4\times10^{-4}$; the PDE residual has fallen from $3.4\times10^{4}$ to $7.8\times10^{-3}$, and the boundary loss to $4.7\times10^{-8}$. Nothing else changed between this run and the previous one: same network, same starting weights, same points, same optimizer, same number of iterations. One constant in one line turned a failure into a solution, which is both the appeal of fixed weighting and its warning.
The same strategy with a smaller boundary weight,
$$ \mathcal{L}_{\mathrm{total}} = \mathcal{L}_{\mathrm{PDE}} + 10^{2}\,\mathcal{L}_{\mathrm{BC}} , $$
is a hundred times stronger than the naive formulation and still two orders of magnitude short of the previous run.
weight_bc = 1e2
net_fixed_bad = new_pinn()
optimizer = torch.optim.Adam(net_fixed_bad.parameters(), lr = lr)
for k in range(n_iter):
optimizer.zero_grad()
u = net_fixed_bad(x_pde)
u_x = derivative(u, x_pde)
u_xx = derivative(u_x, x_pde)
loss_pde = torch.mean((-u_xx - f_pde)**2)
loss_bc = torch.mean(net_fixed_bad(x_bc)**2)
loss = loss_pde + weight_bc*loss_bc
loss.backward()
optimizer.step()
print('Fixed, weight_bc = 1e2 | PDE %.3e | BC %.3e' % (loss_pde.item(), loss_bc.item()))
with torch.no_grad():
u_fixed_bad = net_fixed_bad(x_plot).cpu().numpy().ravel()
plt.figure(figsize = (7, 4))
plt.plot(x_np, u_true, 'k--', linewidth = 2, label = 'Ground Truth', alpha = 0.5)
plt.plot(x_np, u_fixed_bad, linewidth = 2, label = 'Prediction')
plt.title(r'Fixed Weighting, $\omega_{\mathrm{BC}} = 10^2$')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.ylim(-1.2, 1.2)
plt.legend()
plt.show()
The boundary loss is small, $2.9\times10^{-7}$, so the two ends are pinned at zero. Everything between them is wrong: the prediction shoots off the plotting window within a few hundredths of either end, the PDE residual is $1.7\times10^{4}$, and the maximum error is about $10$. A weight of $10^{2}$ was enough to enforce the boundary values, and not enough to keep the interior from the same large-amplitude function that the naive PINN found.
The run did not fail to converge. The optimizer converged, the losses settled, and the boundary condition is satisfied to seven decimal places; by every number available during training this looks like a finished run. It converged to the wrong function, and the only way to know is to compare with the exact solution, which on a real problem is not available. A badly chosen weight does not produce a visibly broken run; it produces a converged-looking run with a wrong answer. That is what makes the choice of weight a matter of care rather than of taste.
The two runs expose two separate difficulties, and it is worth keeping them apart, because they call for different remedies.
The first is that the right value is not known in advance. The weight $10^{4}$ was found by trying values, with the exact solution available to judge them, which is not a procedure available on a problem whose solution is unknown. A sweep over weights, with some independent check of the answer, can address this, at the cost of many training runs.
The second is that even a right value does not stay right. The relative scales of the two gradients change as the network moves through parameter space, so a pair of constants that balances them at one iteration over-corrects at another, and the competition reappears with the roles reversed. The figure below shows the sequence. No sweep can address this, because the quantity being tuned is not a constant of the problem at all. The remedy has to read the gradients as they are, not as they were, which is adaptive weighting.
The same weights at iteration $k+1$. The network has moved, the unweighted gradients have changed, and the fixed weights now make the boundary term dominant. The imbalance is back, in the other direction. Nothing short of re-reading the gradients can keep the balance.
7.1. The Optimizer Sees Gradients, Not Losses¶
The optimizer receives only the combined scalar $\mathcal{L}_{\mathrm{total}}$. It does not know that the first term is a statement about the interior and the second a statement about the boundary, or that a function satisfying one and not the other is not a solution at all. It sees one number and descends. Whatever went wrong in Section 5.2 went wrong inside that one number.
The parameters are updated by gradient descent on the total loss, with learning rate $\eta$,
$$ \theta^{(k+1)} = \theta^{(k)} - \eta\,\nabla_\theta\mathcal{L}_{\mathrm{total}} . $$
Because the objective is a sum, so is its gradient,
$$ \nabla_\theta\mathcal{L}_{\mathrm{total}} = \nabla_\theta\mathcal{L}_{\mathrm{PDE}} + \nabla_\theta\mathcal{L}_{\mathrm{BC}} . $$
The step is the sum of two vectors in parameter space, one per condition, and the optimizer follows the sum. What decides the direction of the step is therefore the gradient contribution of each term, the length and direction of its vector, and not the scalar value of the loss. A small loss can have a long gradient, and a large loss can be nearly flat in $\theta$.
Gradient descent. Left: on a function of one variable, each step is taken against the slope. Right: on a surface, each step follows the negative gradient of the total loss. The optimizer never asks which term that gradient came from.
Adaptive methods compute the coefficients from signals available during training and update them as training proceeds. The schemes differ in which signal they read.
| Method | Signal it reads | Core idea |
|---|---|---|
| Gradient-based | gradient norms $\lVert\nabla_\theta\mathcal{L}_i\rVert$ | balance the contribution of each term to the update |
| Loss-based | loss values and their rate of decrease | upweight terms that are large or decaying slowly |
| Trainable weights | the objective itself | treat the weights as parameters and optimize them |
| NTK-based | the neural tangent kernel | upweight the terms the network learns most slowly |
This notebook uses the first, because the imbalance in question is an imbalance of parameter gradients, and the gradient-based rule acts on exactly that quantity.
$$ \begin{align*} \mathcal{L}_{\mathrm{total}} &= \omega_{\mathrm{PDE}}\mathcal{L}_{\mathrm{PDE}} + \omega_{\mathrm{BC}}\,\mathcal{L}_{\mathrm{BC}} \\\\ \nabla_\theta\mathcal{L}_{\mathrm{total}} &= \omega_{\mathrm{PDE}} \nabla_\theta\mathcal{L}_{\mathrm{PDE}} + \omega_{\mathrm{BC}}\nabla_\theta\mathcal{L}_{\mathrm{BC}} \end{align*} $$
For each loss term, measure the length of its gradient,
$$ g_i = \bigl\lVert \nabla_\theta\mathcal{L}_i \bigr\rVert_2 , $$
and their mean, and set each weight to the ratio of the two,
$$ \bar g = \frac{g_{\mathrm{PDE}} + g_{\mathrm{BC}}}{2} , \qquad \omega_i = \frac{\bar g}{g_i + \varepsilon} , $$
with a small $\varepsilon > 0$ against division by zero. The term with the longer gradient receives the smaller weight and the term with the shorter gradient the larger one, so that every scaled gradient $\omega_i g_i$ has length $\bar g$. The weights are recomputed at every iteration from the current gradients, and so they follow the network as it moves. Both limitations of Section 6.3 are addressed at once: no value has to be guessed, and the value is never stale.
Gradient-based rebalancing. At iteration $k+1$ the weights are recomputed from the current gradient norms, $\omega_i = \bar g / g_i$, and the two scaled vectors are brought back to the same length. The balance of the previous figure is restored without any hand-tuned constant.
Adaptive weighting
The function grad_norm collects the gradient of one loss with respect to every parameter of the network and returns its length. Two of its arguments deserve a word. retain_graph=True keeps the computation graph alive, because the same graph is needed twice, once per loss, and then once more for loss.backward(). And allow_unused=True is needed because the bias of the output layer adds a constant to $u_\theta$, which $u_{xx}$ does not see, so its gradient with respect to $\mathcal{L}_{\mathrm{PDE}}$ does not exist and is returned as None. The weights are plain numbers, computed with .item(), so they scale the losses without being differentiated themselves. The loop records both losses and both weights at every iteration, so that the two conditions can be compared with the naive run of Section 5.2 and the weights with what was chosen by hand in Section 6.
def grad_norm(loss, net):
grads = torch.autograd.grad(loss, net.parameters(), retain_graph = True, allow_unused = True)
total = 0.0
for g in grads:
if g is not None: # the output bias does not enter u_xx
total = total + torch.sum(g**2)
return torch.sqrt(total).item()
net_adapt = new_pinn()
optimizer = torch.optim.Adam(net_adapt.parameters(), lr = lr)
loss_pde_hist = np.zeros(n_iter)
loss_bc_hist = np.zeros(n_iter)
weight_pde_hist = np.zeros(n_iter)
weight_bc_hist = np.zeros(n_iter)
for k in range(n_iter):
optimizer.zero_grad()
u = net_adapt(x_pde)
u_x = derivative(u, x_pde)
u_xx = derivative(u_x, x_pde)
loss_pde = torch.mean((-u_xx - f_pde)**2)
loss_bc = torch.mean(net_adapt(x_bc)**2)
g_pde = grad_norm(loss_pde, net_adapt)
g_bc = grad_norm(loss_bc, net_adapt)
g_mean = (g_pde + g_bc)/2
weight_pde = g_mean/(g_pde + 1e-12)
weight_bc = g_mean/(g_bc + 1e-12)
loss = weight_pde*loss_pde + weight_bc*loss_bc
loss.backward()
optimizer.step()
loss_pde_hist[k] = loss_pde.item()
loss_bc_hist[k] = loss_bc.item()
weight_pde_hist[k] = weight_pde
weight_bc_hist[k] = weight_bc
print('Adaptive | PDE %.3e | BC %.3e' % (loss_pde.item(), loss_bc.item()))
print('Final weights | PDE %.3e | BC %.3e' % (weight_pde, weight_bc))
with torch.no_grad():
u_adapt = net_adapt(x_plot).cpu().numpy().ravel()
plt.figure(figsize = (7, 10))
plt.subplot(3, 1, 1)
plt.plot(x_np, u_true, 'k--', linewidth = 3, label = 'Ground Truth', alpha = 0.5)
plt.plot(x_np, u_adapt, linewidth = 1.5, label = 'Prediction')
plt.title('Adaptive Weighting')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.ylim(-1.2, 1.2)
plt.legend()
plt.subplot(3, 1, 2)
plt.semilogy(iteration, loss_pde_hist, label = r'$\mathcal{L}_{\mathrm{PDE}}$')
plt.semilogy(iteration, loss_bc_hist, label = r'$\mathcal{L}_{\mathrm{BC}}$')
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.legend()
plt.subplot(3, 1, 3)
plt.semilogy(iteration, weight_pde_hist, label = r'$\omega_{\mathrm{PDE}}$')
plt.semilogy(iteration, weight_bc_hist, label = r'$\omega_{\mathrm{BC}}$')
plt.xlabel('Iteration')
plt.ylabel('Weight')
plt.legend()
plt.tight_layout()
plt.show()
The solution is recovered, with a maximum error of $9\times10^{-4}$, a PDE residual of $4.1\times10^{-2}$ and a boundary loss of $2.5\times10^{-12}$. No constant was chosen. The middle panel is the counterpart of the naive run of Section 5.2, and the difference is in the boundary loss: there it was driven up to $60$, here it is held between $10^{-4}$ and $10^{-12}$ from the first iterations, while the PDE loss descends in two stages, a long plateau near $10^{3}$ and then a drop to $10^{-1}$ after iteration $20000$. The boundary condition is never sacrificed for the residual, which is what balancing the two gradients means in practice. The bottom panel shows what the rule chose to achieve this: $\omega_{\mathrm{PDE}}$ sits at $0.5$ throughout, as it must when one gradient is far longer than the other, and $\omega_{\mathrm{BC}}$ rises within the first few hundred iterations from $10^{2}$ to $10^{4}$ and then fluctuates between $10^{4}$ and $10^{6}$ for the rest of the run, ending at $2.7\times10^{4}$.
This panel is the retrospective verdict on Section 6. The fixed weight $10^{4}$ sits at the lower edge of the band the rule chose, which is why it worked; the fixed weight $10^{2}$ is a hundred times below it, which is why it did not. And the band is two decades wide and never settles, which is the second limitation of Section 6.3 made visible: the right weight is not a number but a moving target, and a constant can only ever be right on average.
The rule of Section 7.2 trusts a single measurement. The gradient norms are read at one iteration and the weights are set from them, and the bottom panel of that figure shows the consequence: $\omega_{\mathrm{BC}}$ jumps by a factor of ten or a hundred from one iteration to the next, because the boundary gradient itself does. Each of those jumps is passed straight into the update. The remedy is the one used wherever a noisy measurement drives a control: do not replace the weight by the new measurement, move it part of the way there,
$$ \hat\omega_i^{(k)} = \frac{\bar g^{(k)}}{g_i^{(k)} + \varepsilon}, \qquad \omega_i^{(k)} = (1 - \alpha)\,\omega_i^{(k-1)} + \alpha\,\hat\omega_i^{(k)} , $$
with $0 < \alpha \le 1$. This is an exponential moving average of the instantaneous weights. At $\alpha = 1$ it is Section 7.2; at small $\alpha$ the weight remembers roughly the last $1/\alpha$ measurements and a single spike moves it by a fraction $\alpha$ of the spike. The run below uses $\alpha = 0.01$, a memory of about a hundred iterations, and starts both weights at $1$, so that the first hundred iterations are a ramp from the naive objective to the balanced one.
Moving-average adaptive weighting
The code is that of Section 7.2 with two lines changed, the ones that update the weights, and the two weights initialized before the loop. The instantaneous weight $\hat\omega_{\mathrm{BC}}$ is recorded alongside the averaged one so that the two can be drawn together.
alpha_ma = 0.01
net_ma = new_pinn()
optimizer = torch.optim.Adam(net_ma.parameters(), lr = lr)
weight_pde, weight_bc = 1.0, 1.0
loss_pde_hist = np.zeros(n_iter)
loss_bc_hist = np.zeros(n_iter)
weight_bc_raw = np.zeros(n_iter)
weight_bc_hist = np.zeros(n_iter)
for k in range(n_iter):
optimizer.zero_grad()
u = net_ma(x_pde)
u_x = derivative(u, x_pde)
u_xx = derivative(u_x, x_pde)
loss_pde = torch.mean((-u_xx - f_pde)**2)
loss_bc = torch.mean(net_ma(x_bc)**2)
g_pde = grad_norm(loss_pde, net_ma)
g_bc = grad_norm(loss_bc, net_ma)
g_mean = (g_pde + g_bc)/2
weight_pde_new = g_mean/(g_pde + 1e-12)
weight_bc_new = g_mean/(g_bc + 1e-12)
weight_pde = (1 - alpha_ma)*weight_pde + alpha_ma*weight_pde_new
weight_bc = (1 - alpha_ma)*weight_bc + alpha_ma*weight_bc_new
loss = weight_pde*loss_pde + weight_bc*loss_bc
loss.backward()
optimizer.step()
loss_pde_hist[k] = loss_pde.item()
loss_bc_hist[k] = loss_bc.item()
weight_bc_raw[k] = weight_bc_new
weight_bc_hist[k] = weight_bc
print('Moving average | PDE %.3e | BC %.3e' % (loss_pde.item(), loss_bc.item()))
print('Final weights | PDE %.3e | BC %.3e' % (weight_pde, weight_bc))
with torch.no_grad():
u_ma = net_ma(x_plot).cpu().numpy().ravel()
plt.figure(figsize = (7, 10))
plt.subplot(3, 1, 1)
plt.plot(x_np, u_true, 'k--', linewidth = 3, label = 'Ground Truth', alpha = 0.5)
plt.plot(x_np, u_ma, linewidth = 1.5, label = 'Prediction')
plt.title('Moving-Average Adaptive Weighting')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.ylim(-1.2, 1.2)
plt.legend()
plt.subplot(3, 1, 2)
plt.semilogy(iteration, loss_pde_hist, label = r'$\mathcal{L}_{\mathrm{PDE}}$')
plt.semilogy(iteration, loss_bc_hist, label = r'$\mathcal{L}_{\mathrm{BC}}$')
plt.xlabel('Iteration')
plt.ylabel('Loss')
plt.legend()
plt.subplot(3, 1, 3)
plt.semilogy(iteration, weight_bc_raw, linewidth = 0.8, alpha = 0.4, label = r'$\hat\omega_{\mathrm{BC}}$ (instantaneous)')
plt.semilogy(iteration, weight_bc_hist, 'r', linewidth = 2, label = r'$\omega_{\mathrm{BC}}$ (moving average)')
plt.xlabel('Iteration')
plt.ylabel('Weight')
plt.legend()
plt.tight_layout()
plt.show()
The solution is again recovered, with a maximum error of $1.9\times10^{-3}$, a PDE residual of $9.4\times10^{-2}$ and a boundary loss of $5.2\times10^{-12}$, all within a factor of two of Section 7.2. What has changed is visible in the two lower panels. The weight is now a smooth curve: it ramps up over the first few hundred iterations, settles near $10^{7}$, and descends over the run to $3\times10^{4}$, while the instantaneous values it averages scatter across two decades around it. The losses follow suit; the PDE loss descends in one smooth sweep instead of the plateau-and-drop of Section 7.2, and the boundary loss is held in a narrow band near $10^{-7}$ with none of the excursions seen there.
Two things are worth noticing. First, the averaged weight sits a decade above the band of Section 7.2 for most of the run, not in the middle of it. An arithmetic average of a quantity that spikes upward is pulled up by the spikes, and the gradient ratio is such a quantity; averaging $\log\hat\omega$ instead would give the geometric mean and would sit in the middle. Second, the smoothing cost nothing in accuracy here and bought a far more regular training curve, which is why most practical implementations of adaptive weighting include it, with an update every few iterations rather than every iteration as a further economy.
Adaptive weighting rebalances the competition without ending it. Every condition is still a penalty, satisfied only approximately, and the boundary values of the adaptive run are small, not zero. The moving average of Section 7.3 tames the swings of the measured weights but does not change this: a smoothed penalty is still a penalty. The method removes the need to guess a constant. It does not remove the boundary term from the objective, and for a condition that can be written in closed form there is a more direct option.
8.1. Building the Boundary Condition into the Output¶
A hard constraint attacks the competition from the other side. Instead of adjusting how strongly a condition is penalized, the output of the network is built so that the condition holds for every value of $\theta$, and the corresponding loss term disappears, together with its weight. For the homogeneous Dirichlet condition of this problem, write
$$ u_\theta^{\mathrm{hard}}(x) = 4x(1-x)\,N_\theta(x) , $$
where $N_\theta$ is the network. This is the masking-function construction of Embedding Physical Knowledge into Neural Networks, applied to the boundary condition of a PINN. The multiplier vanishes at both ends, so
$$ u_\theta^{\mathrm{hard}}(0) = u_\theta^{\mathrm{hard}}(1) = 0 $$
identically, whatever the network returns there. The boundary condition is removed from the gradient competition altogether, no weight has to be chosen for it, and training minimizes the PDE residual alone. The price is a change of hypothesis space: the network now represents $u / 4x(1-x)$ rather than $u$, and the multiplier has to be designed so that it enforces the condition without making the interior harder to represent. When the condition has a simple closed form, as here, the construction is almost always preferable to a penalty, for the same reason that eliminating a variable is preferable to constraining it.
Hard constraint. The network output $N_\theta(x)$ is multiplied by a function that vanishes at both ends, so $\hat u(0) = \hat u(1) = 0$ holds by construction and the boundary term drops out of the loss. The factor $4$ used in the code only rescales the network output and changes nothing about the constraint.
Hard-constrained PINN
The same network, points, optimizer, learning rate, iterations and starting weights. The only changes are the line that defines $u$ inside the loop and the loss, which is now the PDE term alone. The second derivative is taken of the product, not of $N_\theta$, so the multiplier enters the residual through the chain rule; automatic differentiation takes care of that. The boundary loss is evaluated once at the end, to confirm that it is zero.
net_hard = new_pinn()
optimizer = torch.optim.Adam(net_hard.parameters(), lr = lr)
for k in range(n_iter):
optimizer.zero_grad()
u = 4*x_pde*(1 - x_pde)*net_hard(x_pde) # zero at both ends, whatever the network says
u_x = derivative(u, x_pde)
u_xx = derivative(u_x, x_pde)
loss_pde = torch.mean((-u_xx - f_pde)**2) # no BC term
loss_pde.backward()
optimizer.step()
with torch.no_grad():
u_hard = (4*x_plot*(1 - x_plot)*net_hard(x_plot)).cpu().numpy().ravel()
loss_bc = torch.mean((4*x_bc*(1 - x_bc)*net_hard(x_bc))**2)
print('Hard constraint | PDE %.3e | BC %.3e' % (loss_pde.item(), loss_bc.item()))
error = np.abs(u_hard - u_true)
plt.figure(figsize = (7, 7))
plt.subplot(2, 1, 1)
plt.plot(x_np, u_true, 'k--', linewidth = 3, label = 'Ground Truth', alpha = 0.5)
plt.plot(x_np, u_hard, linewidth = 1.5, label = 'Prediction')
plt.title('Hard Constraint')
plt.xlabel('$x$')
plt.ylabel('$u(x)$')
plt.ylim(-1.2, 1.2)
plt.legend()
plt.subplot(2, 1, 2)
plt.plot(x_np, error, linewidth = 2)
plt.xlabel('$x$')
plt.ylabel('Absolute Error')
plt.tight_layout()
plt.show()
The hard-constrained model is the most accurate of the four. The PDE residual is $6.8\times10^{-3}$, the smallest in the notebook, and the boundary loss is exactly zero, not because it was minimized but because it was never possible for it to be anything else. The error panel shows the signature of the construction: the error is about $2\times10^{-4}$ across the interior and drops to exactly zero at $x = 0$ and $x = 1$, where the multiplier vanishes. There was no boundary term to weight and nothing to trade the interior against, so every gradient step went into the equation.
The five steps close on a hierarchy that is worth stating plainly. Impose exactly what can be imposed exactly, as the hard constraint does for the boundary condition. Penalize the rest, because a penalty can be written for a requirement that admits no parameterization, an inequality, an integral conservation law, a condition on an unknown interface. And weight the penalties by what their gradients say, not by what their loss values say, letting the gradients choose the weight when the right value is unknown or changes during training. A naive sum is the special case in which none of this has been thought about, and Section 5.2 is what it costs.
Key takeaways
- A penalty term is prior knowledge written into the objective. An overfitted model is one whose coefficients are large and nearly cancelling, so the magnitude of the coefficients is itself worth a cost; ridge regression adds $\lambda\lVert\theta\rVert_2^2$ to the residual sum of squares and trades a little bias for a large reduction in variance.
- Regularization is MAP estimation. Maximum likelihood under Gaussian noise is least squares; a Gaussian prior on the coefficients turns the maximum a posteriori estimate into ridge regression, with $\lambda$ the ratio of the noise variance to the prior variance. The penalty is not a trick but a prior, and the coefficient is a statement of relative trust.
- The physics loss is a soft constraint, not a regularizer in the narrow sense. It carries problem-specific knowledge, a governing equation rather than a preference for small parameters, into the objective by penalty, and a PINN is a MAP estimate whose prior is a physical law.
- A naive PINN can fail while its loss falls. With unit weights the PDE term of the $\alpha = 7$ Poisson problem dominated the update so completely that the boundary condition was never enforced; after $30000$ iterations the residual was $3.4\times10^{4}$, the boundary loss $58$, and the prediction a different function altogether.
- A fixed weight can work, and can fail silently. $\omega_{\mathrm{BC}} = 10^{4}$ recovered the solution to $4\times10^{-4}$, while $\omega_{\mathrm{BC}} = 10^{2}$ converged, with a small boundary loss, to a wrong interior. The right value is not known in advance, and even a right value becomes wrong as the gradients change during training.
- The reason is in the gradients, not the losses. The optimizer follows the sum of the term gradients, and the two differ in length for reasons unrelated to the coefficients. Adaptive weighting reads the gradient norms at every iteration and rescales the terms so that their contributions are equal; it found a boundary weight of order $10^{4}$ to $10^{6}$ by itself and recovered the solution to $9\times10^{-4}$. Smoothing the weights with a moving average removes their iteration-to-iteration swings at no cost in accuracy. The conditions remain penalties.
- A masking function removes the boundary condition from the competition by building it into the output, $u = 4x(1-x)N_\theta(x)$. There is no boundary loss to weight, and the PDE residual alone reached $6.8\times10^{-3}$, the smallest of the four runs.
Looking ahead. Weighting and masking deal with conditions that compete for the parameters. The remaining limitations of a PINN are of a different kind: the collocation points can miss the part of the domain that matters, the solution can have a shape that a smooth network does not easily represent, and a transient problem can be trained without regard to the direction of time. PINN: Limitations and Extensions takes these up.