The conjugate kernel for efficient training of physics-informed deep operator networks

Authors: amanda a howard saad qadeer
andrew w engel
adam  tsou
max  vargas
tony  chiang
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-06 00:51:58

Abstract: Recent work has shown that the empirical Neural Tangent Kernel (NTK) can significantly improve the training of physics-informed Deep Operator Networks (DeepONets). The NTK, however, is costly to calculate, greatly increasing the cost of training such systems. In this paper, we study the performance of the empirical Conjugate Kernel (CK) for physics-informed DeepONets, an efficient approximation to the NTK that has been observed to yield similar results. For physics-informed DeepONets, we show that the CK performance is comparable to the NTK, while significantly reducing the time complexity for training DeepONets with the NTK.

Main paper: 
 introduction: 
 One current research focus in scientific machine learning is physics-informed neural networks (PINNs, (Raissi et al., 2019) ) and deep operator networks (DeepONets, (Lu et al., 2019) ). While PINNs and physics-informed DeepONets (PI-DeepONets, (Wang et al., 2021; Goswami et al., 2022a) ) have shown tremendous promise across a wide range of applications, e.g., Goswami et al. (2022b) ; Hao et al. (2023) ; Koric & Abueidda (2023) ; Karniadakis et al. (2021) , training PI-DeepONets can be difficult and can result in large errors when compared with exact solutions (Wang et al., 2022b) . Recent work has focused on efficient methods for improving accuracy of PINNs and PI-DeepONets, some of which include schemes that implement causality (Wang et al., 2022a) , adaptive point selection methods (Wu et al., 2023) , iterative methods to successively reduce the prediction in errors (Ainsworth & Dong, 2021; 2022; Howard et al., 2023; Aldirany et al., 2023; Wang & Lai, 2023) , and adaptive weighting schemes (McClenny & Braga-Neto, 2023) .One method that has emerged for increasing training accuracy is the use of the Neural Tangent Kernel (NTK) as a weighting scheme in the loss function in PINNs (Wang et al., 2022c) and PI-DeepONets (Wang et al., 2022b) . While the NTK significantly increases the accuracy, it greatly increases the computational cost by approximately a factor of six for each iteration. In this work, we present the use of the conjugate kernel (CK), an efficient approximation of the NTK, for adaptive weighting for training PI-DeepONets. While significantly less expensive to compute, we show that the CK provides almost equally accurate results for PI-DeepONets. 
