Improved Architectures and Training Algorithms for Deep Operator Networks

Abstract

Operator learning techniques have recently emerged as a powerful tool for learning maps between infinite-dimensional Banach spaces. Trained under appropriate constraints, they can also be effective in learning the solution operator of partial differential equations (PDEs) in an entirely self-supervised manner. In this work we analyze the training dynamics of deep operator networks (DeepONets) through the lens of Neural Tangent Kernel theory, and reveal a bias that favors the approximation of functions with larger magnitudes. To correct this bias we propose to adaptively re-weight the importance of each training example, and demonstrate how this procedure can effectively balance the magnitude of back-propagated gradients during training via gradient descent. We also propose a novel network architecture that is more resilient to vanishing gradient pathologies. Taken together, our developments provide new insights into the training of DeepONets and consistently improve their predictive accuracy by a factor of 10-50x, demonstrated in the challenging setting of learning PDE solution operators in the absence of paired input-output observations. All code and data accompanying this manuscript will be made publicly available at https://github.com/PredictiveIntelligenceLab/ImprovedDeepONets.

Appendices

A Nomenclature

Table 8 summarizes the main symbols and notation used in this work.

Table 8 Nomenclature: Summary of the main symbols and notation used in this work

B Hyper-parameter settings

See Table 9.

Table 9 Default hyper-parameter settings for each benchmark employed in this work (unless otherwise stated)

C Computational cost

Training: Table C summarizes the computational cost (hours) of training physics-informed DeepONet models with different network architectures and weighting schemes. All networks are trained using a single NVIDIA RTX A6000 graphics card (Table 10).

Table 10 Computational cost (hours) for training physics-informed a DeepONet model across the different becnhmarks and architectures employed in this work

D Proof of Lemma 3.2

Proof

First we observe that $\varvec{H}(\varvec{\theta })$ is a Gram matrix. Indeed, we define

$$\begin{aligned} v_i = T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)), \quad i=1.2, \dots , N^*. \end{aligned}$$

(D.1)

Then by the definition of $\varvec{H}(\varvec{\theta })$, we have

$$\begin{aligned} \varvec{H}_{ij}(\varvec{\theta }) = \langle v_i, v_j \rangle . \end{aligned}$$

(D.2)

• (a) By the definition of a Gram matrix.

• (b) Let $\Vert \varvec{H}(\varvec{\theta })\Vert _\infty = \langle v_i, v_j \rangle$ for some i, j and $\langle v_k, v_k \rangle = \max _{1 \le k \le N^*} \varvec{H}_{kk}(\varvec{\theta })$/ for some k. Then we have

  $$\begin{aligned} \max _{1 \le k \le N^*} \varvec{H}_{kk}(\varvec{\theta }) = \langle v_k, v_k \rangle \le \Vert \varvec{H}(\varvec{\theta })\Vert _\infty = \langle v_i, v_j \rangle \le \Vert v_i\Vert \Vert v_j\Vert \le \Vert v_k\Vert ^2 = \langle v_k, v_k \rangle . \end{aligned}$$

  (D.3)

  Therefore, we have

  $$\begin{aligned} \Vert \varvec{H}(\varvec{\theta })\Vert _\infty = \max _{1\le k \le N^*} \varvec{H}_{kk}(\varvec{\theta }). \end{aligned}$$

  (D.4)

$\square$

E Proof of Lemma 3.4

Proof

Recall the definition of the loss function 3.6,

$$\begin{aligned} \mathcal {L}(\varvec{\theta }) = \frac{2}{N^*} \sum _{i=1}^{N^*} \left| T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i))\right| ^2. \end{aligned}$$

(E.1)

Now we consider the corresponding gradient flow

$$\begin{aligned} \frac{d \varvec{\theta }}{ dt}&= - \nabla \mathcal {L}(\varvec{\theta }), \end{aligned}$$

(E.2)

$$\begin{aligned}&= - \frac{4}{N^*} \sum _{k=1}^{N^*} \left( T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)) \right) \frac{d T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)) }{d \varvec{\theta }}. \end{aligned}$$

(E.3)

For $1 \le j \le N^*$, note that

$$\begin{aligned}&\frac{d T^{(j)}(u^{(j)}(\varvec{x}_j), G_{\varvec{\theta }}(\varvec{u}^{(j)})(\varvec{y}_j))}{ dt } = \end{aligned}$$

(E.4)

$$\begin{aligned}&= \frac{d T^{(j)}(u^{(j)}(\varvec{x}_j), G_{\varvec{\theta }}(\varvec{u}^{(j)})(\varvec{y}_j))}{ d \varvec{\theta } } \cdot \frac{d \varvec{\theta }}{ dt } \end{aligned}$$

(E.5)

