Inverse Dirichlet weighting enables reliable training of physics informed neural networks

Like most deep neural networks, PINNs are usually trained by first-order parameter update based on (approximate) loss gradients. The update rule can be interpreted as a discretization of the gradient flow, i.e.

$$\begin{align} \boldsymbol{\theta} (\tau + 1) & = \boldsymbol{\theta} (\tau) - \eta (\tau) \sum_{k = 1}^{K}\lambda_k \nabla_{\boldsymbol{\theta}} \mathcal{L}_{k}. \end{align} \tag{ 4 }$$

The learning rate $\eta(\tau)$ depends on the training epoch τ, reflecting the adaptive step size used in optimizers like Adam [38]. The gradients are batch gradients $\nabla_{\boldsymbol{\theta}} \mathcal{L}_k = \frac{1}{\vert \textrm{B} \vert} \sum_{i \in \textrm{B}} \nabla \ell_k (\boldsymbol{u}_{\boldsymbol{\theta}}^{i}, \boldsymbol{u}_i)$, where $\textrm{B}$ is a mini-batch, created by randomly splitting the N training data points into $N/\vert \textrm{B}\vert$ distinct partitions, and $\ell_k$ is one summand (for one data point i) of the loss $\mathcal{L}_k$ corresponding to the kth objective. For batch size $\vert \textrm{B}\vert = 1$, mini-batch gradient descent reduces to stochastic gradient descent algorithm.

For approximating a function $\boldsymbol{u}(\boldsymbol{x})$ from data u _i , adding derivatives $\mathcal{D}_{\boldsymbol{x}}^m$, $m\geqslant 1$, of the state variable into the training loss has been shown to improve data efficiency and generalization power of neural networks [7]. Minimizing the resulting loss

$$\begin{equation} \sum_{i = 1}^{N} \Bigg[ \lambda_0~\vert \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i) - \boldsymbol{u}(\boldsymbol{x}_i) \vert^2 + \sum_{k = 1}^{K} \lambda_k \vert \mathcal{D}_{\boldsymbol{x}}^k \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i) - \mathcal{D}_{\boldsymbol{x}}^k \boldsymbol{u}(\boldsymbol{x}_i) \vert^2~\Bigg ] \end{equation} \tag{ 5 }$$

using the update rule in equation (4) can be interpreted as a finite approximation to the Sobolev gradient flow of the loss functional. Without loss of generality, we drop the dependence of the state variable on time, writing $\boldsymbol{u}_i = \boldsymbol{u}(\boldsymbol{x}_i)$. The neural network is then trained with access to the (K + 2)-tuples $\{(\boldsymbol{x}_i, \boldsymbol{u}(\boldsymbol{x}_i), \mathcal{D}_{\boldsymbol{x}}^1~\boldsymbol{u}(\boldsymbol{x}_i),\ldots,\mathcal{D}_{\boldsymbol{x}}^K \boldsymbol{u}(\boldsymbol{x}_i)) \}_{i = 1}^N$ instead of the usual training set $\{(\boldsymbol{x}_i, \boldsymbol{u}(\boldsymbol{x}_i))\}_{i = 1}^N$. Such Sobolev training is an instance of scalarized multi-objective optimization with each objective based on the approximation of a specific derivative. The weights are usually chosen as $\lambda_0 = \ldots = \lambda_\textrm{K} = 1$ [7].

The additional information from the derivatives introduces an inductive bias, which has been found to improve data efficiency and network generalization [7]. This is akin to PINNs, which use derivatives from the physical equation model to achieve inductive bias. Indeed, the Sobolev loss in equation (5) is a special case of the PINN loss from equation (3) with each objective based on a PDE with right-hand side $\mathcal{F} = \mathcal{D}_{\boldsymbol{x}}^k \boldsymbol{u}$, $k = 1,\ldots ,K$.

The update rule in equation (4) is known to lead to stiff learning dynamics for loss functionals whose Hessian matrix has large positive eigenvalues [34]. This is the case for Sobolev training, even with K = 1, when the function $\boldsymbol{u}(\boldsymbol{x})$ is highly oscillatory, leading us to the following proposition, proven in appendix B.1: Stiff learning dynamics in first-order parameter update

Proposition 2.2.1. For Sobolev training on the tuples $\{\boldsymbol{x}_i, \boldsymbol{u}(\boldsymbol{x}_i), \partial_{\boldsymbol{x}}^m \boldsymbol{u}(\boldsymbol{x}_i) \}_{i = 1}^{N}, m \geqslant 1$, with total loss $\mathcal{L} = \sum_{i = 1}^{N} \left( {\lambda_0} \vert \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_i) - \boldsymbol{u}(\boldsymbol{x}_i) \vert^2~\right. + \left. {\lambda_1} \vert \mathcal{D}_{\boldsymbol{x}}^m \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i) - \mathcal{D}_{\boldsymbol{x}}^m \boldsymbol{u}(\boldsymbol{x}_i)\vert^2\right)$, the rate of change of the residual of the dominant Fourier mode $\widetilde{\boldsymbol{r}}(\boldsymbol{k}_0) = \widetilde{\boldsymbol{u}}_{\boldsymbol{\theta}} (\boldsymbol{k}_0) - \widetilde{\boldsymbol{u}} (\boldsymbol{k}_0)$ using first-order parameter update is:

$$\begin{align*} \left\vert \frac{\mathrm{d}\widetilde{\boldsymbol{r}}({\boldsymbol{k}_0})}{\mathrm{d}\tau} \right\vert & = \frac{\eta({\tau})}{O(\boldsymbol{k}_0) } \left[ {\lambda_{0}} O(1) + {\lambda_{1}} O(\boldsymbol{k}_0^{2m}) \right] \left\vert \frac{\partial \widetilde{\boldsymbol{r}}({\boldsymbol{k}_0})}{\partial \boldsymbol{\theta}} \right\vert , \end{align*}$$