conjugate kernel and the neural tangent kernel: 
 The NTK (often called the empirical NTK) for a deep neural network (DNN) or a DeepONet is defined as the Gram matrix of the Jacobian of the DNN with respect to the network parameters (Jacot et al., 2018) , originally defined for infinitely-wide neural networks (NNs). The NTK after training a finite-width NN is referred to as the after kernel (Long, 2021) . The CK (Fan & Wang, 2020; Wang et al., 2022d ) is a "0-order" approximation to the NTK, defined as the Gram matrix of the parameters of the last layer of a FFNN (Hu & Huang, 2021) . Here, we are only interested in kernels after training, and so both the NTK and CK will denote empirical kernels obtained from finite-width NNs after some amount of training.Mathematically, the NTK is represented as the kernel induced by the mapEQUATIONfor a feedforward neural network f θ NN with trainable parameters θ. The (empirical) NTK is then given by NTK(z, z)EQUATIONDenoting the parameters in the last layer of f θ NN by θ L , the CK is the kernel induced by the mapEQUATION)The CK is then defined as CK(z, z) := Φ CK (z) ⊤ Φ CK ( z).(4) If we define E as the contribution to the NTK for all but the last layer, we can write the relationship between the NTK and CK as NTK(z, z) = CK(z, z) + E(z, z).The NTK has been shown to significantly improve the convergence of physics-informed training, e.g., Wang et al. (2022c; b) ; Howard et al. (2023) ; Bai et al. (2023) , however, it is extremely computationally expensive to compute, thereby greatly increasing the training time of the NN or operator network. The CK has been shown to be an efficient approximation of the NTK, resulting in improved robustness, better conditioning, and increased accuracy in some test cases (Qadeer et al., 2023) .In operator learning (Li et al., 2020a; Lu et al., 2021; Li et al., 2020b; Wen et al., 2022) a map is learned between two Banach spaces. For a general parametric PDE of the form N (u, s) = 0 with boundary conditions B(u, s) = 0 the operator network learns the PDE solution, denoted by G(u) = s(u) for G : U → S, where U is the space of input parameters, e.g., the initial conditions, with u ∈ U, and S is the space of PDE solutions on a domain Ω. s ∈ S is an unknown function governed by the PDE.A standard DeepONet consists of the branch and trunk networks, which for our applications are both feedforward neural networks. The input to the branch network is the initial condition u ∈ U, discretized at a set of M discrete points, or sensor locations. The input to the trunk network is the spatial and time coordinates. The trunk and branch networks are trained simultaneously, and the DeepONet output is given byEQUATIONwhere θ denotes the trainable parameters of the unstacked DeepONet, b k is the k-th component of the branch output and t k is the k-th component of the trunk output (Lu et al., 2019; 2021) . In this work we use "modified" DeepONets (Wang et al., 2022b) , an architecture that introduces encoder layers for the branch and trunk nets and has been shown to increase the accuracy of PI-DeepONet training. PI-DeepONets are trained without data to satisfy the given PDE through the use of automatic differentiation to calculate the derivatives. The loss function is given by (Wang et al., 2022b) :EQUATION)where λ i are weighting terms and where N * = N R, N is the number of samples considered as initial conditions, and R is the number of collocation points at which the residual, initial conditions, and boundary conditions are sampled. T (i) denotes the operators in the loss function, including the boundary condition and differential operator. The weighting terms λ i in Eq. 6 require careful tuning to increase the training accuracy. Recent work has focused on the use of the NTK to successfully increase the accuracy for PI-DeepONets by locally choosing the optimal weights (Wang et al., 2022c; 2021) . The NTK matrix is found asEQUATIONThen, we can define the adaptive weights at a given iteration n byλ i = max 1≤i≤N * H ii (θ n ) H ii (θ n ) α (8)where H ii is calculated using the NTK. The exponent α is generally chosen as α = 0.5 or α = 1, depending on the task. The DeepONet is trained using the ADAM optimizer in Jax Bradbury et al. (2018) , with adaptive weights given by the NTK or CK used at every iteration. 
results: 
 The CK matrix is defined in an identical manner as the NTK, however, the derivatives are calculated only with respect to the final layer of the branch and trunk networks, denoted by θ * :EQUATIONWe can now compute the adaptive weighting given by Eq. 8 by setting H ii as either the NTK or the CK. In each case, we train with α = 0.5 or α = 1, and select the best value of α for this study.To provide a comparison with previous literature Wang et al. (2022b) , we contrast the NTK and CK weighting with the case of using fixed weighting, chosen as λ r = 1 and λ bc = 10.The viscous one-dimensional Burgers equation given in Appendix A.2 with viscosity ν and periodic boundary conditions presents a challenge for DeepONets, as shown in Wang et al. (2022b; 2021) . To provide a direct comparison, we follow the procedure as outlined in Wang et al. (2022b) to generate a training and testing set by selecting initial conditions u(x) from a Gaussian random field, and train using the same network hyperparameters. Results comparing the CK, the NTK, and using predetermined fixed weights are shown in Table 1 . The use of the CK not only leads to comparable or improved test accuracies compared to the NTK but requires only 36.8% of the training time. error are over 500 initial conditions in the test set. Mean run times are computed using an NVIDIA A100 GPU. The weighting scheme refers to the expressions chosen for λ i in Eq. 6.ν Fixed weighting NTK weighting CK weighting 0.01 3.65% ± 6.26% 0.63% ± 0.50%(α = 1) 0.66% ± 0.57%(α = 1) 0.001 9.18% ± 8.03% 2.80% ± 3.00%(α = 0.5) 3.04% ± 3.82%(α = 0.5) 0.0001 25.10% ± 9.53% 11.44% ± 5.91%(α = 0.5) 8.34% ± 3.62%(α = 1) Mean time (h) 1.582 6.096 2.243Table 2 : Comparison of the relative ℓ 2 error and computational time for the wave equation. The mean and SD of the relative ℓ 2 error are over 500 initial conditions in the test set. Run times are computed using an NVIDIA A100 GPU. The weighting scheme refers to the expressions chosen for λ i in Eq. 6.Fixed weighting NTK weighting CK weighting Error 3.45% ± 1.83% 0.93% ± 0.50% (α = 0.5) 0.96% ± 0.74% (α = 1) Run time (h) 1.305 5.949 1.881The wave equation is given byEQUATION)We take u(x) = 5 n=1 b n sin(nπx). Then, the exact solution is given by s(x, t) = In Table 2 we report the relative ℓ 2 errors between the exact solution and results. We note that the CK error is very close to the NTK error, however, it takes only about 32% of the computational time to achieve results of this accuracy. While the CK does not offer additional accuracy over the NTK, the lower computational time is a compelling argument for why one may consider using it.  