$$\begin{aligned}&= \frac{d T^{(j)}(u^{(j)}(\varvec{x}_j), G_{\varvec{\theta }}(\varvec{u}^{(j)})(\varvec{y}_j))}{ d \varvec{\theta } }\nonumber \\&\quad \cdot \left[ - \frac{4}{N^*} \sum _{i=1}^{N^*} \left( T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)) \right) \frac{d T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)) }{d \varvec{\theta }} \right] \end{aligned}$$

(E.6)

$$\begin{aligned}&= - \frac{4}{ N^*} \sum _{i=1}^{N^*} \left( T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)) \right) \nonumber \\&\quad \left\langle \frac{d T^{(i)}(u^{(i)}(\varvec{x}_i), G_{\varvec{\theta }}(\varvec{u}^{(i)})(\varvec{y}_i)) }{d \varvec{\theta }}, \frac{d T^{(j)}(u^{(j)}(\varvec{x}_j), G_{\varvec{\theta }}(\varvec{u}^{(j)})(\varvec{y}_j))}{ d \varvec{\theta } } \right\rangle . \end{aligned}$$

(E.7)

By Definition 3.1, 3.3, we conclude that

$$\begin{aligned} \frac{d \varvec{T}\left( \varvec{U}(\varvec{X}), G_{\varvec{\theta }(t)} (\varvec{U}) (\varvec{Y}) )\right) }{ d t} =- \frac{4}{ N^*} \varvec{K}(\varvec{\theta }) \cdot \varvec{T}(\varvec{U}(\varvec{X}), G_{\varvec{\theta }(t)} (\varvec{U}) (\varvec{Y}) )), \end{aligned}$$

(E.8)

where $\varvec{K}$ is a $N^* \times N^*$ matrix whose entries are given by

$$\begin{aligned} \varvec{K}_{ij}(\varvec{\theta }) = \left\langle \frac{d T^{(i)} \left( u^{(i)} (\varvec{x}_i), G_{\varvec{\theta }} (\varvec{u}^{(i)})(\varvec{y}_i) \right) }{d \varvec{\theta }} , \frac{d T^{(j)} \left( u^{(j)} (\varvec{x}_j), G_{\varvec{\theta }} (\varvec{u}^{(j)})(\varvec{y}_j) \right) }{d \varvec{\theta }} \right\rangle . \end{aligned}$$

(E.9)

$\square$

F Antiderivative

See Fig. 15.

Fig. 15

Anti-derivative operator: Training loss convergence of DeepONet models using different loss schemes for $4 \times 10^4$ 40,000 iterations of gradient descent using the Adam optimizer. Here we remark that all losses are weighted, and, therefore, the magnitude of these losses is not informative

G Advection equation

See Figs. 16, 17 and 18.

Fig. 16

Advection equation: Training loss convergence of a DeepONet model using different DeepONet architectures and weighting schemes for $3 \times 10^5$ iterations of gradient descent using the Adam optimizer. Here we remark that all losses are unweighted

Fig. 17

Advection equation: Relative $L^2$ error of physics-informed DeepONets represented by different architectures and trained using different fixed weights $\lambda _{\text {bc}} = \lambda _{\text {ic}} = \lambda \in [10^{-2}, 10^{2}]$, averaged over the test data-set

Fig. 18

Advection equation: Predicted solution of the best trained physics-informed DeepONet for three different examples in the test data-set

H Burger’s equation

See Figs. 19, 20, 21, 22 and 23.

Fig. 19

Burger’s equation: Training loss convergence of a DeepONet models using different DeepONet architectures and weighting schemes for $2 \times 10^5$ iterations of gradient descent using the Adam optimizer. Here we remark that all losses are unweighted

Fig. 20

Burger’s equation: Relative $L^2$ error of physics-informed DeepONets represented by different architectures and trained using different fixed weights $\lambda _{\text {bc}} = \lambda _{\text {ic}} = \lambda \in [10^{-2}, 10^{2}]$, averaged over the test data-set

Fig. 21

Burger’s equation ($\nu = 0.01$): Predicted solution of the best trained physics-informed DeepONet for three different examples in the test data-set

Fig. 22

Burger’s equation ($\nu = 0.001$): Predicted solution of the best trained physics-informed DeepONet for three different examples in the test data-set

Fig. 23

Burger’s equation ($\nu = 0.0001$): Predicted solution of the best trained physics-informed DeepONet for three different examples in the test data-set

I Stokes flow

See Figs. 24, 25, 26 and 27.

Fig. 24

Stokes equation: Training loss convergence of a DeepONet models using different DeepONet architectures and weighting schemes for $2 \times 10^5$ iterations of gradient descent using the Adam optimizer. Here we remark that all losses are unweighted

Fig. 25

Stokes equation: Predicted solution of the best trained physics-informed DeepONet for one example in the test data-set

Fig. 26

Stokes equation: Predicted solution of the best trained physics-informed DeepONet for one example in the test data-set

Fig. 27

Stokes equation: Predicted solution of the best trained physics-informed DeepONet for one example in the test data-set