where $O(\cdot)$ is the Bachmann-Landau big-O symbol. For ${\lambda_0} = {\lambda_1} = 1$, this leads to training dynamics with slow time scale $\mathrm{d}/\mathrm{d}t \in O(\boldsymbol{k}_0^{-1})$ and fast time scale $\mathrm{d}/\mathrm{d}t \in O(\boldsymbol{k}_0^{2m-1})$ for $\widetilde{\boldsymbol{r}}({\boldsymbol{k}_0}) \in O(1)$.

This proposition suggests that losses containing high orders of derivatives m, dominant high frequencies k ₀, or a combination of both exhibit dominant gradient statistics. This can lead to biased optimization of the dominant objectives with large gradients at the expense of the other objectives with relatively smaller gradients. This is akin to the well-known vanishing gradients phenomenon across the layers of a neural network [39], albeit now across different objectives of a multi-objective optimization problem. We therefore call this phenomenon vanishing task-specific gradients.

Proposition 2.2.1 also admits a qualitative comparison with the phenomenon of numerical stiffness in PDEs of the form

$$\begin{equation} \frac{\partial {\widetilde{\boldsymbol{u}}}}{\partial t} = \mathbf{L} \widetilde{\boldsymbol{u}} + \mathbf{F}(\widetilde{\boldsymbol{u}}, t) . \end{equation} \tag{ 6 }$$

If the real parts of the eigenvalues of L, $c = \mathrm{Re}(\lambda(\mathbf{L})) \gg 1$, then the solution of the above PDE has fast modes $\widetilde{\boldsymbol{u}} \in O(1)$ with $\mathrm{d}/\mathrm{d}t \in O(c)$ and slow modes $\widetilde{\boldsymbol{u}} \in O(1/c)$ with $\mathrm{d}/\mathrm{d}t \in O(1)$. Therefore, high frequency modes evolve on shorter time scales than low frequency modes. The rate amplitude for spatial derivatives of order m is $\mathrm{d}/\mathrm{d}t \in O(\boldsymbol{k}^{-m})$ for the wavenumber k [40]. Due to this analogy, we call the learning dynamics of a neural network with vanishing task-specific gradients stiff. Stiff learning dynamics leads to discrepancy in the convergence rates of different objectives in a loss function.

Taken together, vanishing task-specific gradients constitute a likely failure mode of PINNs, since their training amounts to a generalized version of Sobolev training.

While ad hoc manual tuning is still commonplace, it has recently been proposed to determine the weights as moving averages over the inverse gradient magnitude statistics of a given objective as [34]:

$$\begin{align} \hat{\lambda}_k (\tau) & = \frac{\max \{\vert \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_\textrm{1} (\tau) \vert \} }{\lambda_k (\tau) \overline{ \vert \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\tau) \vert}},\nonumber\\ \lambda_k (\tau + 1) & = \alpha \lambda_k(\tau) + (1-\alpha) \hat{\lambda}_k(\tau), \end{align} \tag{ 7 }$$

where $\nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_\textrm{1}$ is the gradient of the residual, i.e. of the first objective in equation (3), and α is a user-defined learning rate (here always α = 0.5). All gradients are taken with respect to the shared parameters across tasks, and $| \cdot |$ is the component-wise absolute value. The overbar signifies the algebraic mean over the vector components. The maximum in the numerator is over all components of the vector. While this max/avg weighting has been shown to improve PINN performance [34], it does not completely alleviate the problem of vanishing task-specific gradients.

Instead of determining the loss weights proportional to the inverse average gradient magnitude, we here propose to use weights based on the gradient variance. In particular, we propose to use weights for which the variances over the components of the back-propagated weighted gradients $\lambda_k \nabla_{\boldsymbol{\theta}} \mathcal{L}_k$ become equal across all objectives, thus directly preventing vanishing task-specific gradients. This can be achieved in any (stochastic) gradient-descent optimizer by using the update rule in equation (4) with weights

$$\begin{align} \hat{\lambda}_k (\tau) & = {\frac{\gamma(\tau)}{\textrm{std}\{\nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\tau) \}}}, \nonumber\\ \lambda_k (\tau + 1) & = \alpha \lambda_k(\tau) + (1-\alpha) \hat{\lambda}_k(\tau), \end{align} \tag{ 8 }$$

where $\gamma(\tau) = \max_{t = 1,\ldots ,K} \left( \textrm{std}\{\nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_t (\tau) \} \right)$, and $\textrm{std}\{\cdot\}$ is the empirical (sample) standard deviation over the vector components. Under the assumption that the back-propagated gradients are normally distributed, the weighting strategy in equation (8) leads to balanced gradient distributions with variance $\gamma(\tau)^2$, i.e.

$$\begin{equation*} \frac{\gamma(\tau)}{\textrm{std}\{\nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\tau) \}} \mathcal{N} \left( \mu_k, {\textrm{Var}}\{ \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\tau) \} \right) = \mathcal{N} \left( \mu_k, \gamma(\tau)^2~\right), \end{equation*}$$

$\mathcal{N}\left( \cdot \right)$ is the normal distribution with mean $\mu_k, k = {1,\ldots ,K}$, and given variance.

For optimizers that are invariant to diagonal scaling of the gradients, such as the popular Adam optimizer [38], the weights in equation (8) can be efficiently computed as the inverse of the square root of the Dirichlet energy of each objective, i.e.

$$\begin{equation*} \textrm{std}\{\nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\tau) \} \propto \sqrt{\int_{\Omega_{\boldsymbol{\theta}}} \vert \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\tau) \vert^2 \,\mathrm{d}\boldsymbol{\theta} \, .} \end{equation*}$$

We therefore refer to this weighting strategy as inverse Dirichlet weighting.

