On the impact of larger batch size in the training of Physics Informed Neural Networks
Abstract
Physics Informed Neural Networks (PINNs) have demonstrated remarkable success in learning complex physical processes such as shocks and turbulence, but their applicability has been limited due to long training times.
In this work, we explore the potential of large batch size training to save training time and improve final accuracy in PINNs.
We show that conclusions about generalization gap brought by large batch size training on image classification tasks may not be compatible with
PINNs.
We conclude that larger batch sizes always beneficial to training PINNs.
Introduction
The advance of computation hardware technologies in the last decade has enabled researchers to explore the possibility of extremely large scale training of neural networks for image based tasks such as classification and segmentation.
In the course of 3 years, from 2016 to 2019, the training time of the ResNet50 [1] architecture for ImageNet[2] classification task has been reduced fromhours[1] in the original paper to 1 minute [3], while maintaining comparable Top-1 testing accuracy, by using various techniques such as learning rate warm up and layer-wise gradient normalization[4, 5].
However, other studies[6, 7] suggest that training with extremely large batch size often suffer from poor generalization performance on the test set, possibly due to the discrepancy between the loss landscapes on the train data and test data.
Keskar et al.
[6] numerically show that the training loss landscape, linearly parameterized around the local minima, is often “flatter” and matches that of the test loss, when using smaller batch size.
They, therefore, pose a conjecture that a larger batch size leads to “sharper” local minimum in the training loss landscape that results in poor generalization accuracy.
In [8], the authors empirically validate on ImageNet the assumption that decreasing learning rate and increasing batch size have similar effects on training and suggest that under the current parallelization hardware framework, larger batch sizes should be preferable due to its superior efficiency.
While numerous works[6, 9, 8] have studied the effect of batch size on Convolutional Neural Networks (CNNs) and image-based classification tasks, no such studies have been carried out for physicsinformed machine learning [10].
In this paper, we extend such analysis onto Physics Informed Neural
Networks (PINNs)[11], a framework for learning representations of the solutions to non-linear partial differential equations (PDEs).
Unlike the ImageNet classification task, where the train and test data are finite, the input data for PINNs is often sampled uniformly within the interior and on the boundary of the domain.
Hence, we don’t have a limitation on the amount of data we can use during training and each training iteration is carried out using a set of unique points.
Additionally, PINNs are trained in a self-supervised manner where we the train the network to satisfy certain boundary conditions and an underlying PDE, which directly contrasts with image classification problems.
Furthermore, in
PINNs, the evaluation metric is often calculated on test points that are also evenly sampled across the domain, leading to the same train and test data distributions.
This means that the generalization gap as described in [6] does not directly apply when dealing with PINNs.
Based on these observations, we conclude that the literature surrounding batch size on vision tasks may not generalize to PINNs and that this area demands more research.
Thus, in this paper, we explore the impact of larger batch size under the scientific computing scenario with PINNs.
The structure of the remaining of this paper is as follows: in Section 2, we briefly describe Physics
Informed Neural Networks (PINNs); in Section 3 we carry out experiments demonstrate the benefit of utilizing larger batch-sizes for PINNs through two different PDEs, namely the Poisson equation and the Navier-Stokes equation; finally, in Section 4 we summarize our findings and outline possible future areas of research.
Physics Informed Neural Networks (PINNs)
We begin with a brief overview of physics-informed neural networks (PINNs)[12] which is a method for approximating the solution to differential equations using neural networks.
Consider a PDE with a Dirichlet boundary condition
[Equation]
where N is some nonlinear operator.
To approximate this PDE using PINNs, we define a neural network u θ and tune the parameters θ of the network to minimize the following loss function
[Equation]
D Ω , D ∂
Ω
are uniform distributions supported on the interior and the boundary of the domain Ω , and α is the weight for the boundary loss term used to balance the residual and the boundary losses.
The weight α may be fixed or adaptive [13].
By minimizing the loss function above, the networks learns a solution to the desired PDE in (1).
Experiments on PINNs
For the experiments carried out in this section, we consider the following PDEs
Poisson Equation
The 2 − dimensional Poisson PDE consists of finding u : Ω → R satisfying
[Equation]
where u xx , u yy are the second-order partial derivatives with respect to input x, y and Ω ⊂ R 2 .
For our experiments, we choose Ω =
[ − 1 , 1] 2 and choose a model solution u and corresponding forcing function f as
[Equation]
Navier-Stokes Equation We consider a decaying turbulence simulation in a periodic square domain Ω = [0 , 2 π ] 2 with Reynolds number Re = 100 and T = 0 .
We pose this problem using the velocity-vorticity formulation:
[Equation]
3.1 Sharpness of PINNs local minimum
First, utilizing the 2D Poisson Equation, we demonstrate that training with larger batch size leads to sharper local minimum.
The specific architecture we utilize along with optimization hyper-parameters have been listed in Appendix B.4.
We visualize the sharpness of the loss function qualitatively at the parameters obtained by the end of the training from singular values of the Hessian matrix of the loss function evaluated at the resulting parameters.
Given the O ( N 3 ) complexity in performing the
SVD, we reduce the neural network to 3 hidden layers, each with 16 neurons, and choose k = 4 in equation (4).
Figure 1 shows the sorted singular values of parameters obtained from training with different batch size, as normalized by the L 2 norm of the singular value vectors.
The singular values obtained from larger batch sizes decay consistently faster than that obtained from smaller batch sizes, suggesting that sharper loss landscapes.
While Keskar et al.
[6] propose that sharper local minima exacerbates the discrepancy between the validation and training metrics, in the case of PINNs the continuous uniform sampling fills the gap and consequently yields improved errors at sharper local minima.
3.2 Learning rate scaling for PINNs
Each curve is evaluated on one set of parameters returned by an instance of training with size the with its L
While the example above have shown the efficacy of increasing batch size, in this section we further demonstrate that an optimal learning rate choice can significantly improve prediction accuracy and accelerate training.
While convex optimization suggests that less noise in the gradient estimate in each step allows the use of larger learning rates and consequently leads to faster convergence[14], it is not sufficiently studied how the learning rate should be scaled in proportion to the batch size increase in neural networks.
Given the scaling scheme
[Equation]
two commonly used schemes are linear scaling[8, 15], where κ = 1 , and square root scaling[16], where κ = 0 .
In this experiment, we will use PINNs to solve the 2 -dimensional Poisson problem as using k = 12 in equation (4) and a more challenging Navier-Stokes (NS) equation to explore the optimal learning rate scaling schemes.
For both these cases, we perform a grid search to determine the optimal learning rate for a batch size.
The specific architectures and hyperparameter choices we utilized for our experiments have been discussed in the Appendix.
Figure 2 summarizes our findings through heat maps of the final mean ensemble errors for combinations of batch size and learning rates on the 2 − D Poisson and turbulence problem.
Additionally, we plot the number of iterations required to achieve a particular error tolerance.
In summary, we infer the following from our results when utilizing a large batch size for training
PINNs:
• Consistent improvements in the accuracy of the solution obtained.
• Higher tolerance to choice of learning rate; large batch sizes allows using larger learning rates without compromising the stability of training.
However, we do note that beyond a certain threshold, higher learning rates would always lead to divergence of training.
• Quicker convergence to a desired error tolerance.
Finally, the inferred slope κ for the scaling of optimal learning rate were 0 .
7636 and − 0 .
respectively for the 2 − D Poisson and turbulence.
This shows that an optimal scaling scheme is likely dependent on properties of each specific problem and that a general learning rate scheme may not exist.
Discussion
We demonstrate that PINNs greatly benefit from higher batch sizes, which not only improve the accuracy but also speed up the training process.
While the choice of optimal learning rate helps improve accuracy, we see that the improvement gained from learning rate tuning is not as impactful [Figure] as that from increasing batch size.
We are aware that PINNs still are significantly more expensive than classical numerical solvers.
However, our work demonstrates that PINNs can become more competitive when worked with larger batch sizes.
Given that they can achieve a required tolerance faster with higher batch sizes, a particularly effective route to this accelerate training would be employing data-parallelism to reduce the time per iteration.
Future lines of research regarding this topic include analysing the behaviour of batch size on more complicated PDEs.
While this study does not consider higher dimensional PDEs, we conjecture that higher batch sizes could significantly help in that space as well, where utilizing classic numerical methods is often very computationally expensive or impossible in practice.
Finally, while PINNs solve a single instance of a PDE, a possible future direction of study would involve the study of batch size in the context of operator learning architectures [17, 18, 19, 20, 21].
Using these architectures, it is possible to evaluate the solution of PDEs under different parameters almost instantaneously, after spending time training the model prior to deployment.
Increasing batch sizes and reducing training time of operator learning frameworks has the potential of improving the practicality of these methods in both industry and academia.  Acknowledgements
We would like to acknowledge support from the US Department of Energy under the Advanced
Scientific Computing Research program (grant DE-SC0019116), the US Air Force (grant AFOSR FA9550-20-1-0060), and US Department of Energy/Advanced Research Projects Agency (grant DE-AR0001201). This research used resources of the Argonne Leadership Computing Facility, which is a DOE Office of Science User Facility supported under Contract DE-AC02-06CH11357.
We additionally thank Corey Adams and Bethany Lusch for helpful discussions
A Architecture A.1 Multilayer Pecerptron (MLP)
A Multilayer Pecerptron (MLP) consists of compositions of layers, which are affine transformations followed by an activation function σ : R → R .
That is, a layer computes the following transformation from R d in to R d out x (cid:55)→ σ ( Wx + b ) , where W ∈ R d in × d out and b ∈ R d out are trainable parameters and σ is applied element-wise.
The MLP architecture then consists of multiple iterations of the transformation above being carried out, each with independent trainable parameters.
A.2 Positional Encoding
The harmonic encoding given input vector x can be expressed as the concatenation of feature vectors [sin(2 k π x ) , cos( [Equation] x [Equation]
[Equation]
= [Equation] [Equation] [Equation]
[Equation] [Equation] A.3 Modified MLP A modified MLP architecture for PINNs was described in [22] and found to consistently yield better results for PINNs.
An L -layer architecture of the modified MLP is defined as follows
[Equation] [Equation]
[Equation] [Equation]
[Equation]
14) where σ denotes a nonlinear activation function, ⊙ denotes a point-wise multiplication, and X denotes an batch of input coordinates.
B Hyper Parameter and Architecture Choices for Numerical Experiments B.1 Optimizer Choice
We train using Adam[23] and learning rate schedule consisting of a linear warmup phase of 1 ,steps followed by an exponential decay schedule with rate 0 .
1 per 10 , 000 transition steps.
The total number of iterations is 20 , 000 .
At the end of training, the learning rate is reduced to 0 .
01 of the peak value 1 .
B.2 Singular Value Spectrum for 2D Poisson Problem
The weight on the boundary loss term in equation (2) is 1000 for the evaluation of singular value spectrum.
We set the MLP following the positional encoding to be 3 hidden layers, each withneurons.
B.3 Grid Search for 2D Poisson
We utilize a 3-layer MLP of width 256 with a positional encoding layer consisting of 5 frequencies, and utilize a boundary loss weight α = 100 in Equation 2.
The results presented were averaged over an ensemble size of 32.
B.4 Navier Stokes
The architecture we utilize is largely based on the architecture utilized in [24]: a modified MLP architecture consisting of 3 hidden layers of width 256 .
with harmonic expansion layer at the start to ensure that the network implicitly satisfies the periodic boundary conditions.
Additionally, we parameterize the model using two separate networks as
[Equation]
Before, we carry out our grid search, we train the network ˜ u
θ 2 to capture the initial condition with high accuracy (0.01% relative L 2 error) and keep this fixed for all the experiments we perform.
Hence, we’re only evolving θ 1 during the process of training our model and do not have losses for the initial condition competing during the process of training.
However, we assign a weight of 500 to the continuity residual equation over the vorticity residual equation.
Finally, we train the network utilizing the causality enforcing training process described in [24].
C Larger batch size helps overcome the spectral bias
We use a one dimensional Poisson problem to demonstrate that larger batch size helps to stabilize training.
Poisson Equation The 1 − dimensional Poisson PDE consists of finding u : Ω → R satisfying [Equation] (16) where u xx are the second-order partial derivatives with respect to input x, y and Ω ⊂ R .
For our experiments, we choose Ω =
[ − 1 , 1] .
We choose the solution u =
[Equation]
In this case the weight on the boundary loss term in equation (2) is 10 , and we use the same architecture as in section B.3.
Figure 3 shows that larger batch size stabilizes the training, reduce the uncertainty, and shift the predictive power spectrum towards higher frequency.
[Figure]