discussion: 
 We have shown that the CK is an efficient approximation of the NTK, and can result in similar training accuracy for physics-informed DeepONets, while significantly reducing the computational time. In the cases presented, the CK results have errors quite similar to the NTK results, including outperforming the NTK for the smallest viscosity for the viscous Burgers equation. While the accuracy is similar between the two, using the CK rather than the NTK cuts the training time by roughly two-thirds, and for this reason, the CK should be considered as an alternative for the NTK when training physics-informed DeepONets. Testing the CK for two cases limits the generalization in this study, however the use of the CK will be further explored in future work, including tuning for the optimal value of the exponent α to ensure the robustness of results. where ν is the viscosity.Learning rate (10 -3 , 5000, 0.9) Branch network size [101, 100, 100, 100, 100, 100, 100,  
 
 Appendix:  


Decision: The authors present a neural tangent kernel (NTK) approximation for training physics-informed deepOnet more accurately and faster, by computing a conjugate NTK with only the last layer of the network. The proposed method achieves similar performance as full NTK computations but allows for much faster training. Whilst the work is promising it is lacking in rigour in that the results are not put in context with a broader range of other approaches and it is not clear how well this method can scale. Nevertheless it is a good contribution to the community and should be accepted for the poster session.

Reviews: review: title: title: Review
main_review: main_review: **Summary of the paper:** The authors present a neural tangent kernel (NTK) approximation for training physics-informed deepOnet more accurately and faster, by computing a conjugate NTK with only the last layer of the network. The proposed method achieves similar performance as full NTK computations but allows for much faster training.

**Strength of the paper:** The paper does a good job at introducing the main differences between NTK and conjugate kernel, and the motivation behind developing fast NTK approximations. It also describes well key aspects of physics-informed deepOnet, and how they fit within NTK computations. The results are sound: the conjugate gradient achieves similar accuracy as NTK, with a significant training time gain.

**Weakness of the paper:** My understanding is that NTK/CK are only used to compute weighing terms in the loss function, but the deepOnet weights are updated using a standard gradient descent. For someone not familiar with NTK, this is somewhat confusing, as equation 1 and 3 may intuitively lead to think that NTK updates neural network weights in a different fashion than standard gradient descent. In table 1, it is also unfortunate to use the term "weight" to refer to $\lambda_i$, as this may be mistaken with the network weights/parameters. 

I wish the authors had expanded on the accuracy results for CK and low viscosity Burgers equation. It is surprising that CK achieves better performance than NTK, given that CK is an approximation of NTK. Do the authors have an idea why this happens?
title: title: uses conjugate kernel (CK) for updating the loss term coefficients in physics-informed deeponets
main_review: main_review: This works uses conjugate kernel (CK) as an efficient approximation of neural tangent kernel or NTK when the latter is used to increase the accuracy of physics-informed deeopnets in operator learning. Using two examples on Burgers’ and wave equations, the authors show that using CK, while insignificantly affecting the accuracy of the model, can reduce the training times (compared to cases when NTK is used). 

The main idea of the paper is to use the Gram matrix of the Jacobian of the DNN with respect to its last layers’ parameters. This matrix is used to tune the coefficients of the different terms that appear in the loss function (data loss terms and PDE residuals) of a PI-DeepONet. The contribution of this paper is to use eq. 8 instead of eq. 7 for updating the weights. 

I do not believe the paper makes sufficient contributions to the field to warrant its presentation in this workshop. This is especially the case since the results section lacks rigor. It is unclear if this method can be scaled up to cases where, e.g., the PDE has 3D spatial variables or has multiple outputs (e.g., the navier-stokes equations). It is also unclear why the authors are comparing their approach to a fixed-weight, i.e., to be fair, you can compare your approach to an adaptive weight scheme that is not based on NTK or CK (e.g., based on the moving average of the gradient magnitudes which are already available during training).