The weighting strategies in equations (7) and (8) are fundamentally different. Equation (7) weights objectives in inverse proportion to certainty as quantified by the average gradient magnitude. Equation (8) uses weights that are inversely proportional to uncertainty as quantified by the variance of the loss gradients. Equation (8) is also different from previous uncertainty-weighted approaches [31] because it measures the training uncertainty stemming from variance in the loss gradients rather than the uncertainty from the model's observational noise.

We compare weighting-based scalarization approaches with gradient-based multi-objective approaches using Karush-Kuhn-Tucker (KKT) local optimality conditions to find a descent direction that decreases all objectives. Specifically, we adapt the multiple gradient descent algorithm (MGDA) [33] to PINNs. This algorithm leverages the KKT conditions

$$\begin{align} \exists \boldsymbol{\lambda} & = \{\lambda_k \} \in \mathbb{R}_+^K \nonumber \\ \textrm{s.t. } \sum_{k = 1}^{K} \lambda_k \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k (\boldsymbol{\theta}^\textrm{sh}, \boldsymbol{\theta}^k) &= 0~\textrm{~~and } \: \sum_{k = 1}^{K} \lambda_k = 1 \end{align} \tag{ 9 }$$

to solve an MTL problem. The MGDA is guaranteed to converge to a Pareto-stationary point [41]. Given the gradients $\nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_k$ of all objectives $k = 1,\ldots ,K$ with respect to the shared parameters, we find weights λ_k that satisfy the above KKT conditions. The resulting λ leads to a descent direction that improves all objectives [33], see also figure A.1. We adapt the MGDA to PINNs as described in appendix A.2. We refer to this adapted algorithm as pinn-MGDA.

Zoom In Zoom Out Reset image size

We characterize the phenomenon of vanishing task-specific gradients, demonstrate its impact on the accuracy of a learned function approximation, and empirically validate Proposition 2.2.1. For this, we consider a generic neural network with five layers and 64 neurons per layer, tasked with learning a 2D function that contains a wide spectrum of frequencies and unknown coefficients $\boldsymbol{\hat{\xi}} = (\hat{\xi}_1,\hat{\xi}_2,\hat{\xi}_3,\hat{\xi}_4)$ using Sobolev training with the loss from equation (10). The details of the test problem and the training setup are given in appendix B.

We first confirm that in this test case vanishing task-specific gradients occur when using uniform weights, i.e. when setting $\lambda_k = 1$ for all $k = 0,\ldots , 4$. For this, we plot the gradient histogram of the five loss terms in figure 1(A). The histograms clearly show that training is biased in this case to mainly optimize the objective containing the highest derivative k = 4. All other objectives are neglected, as predicted by proposition 2.2.1. This leads to a uniformly high approximation error on the test data (training vs. test data split 50:50, see appendix B), as shown in figure 2(B), as well as a uniformly high error in the estimated coefficients ξ (figure 2(C)). When determining the weights using the Pareto-seeking pinn-MGDA, the accuracy considerably improves (figures 2(B) and (C)). However, the phenomenon of vanishing task-specific gradients is still present, as seen in figure 1(B), this time favoring the objective for k = 0 and neglecting the derivatives. Max/avg weighting as proposed in reference [34] and given in equation (7) improves the accuracy for higher derivatives (figure 2(B)), but still suffers from unbalanced gradient distributions (figure 1(C)), which leads to a uniformly high error in the estimated coefficients ξ (figure 2(C)). The inverse Dirichlet weighting proposed here and detailed in equation (8) shows balanced gradient histograms (figure 1(D)) closest to those observed when using ε-optimal analytical weights (figure 1(E)). This leads to orders of magnitude better accuracy in both the function approximation and the coefficient estimates, as shown in figures 2(B) and (C), reaching the performance achieved when using ε-optimal weights.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

This is confirmed when looking at the power spectrum of the function approximation $\boldsymbol{u}_{\boldsymbol{\theta}}$ learned by the neural network in figure 2(D). Inverse Dirichlet weighting leads to results that are as good as those achieved by optimal weights in this test case, in particular when approximating high frequencies, as predicted by proposition 2.2.1. A closer look at the power spectra learned by the different weighting strategies also reveals that the max/avg strategy fails to capture high frequencies, whereas the pinn-MDGA performs surprisingly well (figure 2(D)) despite its unbalanced gradients (figure 1(B)). This may be problem-specific, as the weight trajectories taken by pinn-MDGA during training (figure B.1) are fundamentally different from the behavior of the other methods. This could be because the pinn-MGDA uses a fundamentally different strategy, purely based on satisfying Pareto-stationarity conditions (equation (9)). As shown in figure B.2, this leads to a rapid attenuation of the Fourier spectrum of the residual $\vert \tilde{\boldsymbol{r}}(\boldsymbol{k})\vert$ over the first few learning epochs in this test case, allowing pinn-MGDA to escape the stiffness defined in proposition 2.2.1.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

In summary, we observe stark differences between different weighting strategies. Of the strategies compared, only the inverse Dirichlet weighting proposed here is able to avoid vanishing task-specific gradients, performing on par with the ε-optimal solution. The pinn-MGDA seems to be able to escape, rather than avoid, the stiffness caused by vanishing task-specific gradients. We also observe a clear correlation between having balanced task-specific gradient distributions and achieving good approximation and estimation accuracy.

We investigate the use of PINNs in modeling multi-scale dynamics in space and time where vanishing task-specific gradients are a common problem. As a real-world test problem, we consider meso-scale active turbulence as it occurs in living fluids, such as bacterial suspensions [43, 44] and multi-cellular tissues [45]. Unlike inertial turbulence, active turbulence has energy injected at small length scales comparable to the size of the actively moving entities, like bacteria or individual motor proteins. The spatio-temporal dynamics of active turbulence can be modeled using the generalized incompressible Navier–Stokes equation [43, 44]:

$$\begin{align} \frac{\partial \boldsymbol{u}}{\partial t} + \xi_0 (\boldsymbol{u} \cdot \nabla \boldsymbol{u}) & = - \nabla p + \xi_1~\nabla \vert \boldsymbol{u}\vert^2 - \alpha \boldsymbol{u} - \beta \vert \boldsymbol{u}\vert^2~\boldsymbol{u} \nonumber \\ &\qquad +\Gamma_0~\Delta \boldsymbol{u} - \Gamma_2~\Delta^2~\boldsymbol{u}\\ \nabla \cdot \boldsymbol{u} & = 0~,\nonumber \end{align} \tag{ 11 }$$

where $\boldsymbol{u}(t,\boldsymbol{x})$ is the mean-field flow velocity and $p(t,\boldsymbol{x})$ the pressure field. The coefficient ξ₀ scales the contribution from advection and ξ₁ the active pressure. The α and β terms correspond to a quartic Landau-type velocity potential. The higher-order terms with coefficients Γ₀ and Γ₂ model passive and active stresses arising from hydrodynamic and steric interactions, respectively [44]. Figure C.2 illustrates the rich multi-scale dynamics observed when numerically solving this equation on a square for the parameter regime $\Gamma_0 \lt 0$ and $\Gamma_2 \gt 0$ [43] (see appendix C for details on the numerical methods used). An important characteristic of equation (11) is the presence of disparate length and time scales in addition to fourth-order derivatives and nonlinearities in u , which jointly lead to considerable numerical stiffness [40]. Given such large dominant wavenumbers k ₀, any data-driven approach, forward or inverse, is expected to encounter significant learning pathologies as per proposition 2.2.1.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

We demonstrate these pathologies by applying PINNs to approximate the forward solution of equation (11) in 2D for the parameter regime explored in reference [43] in both square (figure C.2) and annular (figure 3(E)) domains. The details of the simulation and training setups are given in appendix C. The loss function includes the terms as given in equation (3), including the PDE residual, initial and boundary conditions, and the divergence-freeness constraint.

Zoom In Zoom Out Reset image size

Figures 3(A) and (B) show the average (over different initializations of the neural network weights) relative L2 errors of the vorticity field over training epochs when using different PINN weighting strategies. As predicted by proposition 2.2.1, the uniformly weighted PINN fails completely both in the square (figure 3(A)) and in the annular geometry (figure 3(B)). Optimization is entirely biased towards minimizing the PDE residual, neglecting the initial, boundary, and incompressibility conditions. This imbalance is also clearly visible in the loss gradient distributions shown in figure C.5(A). Comparing the approximation accuracy after 7000 training epochs confirms the findings from the previous section with inverse Dirichlet weighting outperforming the other methods, followed by pinn-MGDA and max/avg weighting. Inverse Dirichlet weighting is found to assign adequate weights for the initial and boundary conditions relative to the equation residual (see figures C.3(A) and (D)).

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

In figure 3(C), we compare the convergence of the learned approximation $\boldsymbol{u}_{\boldsymbol{\theta}}$ to a ground-truth numerical solution for different spatial resolutions h of the discretization grid. The errors scale as $O(h^r)$ with convergence orders r as identified by the dashed lines. While the pinn-MGDA and max/avg weighting achieve linear convergence (r = 1), the inverse Dirichlet weighted PINN converges with order r = 4. The uniformly weighted PINN does not converge at all and is therefore not shown in the plot. The evolution of the Fourier spectra $| \widetilde{r}(\boldsymbol{k}) |$ of the residuals of both velocity components $(u,v) = \boldsymbol{u}$ shown in figure C.4 confirm the rapid solution convergence achieved by inverse Dirichlet weighted PINNs.

Inverse Dirichlet weighting also perfectly recovers the power spectrum E(k) of the active turbulence velocity field, as shown in figure 3(D). This confirms that inverse Dirichlet weighting enables PINNs to capture high frequencies in the PDE solution, corresponding to fine structures in the fields, by avoiding the stiffness problem from proposition 2.2.1. Indeed, the PINNs trained with the other weighting methods fail to predict small structures in the flow field, as shown by the spatial distributions of the prediction errors in figures 3(F)–(I) compared to the ground-truth numerical solution in figure 3(E).

In Appendices D and E, we provide results for two additional examples: approximating the solution of the 2D Poisson equation (see figure D.1) on a square with a large spectrum of frequencies, and learning the spatio-temporal dynamics of curvature-driven level-set flow (see figures E.1 and E.2) developing small structures. In both of these additional multi-scale examples we notice similar behavior, with inverse Dirichlet weighting outperforming the other approaches either in terms of accuracy or convergence rate.

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

Zoom In Zoom Out Reset image size

In summary, these results suggest that avoiding vanishing task-specific gradients can enable the use of PINNs in multi-scale problems previously not amenable to neural-network modeling. Moreover, we observe that seemingly small changes in the weighting strategy can lead to significant differences in the convergence order of the learned approximation.

Artificial neural networks tend to 'forget' their parameterization when sequentially training for multiple tasks [46], a phenomenon referred to as catastrophic forgetting. Sequential training is often used, e.g. when learning to de-noise data and estimate parameters at the same time, like combined image segmentation and denoising [47], or when learning solutions of fluid mechanics models with additional constraints like divergence-freeness of the flow velocity field [16].

In PINNs, catastrophic forgetting can occur, e.g. in the frequently practiced procedure of first training on measurement data alone and adding the PDE model only at later epochs [48]. Such sequential training can be motivated by the computational cost of evaluating the derivatives in the PDE using automatic differentiation over the network, so it seems prudent to pre-condition the network for data fitting and introduce the residual later in order to regress for the missing network parameters.

Using the active turbulence test case described above, we investigate how the different weighting schemes compared here influence the phenomenon of catastrophic forgetting. For this, we consider the inverse problem of inferring unknown model parameters and latent fields like the effective pressure $p^*$ from flow velocity data u _i . We train the PINNs in three sequential steps: (1) first fitting flow velocity data, (2) later including the additional objective for the incompressibility constraint $\nabla \cdot \boldsymbol{u} = 0$, (3) finally including also the objective for the PDE model residual. This creates a setup in which catastrophic forgetting occurs when using uniform PINN weights, as confirmed in figure 4(A) (arrows). As soon as additional objectives (divergence-freeness at epoch 1000 and equation residual at epoch 2000) are introduced, the PINN loses the parameterization learned on the previous objectives, leading to large and increasing test error for the field u .

Zoom In Zoom Out Reset image size

Dynamic weighting strategies can protect PINNs against catastrophic forgetting by adequately adjusting the weights whenever an additional objective is introduced. This is confirmed in figure 4(A) with inverse Dirichlet and pinn-MGDA only minimally affected by catastrophic forgetting. The max/avg weighting strategy [34], however, does not protect the network quite as well. When looking at the test errors for the learned model coefficients (figure 4(B)) and the accuracy of the estimated latent pressure field $p(t,\boldsymbol{x})$ (figure 4(C)), for which no training data is given, the inverse Dirichlet weighting proposed here outperforms the other approaches. Snapshots of the reconstructed pressure fields using different strategies are shown in figure C.7. Results for inference from noisy data are shown in figures C.6 and C.8.

In summary, inverse Dirichlet weighting can protect a PINN against catastrophic forgetting, making sequential training a viable option. In our tests, it offered the best protection of all compared weighting schemes.

For Sobolev loss functions based on elliptic PDEs, optimal weights λ_k that minimize all objectives in an unbiased way can be determined by ε-closeness analysis [42]. ε -closeness

Definition 1. A candidate solution u is called ε-close to the true solution $\hat{\boldsymbol{u}}$ if it satisfies

$$\begin{equation} \frac{\vert \mathcal{D}^{\boldsymbol{\beta}} \boldsymbol{u} ({\boldsymbol{x}}) -\mathcal{D}^{\boldsymbol{\beta}} \hat{\boldsymbol{u}} ({\boldsymbol{x}})\vert}{\vert \mathcal{D}^{\boldsymbol{\beta}} \hat{\boldsymbol{u}}({\boldsymbol{x}}) \vert} \leqslant \varepsilon \end{equation} \tag{ A1 }$$

for any multi-index $\boldsymbol{\beta} \in \mathbb{N}^{{d}}$ and $\boldsymbol{x} = (x_1,\ldots ,x_d) \in \Omega \subset \mathbb{R}^d$, thus $\mathcal{D}^{\boldsymbol{\beta}} \boldsymbol{u} (\boldsymbol{x}) = \frac{\partial^{\vert \boldsymbol{\beta} \vert} \boldsymbol{u}(\boldsymbol{x})}{\partial x_1^{\beta_1}\partial x_2^{\beta_2}\ldots \partial x_d^{\beta_d}}$ with $\vert \boldsymbol{\beta}\vert = \sum_{i = 1}^{d} \beta_i$.

For the loss function typically used in Sobolev training of neural networks, i.e.

$$\begin{equation*} \mathcal{L}_k( \cdot ) = \int_{\Omega} \left( \mathcal{D}_{{\boldsymbol{x}}}^{k} \hat{\boldsymbol{u}} - \mathcal{D}_{{\boldsymbol{x}}}^{k} \boldsymbol{u}\right)^2~\mathrm{d} \boldsymbol{x}. \end{equation*}$$

We can use ε-closeness analysis to show that

$$\begin{equation} \mathcal{L}_{k} \leqslant \varepsilon^2~\int_{\Omega} \left( \vert \xi_k \mathcal{D}_{\boldsymbol{x}}^k (\boldsymbol{\hat{u}}) \vert \right)^2~\mathrm{d} \boldsymbol{x} \, . \end{equation} \tag{ A2 }$$

Using this inequality, we can bound the total loss as

$$\begin{equation} \mathcal{L}(\boldsymbol{u}) \leqslant \varepsilon^2~\sum_{k = 1}^{K} \lambda_k \int_{\Omega} \left( \vert \xi_k \mathcal{D}_{\boldsymbol{x}}^k (\boldsymbol{\hat{u}}) \vert \right)^2~\mathrm{d} \boldsymbol{x} = \varepsilon^2~\sum_{k = 1}^{K} \lambda_k \mathcal{I}_k, \end{equation} \tag{ A3 }$$

with $\mathcal{I}_k = \int_{\Omega} \left( \vert \xi_k \mathcal{D}_{\boldsymbol{x}}^k (\boldsymbol{\hat{u}}) \vert \right)^2\! \mathrm{d} {\boldsymbol{x}}$. This bound can be used to determine the weights $\{\lambda _k \} _{k = 1}^K$ that result in the smallest ε:

$$\begin{align} \mathcal{L}_k(\boldsymbol{u}) & \leqslant \frac{1}{\lambda_k} \mathcal{L}(\boldsymbol{u}) \leqslant \varepsilon^2~\lambda_k \mathcal{I}_k \:\: \forall k = {1,\ldots ,K}, \nonumber\\ \mathcal{L}(\boldsymbol{u}) & \leqslant \varepsilon^2~\min \{\lambda_1~\mathcal{I}_1,\ldots ,\lambda_K \mathcal{I}_K \}. \end{align} \tag{ A4 }$$

The weights $\lambda_k, \forall k \in {1,\ldots ,K}$, that lead to the tightest upper bound on the total loss in equation (A4) can be found by solving the maximization-minimization problem described in following corollary: A.1.1.

Corollary A maximization-minimization problem can be transformed into a piecewise linear convex optimization problem by introducing an auxiliary variable z, such that

$$\begin{align*} \textrm{maximize} \:\:z \qquad\quad & \quad \\ \textrm{subject~to } \mathbf{1}^\top \boldsymbol{\lambda}^{*} = 1; \\ z \leqslant \lambda_k \mathcal{I}_k\,, \quad k = {1,\ldots ,K} . \end{align*}$$

The analytical solution to this problem exists and is given by $\lambda_k^* = \frac{\Pi_{j\neq k} \mathcal{I}_j}{\sum_{k = 1}^{K} \Pi_{j\neq k} \mathcal{I}_j}, \quad k = 1,\ldots ,K$.

For a Sobolev training problem with loss given by equation (5), weights that lead to ε-optimal solutions are computed using

$$\begin{equation*} \mathcal{I}_k = \int_{\Omega} \left(\vert \xi_k \mathcal{D}_{\boldsymbol{x}}^k (\boldsymbol{\hat{u}}) \vert \right)^2\! \mathrm{d} \boldsymbol{x} .\end{equation*}$$

We adapt the MGDA to PINNs as illustrated in figure A.1 by formulating the conditions for Pareto stationary as stated in equation (9) as a quadratic program:

$$\begin{equation} \underset{\boldsymbol{\lambda} \in \mathbb{R}^K}{\textrm{minimize }} \frac{1}{2} \boldsymbol{\lambda}^\top \boldsymbol{Q} \boldsymbol{\lambda} \,\,\,\textrm{ s.t. } \,\,\,\lambda_k \geqslant 0 \: \forall k \in [K] \,\,\,\textrm{ and } \,\,\,\sum_{k = 1}^K \lambda_k = 1, \end{equation} \tag{ A5 }$$

where $\boldsymbol{Q} = \boldsymbol{U}^\top\boldsymbol{U}$. The matrix $\boldsymbol{U} \in \mathbb{R}^{\vert \textrm{B} \vert \times K}$ is given by

$$\begin{equation*} \boldsymbol{U} = \begin{bmatrix} \vdots \:\: \qquad \qquad \vdots \qquad \qquad \:\: \vdots \:\: \\ \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_1 \:\: \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_2 \:\: \ldots \:\: \nabla_{\boldsymbol{\theta}^\textrm{sh}} \mathcal{L}_K \\ \vdots \:\: \qquad \qquad \vdots \qquad \qquad \:\: \vdots \:\: \\ \end{bmatrix}. \end{equation*}$$

We solve the quadratic program in equation (A5) using the Frank-Wolfe algorithm.

Let us assume the Fourier spectrum of the signal is a delta function, i.e. $\widetilde{u}(k) = \Gamma_0~\delta (k - k_0)$, with k₀ being the dominant wavenumber. The residual at training time (t) is given as $r(\cdot,t) = u_\theta - u$, and its corresponding Fourier coefficients for the wavenumber k is given by $\widetilde{r}(k,t) = \widetilde{u}_{\theta}(k,t) - \widetilde{u}(k)$. The learning rate of the residual for the dominant frequency is given as,

$$\begin{align} \Big \vert \frac{\partial \widetilde{r}(k_0,\tau)}{\partial \tau} \Big \vert & = \Big \vert \frac{\partial \widetilde{u}_{\theta}(k_0,\tau)}{\partial \theta} \Big \vert \Big \vert \frac{\partial \theta}{\partial \tau} \Big \vert \nonumber \\ & = \Big \vert \eta \frac{\partial \widetilde{u}_{\theta}(k_0)}{\partial \theta} \Big \vert \Big \vert \frac{\partial \mathcal{L}}{\partial \theta} \Big \vert \nonumber \\ & = \Big \vert \eta \frac{\partial \widetilde{u}_{\theta}(k_0)}{\partial \theta} \Big \vert \Big \vert \lambda_{{0}} \frac{\partial \mathcal{L}_{{0}}}{\partial \theta} + \lambda_{{1}} \frac{\partial \mathcal{L}_{{1}}}{\partial \theta} \Big \vert , \end{align} \tag{ B1 }$$

where $\mathcal{L}_0~{ = \sum_{i = 1}^{N} \vert \boldsymbol{u}_{\boldsymbol{\theta}}(\boldsymbol{x}_i) - \boldsymbol{u}(\boldsymbol{x}_i) \vert^2}$ and $\mathcal{L}_1~{ = \sum_{i = 1}^{N} \vert \mathcal{D}_{\boldsymbol{x}}^m \boldsymbol{u}_{\boldsymbol{\theta}} (\boldsymbol{x}_i) - \mathcal{D}_{\boldsymbol{x}}^m \boldsymbol{u}(\boldsymbol{x}_i)\vert^2}$. Equation (B1) follows from the gradient flow of the parameters $\partial_\tau \theta \approx \eta \partial_\theta \mathcal{L}$ as described in equation (4). From Parseval's Theorem [50], we can compute the magnitude of the gradients for both objectives $\mathcal{L}_{{0}}$ and $\mathcal{L}_{{1}}$ with respect to the network parameters θ.

$$\begin{align} \frac{\partial \mathcal{L}_{{0}}}{\partial \theta} & = 2~\sum_{k = -N/2}^{N/2-1} \Big \vert \textbf{Re}\Big [ \widetilde{u}_{\theta} (k) - \widetilde{u}(k) \Big ] \frac{\partial \widetilde{u_\theta} (k) }{\partial \theta} \Big \vert \nonumber \\ & \leqslant 2~\Big \vert \Gamma_0~\frac{\partial \widetilde{u}_{\theta}(k_0)}{\partial \theta} \Big \vert + 2~\sum_{k = -N/2}^{N/2-1} \Big \vert \widetilde{u}_{\theta}(k) \frac{\partial \widetilde{u}_{\theta} (k)}{\partial \theta } \Big \vert \nonumber \\ &\approx O(1) \Big \vert \frac{\partial \widetilde{r}(k_0)}{\partial \theta } \Big \vert . \end{align} \tag{ B2 }$$

The second step in the above calculation involves triangular inequality and the assumption that the Fourier spectrum of the signal is a delta function. In a similar fashion, we can repeat the same analysis for the Sobolev part of the objective function as follows,

$$\begin{align} \frac{\partial \mathcal{L}_{{1}}}{\partial \theta} & = 2~\sum_{k = -N/2}^{N/2-1} \Big \vert \textbf{Re} \Big [ (ik)^m \big (\widetilde{u}_{\theta} (k) - \widetilde{u}(k) \big ) \Big ] \frac{\partial (ik)^m \widetilde{u}_{\theta} (k) }{\partial \theta} \Big \vert \nonumber \\ & = 2~\sum_{k = -N/2}^{N/2-1} \Big \vert \textbf{Re} \Big [ (ik)^{2m} \big (\widetilde{u}_{\theta} (k) - \widetilde{u}(k) \big) \Big ] \frac{\partial \widetilde{u}_{\theta} (k) }{\partial \theta} \Big \vert \nonumber \\ & \leqslant 2~\Big \vert \Gamma_0~\textbf{Re} \big [ (ik_0)^{2m} \big ] \frac{\partial \widetilde{u}_{\theta}(k_0)}{\partial \theta} \Big \vert + 2~\sum_{k = -N/2}^{N/2-1} \Big \vert \textbf{Re} \big [(ik)^{2m} \big ] \widetilde{u}_{\theta}(k) \frac{\partial \widetilde{u}_{\theta} (k)}{\partial \theta} \Big \vert \nonumber \\ & \approx O (k_0^{2m}) \Big \vert \frac{\partial \widetilde{r}(k_0)}{\partial \theta} \Big \vert . \end{align} \tag{ B3 }$$

By substituting equations (B2) and (B3) into equation (B1), we get,

$$\begin{align*} \Big \vert \frac{\partial \widetilde{r}(k_0,\tau)}{\partial \tau} \Big \vert & = \Big \vert \eta \frac{\partial \widetilde{u}_{\theta}(k_0)}{\partial \theta} \Big \vert \left( \lambda_{{0}} O(1) \Big \vert \frac{\partial \widetilde{r}(k_0)}{\partial \theta} \Big \vert + \lambda_{{1}} O (k_0^{2m}) \Big \vert \frac{\partial \widetilde{r}(k_0)}{\partial \theta} \Big \vert \right) \\ & = \frac{\eta}{O(k_0)} \Big[ \lambda_{{0}} O(1) + \lambda_{{1}} O (k_0^{2m}) \Big] \Big \vert \frac{\partial \widetilde{r}(k_0)}{\partial \theta} \Big \vert. \end{align*}$$

We use the fact that the spectral rate of the network function residual is inversely proportional to the wavenumber, i.e. $\partial \widetilde{u}_\theta (k)/\partial \theta = O\left(k^{-1}\right)$ [26].

We use a neural network $f_{\theta}: (x,y) \rightarrow (u)$ with five layers and 64 neurons per layer and $\sin(\cdot)$ activation function. The target function is a sinusoidal forcing term of the form,

$$\begin{equation} u(x,y) = \sum_{i = 1}^{M} A_x^i \cos(2\pi l_x^i x/L + \phi_x^i) A_y^i \sin(2\pi l_y^i y/L + \phi_y^i), \end{equation} \tag{ B4 }$$

where M = 20 and $L = 2\pi$ is the size of the domain. The parameters $A_x, A_y$ are drawn independently and uniformly from the interval $[-5,5]$, and similarly $\phi_x, \phi_y$ are drawn from the interval $[0, 2\pi]$. The parameters $l_x, l_y$ which set the local length scales are sampled uniformly from the set $\{1,2,3,4,5\}$. The target function $u(x,y)$ is given on a 2D spatial grid of resolution $128 \times 128$ covering the domain $[0, 2\pi] \times [0, 2\pi]$. The grid data is then divided into a $50/50$ train and test split which corresponds to 8192 training points used. We train the neural network for 20 000 epochs using the Adam optimizer [38]. Each epoch is composed of two mini-batch updates of batch size $\vert \textrm{B} \vert = 4096$. The initial leaning rate is chosen as $\eta = 10^{-3}$ and decreased by a factor of 10 after 10 000 and 15 000 epochs.

We study the forward solution of the active turbulence problem using PINNs in the $(\omega-\psi)$ formulation. We employ a neural network $f_{\theta}: (x,y,t) \rightarrow (u,v)$ with seven layers and 100 neurons per layer and ${\sin}$-activation functions. For choosing the activation function we conducted extensive numerical experiments using different activation functions for a simple function reconstruction problem with the objective

$$\begin{equation*} \mathcal{L}(\boldsymbol{u}, \boldsymbol{u}_\theta, \theta) = \frac{1}{N_\textrm{int}} \sum_{i = 1}^{N_\textrm{int}} \Big \vert \boldsymbol{u}_{\boldsymbol{\theta}}(t_i, \boldsymbol{x}_i^{\Omega}) - \boldsymbol{u}_{i} ) \Big \vert^2, \end{equation*}$$

where $\boldsymbol{u}_{i} = \boldsymbol{u}(t_i, \boldsymbol{x}_i^{\Omega})$. In figure C.1 we clearly see that the sin activation function leads to the best approximation accuracy for the velocities and their associated derivatives. The vorticity ω is directly computed from the network outputs using automatic differentiation. For the solution in the square domain, we select a sub-domain $[\frac{\pi}{4},\frac{3\pi}{4}] \times [\frac{\pi}{4},\frac{3\pi}{4}]$ with 50 time steps (dt = 0.01) accounting to a total time of 0.5 units in simulation time. We choose $N_\mathrm{I} = 4096$ points for the initial conditions, $N_\mathrm B = 12\,800$ boundary points and a total of $N_\textrm{int} = 192200$ residual points in the interior of the square domain. The network is trained for 7000 epochs using the Adam optimizer with each epoch consisting of 49 mini-batch updates with a batch size $\vert\textrm{B}\vert = 4000$. The velocities are provided at the boundaries of the training domain.

For the forward simulation in the annular geometry, the domain is chosen with the inner radius of $r_\textrm{inner} \approx 0.245$ units and the outer radius $r_\textrm{outer}\approx 0.85$ units. We choose $N_\mathrm{I} = 3536$ points for the initial conditions, $N_\mathrm B = 14\,400$ boundary points and a total of $N_\textrm{int} = 162\,400$ residual points in the interior of the domain. The network training is again performed over 7000 epochs using Adam with each epoch consisting of 41 mini-batch updates with a batch size $\vert\textrm{B}\vert = 4000$. Dirichlet boundary conditions are enforced at the inner and outer rims of the annular geometry by providing the velocities. The network is optimized subject to the loss function

$$\begin{equation*} \mathcal{L}(\boldsymbol{u}, \theta) = \lambda_1~\mathcal{L}_\textrm{boundary} + \lambda_2~\mathcal{L}_\textrm{res} + \lambda_3~\mathcal{L}_\textrm{div} + \lambda_4~\mathcal{L}_\textrm{init}. \end{equation*}$$

Here, $\lambda_1, \lambda_2, \lambda_3, \lambda_4$ correspond to the weights of the objectives for boundary condition, equation residual, divergence, and initial condition, respectively.

For studying the inverse problem we resort to the primitive formulation $(u,v,p^{*})$ of the active turbulence model, with $p^{*} = -\nabla p+\xi_1\nabla |\boldsymbol{u}|^2$ as the effective pressure. We used a neural network $f_\theta:(x,y,t)\rightarrow (u,v,p^*)$ with five layers and 100 neurons per layer and sin-activation functions. Given the measurement data of flow velocities $\{u_i, v_i \}_{i = 1}^{N}$ at discrete points, we infer the effective pressure $p^{*}$ along with the parameters $\Sigma = \{\xi_0, \alpha, \beta, \Gamma_0, \Gamma_2 \}$. The total of $204\,800$ points, sampled from the same domain as for the forward solution in the square domain, are utilized with a $70/30$ train and test split, which corresponds to $143\,360$ training points used. The network is trained for a total of 7000 epochs using the Adam optimizer with each epoch consisting of 35 mini-batch updates with a batch-size $\vert\textrm{B}\vert = 4096$. We optimize the model subject to the loss function

$$\begin{equation*} \mathcal{L}(\boldsymbol{u}, \theta) = \lambda_1~\mathcal{L}_\textrm{fit} + \lambda_2~\mathcal{L}_\textrm{res} + \lambda_3~\mathcal{L}_\textrm{div}, \end{equation*}$$

with $\lambda_1, \lambda_2, \lambda_3$ corresponding to the data-fidelity, equation residual and the objective that penalises non-zero divergence in velocities, respectively. The training is done in a three-step approach, with the first 1000 epochs used for fitting the measurement data (pre-training, $\lambda_1 = 1, \lambda_2 = \lambda_3 = 0$), the next 1000 epochs training with an additional constraint on the divergence ($\lambda_1~\neq 0, \lambda_2~\neq 0 , \lambda_3 = 0$), and finally followed by the introduction of the equation residual ($\lambda_1~\neq 0, \lambda_2~\neq 0 , \lambda_3~\neq 0$). At every checkpoint introducing a new task, we reset all λ_k for the computation of $\hat{\lambda}_k$ (see equation (7) main text) of the max/avg variant to $\lambda_k = 1$. For both, forward and inverse problems, we choose an initial learning rate $\eta = 10^{-3}$ and decrease it by a factor of 10 after 3000 and 6000 epochs (see figure C.3).

The training data is generated from uniformly sampling in the domain $[-1.5,1.5] \times [-1.5,1.5]$ with a resolution of $100 \times 100$ in space and 80 time snapshots separated with time-step dt = 0.001. This lead to $N_\mathrm I = 10\,000, N_\mathrm B = 32\,000$ and $N_\textrm{int} = 768\,000$ points in the interior of the space and time domain. The network is trained for 2500 epochs using the Adam optimizer with each epoch consisting of 187 mini-batch updates with batch size $\vert\textrm{B}\vert = 4096$. The neural network $f_\theta: (x,y,t) \rightarrow (\phi)$ with 5 layers and 64 nodes is chosen for this problem. We found that using the ELU activation function resulted in the best performance. The total loss is given as,

$$\begin{equation} \mathcal{L}(\phi) = \lambda_1~\mathcal{L}_\textrm{boundary} + \lambda_2~\mathcal{L}_\textrm{res} + \lambda_3~\mathcal{L}_\textrm{init}{} .\end{equation} \tag{ E2 }$$

We choose the initial learning rate as $\eta = 10^{-3}$ and decrease it after 1500 and 2000 epochs by a factor of 10.

We initialize all weights to $\lambda_k = 1,k = 1,\dots,K$. For all experiments we update the dynamic weights at the first batch of every fifth epoch. The moving average is then performed with α = 0.5. The mini-batches $\boldsymbol{X}_\textrm{B}$ are randomized after each epoch, such that for every update of the weights, λ_k will be computed with respect to a different subset $\boldsymbol{X}_\textrm{B}$ of all training points $\mathcal{X}$. Furthermore, batching is only performed over the interior points N_int. As $N_\mathrm{I},N_\mathrm{B}\ll N_{\textrm{int}}$, the points for initial and boundary conditions do not require batching and remain fixed throughout the training.