# Hutchinson Trace Estimation for High-Dimensional and High-Order Physics-
Informed Neural Networks

Zheyuan Hu Equal ContributionDepartment of Computer Science, National
University of Singapore, Singapore, 119077
(e0792494<EMAIL_ADDRESS>Zekun Shi11footnotemark: 1
22footnotemark: 2 George Em Karniadakis Division of Applied Mathematics, Brown
University, Providence, RI 02912, USA (george_karniadakis@brown.edu)Advanced
Computing, Mathematics and Data Division, Pacific Northwest National
Laboratory, Richland, WA, United States Kenji Kawaguchi22footnotemark: 2

###### Abstract

Physics-Informed Neural Networks (PINNs) have proven effective in solving
partial differential equations (PDEs), especially when some data are available
by blending seamlessly data and physics. However, extending PINNs to high-
dimensional and even high-order PDEs encounters significant challenges due to
the computational cost associated with automatic differentiation in the
residual loss function calculation. Herein, we address the limitations of
PINNs in handling high-dimensional and high-order PDEs by introducing the
Hutchinson Trace Estimation (HTE) method. Starting with the second-order high-
dimensional PDEs, which are ubiquitous in scientific computing, HTE is applied
to transform the calculation of the entire Hessian matrix into a Hessian
vector product (HVP). This approach not only alleviates the computational
bottleneck via Taylor-mode automatic differentiation but also significantly
reduces memory consumption from the Hessian matrix to an HVP’s scalar output.
We further showcase HTE’s convergence to the original PINN loss and its
unbiased behavior under specific conditions. Comparisons with the Stochastic
Dimension Gradient Descent (SDGD) highlight the distinct advantages of HTE,
particularly in scenarios with significant variability and variance among
dimensions. We further extend the application of HTE to higher-order and
higher-dimensional PDEs, specifically addressing the biharmonic equation. By
employing tensor-vector products (TVP), HTE efficiently computes the colossal
tensor associated with the fourth-order high-dimensional biharmonic equation,
saving memory and enabling rapid computation. The effectiveness of HTE is
illustrated through experimental setups, demonstrating comparable convergence
rates with SDGD under memory and speed constraints. Additionally, HTE proves
valuable in accelerating the Gradient-Enhanced PINN (gPINN) version as well as
the Biharmonic equation. Overall, HTE opens up a new capability in scientific
machine learning for tackling high-order and high-dimensional PDEs.

##  1 Introduction

Physics-Informed Neural Networks (PINNs) [44] have made significant strides in
solving partial differential equation (PDE) problems in scientific computing,
particularly for low-dimensional equations. This is attributed to their
powerful fitting capability [34], generalization [36], and stability in
optimization [37]. However, despite the demonstrated capacity of neural
networks in modeling high-dimensional data in image and text domains, PINNs’
exploration of high-dimensional and high-order PDEs has been relatively
limited. The primary challenge arises from the enormous computational cost
involved in the automatic differentiation process of calculating the residual
loss function steps for high-dimensional and high-order PDEs under the PINNs
framework.

Solving problems related to high-dimensional PDEs is crucial, considering the
curse-of-dimensionality, which is prevalent in various fields such as the
Black-Scholes equation in mathematical finance for option pricing, the
Hamilton-Jacobi-Bellman (HJB) equation in optimal control, and the Schrödinger
equation in quantum physics. Despite the availability of specialized high-
dimensional PDE solvers for different domains, such as DeepBSDE [16] and Deep
Splitting Method [1] for high-dimensional parabolic equations and specific
algorithms for the Hamilton-Jacobi equation [13, 14], a unified and effective
approach for solving general high-dimensional PDEs is lacking.

PINNs is a versatile method capable of handling arbitrary PDEs and being mesh-
free, in principle it is not susceptible to the curse-of-dimensionality.
Despite recent efforts to use PINNs for solving high-dimensional PDEs, such as
Stochastic Dimension Gradient Descent (SDGD) [22] and random smoothing [18,
23], these approaches have their limitations. SDGD attempts to reduce each
stochastic gradient’s computational cost and memory consumption by sampling
dimensions. However, when significant variability and variance among
dimensions exist, SDGD’s stochastic gradients may suffer from substantial
variance, hindering convergence. As for random smoothing methods, they smooth
the original neural network with Gaussian noise, impacting the network’s
expressive power and making it less effective in handling non-smooth PDE
solutions.

To address the limitations of PINN in tackling high-dimensional and even high-
order PDEs, herein, we propose the use of Hutchinson Trace Estimation (HTE)
[25] to accelerate PINNs and reduce memory cost. HTE can be applied to various
common PDEs, such as second-order parabolic equations and the biharmonic
equation. Firstly, for second-order parabolic PDEs, the computational
bottleneck lies in the need to calculate the entire Hessian, with its
computational cost increasing quadratically with the dimensionality. HTE is
seamlessly integrated into the PINNs framework to alleviate this central
computational bottleneck. Specifically, HTE transforms the calculation of the
entire Hessian matrix into a Hessian vector product (HVP). Unlike the full
Hessian, which is a large matrix of dimensions $d\times d$, where $d$ is the
dimensionality of the PDE problem, the output of HVP is a scalar,
significantly reducing memory consumption. Moreover, HVP can efficiently
utilize Taylor-mode automatic differentiation [7] in JAX [8] for rapid
computation, achieving both acceleration and reduced memory usage.

Furthermore, we demonstrate that the loss of PINNs with HTE converges to the
original version of PINN loss under certain conditions and is unbiased,
providing a solid foundation for its convergence properties. We also discuss
the distinctions, similarities, and applicability between HTE and SDGD, the
most related methodologies. We present examples illustrating when HTE or SDGD
is more suitable and how their variances depend on the PDE problem. After
addressing second-order parabolic equations, we explore how HTE is even more
efficient in higher-order and higher-dimensional scenarios, motivating the
extension of HTE to the biharmonic equation.

Given that the biharmonic equation is a fourth-order high-dimensional
equation, attempting to solve it directly with conventional PINNs would lead
to a massive tensor of size $d^{4}$, where $d$ represents the dimensionality
of the PDE problem, and 4 denotes the order of the equation. However, by
extending the Hessian Vector Product (HVP) in HTE to a tensor-vector product
(TVP), we can transform the computation of the entire colossal tensor into a
scalar which is the TVP output. This not only saves a considerable amount of
memory but also facilitates rapid computation using Taylor-mode
differentiation. We also demonstrate that the extended HTE method provides an
unbiased estimate for the biharmonic operator, and we analyze the sources of
its variance. To emphasize the importance of properly implementing the HVP and
TVP in the proposed HTE method in JAX, we present a pseudocode to familiarize
readers with HTE and its efficient implementation.

Ultimately, we showcase the effectiveness of HTE in various experimental
setups. Firstly, we demonstrate that HTE and SDGD exhibit similar convergence
rates under the same memory constraints and speed limitations on a second-
order nonlinear anisotropic parabolic equation. Notably, they successfully
scale PINNs to a very high dimension (100,000 dimensions). Secondly, we
illustrate how HTE can accelerate the more accurate Gradient-Enhanced PINN
(gPINN). Since gPINN requires additional first-order derivatives of the
original PINN residual, the computational load is substantial, but HTE
mitigates this burden while keeping the improvement brought by gPINN. Lastly,
we highlight the advantages of HTE in the context of higher-order, high-
dimensional biharmonic equations.

This paper is arranged as follows. We present related work in Section 2. Then,
we introduce the proposed HTE method in Section 3. We present computational
experiments in Section 4 and summarize the paper in Section 5.

##  2 Related Work

Hutchinson Trace Estimation. Hutchinson Trace Estimation (HTE) [25] is an
unbiased estimator of the trace of a matrix, which was first proposed for the
influence matrix related to the Laplacian smoothing splines computation. Since
its initial publication, progress has been made for improvement. Hutch++ [39]
reduces the variance of HTE by employing a low-rank approximation methodology,
whose complexity is proved to be optimal among all matrix-vector query-based
methods for trace estimation. Later, Persson et al. [42] adopted the Nyström
approximation to improve Hutch++ further; [48] provided a modern analysis of
the HTE’s error. Roosta et al. [46] proved an improved estimation for the
sample efficiency in HTE. In modern deep learning, HTE has also been applied
to the Diffusion model [49] to infer the probability density function of the
model, where HTE can reduce the computation cost when dealing with the
Jacobian of the Diffusion model, which can be a convolutional neural network
or even a Transformer [50], whose derivative is extremely costly. To the best
of our knowledge, we are the first to apply HTE for efficient physics-informed
neural networks (PINNs) computations for high-order PDEs and in high-
dimensions.

Efficient Automatic Differentiation. To implement HTE efficiently, we adopt
the Taylor-mode automatic differentiation (AD) [7] in JAX [8], which is
tailored for high-order AD. Indeed, the type of AD affects the speed and
memory cost significantly, and the commonly used stacked forward/backward
modes of AD are much slower. In addition to the superior Taylor-mode AD, other
attempts exist to avoid the costly high-order derivatives. Randomized
smoothing [18, 23] randomly smooths the model using Gaussian noises for
inference, and its derivatives can be simulated via Monte Carlo simulation.
The generalization ability of a randomly smoothed model can be interpreted by
the information bottleneck theory [35]. Randomized AD [41] incorporated
stochastic gradient descent (SGD) to accelerate and avoid the need for full
AD.

High-Dimensional PDE Solver. Several works have considered high-dimensional
PDE solvers. In [51], the significance of $L^{\infty}$ loss and adversarial
training in addressing high-dimensional Hamilton-Jacobi-Bellman equations was
demonstrated. The separable PINNs approach in [12] employs a structure that
allows the residual points to be a tensor product of per-dimension points,
thereby increasing the batch size. Nevertheless, when confronted with problems
surpassing ten dimensions, memory usage becomes a serious bottleneck. DeepBSDE
[16, 17] and its extensions [2, 10, 19, 24, 31, 6] are grounded in the
classical BSDE interpretation for specific high-dimensional parabolic PDEs,
leveraging deep learning models to approximate the unknowns in the BSDE
formulation. The deep splitting method [1] unifies the splitting method with
deep neural networks. Chen et al. [11] solved forward and inverse problems of
Fokker-Planck equations, including Brownian noise and Levy noise in high
dimensions using PINNs. FBSNN [44] established a connection between high-
dimensional parabolic PDEs and forward-backward stochastic differential
equations, employing deep neural networks to learn the unknown solution. The
multilevel Picard methods [3, 4, 5, 26, 27] represent a nonlinear extension of
Monte Carlo capable of solving parabolic PDEs under specific constraints. In
[52, 53], tensor neural networks were proposed, adopting a separable structure
for efficient numerical integration in solving high-dimensional Schr”odinger
equations. More recently, SDGD [22] emerged as a method designed to sample
dimensions in PDEs, aiming to scale up and accelerate high-dimensional PINNs.
Randomized smoothing [18, 23] adopts the smoothed model with Gaussian noise so
that Monte Carlo can simulate its inference and derivatives to combat the
curse of dimensionality and avoid the costly automatic differentiation in
high-order and high dimensions.

Physics-Informed Machine Learning. The methodology developed in this paper is
based on the concept of Physics-Informed Machine Learning [33], especially
Physics-Informed Neural Networks (PINNs) [45]. PINNs model and approximate the
PDE solution by neural networks as surrogate models, which are trained through
the boundary and residual losses. These approaches are shown theoretically to
discover the underlying solutions governed by the PDEs [21, 40, 47]. PINNs
have been successful in numerous fields in science and engineering [9, 15, 54,
30], and effective PINNs variants have been proposed to deal with different
applications and problem settings [20, 28, 29, 32, 43].

##  3 Proposed Method

###  3.1 Preliminaries

This paper focuses on employing Physics-Informed Neural Networks (PINNs) to
address high-dimensional and high-order Partial Differential Equation (PDE)
problems. Additionally, it introduces the Hutchinson Trace Estimation (HTE) to
accelerate PINN and reduce its memory consumption. Below, we provide an
introduction to PINNs and HTE.

Physics-Informed Neural Networks (PINNs). This paper focuses on employing
PINNs [45] solving partial differential equations (PDEs) defined on the domain
$\Omega\subset\mathbb{R}^{d}$ with the boundary/initial condition on $\Gamma$
and the residual condition within $\Omega$:

$\displaystyle\mathcal{B}u(\boldsymbol{x})=B(\boldsymbol{x})\ \text{on}\
\Gamma,\qquad\mathcal{L}u(\boldsymbol{x})=g(\boldsymbol{x})\ \text{in}\
\Omega,$ (1)

Given the boundary points
$\\{\boldsymbol{x}_{b,i}\\}_{i=1}^{n_{b}}\subset\Gamma$ and the residual
points $\\{\boldsymbol{x}_{r,i}\\}_{i=1}^{n_{r}}\subset\Omega$, PINNs minimize
the discrepancy in the residual and on the boundary:

$\displaystyle\mathcal{L}(\theta)$
$\displaystyle=\lambda_{b}\mathcal{L}_{b}(\theta)+\lambda_{r}\mathcal{L}_{r}(\theta)$
(2)
$\displaystyle=\frac{\lambda_{b}}{n_{b}}\sum_{i=1}^{n_{b}}{|\mathcal{B}u_{\theta}(\boldsymbol{x}_{b,i})-B(\boldsymbol{x}_{b,i})|}^{2}+\frac{\lambda_{r}}{n_{r}}\sum_{i=1}^{n_{r}}{|\mathcal{L}u_{\theta}(\boldsymbol{x}_{r,i})-g(\boldsymbol{x}_{r,i})|}^{2},$

where $\lambda_{b}$ and $\lambda_{r}$ are the boundary and residual loss
weights, respectively.

Hutchinson Trace Estimation (HTE). The trace of a matrix
$A\in\mathbb{R}^{d\times d}$ can be randomly estimated as follows:

$\displaystyle\operatorname{Tr}(A)=\mathbb{E}_{\boldsymbol{v}\sim
p(\boldsymbol{v})}\left[\boldsymbol{v}^{\mathrm{T}}A\boldsymbol{v}\right],$
(3)

for all random variable $\boldsymbol{v}\in\mathbb{R}^{d}$ such that
$\mathbb{E}_{\boldsymbol{v}\sim
p(\boldsymbol{v})}[\boldsymbol{v}\boldsymbol{v}^{T}]=I$. Therefore, the trace
can be estimated by Monte Carlo:

$\operatorname{Tr}(A)\approx\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}A\boldsymbol{v}_{i},$
(4)

where each $\boldsymbol{v}_{i}\in\mathbb{R}^{d}$ are $i.i.d.$ samples from
$p(\boldsymbol{v})$.

There are several viable choices for the distribution $p(\boldsymbol{v})$ in
HTE, such as the most common standard normal distribution. However, to
minimize the variance of HTE, we opt for the Rademacher distribution as
follows: for each dimension of the vector $\boldsymbol{v}\sim
p(\boldsymbol{v})$, it is a discrete probability distribution that has a 50%
chance of getting +1 and a 50% chance of getting -1. The proof for the minimal
variance is given in [48].

###  3.2 HTE for High-Dimensional Second-Order Parabolic PDEs

In this subsection, we introduce the second-order parabolic equation under
consideration. Subsequently, we elaborate on how the HTE technique is
seamlessly incorporated into the PINNs framework, facilitating accelerated
convergence and reduced memory consumption. Following this, we delve into the
theoretical properties of the HTE loss functions. Lastly, we discuss the
efficient implementation of HTE, providing a JAX [8] pseudocode for HTE to aid
readers in realizing its efficiency.

We focus on a class of second-order parabolic equations, which include the
Fokker-Planck equation in statistical mechanics, the Black-Scholes equation in
mathematical finance, the Hamilton-Jacobi-Bellman equation in optimal control,
Schrödinger equation in quantum physics, etc., all ubiquitous in science and
engineering:

$\displaystyle\partial_{t}u(\boldsymbol{x},t)+{\operatorname{Tr}\left(\sigma\sigma^{\mathrm{T}}(\boldsymbol{x},t)\left(\operatorname{Hess}_{\boldsymbol{x}}u\right)(\boldsymbol{x},t)\right)}+f(\boldsymbol{x},t,u,\nabla_{\boldsymbol{x}}u)=0,\quad\boldsymbol{x}\in\mathbb{R}^{d},t\in[0,T],$
(5)

where $u(\boldsymbol{x},t)$ is the unknown exact solution we wish to solve,
$\operatorname{Hess}_{\boldsymbol{x}}u$ denotes the Hessian matrix of $u$,
$\sigma(\boldsymbol{x},t)\in\mathbb{R}^{d\times d}$ is a known matrix-valued
function, and $f(\boldsymbol{x},t,u,\nabla_{\boldsymbol{x}}u)$ is also a known
scalar function. PDEs with this form are of great interest in high-dimensions
addressing the curse-of-dimensionality; see [1, 16, 22, 23, 44]. PINN’s memory
and speed bottleneck are traces of the second-order Hessian part, which is
high-dimensional and high-order. Concretely, the computational cost of
automatic differentiation in PINNs increases exponentially with the order [7].
Regarding dimensionality, for second-order parabolic equations, the size of
the Hessian matrix increases quadratically with the dimension, making high
dimensionality and high-order the primary computational bottlenecks.

To overcome the speed and memory bottleneck due to the Hessian trace, we can
use the Hutchinson Trace Estimation (HTE) to estimate the trace in the
equation for efficient PINNs since computing Hessian-Vector Product (HVP) is
much faster and more memory-efficient than the full Hessian. The memory cost
of HTE is significantly smaller than that of full PINNs and full Hessian since
the output of HTE is a scalar, which reduces to $O(1)$ memory cost compared to
the $O(d^{2})$ in the full Hessian.

More specifically, the residual loss of the original PINN [44] on a point
$(\boldsymbol{x},t)$ is

$L_{\text{PINN}}(\theta)=\frac{1}{2}\left[\operatorname{Tr}\left(A_{\theta}(\boldsymbol{x},t)\right)+B_{\theta}(\boldsymbol{x},t)\right]^{2},$
(6)

where
$A_{\theta}(\boldsymbol{x},t):=\sigma\sigma^{\mathrm{T}}(\boldsymbol{x},t)\left(\operatorname{Hess}_{\boldsymbol{x}}u\right)(\boldsymbol{x},t)\in\mathbb{R}^{d\times
d}$ and
$B_{\theta}(\boldsymbol{x},t):=\partial_{t}u(\boldsymbol{x},t)+f(\boldsymbol{x},t,u,\nabla_{\boldsymbol{x}}u)\in\mathbb{R}.$

The loss of the HTE replaces the Hessian trace part with a stochastic trace
estimator:

$L_{\text{HTE}}(\theta;\\{\boldsymbol{v}_{i}\\}_{i=1}^{V})=\frac{1}{2}\left(\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\boldsymbol{v}_{i}+B_{\theta}(\boldsymbol{x},t)\right)^{2},$
(7)

where we adopt an HTE with batch size $V$, i.e., we sample
$\\{\boldsymbol{v}_{i}\\}_{i=1}^{V}$ which are $i.i.d.$ samples from the
distribution $p(\boldsymbol{v})$ such that $\mathbb{E}_{\boldsymbol{v}\sim
p(\boldsymbol{v})}[\boldsymbol{v}\boldsymbol{v}^{\mathrm{T}}]=I$, which we
consider the Rademacher distribution to minimize its variance [48].

Although HTE is an unbiased trace estimator, the HTE loss in equation (7) is
biased due to the nonlinear mean square error loss function, which breaks the
linearity of mathematical expectation, i.e., the HTE loss in equation (7) is
biased while it converges to the exact PINN loss almost surely (a.s.) as
$V\rightarrow\infty$. To correct the bias, similar to the techniques in Hu et
al. [23], we are required to sample two sets of samples:

$\displaystyle L_{\text{HTE,
unbiased}}(\theta;\\{\boldsymbol{v}_{i},\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V})$
$\displaystyle=\frac{1}{2}\left(\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\boldsymbol{v}_{i}+B_{\theta}(\boldsymbol{x},t)\right)\left(\frac{1}{V}\sum_{i=1}^{V}\hat{\boldsymbol{v}}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\hat{\boldsymbol{v}}_{i}+B_{\theta}(\boldsymbol{x},t)\right),$
(8)

where $\\{\boldsymbol{v}_{i}\\}_{i=1}^{V}$ and
$\\{\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V}$ are 2$V$ $i.i.d.$ samples from the
distribution $p(\boldsymbol{v})$. The properties of the two loss functions are
summarized in the following theorem.

######  Theorem 3.1.

The loss $L_{\text{HTE}}(\theta)$ in equation (7) converges almost surely
(a.s.) to the exact PINN loss $L_{\text{PINN}}(\theta)$ in equation (6), as
$V\rightarrow\infty$, i.e.,

$\displaystyle\mathbb{P}\left(\lim_{V\rightarrow\infty}L_{\text{HTE}}(\theta;\\{\boldsymbol{v}_{i}\\}_{i=1}^{V})=L_{\text{PINN}}(\theta)\right)=1.$
(9)

The loss $L_{\text{HTE, unbiased}}(\theta)$ in equation (8) is an unbiased
estimator for the exact PINN loss $L_{\text{PINN}}(\theta)$ in equation (6),
i.e.,

$\displaystyle\mathbb{E}_{\\{\boldsymbol{v}_{i},\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V}}\left[L_{\text{HTE,
unbiased}}(\theta;\\{\boldsymbol{v}_{i},\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V})\right]=L_{\text{PINN}}(\theta).$
(10)

###### Proof.

The proof is presented in Appendix A.1. ∎

In HTE, there exists a bias-variance tradeoff as in the Randomized Smoothing
PINNs in [23], where the biased version (equation (7)) involves fewer samples,
resulting in lower variance, while the unbiased version (equation (8)) employs
more samples, leading to excessive variance. In practice, we utilize the
biased version as it proves sufficient. This is attributed to the negligible
bias of HTE in practical applications if $V$ is large enough, whose bias is
much less obvious than the bias in Randomized Smoothing PINN considered in
[23]. HTE’s convergence results are already satisfactory. Indeed, HTE loss
converges to the exact PINN loss, so its bias is negligible if $V$ is large
enough.

To conclude this subsection, we provide an implementation of HTE in a few
lines of code using JAX [8]. Although HTE is mathematically formulated as an
HVP, the implementation details significantly impact its speed and memory
usage. For instance, a straightforward computation of the entire Hessian
followed by a vector multiplication will not yield any efficiency gains.
Below, we show a sample code for highly efficient HTE in JAX Taylor-mode
autodiff, based on fully forward mode [7], which is super fast due to the
absence of backward mode autodiff.

Specifically, given an input function fn, we return its HTE,
$\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}(\text{Hess
fn})\boldsymbol{v}_{i}$, to estimate its Hessian trace.

###  3.3 Comparison with SDGD: Variance and Applicability

Recently, the related work of SDGD [22] has demonstrated robust capabilities
in tackling high-dimensional PDEs and the curse-of-dimensionality. SDGD
primarily relies on sampling dimensions to mitigate the computational and
memory burdens associated with increasing dimensions. In contrast, HTE employs
HVP to avoid the challenges posed by computing the full Hessian, whose cost
grows quadratically with dimensionality. In this subsection, we will compare
their variances and discuss their similarity and applicability.

####  3.3.1 SDGD as a Special Case of HTE

We show that SDGD can be regarded as a special case of HTE. Consider the
discrete distribution $p(\boldsymbol{v})$ such that
$\boldsymbol{v}=\sqrt{d}\boldsymbol{e}_{i}$ for $i=1,2,\cdots,d$ with
probability $1/d$, where $\boldsymbol{e}_{i}$ denotes the standard basis
vector with a one in the $i$th coordinate and zero elsewhere. Then,
$\mathbb{E}_{\boldsymbol{v}\sim
p(\boldsymbol{v})}[\boldsymbol{v}\boldsymbol{v}^{\mathrm{T}}]=I$, and the
corresponding estimator is $\operatorname{Tr}(A)\approx\frac{d}{|I|}\sum_{i\in
I}A_{ii}$ where $I\subset\\{1,2,\cdots,d\\}$ is a multiset and $|I|$ denotes
its cardinality, which is exactly sampling dimensionality as SDGD. Despite the
theoretical similarity, SDGD differs in the following aspects:

  * •

To minimize gradient variance, SDGD employs sampling without replacement,
meaning each dimension is selected at most once. On the other hand, if we
consider SDGD as a special case of HTE, it uses sampling with resampling,
i.e., a dimension can be sampled several times; thus, the sampled dimension
indices, denoted as $I$, form a multiset instead of a set in the original SDGD
formulation.

  * •

While SDGD can be considered a special case of HTE, implementing SDGD using
HTE is inefficient. This is because the probability distribution in this
context takes values in a sparse vector, where most elements are zero.
Optimization for sparse vectors can be applied in this scenario. However, HTE
is generally implemented using non-sparse vectors, and this advantage cannot
be fully exploited in this specific case.

####  3.3.2 Variance Comparison between HTE and SDGD

Both SDGD and HTE can speed up and scale up PINNs in high-dimensional and
high-order PDEs. However, their source of variance differs due to various
sampling and estimations. Specifically, we focus on the second-order parabolic
PDE given in equation (5), and SDGD and HTE estimate the trace of 2nd-order
derivatives $\operatorname{Tr}(A)=\sum_{i=1}^{d}A_{ii},$ where
$A=\sigma\sigma^{\mathrm{T}}(\boldsymbol{x},t)\left(\operatorname{Hess}_{\boldsymbol{x}}u\right)(\boldsymbol{x},t)$.
For simplicity, we drop $A$’s dependency on $\theta,\boldsymbol{x},t$. In the
following theorems and their corresponding remarks, we prove the variance of
SDGD and HTE and explain their various sources.

######  Theorem 3.2.

Given the SDGD estimator of the trace
$\operatorname{Tr}(A)\approx\frac{d}{B}\sum_{i\in I}A_{ii}$ given an index set
$I\subset\\{1,2,\cdots,d\\}$ whose cardinality $|I|=B$, where $B$ is the
SDGD’s batch size for the dimension, its variance is

$\displaystyle\mathbb{V}\left[\frac{d}{B}\sum_{i\in
I}A_{ii}\right]=\frac{1}{\binom{d}{B}}\left(\sum_{I:|I|=B}\frac{d}{B}\sum_{i\in
I}A_{ii}-\sum_{i=1}^{d}A_{ii}\right)^{2},$ (11)

where $\binom{d}{B}=\frac{d(d-1)\cdots(d-B+1)}{B(B-1)\cdots 1}$.

###### Proof.

The proof is presented in Appendix A.3. ∎

######  Remark 1.

SDGD’s variance comes from the variance between diagonal elements across
different dimensions.

######  Theorem 3.3.

Given the HTE estimator of the trace
$\operatorname{Tr}(A)\approx\frac{1}{V}\sum_{k=1}^{V}\boldsymbol{v}_{k}^{\mathrm{T}}A\boldsymbol{v}_{k},$
where $V$ is the HTE batch size and each dimension of
$\boldsymbol{v}_{k}\in\mathbb{R}^{d}$ is an $i.i.d.$ sample from the
Rademacher distribution, then its variance is $\frac{1}{V}\sum_{i\neq
j}A_{ij}^{2}.$

###### Proof.

The proof is presented in Appendix A.4. ∎

######  Remark 2.

HTE’s variance only comes from the off-diagonal elements.

Based on this observation, the following statements can be inferred:

  * •

If the diagonal elements are similar, i.e., the PDE exact solution is
symmetric, then SDGD has low variance.

  * •

If the off-diagonal elements of the Hessian are zero, i.e., different
dimensions do not interact, then HTE is exact.

  * •

If the scales of off-diagonal elements are much larger than the diagonal ones,
then HTE suffers from huge variance.

Hence, we propose the following examples to showcase when HTE or SDGD can
outperform each other and when they perform similarly:

  * •

SDGD fails but HTE is accurate. Consider the 2D exact solution
$f(x,y)=-kx^{2}+ky^{2}$ and its Laplacian $\Delta f(x,y)=0$, where
$k\in\mathbb{R}^{+}$ is large enough. For SDGD, we choose the batch size of
dimension as 1. Thus, the estimator of SDGD will either be
$\frac{\partial^{2}f(x,y)}{\partial x^{2}}=-2k$ or
$\frac{\partial^{2}f(x,y)}{\partial y^{2}}=2k$, whose variance is $4k^{2}$.
For HTE, since the off-diagonal elements in the Hessian matrix of $f$ are
zero, HTE is exact, i.e., has zero variance. In brief, for this case, SDGD
suffers from the variance $4k^{2}$ when $k$ is large, while HTE’s variance is
zero.

  * •

HTE fails but SDGD is accurate. Consider the 2D exact solution $f(x,y)=kxy$
and its Laplacian $\Delta f(x,y)=0$, where $k\in\mathbb{R}^{+}$ is large
enough. Since the diagonal of its Hessian matrix is all zero, SDGD is exact,
i.e., SDGD has zero variance since $\frac{\partial^{2}f(x,y)}{\partial
x^{2}}=\frac{\partial^{2}f(x,y)}{\partial y^{2}}=0$ is exact. For HTE, suppose
that its batch size is $V=1$. Then HTE’s estimator is
$2v_{1}v_{2}\frac{\partial^{2}f(x,y)}{\partial x\partial y}=2v_{1}v_{2}k$
where $v_{1}$ and $v_{2}$ follows the Rademacher distribution. Thus, HTE’s
variance is $4k^{2}$. In brief, for this case, HTE suffers from the variance
$4k^{2}$ when $k$ is large, while SDGD’s variance is zero and SDGD is exact.

  * •

HTE and SDGD have the same nonzero variance. Consider the 2D exact solution
$f(x,y)=k(-x^{2}+y^{2}+xy)$ and its Laplacian $\Delta f(x,y)=0$. For SDGD, we
choose the batch size of dimension as 1. Thus, the estimator of SDGD will
either be $\frac{\partial^{2}f(x,y)}{\partial x^{2}}=-2k$ or
$\frac{\partial^{2}f(x,y)}{\partial y^{2}}=2k$, whose variance is $4k^{2}$.
For HTE, suppose that its batch size is $V=1$. Then HTE’s estimator is
$2v_{1}v_{2}\frac{\partial^{2}f(x,y)}{\partial x\partial y}=2v_{1}v_{2}k$,
where $v_{1}$ and $v_{2}$ follows the Rademacher distribution. Thus, HTE’s
variance is $4k^{2}$. In brief, HTE and SDGD have the same nonzero variance
$4k^{2}$.

However, we emphasize that the specific use case can be much more complex.
Since we are applying SDGD or HTE to the neural network at each iteration, the
characteristics of the function represented by the network determine which
method is superior. This, in turn, is challenging to observe directly.
Therefore, we provide these insights to readers as a basis for consideration.
In practical applications, SDGD and HTE algorithms can be judiciously chosen
based on some a priori knowledge of the PDE problem and the true solution. One
can explore the simultaneous application of two methods and observe the rate
of training loss reduction for each. Given the strong correlation between the
error in approximating the true solution by PINN and the training loss [21],
algorithms with lower variance are more likely to converge faster.

####  3.3.3 Applicability Comparison between HTE and SDGD

Next, we move to discuss the applicability of HTE compared to SDGD. HTE is
more beneficial in higher dimensions and higher orders; otherwise, regular
autograd methods may already be sufficient. For instance, consider the extreme
case of one dimension. In this scenario, the Hessian is just a scalar.
Consequently, the computational cost and memory consumption for both HVP and
the full Hessian are exactly the same, hence providing no advantage for HTE.
Another example will be high-dimensional but low-order; while the HTE approach
can be applied to first-order PDEs, i.e., the norm square of the Jacobian can
be unbiasedly estimated by Jacobian-vector product (JVP):

$\|\nabla_{\boldsymbol{x}}u(\boldsymbol{x})\|^{2}=\mathbb{E}_{v\sim
p(v)}|v^{\mathrm{T}}\nabla_{\boldsymbol{x}}u(\boldsymbol{x})|^{2},$ (12)

experiments indicate that memory savings are insignificant. This term
$\|\nabla_{\boldsymbol{x}}u(\boldsymbol{x})\|^{2}$ is ubiquitous in scientific
computing, e.g., the first-order Hamilton-Jacobi equation, and the Schrödinger
equation in the variation form. (see section 5.4 of SDGD [22] where SDGD
scales up tensor neural networks for the Schrödinger equation).

More generally, from a mathematical perspective, for a $d$-dimensional $n$th-
order PDE, the pure autograd computation results in a tensor of dimensions
$d^{n}$. For example, the Jacobian is of dimension $d$, and the Hessian is of
dimension $d^{2}$. In contrast, the output of HTE are $V$ scalars when using
$V$ as the batch size; thus, the output dimension is $V$. Consequently, as the
dimension and order increase, the advantages of HTE become more pronounced.

###  3.4 Applications to Biharmonic Equations

Given that HTE demonstrates superior performance for high-order and high-
dimensional PDEs, we illustrate its application to the biharmonic equation in
this subsection. The biharmonic equation is a fourth-order PDE that can be
defined in any dimension. We will demonstrate how to extend HTE, involving the
generalization of HVP to a tensor-vector product, and provide a JAX pseudo-
code for efficient implementation.

Concretely, the fourth-order and $d$-dimensional Biharmonic operator is given
by

$\displaystyle\Delta^{2}u(\boldsymbol{x})=\sum_{i=1}^{d}\sum_{j=1}^{d}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}^{2}\partial\boldsymbol{x}_{j}^{2}}u(\boldsymbol{x}).$
(13)

The biharmonic equation is crucial in understanding and simulating complex
behaviors. It is a cornerstone in analyzing systems governed by higher-order
differential equations, including modeling elastic membranes, thin plates, and
stream functions in fluid dynamics. We will attempt to generalize HTE,
starting with the extension of HVP to a tensor-vector product (TVP), to solve
high-dimensional and high-order Biharmonic equations efficiently.

######  Theorem 3.4.

The biharmonic operator can be unbiasedly estimated using the following TVPs:

$\displaystyle\Delta^{2}u(\boldsymbol{x})$
$\displaystyle=\frac{1}{3}\mathbb{E}_{\boldsymbol{v}\sim\mathcal{N}(0,I)}\left[\frac{\partial^{4}}{\partial\boldsymbol{x}^{4}}u(\boldsymbol{x})[\boldsymbol{v},\boldsymbol{v},\boldsymbol{v},\boldsymbol{v}]\right],$
(14)

where $\mathcal{N}(0,I)$ is the $d$-dimensional unit Gaussian, and

$\frac{\partial^{4}}{\partial\boldsymbol{x}^{4}}u(\boldsymbol{x})\in\mathbb{R}^{d\times
d\times d\times d},\quad\text{where
}\left[\frac{\partial^{4}}{\partial\boldsymbol{x}^{4}}u(\boldsymbol{x})\right]_{ijkl}=\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}\partial\boldsymbol{x}_{j}\partial\boldsymbol{x}_{k}\partial\boldsymbol{x}_{l}}u(\boldsymbol{x}),$
(15)

and the TVP is defined as

$\frac{\partial^{4}}{\partial\boldsymbol{x}^{4}}u(\boldsymbol{x})[\boldsymbol{v},\boldsymbol{v},\boldsymbol{v},\boldsymbol{v}]=\sum_{i,j,k,l=1}^{d}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}\partial\boldsymbol{x}_{j}\partial\boldsymbol{x}_{k}\partial\boldsymbol{x}_{l}}u(\boldsymbol{x})\boldsymbol{v}_{i}\boldsymbol{v}_{j}\boldsymbol{v}_{k}\boldsymbol{v}_{l},$
(16)

###### Proof.

The proof is presented in Appendix A.2. ∎

Therefore, we can also obtain an unbiased estimate of the biharmonic operator
through sampling. To help the reader understand the implementation of the
tensor vector product, we provide a pseudocode implementation based on JAX [8]
Taylor Mode automatic differentiation [7]. We want to emphasize again that the
implementation details significantly impact the speed and memory usage. A
straightforward computation of the entire fourth-order derivative tensor
followed by the vector multiplication will not yield any efficiency gains.
Rather, we should leverage Taylor-mode automatic differentiation to allow JAX
to implicitly compute the Tensor-Vector Product (TVP), thereby avoiding the
massive memory consumption associated with explicit calculations. Furthermore,
internal optimizations within JAX can be exploited to enhance speed by
directly invoking JAX’s TVP functionality and extracting only the output TVP,
as opposed to the entire fourth-order derivative tensor.

Similar to our previous treatment of the HTE for Hessian trace estimation, we
analyze the HTE’s variance when estimating the biharmonic operator. Here, we
employ a Gaussian distribution instead of a Rademacher distribution, and thus
both diagonal and off-diagonal elements contribute to a certain extent of
variance. Consequently, a larger HTE batch size, denoted as $V$, is required
when estimating the biharmonic operator.

##  4 Computational Experiments

###  4.1 Nonlinear Second-Order PDEs with Anisotropic and Inseparable
Solutions

This subsection considers nonlinear second-order Sine-Gordon PDEs with
anisotropic and inseparable solutions. The following complicated exact
solution is considered in related work [22, 23] to test the convergence of
high-dimensional PDE solvers:

$u_{\text{exact}}(\boldsymbol{x})=\left(1-\|\boldsymbol{x}\|_{2}^{2}\right)\left(\sum_{i=1}^{d-1}c_{i}\sin(\boldsymbol{x}_{i}+\cos(\boldsymbol{x}_{i+1})+\boldsymbol{x}_{i+1}\cos(\boldsymbol{x}_{i}))\right),$
(17)

where $c_{i}\sim\mathcal{N}(0,1)$. The term $1-\|\boldsymbol{x}\|_{2}^{2}$ is
multiplied for a zero boundary condition on the unit ball to prevent the
boundary loss from leaking the solution’s information. The solution’s $i$th
and $(i+1)$th dimension correlate, also known as the two-body interaction.

Based on the idea, we further introduce the following complicated exact
solutions to test various algorithms:

$u_{\text{exact}}(\boldsymbol{x})=\left(1-\|\boldsymbol{x}\|_{2}^{2}\right)\left(\sum_{i=1}^{d-2}c_{i}\exp(\boldsymbol{x}_{i}\boldsymbol{x}_{i+1}\boldsymbol{x}_{i+2})\right),$
(18)

where the three-body interaction between each pair
$(\boldsymbol{x}_{i},\boldsymbol{x}_{i+1},\boldsymbol{x}_{i+2})$ is
considered.

As for the nonlinear PDE, the Sine-Gordon equation with zero boundary
condition is considered

$\displaystyle\Delta
u(\boldsymbol{x})+\sin\left(u(\boldsymbol{x})\right)=g(\boldsymbol{x}),\quad\boldsymbol{x}\in\mathbb{B}^{d}=\\{\boldsymbol{x}\
|\ \|\boldsymbol{x}\|_{2}<1\\},$ (19) $\displaystyle
u(\boldsymbol{x})=0,\quad\boldsymbol{x}\in\mathbb{S}^{d-1}=\\{\boldsymbol{x}\
|\ \|\boldsymbol{x}\|_{2}=1\\},$ (20)

where $g(\boldsymbol{x})=\Delta
u_{\text{exact}}(\boldsymbol{x})+\sin\left(u_{\text{exact}}(\boldsymbol{x})\right)$,
since its nonlinearity order is infinity, i.e., its nonlinearity is strong.

Here are the implementation details. The model is a 4-layer fully connected
network with 128 hidden units, which is trained via Adam [37] for 10K epochs
in the two-body case and 20K epochs in the three-body case, respectively, with
an initial learning rate 1e-3, which linearly decays to zero at the end of the
optimization. We select 100 random residual points at each Adam epoch and 20K
fixed testing points uniformly from the unit ball. We adopt the following
model structure to satisfy the zero boundary condition with hard constraint
and to avoid the boundary loss [38]
$(1-\|\boldsymbol{x}\|_{2}^{2})u_{\theta}(\boldsymbol{x}),$ where
$u_{\theta}(\boldsymbol{x})$ is the neural neural network model. We repeat our
experiment five times with five independent random seeds and report the mean
and standard deviation of the errors. For HTE, we choose $V=16$. For SDGD, we
chose the batch size of dimension 16.

Method | Metric | 100 D | 1,000 D | 5,000D | 10,000 D | 100,000 D
---|---|---|---|---|---|---
PINNs [45] | Speed | 1367.61it/s | 141.45it/s | 8.43it/s | N.A. | N.A.
Memory | 1123MB | 2915MB | 14283MB |  $>$80GB |  $>$80GB
Error_1 |  6.24E-3$\pm$2.83E-3 |  1.20E-3$\pm$6.47E-4 |  2.64E-3$\pm$1.33E-3 | N.A. | N.A.
Error_2 |  7.43E-3$\pm$4.56E-4 |  1.00E-3$\pm$1.45E-5 |  2.02E-4$\pm$2.45E-6 | N.A. | N.A.
SDGD [22] | Speed | 1853.67it/s | 1444.00it/s | 1243.91it/s | 1027.38it/s | 356.62it/s
Memory | 883MB | 885MB | 909MB | 929MB | 1173MB
Error_1 |  6.28E-3$\pm$2.55E-3 |  1.55E-03$\pm$6.72E-4 |  2.63E-3$\pm$1.21E-3 |  1.93E-3$\pm$7.81E-4 |  2.30E-3$\pm$1.33E-3
| Error_2 |  7.60E-3$\pm$6.00E-4 |  1.15E-3$\pm$2.80E-4 |  2.05E-4$\pm$3.98E-6 |  1.04E-4$\pm$4.30E-6 |  1.56E-5$\pm$9.57E-6
HTE (Ours) | Speed | 1792.03it/s | 1566.19it/s | 1275.81it/s | 1066.10it/s | 345.10it/s
Memory | 869MB | 890MB | 913MB | 927MB | 1089MB
Error_1 |  6.30E-3$\pm$2.88E-3 |  1.25E-3$\pm$3.40E-4 |  2.61E-3$\pm$1.31E-3 |  1.84E-3$\pm$9.56E-4 |  2.38E-3$\pm$1.72E-3
| Error_2 |  7.58E-3$\pm$5.39E-4 |  9.94E-4$\pm$3.23E-6 |  2.01E-4$\pm$2.27E-6 |  1.02E-4$\pm$3.88E-6 |  1.59E-5$\pm$9.29E-6
Table 1: Computational results for Sine-Gordon equations with the two-body
(Error_1) and three-body (Error_2) exact solutions.

The computational results for the two-body and three-body cases are presented
in Table 1, where we compare vanilla PINN [45], SDGD [22], and HTE (Ours), in
terms of speed (iteration per second), memory cost, and relative $L_{2}$ error
for the two-body case (Error_1) and that for the three body case (Error_2).
Note that when solving for the two PDEs with different solutions, the speed
and the memory costs are similar, and only the error differs, so we only
report speed and memory once.

Regarding the regular PINNs, we observe a significant slowdown in speed and a
rapid increase in memory consumption as the dimensionality rises. This is
because the regular PINNs require the computation of the full Hessian, and the
computational cost quadratically increases with dimension. As a result,
regular PINNs exceed the memory limit of an A100 GPU in the case of 10K
dimensions. For HTE and SDGD, we observe similar convergence speeds, results
and memory consumption. In scenarios with 100D, 1K D, and 5K D, the final
errors are comparable to regular PINNs, demonstrating that HTE can overcome
the curse-of-dimensionality. It accelerates regular PINNs without compromising
its effectiveness. Moreover, the PDEs we used cannot be simplified into lower-
dimensional subproblems as they are inherently high-dimensional. We also
conducted tests on the PDE exact solutions of two-body and three-body
interactions to demonstrate the efficacy of HTE across various high-
dimensional problems.

###  4.2 High-Dimensional Gradient-Enhanced PINN

The Gradient-Enhanced Physics-Informed Neural Network (gPINN) [55] is a
technique designed to enhance the performance of PINNs. It achieves this by
incorporating additional regularization through the computation of an extra
first-order derivative for the PINN’s residual loss. However, due to the extra
derivatives, gPINN becomes computationally expensive, particularly in higher
dimensions, where the required additional derivatives are proportional to the
dimensionality of the PDE. In this subsection, we illustrate how to leverage
the proposed HTE method to expedite the training of gPINN. To illustrate, we
continue to consider the Sine-Gordon equation with a two-body interactive
solution discussed in the preceding section in equation (17).

Given the Sine-Gordon equation

$\displaystyle\Delta
u(\boldsymbol{x})+\sin\left(u(\boldsymbol{x})\right)=g(\boldsymbol{x}),\quad\boldsymbol{x}\in\mathbb{B}^{d}=\\{\boldsymbol{x}\
|\ \|\boldsymbol{x}\|_{2}<1\\},$ (21) $\displaystyle
u(\boldsymbol{x})=0,\quad\boldsymbol{x}\in\mathbb{S}^{d-1}=\\{\boldsymbol{x}\
|\ \|\boldsymbol{x}\|_{2}=1\\},$ (22)

where $g(\boldsymbol{x})=\Delta
u_{\text{exact}}(\boldsymbol{x})+\sin\left(u_{\text{exact}}(\boldsymbol{x})\right)$,
the PINN loss is given by

$\displaystyle
L_{\text{PINN}}(\theta)=\frac{1}{2}\left(u_{\theta}(\boldsymbol{x})+\sin\left(u_{\theta}(\boldsymbol{x})\right)-g(\boldsymbol{x})\right)^{2}:=\frac{1}{2}r_{\theta}(\boldsymbol{x})^{2},$
(23)

where $r_{\theta}(\boldsymbol{x})=\Delta
u_{\theta}(\boldsymbol{x})+\sin\left(u_{\theta}(\boldsymbol{x})\right)-g(\boldsymbol{x})$
is the residual prediction. The gPINN loss is given by

$\displaystyle
L_{\text{gPINN}}(\theta)=\frac{1}{2}r_{\theta}(\boldsymbol{x})^{2}+\frac{1}{2}\lambda_{\text{gPINN}}\left\|\nabla_{\boldsymbol{x}}r_{\theta}(\boldsymbol{x})\right\|^{2},$
(24)

i.e., gPINN optimizes both the residual part and its derivative to be zero,
where $\lambda_{\text{gPINN}}\in\mathbb{R}^{+}$ is the weight for the gPINN
regularization. To accelerate gPINN, we use the HTE-based gPINN loss

$\displaystyle L_{\text{HTE,
gPINN}}(\theta)=\frac{1}{2}\hat{r}_{\theta}(\boldsymbol{x})^{2}+\frac{1}{2}\lambda_{\text{gPINN}}\left\|\nabla_{\boldsymbol{x}}\hat{r}_{\theta}(\boldsymbol{x})\right\|^{2},$
(25)

where
$\hat{r}_{\theta}(\boldsymbol{x})=\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}\left(\operatorname{Hess}u_{\theta}(\boldsymbol{x})\right)\boldsymbol{v}_{i}+\sin\left(u_{\theta}(\boldsymbol{x})\right)-g(\boldsymbol{x})$
is the HTE-based residual prediction. Thus, we are only required to take an
additional derivative to the HVP of the neural network model, which is much
more efficient than that of the full Hessian.

The implementation details are given as follows. The model is a 4-layer fully
connected network with 128 hidden units, which is trained via Adam [37] for
10K epochs with an initial learning rate 1e-3, which linearly decays to zero
at the end of the optimization. We select 100 random residual points at each
Adam epoch and 20K fixed testing points uniformly from the unit ball. We adopt
the following model structure to satisfy the zero boundary condition with hard
constraint and to avoid the boundary loss [38]
$(1-\|\boldsymbol{x}\|_{2}^{2})u_{\theta}(\boldsymbol{x}),$ where
$u_{\theta}(\boldsymbol{x})$ is the neural neural network model. We repeat our
experiment five times with five independent random seeds and report the mean
and standard deviation of the errors. For HTE, we choose $V=16$ in all
settings. We choose $\lambda_{\text{gPINN}}$ as ten in 100D and 1K D cases and
choose it as ten thousand in 10K D and 100K D cases. The gPINN loss weight
$\lambda_{\text{gPINN}}$ can be chosen to make the scale of the PINN loss and
the scale of the gPINN loss similar during initialization, i.e., choose
$\lambda_{\text{gPINN}}$ such that
${r}_{\theta}(\boldsymbol{x})^{2}\approx\lambda_{\text{gPINN}}\left\|\nabla_{\boldsymbol{x}}{r}_{\theta}(\boldsymbol{x})\right\|^{2}$
and
$\hat{r}_{\theta}(\boldsymbol{x})^{2}\approx\lambda_{\text{gPINN}}\left\|\nabla_{\boldsymbol{x}}\hat{r}_{\theta}(\boldsymbol{x})\right\|^{2}$
at initialization.

Method | Metric | 100D | 1K D | 10K D | 100K D
---|---|---|---|---|---
PINN [45] | Speed | 1367.61it/s | 141.45it/s | N.A. | N.A.
Error | 6.24E-3$\pm$2.83E-3 | 1.20E-3$\pm$6.47E-4 | N.A. | N.A.
gPINN [55] | Speed | 458.84it/s | 51.42it/s | N.A. | N.A.
Error | 4.20E-3$\pm$1.21E-2 | 4.42E-4$\pm$1.84E-4 | N.A. | N.A.
HTE PINN (Ours) | Speed | 1792.03it/s | 1566.19it/s | 1066.10it/s | 345.10it/s
Error | 6.30E-3$\pm$2.88E-3 | 1.25E-3$\pm$3.40E-4 | 1.84E-3$\pm$9.56E-4 | 2.38E-3$\pm$1.72E-3
HTE gPINN (Ours) | Speed | 616.79it/s | 461.81it/s | 229.18it/s | 153.35it/s
Error | 5.08E-3$\pm$2.00E-3 | 4.88E-4$\pm$1.59E-4 | 4.68E-4$\pm$4.61E-4 | 4.77E-5$\pm$4.70E-5
Table 2: Computation results for gPINN.

The computational results are presented in Table 2, where three algorithms are
compared: vanilla PINN [45], vanilla gPINN [55], HTE-based PINN, and HTE-based
gPINN. In the preliminary, we observe that both PINN and gPINN encounter
failure when exceeding the 80GB memory limit of the A100 GPU in dimensions
surpassing 1000D. Conversely, PINN and gPINN based on HTE successfully operate
in dimensions as high as 100K, demonstrating swift performance. Comparing PINN
and gPINN, we note an enhancement in the latter’s performance improvement,
albeit at a slightly slower convergence pace, due to the additional first-
order derivatives of the residual. The advantages of gPINN become more
pronounced in higher dimensions (exceeding 1000D), leading to an additional
order of magnitude reduction in PINN error. However, in scenarios with 100D
and 1000D, HTE-based gPINN performs slightly inferior to conventional gPINN.
This arises because HTE is ultimately an approximation, and in high-
dimensional gPINN, it introduces $d$ additional losses, where $d$ is the
dimensionality. Hence, the approximation with HTE tends to incur more
significant errors. Nevertheless, the performance of HTE-based gPINN remains
comparable to conventional gPINN. In summary, we verify that HTE can
accelerate and scale up gPINN to very high dimensions. Moreover, gPINN
continues to enhance the performance of PINN in high-dimensional PDEs. To the
best of our knowledge, previous work [55] only demonstrated gPINN’s
effectiveness in low-dimensional PDEs, while our work is the first to show
that in high dimensions.

###  4.3 Biharmonic Equation

In this subsection, we further validate the proposed HTE method, demonstrating
its capability to accelerate the convergence of PINNs in high-dimensional and
high-order biharmonic equations while significantly reducing its memory
consumption. The following complicated exact solution is considered to test
the convergence of high-dimensional PDE solvers:

$u_{\text{exact}}(\boldsymbol{x})=\left(1-\|\boldsymbol{x}\|_{2}^{2}\right)\left(4-\|\boldsymbol{x}\|_{2}^{2}\right)\left(\sum_{i=1}^{d-2}c_{i}\exp(\boldsymbol{x}_{i}\boldsymbol{x}_{i+1}\boldsymbol{x}_{i+2})\right),$
(26)

where $c_{i}\sim\mathcal{N}(0,1)$. The term
$(1-\|\boldsymbol{x}\|_{2}^{2})(4-\|\boldsymbol{x}\|_{2}^{2})$ is multiplied
for zero boundary condition to prevent the boundary loss from leaking the
solution’s information.

The Biharmonic equation is given as below

$\displaystyle\Delta^{2}u(\boldsymbol{x})=g(\boldsymbol{x}),\quad\boldsymbol{x}\in\\{\boldsymbol{x}\
|\ 1<\|\boldsymbol{x}\|_{2}<2\\},$ (27) $\displaystyle
u(\boldsymbol{x})=0,\quad\boldsymbol{x}\in\\{\boldsymbol{x}\ |\
\|\boldsymbol{x}\|_{2}=1\\}\cup\\{\boldsymbol{x}\ |\
\|\boldsymbol{x}\|_{2}=2\\},$ (28)

where $g(\boldsymbol{x})=\Delta^{2}u_{\text{exact}}(\boldsymbol{x})$ is given
by the exact solution.

Here are the implementation details. The model is a 4-layer fully connected
network with 128 hidden units, which is trained via Adam [37] for 10K epochs
in 100D, 150D, and 200D cases, and 20K epochs in the 50D case, with an initial
learning rate 1e-3, which linearly decays to zero at the end of the
optimization. We select 100 random residual points at each Adam epoch and 20K
fixed testing points within the domain $\boldsymbol{x}\in\\{\boldsymbol{x}\ |\
1\leq\|\boldsymbol{x}\|_{2}\leq 2\\}$. We adopt the following model structure
to satisfy the zero boundary condition with hard constraint and to avoid the
boundary loss [38]:
$(1-\|\boldsymbol{x}\|_{2}^{2})(4-\|\boldsymbol{x}\|_{2}^{2})u_{\theta}(\boldsymbol{x}),$
where $u_{\theta}(\boldsymbol{x})$ is the neural neural network model. We
repeat our experiment five times with five independent random seeds and report
the mean and standard deviation of the errors. For HTE, we control the batch
size $V=16/512/1024$ to see how $V$ affects convergence.

Method | Metric | 50D | 100D | 150D | 200D
---|---|---|---|---|---
PINN [45] | Speed | 19.82it/s | 6.56it/s | 3.50it/s | N.A.
Memory | 6199MB | 21453MB | 44631MB |  $>$80GB
Error | 2.34E-2$\pm$1.91E-2 | 6.01E-3$\pm$2.32E-3 | 3.11E-3$\pm$1.79E-3 | N.A.
HTE ($V=16$) | Speed | 438.79it/s | 109.57it/s | 54.86it/s | 25.89it/s
Memory | 991MB | 2085MB | 3589MB | 7225MB
Error | 2.36E-2$\pm$1.97E-2 | 6.83E-3$\pm$1.63E-3 | 3.91E-3$\pm$1.47E-3 | 4.40E-3$\pm$5.65E-3
HTE ($V=512$) | Speed | 123.98it/s | 54.07it/s | 42.44it/s | 23.50it/s
Memory | 1861MB | 2089MB | 3593MB | 7229MB
Error | 2.34E-2$\pm$1.96E-2 | 6.08E-3$\pm$2.28E-3 | 3.39E-3$\pm$1.99E-3 | 4.03E-3$\pm$4.85E-3
HTE ($V=1024$) | Speed | 78.50it/s | 34.72it/s | 33.77it/s | 20.36it/s
Memory | 2535MB | 3037MB | 3595MB | 7231MB
Error | 2.34E-2$\pm$1.96E-2 | 6.01E-3$\pm$2.35E-3 | 3.13E-3$\pm$1.82E-3 | 3.97E-3$\pm$4.91E-3
Table 3: Computational results for the Biharmonic equation.

The computational results for the biharmonic equation are presented in Table
3, where we report the speed, memory, and relative $L_{2}$ error of the
baseline PINN and our proposed HTE method. Regarding the regular PINN, we
observe a significant slowdown in speed and a rapid increase in memory
consumption as the dimensionality rises. This is because the regular PINNs
require the computation of the full biharmonic operator, and the computational
cost increases to the fourth power with respect to the dimension. As a result,
regular PINNs exceed the memory limit of an A100 GPU in the case of 200
dimensions. In contrast, the second-order PDE case requires 10K dimensions for
PINN to go out-of-memory, signifying that the biharmonic operator is much more
costly than the Hessian due to the high-order derivative. For HTE, we opted
for a batch size of $V=512/1024$ to make HTE’s convergence results similar to
those of PINN, which is in contrast to our earlier example with a second-order
PDE, where we used a batch size of $V=16$. If we just use a small batch size
$V=16$ in the biharmonic equation, then HTE will perform worse than PINN. The
reason behind this is our variance analysis, indicating that when estimating
the biharmonic operator, both diagonal and off-diagonal elements contribute to
the variance, resulting in relatively larger errors than estimating the
Hessian trace. Nevertheless, HTE demonstrates substantial improvements in
terms of speed and memory consumption. In scenarios exceeding 50 dimensions,
its speed is nearly ten times that of regular PINN, and the rate of memory
consumption growth is significantly lower than that of PINN.

##  5 Summary

In this paper, we introduced the Hutchinson Trace Estimation (HTE) method to
enhance the range of capabilities of Physics-Informed Neural Networks (PINNs)
in tackling high-dimensional and high-order partial differential equations
(PDEs). While PINNs have succeeded in low-dimensional problems, their
application to complex PDEs has faced challenges due to computational
bottlenecks in automatic differentiation. HTE, which is applied to calculate
Hessian vector products (HVPs), has emerged as an effective solution to
accelerate PINNs and reduce memory costs.

We first validated our approach on second-order parabolic equations,
demonstrating the convergence of the PINN loss with HTE to the original PINNs
loss under certain conditions, guaranteeing the effectiveness of the HTE
approximation. Comparisons with Stochastic Dimension Gradient Descent (SDGD)
highlighted the superior performance of HTE in scenarios with significant
dimensionality variability and variance, where we delved into the theoretical
comparison between the two methods by proving their corresponding variance.
The extension of HTE to the biharmonic equation showcased its efficiency in
handling higher-order, high-dimensional PDEs by employing tensor-vector
products (TVPs) to transform computations.

Experimental setups confirmed the effectiveness of HTE, revealing comparable
convergence effects with SDGD under memory and speed constraints. Furthermore,
HTE proved valuable in accelerating Gradient-Enhanced PINN (gPINN), which
requires an additional first-order derivative over the PINNs residual and the
fourth-order high dimensional biharmonic equation for better computational
efficiency.

In conclusion, our proposed HTE method addresses the limitations of PINNs in
handling high-dimensional and high-order PDEs, providing a versatile and
efficient approach. Our work contributes to the scientific machine learning
field by presenting a unified framework for solving general high-dimensional
PDEs, demonstrating the effectiveness of HTE, and discussing its advantages
over existing methods.

## Acknowledgement

This research of ZH, ZS, and KK is partially supported by the National
Research Foundation Singapore under the AI Singapore Programme (AISG Award No:
AISG2-TC-2023-010-SGIL) and the Singapore Ministry of Education Academic
Research Fund Tier 1 (Award No: T1 251RES2207). The work of GEK was supported
by the MURI-AFOSR FA9550-20-1-0358 projects and by the DOE SEA-CROGS project
(DE-SC0023191). GEK was also supported by the ONR Vannevar Bush Faculty
Fellowship (N00014-22-1-2795).

##  Appendix A Proof

###  A.1 Proof of Theorem 3.1

Convergence of $L_{\text{HTE}}(\theta;\\{\boldsymbol{v}_{i}\\}_{i=1}^{V})$.
Due to the strong law of large numbers,

$\mathbb{P}\left(\lim_{V\rightarrow\infty}\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\boldsymbol{v}_{i}=\operatorname{Tr}\left(A_{\theta}(\boldsymbol{x},t)\right)\right)=1.$
(29)

Thus,

$\mathbb{P}\left(\lim_{V\rightarrow\infty}\left(\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\boldsymbol{v}_{i}+B_{\theta}(\boldsymbol{x},t)\right)^{2}=\left(\operatorname{Tr}\left(A_{\theta}(\boldsymbol{x},t)\right)+B_{\theta}(\boldsymbol{x},t)\right)^{2}\right)=1.$
(30)

Unbiasedness of $L_{\text{HTE,
unbiased}}(\theta;\\{\boldsymbol{v}_{i},\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V})$.
Since HTE is an unbiased trace estimator,

$\displaystyle\quad\mathbb{E}_{\\{\boldsymbol{v}_{i},\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V}}\left[L_{\text{HTE,
unbiased}}(\theta;\\{\boldsymbol{v}_{i},\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V})\right]$
(31)
$\displaystyle=\frac{1}{2}\mathbb{E}_{\\{\boldsymbol{v}_{i}\\}_{i=1}^{V}}\left[\frac{1}{V}\sum_{i=1}^{V}\boldsymbol{v}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\boldsymbol{v}_{i}+B_{\theta}(\boldsymbol{x},t)\right]\mathbb{E}_{\\{\hat{\boldsymbol{v}}_{i}\\}_{i=1}^{V}}\left[\frac{1}{V}\sum_{i=1}^{V}\hat{\boldsymbol{v}}_{i}^{\mathrm{T}}A_{\theta}(\boldsymbol{x},t)\hat{\boldsymbol{v}}_{i}+B_{\theta}(\boldsymbol{x},t)\right]$
$\displaystyle=\frac{1}{2}\left[\operatorname{Tr}\left(A_{\theta}(\boldsymbol{x},t)\right)+B_{\theta}(\boldsymbol{x},t)\right]^{2}$
$\displaystyle=L_{\text{PINN}}(\theta).$

###  A.2 Proof of Theorem 3.4

For the unit Gaussian distribution, we have that

$\displaystyle\mathbb{E}_{\boldsymbol{v}\sim\mathcal{N}(0,I)}\left[\frac{\partial^{4}}{\partial\boldsymbol{x}^{4}}u(\boldsymbol{x})[\boldsymbol{v},\boldsymbol{v},\boldsymbol{v},\boldsymbol{v}]\right]$
$\displaystyle=\mathbb{E}_{\boldsymbol{v}\sim\mathcal{N}(0,I)}\left[\sum_{i,j,k,l=1}^{d}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}\partial\boldsymbol{x}_{j}\partial\boldsymbol{x}_{k}\partial\boldsymbol{x}_{l}}u(\boldsymbol{x})\boldsymbol{v}_{i}\boldsymbol{v}_{j}\boldsymbol{v}_{k}\boldsymbol{v}_{l}\right]$
(32)
$\displaystyle=\mathbb{E}_{\boldsymbol{v}\sim\mathcal{N}(0,I)}\left[\sum_{i=1}^{d}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}^{4}}u(\boldsymbol{x})\boldsymbol{v}_{i}^{4}\right]+6\mathbb{E}_{\boldsymbol{v}\sim\mathcal{N}(0,I)}\left[\sum_{i\neq
j}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}^{2}\partial\boldsymbol{x}_{j}^{2}}u(\boldsymbol{x})\boldsymbol{v}_{i}^{2}\boldsymbol{v}_{j}^{2}\right]$
$\displaystyle=3\sum_{i=1}^{d}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}^{4}}u(\boldsymbol{x})+6\sum_{i\neq
j}\frac{\partial^{4}}{\partial\boldsymbol{x}_{i}^{2}\partial\boldsymbol{x}_{j}^{2}}u(\boldsymbol{x})$
$\displaystyle=3\Delta^{2}u(\boldsymbol{x}).$

###  A.3 Proof of Theorem 3.2

SDGD’s trace estimator is $\operatorname{Tr}(A)\approx\frac{d}{|I|}\sum_{i\in
I}A_{ii}.$ Suppose that the batch size of the sampled dimension in SDGD is
$B$; then the variance can be directly computed:

$\displaystyle\mathbb{V}\left[\frac{d}{B}\sum_{i\in
I}A_{ii}\right]=\frac{1}{\binom{d}{B}}\left(\sum_{I:|I|=B}\frac{d}{B}\sum_{i\in
I}A_{ii}-\sum_{i=1}^{d}A_{ii}\right)^{2},$ (33)

where $\binom{d}{B}=\frac{d(d-1)\cdots(d-B+1)}{B(B-1)\cdots 1}$.

###  A.4 Proof of Theorem 3.3

The estimator of HTE is
$\operatorname{Tr}(A)\approx\frac{1}{V}\sum_{k=1}^{V}\boldsymbol{v}_{k}^{\mathrm{T}}A\boldsymbol{v}_{k}=\sum_{i=1}^{d}A_{ii}+\frac{1}{V}\sum_{i\neq
j}A_{ij}\sum_{k=1}^{V}\boldsymbol{v}_{k,i}\boldsymbol{v}_{k,j},$ where
$\boldsymbol{v}_{k,i}$ denotes the $i$th dimension of the vector
$\boldsymbol{v}_{k}\in\mathbb{R}^{d}$ and we are using $p(\boldsymbol{v})$ as
the Rademacher distribution.

HTE’s variance can be computed as follows:

$\displaystyle\left(\boldsymbol{v}^{\mathrm{T}}A\boldsymbol{v}-\operatorname{Tr}(A)\right)^{2}$
$\displaystyle=\left(\sum_{i\neq
j}A_{ij}\boldsymbol{v}_{i}\boldsymbol{v}_{j}\right)^{2}$ (34)
$\displaystyle=\sum_{i\neq j}\sum_{k\neq
l}A_{ij}A_{kl}\boldsymbol{v}_{i}\boldsymbol{v}_{j}\boldsymbol{v}_{k}\boldsymbol{v}_{l}.$

If $i=k$ and $j=l$, then
$\mathbb{E}[\boldsymbol{v}_{i}\boldsymbol{v}_{j}\boldsymbol{v}_{k}\boldsymbol{v}_{l}]=\mathbb{E}[v_{i}^{2}]\mathbb{E}[v_{j}^{2}]=1$,
else equals zero. Therefore,

$\displaystyle\mathbb{V}\left[\boldsymbol{v}^{\mathrm{T}}A\boldsymbol{v}\right]$
$\displaystyle=\mathbb{E}\left[\left(\boldsymbol{v}^{\mathrm{T}}A\boldsymbol{v}-\operatorname{Tr}(A)\right)^{2}\right]=\sum_{i\neq
j}A_{ij}^{2}.$ (35)

With the HTE estimator of batch size $V$, the variance is

$\mathbb{V}\left[\frac{1}{V}\sum_{k=1}^{V}\boldsymbol{v}_{k}^{\mathrm{T}}A\boldsymbol{v}_{k}\right]=\frac{1}{V}\sum_{i\neq
j}A_{ij}^{2}.$ (36)

## References

  * [1] Christian Beck, Sebastian Becker, Patrick Cheridito, Arnulf Jentzen, and Ariel Neufeld.  Deep splitting method for parabolic pdes.  SIAM Journal on Scientific Computing, 43(5):A3135–A3154, 2021.
  * [2] Christian Beck, Weinan E, and Arnulf Jentzen.  Machine learning approximation algorithms for high-dimensional fully nonlinear partial differential equations and second-order backward stochastic differential equations.  Journal of Nonlinear Science, 29:1563–1619, 2019.
  * [3] Christian Beck, Lukas Gonon, and Arnulf Jentzen.  Overcoming the curse of dimensionality in the numerical approximation of high-dimensional semilinear elliptic partial differential equations.  arXiv preprint arXiv:2003.00596, 2020.
  * [4] Christian Beck, Fabian Hornung, Martin Hutzenthaler, Arnulf Jentzen, and Thomas Kruse.  Overcoming the curse of dimensionality in the numerical approximation of allen–cahn partial differential equations via truncated full-history recursive multilevel picard approximations.  Journal of Numerical Mathematics, 28(4):197–222, 2020.
  * [5] Sebastian Becker, Ramon Braunwarth, Martin Hutzenthaler, Arnulf Jentzen, and Philippe von Wurstemberger.  Numerical simulations for full history recursive multilevel picard approximations for systems of high-dimensional partial differential equations.  arXiv preprint arXiv:2005.10206, 2020.
  * [6] Sebastian Becker, Patrick Cheridito, Arnulf Jentzen, and Timo Welti.  Solving high-dimensional optimal stopping problems using deep learning.  European Journal of Applied Mathematics, 32(3):470–514, 2021.
  * [7] Jesse Bettencourt, Matthew J. Johnson, and David Duvenaud.  Taylor-mode automatic differentiation for higher-order derivatives in JAX.  In Program Transformations for ML Workshop at NeurIPS 2019, 2019.
  * [8] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018.
  * [9] Shengze Cai, Zhiping Mao, Zhicheng Wang, Minglang Yin, and George Em Karniadakis.  Physics-informed neural networks (pinns) for fluid mechanics: A review.  Acta Mechanica Sinica, 37(12):1727–1738, 2021.
  * [10] Quentin Chan-Wai-Nam, Joseph Mikael, and Xavier Warin.  Machine learning for semi linear pdes.  Journal of scientific computing, 79(3):1667–1712, 2019.
  * [11] Xiaoli Chen, Liu Yang, Jinqiao Duan, and George Em Karniadakis.  Solving inverse stochastic problems from discrete particle observations using the fokker–planck equation and physics-informed neural networks.  SIAM Journal on Scientific Computing, 43(3):B811–B830, 2021.
  * [12] Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, and Eunbyung Park.  Separable pinn: Mitigating the curse of dimensionality in physics-informed neural networks.  arXiv preprint arXiv:2211.08761, 2022.
  * [13] Jérôme Darbon, Gabriel P Langlois, and Tingwei Meng.  Overcoming the curse of dimensionality for some hamilton–jacobi partial differential equations via neural network architectures.  Research in the Mathematical Sciences, 7:1–50, 2020.
  * [14] Jérôme Darbon and Stanley Osher.  Algorithms for overcoming the curse of dimensionality for certain hamilton–jacobi equations arising in control theory and elsewhere.  Research in the Mathematical Sciences, 3(1):19, 2016.
  * [15] Ehsan Haghighat, Maziar Raissi, Adrian Moure, Hector Gomez, and Ruben Juanes.  A physics-informed deep learning framework for inversion and surrogate modeling in solid mechanics.  Computer Methods in Applied Mechanics and Engineering, 379:113741, 2021.
  * [16] Jiequn Han, Arnulf Jentzen, and Weinan E.  Solving high-dimensional partial differential equations using deep learning.  Proceedings of the National Academy of Sciences, 115(34):8505–8510, 2018.
  * [17] Jiequn Han, Arnulf Jentzen, et al.  Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations.  Communications in mathematics and statistics, 5(4):349–380, 2017.
  * [18] Di He, Shanda Li, Wenlei Shi, Xiaotian Gao, Jia Zhang, Jiang Bian, Liwei Wang, and Tie-Yan Liu.  Learning physics-informed neural networks without stacked back-propagation.  In International Conference on Artificial Intelligence and Statistics, pages 3034–3047. PMLR, 2023.
  * [19] Pierre Henry-Labordere.  Deep primal-dual algorithm for bsdes: Applications of machine learning to cva and im.  Available at SSRN 3071506, 2017.
  * [20] Zheyuan Hu, Ameya D Jagtap, George Em Karniadakis, and Kenji Kawaguchi.  Augmented physics-informed neural networks (apinns): A gating network-based soft domain decomposition methodology.  arXiv preprint arXiv:2211.08939, 2022.
  * [21] Zheyuan Hu, Ameya D. Jagtap, George Em Karniadakis, and Kenji Kawaguchi.  When do extended physics-informed neural networks (xpinns) improve generalization?  SIAM Journal on Scientific Computing, 44(5):A3158–A3182, 2022.
  * [22] Zheyuan Hu, Khemraj Shukla, George Em Karniadakis, and Kenji Kawaguchi.  Tackling the curse of dimensionality with physics-informed neural networks.  arXiv preprint arXiv:2307.12306, 2023.
  * [23] Zheyuan Hu, Zhouhao Yang, Yezhen Wang, George Em Karniadakis, and Kenji Kawaguchi.  Bias-variance trade-off in physics-informed neural networks with randomized smoothing for high-dimensional pdes.  arXiv preprint arXiv:2311.15283, 2023.
  * [24] Côme Huré, Huyên Pham, and Xavier Warin.  Deep backward schemes for high-dimensional nonlinear pdes.  Mathematics of Computation, 89(324):1547–1579, 2020.
  * [25] Michael F Hutchinson.  A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines.  Communications in Statistics-Simulation and Computation, 18(3):1059–1076, 1989.
  * [26] Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, Tuan Anh Nguyen, and Philippe von Wurstemberger.  Overcoming the curse of dimensionality in the numerical approximation of semilinear parabolic partial differential equations.  Proceedings of the Royal Society A, 476(2244):20190630, 2020.
  * [27] Martin Hutzenthaler, Arnulf Jentzen, Thomas Kruse, et al.  Multilevel picard iterations for solving smooth semilinear parabolic heat equations.  Partial Differential Equations and Applications, 2(6):1–31, 2021.
  * [28] Ameya D Jagtap and George Em Karniadakis.  Extended physics-informed neural networks (xpinns): A generalized space-time domain decomposition based deep learning framework for nonlinear partial differential equations.  Communications in Computational Physics, 28(5):2002–2041, 2020.
  * [29] Ameya D Jagtap, Kenji Kawaguchi, and George Em Karniadakis.  Adaptive activation functions accelerate convergence in deep and physics-informed neural networks.  Journal of Computational Physics, 404:109136, 2020.
  * [30] Ameya D Jagtap, Dimitrios Mitsotakis, and George Em Karniadakis.  Deep learning of inverse water waves problems using multi-fidelity data: Application to serre–green–naghdi equations.  Ocean Engineering, 248:110775, 2022.
  * [31] Shaolin Ji, Shige Peng, Ying Peng, and Xichuan Zhang.  Three algorithms for solving high-dimensional fully coupled fbsdes through deep learning.  IEEE Intelligent Systems, 35(3):71–84, 2020.
  * [32] Xiaowei Jin, Shengze Cai, Hui Li, and George Em Karniadakis.  Nsfnets (navier-stokes flow nets): Physics-informed neural networks for the incompressible navier-stokes equations.  Journal of Computational Physics, 426:109951, 2021.
  * [33] George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang.  Physics-informed machine learning.  Nature Reviews Physics, 3(6):422–440, 2021.
  * [34] Kenji Kawaguchi.  Deep learning without poor local minima.  In Advances in neural information processing systems (NeurIPS), pages 586–594, 2016.
  * [35] Kenji Kawaguchi, Zhun Deng, Xu Ji, and Jiaoyang Huang.  How does information bottleneck help deep learning?  arXiv preprint arXiv:2305.18887, 2023.
  * [36] Kenji Kawaguchi, Leslie Pack Kaelbling, and Yoshua Bengio.  Generalization in deep learning.  Cambridge University Press, 2022.
  * [37] Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  ICLR, 2015.
  * [38] Lu Lu, Raphael Pestourie, Wenjie Yao, Zhicheng Wang, Francesc Verdugo, and Steven G Johnson.  Physics-informed neural networks with hard constraints for inverse design.  SIAM Journal on Scientific Computing, 43(6):B1105–B1132, 2021.
  * [39] Raphael A Meyer, Cameron Musco, Christopher Musco, and David P Woodruff.  Hutch++: Optimal stochastic trace estimation.  In Symposium on Simplicity in Algorithms (SOSA), pages 142–155. SIAM, 2021.
  * [40] Siddhartha Mishra and Roberto Molinaro.  Estimates on the generalization error of physics informed neural networks (pinns) for approximating pdes.  arXiv preprint arXiv:2006.16144, 2020.
  * [41] Deniz Oktay, Nick McGreivy, Joshua Aduol, Alex Beatson, and Ryan P Adams.  Randomized automatic differentiation.  In International Conference on Learning Representations, 2021.
  * [42] David Persson, Alice Cortinovis, and Daniel Kressner.  Improved variants of the hutch++ algorithm for trace estimation.  SIAM Journal on Matrix Analysis and Applications, 43(3):1162–1185, 2022.
  * [43] Apostolos F Psaros, Kenji Kawaguchi, and George Em Karniadakis.  Meta-learning pinn loss functions.  Journal of computational physics, 458:111121, 2022.
  * [44] Maziar Raissi.  Forward-backward stochastic neural networks: Deep learning of high-dimensional partial differential equations.  arXiv preprint arXiv:1804.07010, 2018.
  * [45] Maziar Raissi, Paris Perdikaris, and George E Karniadakis.  Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations.  Journal of Computational Physics, 378:686–707, 2019.
  * [46] Farbod Roosta-Khorasani and Uri Ascher.  Improved bounds on sample size for implicit matrix trace estimators.  Foundations of Computational Mathematics, 15(5):1187–1212, 2015.
  * [47] Yeonjong Shin, Jerome Darbon, and George Em Karniadakis.  On the convergence of physics informed neural networks for linear second-order elliptic and parabolic type pdes.  arXiv preprint arXiv:2004.01806, 2020.
  * [48] Maciej Skorski.  Modern analysis of hutchinson’s trace estimator.  In 2021 55th Annual Conference on Information Sciences and Systems (CISS), pages 1–5. IEEE, 2021.
  * [49] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.  Score-based generative modeling through stochastic differential equations.  In International Conference on Learning Representations, 2021.
  * [50] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [51] Chuwei Wang, Shanda Li, Di He, and Liwei Wang.  Is $l^2$ physics informed loss always suitable for training physics informed neural network?  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  * [52] Yifan Wang, Pengzhan Jin, and Hehu Xie.  Tensor neural network and its numerical integration.  arXiv preprint arXiv:2207.02754, 2022.
  * [53] Yifan Wang, Yangfei Liao, and Hehu Xie.  Solving schr$\backslash$”$\\{$o$\\}$ dinger equation using tensor neural network.  arXiv preprint arXiv:2209.12572, 2022.
  * [54] Yibo Yang and Paris Perdikaris.  Adversarial uncertainty quantification in physics-informed neural networks.  Journal of Computational Physics, 394:136–152, 2019.
  * [55] Jeremy Yu, Lu Lu, Xuhui Meng, and George Em Karniadakis.  Gradient-enhanced physics-informed neural networks for forward and inverse pde problems.  Computer Methods in Applied Mechanics and Engineering, 393:114823, 2022.


# Towards energy discretization for muon scattering tomography in GEANT4
simulations: A discrete probabilistic approach

A. Ilker Topuz1,2 Madis Kiisk1,3 1 Institute of Physics, University of Tartu,
W. Ostwaldi 1, 50411, Tartu, Estonia  2 Centre for Cosmology, Particle Physics
and Phenomenology, Université catholique de Louvain, Chemin du Cyclotron 2,
B-1348 Louvain-la-Neuve, Belgium  3 GScan OU, Maealuse 2/1, 12618 Tallinn,
Estonia<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

In this study, by attempting to eliminate the disadvantageous complexity of
the existing particle generators, we present a discrete probabilistic scheme
adapted for the discrete energy spectra in the GEANT4 simulations. In our
multi-binned approach, we initially compute the discrete probabilities for
each discrete energy bin, the number of which is flexible depending on the
computational goal, and we solely satisfy the imperative condition that
requires the sum of the discrete probabilities to be the unity. Regarding the
implementation in the GEANT4 code, we construct a one-dimensional probability
grid that consists of sub-cells equaling the number of the energy bin, and
each cell represents the discrete probability of each energy bin by fulfilling
the unity condition. Through uniformly generating random numbers between 0 and
1, we assign the discrete energy in accordance with the associated generated
random number that corresponds to a specific cell in the probability grid.
This probabilistic methodology does not only permits us to discretize the
continuous energy spectra based on the Monte Carlo generators, but it also
gives a unique access to utilize the experimental energy spectra measured at
the distinct particle flux values. Ergo, we initially perform our simulations
by discretizing the muon energy spectrum acquired via the CRY generator over
the energy interval between 0 and 8 GeV along with the measurements from the
BESS spectrometer and we determine the average scattering angle, the root-
mean-square of the scattering angle, and the number of the muon absorption by
using a series of slabs consisting of aluminum, copper, iron, lead, and
uranium. Eventually, we express a computational strategy in the GEANT4
simulations that grants us the ability to verify as well as to modify the
energy spectrum depending on the nature of the information source in addition
to the exceptional tracking speed.

Keywords:  Muon tomography; Characteristic parameters; Energy discretization;
Multi-group approximation; Discrete probabilities; Non-analogue Monte Carlo
simulations; GEANT4

##  1 Introduction

The emerging applications of cosmic ray muon tomography [1, 2, 3] lead to a
significant rise in the utilization of the cosmic particle generators, e.g.
CRY [4], CORSIKA [5], or CMSCGEN [6], where the fundamental parameters such as
the energy spectrum and the angular distribution associated with the generated
muons are represented in the continuous forms routinely governed by the
probability density functions over the corresponding vast intervals. Despite
this perceptible increase in the diversity of the muon generators, the common
difficulties in the hands-on applications that might be relatively exemplified
as perplexing coupling with the Monte Carlo codes, unnecessarily broad and
occasionally unmodifiable parametric intervals for the specific applications,
extensive execution times, and complications in the particle tracking partly
remain steady. Contrary to the continuous mode, the discretized energy
spectra, i.e. multi-group energy approximations to put it another way, have
been ubiquitously employed in the neighboring fields such as nuclear
engineering [7, 8] and medical physics [9, 10] under the umbrella of the non-
analogue Monte Carlo simulations [11] on and on. Along with the discretization
schemes based on the theoretical assumptions, a number of notable empirical
studies founded on the advanced particle detectors such as the BESS
spectrometer [12] represent the experimental energy spectra in the discrete
format. While MNCP6 [13] includes the necessary algorithms to utilize the
discrete energy distributions in the black box paradigm, the general particle
source (GPS) in GEANT4 [14] is the existing pre-configured module that
provides the opportunity of this discrete approach through a macro file
without detailing the algebraic/algorithmic phase.

In the present study, by aiming at obtaining a fast, verifiable, and
modifiable muon source in the GEANT4 simulations, we initially exhibit our
procedure that is based on the multi-binned approximation of the CRY muon
energy spectrum. In the latter instance, we incorporate our discrete spectra
in the GEANT4 code by using a one-dimensional probability grid that operates
under G4ParticleGun. Following this step, we furthermore gain the capability
of utilizing the existing experimental spectra. Then, we test our methodology
over our tomographic setup consisting of plastic scintillators fabricated from
polyvinyl toluene with the dimensions of 100$\times$0.4$\times$100 $\rm
cm^{3}$ and we determine the characteristic parameters such the average
scattering angle, the root-mean-square of the scattering angle, and the number
of the muon absorption by using a set of slabs composed of aluminum, copper,
iron, lead, and uranium with the dimensions of 40$\times$10$\times$40 $\rm
cm^{3}$. Finally, we contrast our simulation outcomes by means of both the CRY
discrete spectrum and the BESS muon spectrum. This study is outlined as
follows. In section 2, we compute the discrete probabilities for the
corresponding discrete energies, and section 3 describes the implementation in
the GEANT4 code. While we express our characteristic parameters as well as our
simulation setup in section 4, we expose our simulation outcomes in section 5.
Finally, our conclusions are stated in section 6.

##  2 Discrete energies and discrete probabilities

At the outset, we strive for the energy discretization based on the extracted
energy list from the CRY muon generator [4] between 0 and 8 GeV. To achieve
this aim, we first set out our constant bin length that is selected as 0.1
GeV. Thus, the number of non-zero energy bins is evidently 80 as written in

$\\#_{\rm Bins}=\frac{E_{\rm Max}-E_{\rm Min}}{L_{\rm
Bins}}=\frac{8.0-0}{0.1}=80$ (1)

Then, in agreement with the energy dataset acquired through the CRY generator,
the number of the counts in the specific energy bin denoted by $E_{i}$ in GeV
is computed by incorporating any $E_{x}\in(E_{i-1},E_{i}]$ under the condition
of $m_{0}=0$ for $E_{0}=0$ as described in

$m_{i}=\sum\limits_{k=1}1\mbox{ if $E_{i-1}<E_{x}\leq E_{i}$ for i={1, 2,
3,.., 80}}$ (2)

Basically, Eq. (2) might be performed by using the existing tools in Python,
e.g. NumPy. While this operation is an approximation by definition, we do not
significantly neglect any fundamental information in the case of the fine
energy bins as accomplished in the present study, which also implies that the
simulation outcomes by means of the finely binned histogram do not yield
consequential differences in comparison with the continuous energy spectra.
Whereas the count numbers in the particular energy bins already provides an
opportunity to employ the discrete energy distributions in the wake of
renormalization, we favor to determine the discrete probabilities that serve
to constitute a probability grid; hence, we first calculate the total count
over 80 energy bins as shown in

$\sum\limits_{i=0}^{\\#_{\rm Bins}}m_{i}=\sum\limits_{i=0}^{80}m_{i}$ (3)

In the end, the discrete probability, i.e. the discrete normalized frequency
to rephrase it, at a given energy bin indicated by $E_{i}$ is the ratio
between the specific count in $E_{i}$ and the total counts over 80 bins by
satisfying the unity condition as noted in

$p_{i}=\frac{m_{i}}{\sum\limits_{i=0}^{80}m_{i}}~{}~{}~{}{\rm
with}~{}~{}~{}\sum\limits_{i=0}^{80}p_{i}=1$ (4)

Finally, we tabulate the discrete probabilities along with the discrete muon
energies obtained via the CRY muon source in Table 1 by entitling D$\neg$CRY.
As indicated in Table 1, we view that the highest discrete probability
appertains to the energy bin of 0.5 GeV, i.e. the mode of our discrete energy
distribution, which is reduced approximately by one order of magnitude at 8
GeV, and the trend observed in D$\neg$CRY resembles to a log-normal
distribution with a positive skewness.

Table 1: D$\neg$CRY discrete probabilities between 0 and 8 GeV.

$E_{i}$ [GeV] | $p_{i}$
---|---
0.0 | 0.000000
0.1 | 0.012536
0.2 | 0.025745
0.3 | 0.028020
0.4 | 0.027066
0.5 | 0.035285
0.6 | 0.028265
0.7 | 0.031579
0.8 | 0.030784
0.9 | 0.027776
1.0 | 0.025464
1.1 | 0.031506
1.2 | 0.028155
1.3 | 0.025807
1.4 | 0.023642
1.5 | 0.021709
1.6 | 0.021526
1.7 | 0.023483
1.8 | 0.021342
1.9 | 0.019691
2.0 | 0.020364
2.1 | 0.018419
2.2 | 0.017184
2.3 | 0.017001
2.4 | 0.016242
2.5 | 0.015398
2.6 | 0.015362

$E_{i}$ [GeV] | $p_{i}$
---|---
2.7 | 0.014713
2.8 | 0.014224
2.9 | 0.014126
3.0 | 0.012842
3.1 | 0.012610
3.2 | 0.012133
3.3 | 0.012903
3.4 | 0.012487
3.5 | 0.011962
3.6 | 0.010641
3.7 | 0.010579
3.8 | 0.009626
3.9 | 0.010384
4.0 | 0.009283
4.1 | 0.008794
4.2 | 0.008843
4.3 | 0.007938
4.4 | 0.007864
4.5 | 0.007693
4.6 | 0.007094
4.7 | 0.007363
4.8 | 0.007192
4.9 | 0.007216
5.0 | 0.006923
5.1 | 0.006433
5.2 | 0.006788
5.3 | 0.006739

$E_{i}$ [GeV] | $p_{i}$
---|---
5.4 | 0.006189
5.5 | 0.006348
5.6 | 0.006653
5.7 | 0.006507
5.8 | 0.005614
5.9 | 0.005895
6.0 | 0.005895
6.1 | 0.005785
6.2 | 0.005577
6.3 | 0.005504
6.4 | 0.004342
6.5 | 0.004354
6.6 | 0.004085
6.7 | 0.003645
6.8 | 0.003999
6.9 | 0.003889
7.0 | 0.003963
7.1 | 0.004317
7.2 | 0.003681
7.3 | 0.003632
7.4 | 0.003620
7.5 | 0.004109
7.6 | 0.003363
7.7 | 0.003584
7.8 | 0.003620
7.9 | 0.003486
8.0 | 0.003596

In view of the fact that we aim at carrying out a multi-group approach, we are
furthermore capable of utilizing the empirical muon energy spectrum as in the
case of the BESS spectrometer [12] where the discrete energies are measured at
the distinct particle fluxes. We adopt the first 36 non-zero energy bins
between 0.598 and 8.1 GeV and their corresponding separate fluxes by labeling
as EXP-BESS, and the discrete probabilities are determined by the quotient
between the particular flux value at a specific energy and the total flux as
follows

$p_{i}=\frac{\phi_{i}}{\sum\limits_{i=0}^{36}\phi_{i}}~{}~{}~{}{\rm
with}~{}~{}~{}\sum\limits_{i=0}^{36}p_{i}=1$ (5)

Since the energy list is already discretized, we exert neither a further
approximation nor a simplifying assumption in Eq. (5). The experimental energy
values obtained through the BESS spectrometer as well as the computed discrete
probabilities are listed in Table 2.

Table 2: Discrete EXP-BESS probabilities between 0 and 8.1 GeV.

$E_{i}$ [GeV] |  $p_{i}$
---|---
0.000 |  0.000000
0.598 |  0.058310
0.644 |  0.056910
0.694 |  0.055511
0.748 |  0.053645
0.806 |  0.052712
0.868 |  0.050379
0.936 |  0.049447
1.010 |  0.046648
1.080 |  0.044782
1.170 |  0.042449
1.260 |  0.041050
1.360 |  0.039790
1.460 |  0.037365
1.570 |  0.035126
1.700 |  0.033680
1.830 |  0.031347
1.970 |  0.029295
2.120 |  0.026636

$E_{i}$ [GeV] |  $p_{i}$
---|---
2.290 |  0.025003
2.460 |  0.022951
2.650 |  0.021551
2.860 |  0.019079
3.080 |  0.017213
3.310 |  0.015440
3.570 |  0.014274
3.850 |  0.012595
4.140 |  0.011055
4.470 |  0.009843
4.810 |  0.008723
5.180 |  0.007744
5.580 |  0.006811
6.010 |  0.005784
6.480 |  0.005271
6.980 |  0.004478
7.520 |  0.003872
8.100 |  0.003233

As can be noticed from Table 2, the EXP-BESS muon spectrum exhibits an
exponentially decreasing trend where the lower bound is associated with the
highest discrete probability, whereas the upper bound is identified with the
minimum discrete normalized frequency. In order to qualitatively show the
variation of both D$\neg$CRY and EXP-BESS, Fig. 1 displays the energy
histograms where (a) indicates D$\neg$CRY in comparison with the discrete
CMSCGEN data reproduced from another study [15] that is dedicated to the CMS
strip tracker, while (b) points out EXP-BESS. From Fig. 1(a), we experience
that our D$\neg$CRY tends to show strong similarities when compared to the
existing qualitative data of the CMSCGEN muon generator, thereby relatively
authenticating our discretized energy spectrum hinged on the CRY muon
generator.

Figure 1: Comparison between discrete energy histograms (a) 80-bin D$\neg$CRY
compared to CMSCGEN and (b) 36-bin EXP-BESS.

##  3 Implementation via G4ParticleGun in GEANT4

On account of attaining the discrete energy spectra, we subsequently integrate
our strategy to inject the incoming muons by means of G4ParticleGun. By
reminding the unity condition, we build a grid by adding up the discrete
probabilities, the inverval of which starts with 0 and ends in 1 as
illustrated in Fig. 2. Thus, each cell in this grid, i.e. the difference
between two points on the probability grid, specifies a discrete probability.
Then, we generate a random number denoted by $\xi$ between 0 and 1 by using
the pre-defined uniform number generator called G4UniformRand(). Finally, we
scan this random number on the probability grid by checking the difference
between the grid points and we assign the particular discrete energy when the
random number matches with the associated cell.

Figure 2: Scheme of probability grid implemented via G4ParticleGun in GEANT4.

Intending to reveal out the performance of our procedure, the energy
histrograms before and after the probability grid are put on view in Fig. 3
where the term called input indicates the discrete energy values as well as
the discrete probabilities that are stored in an array before the activation
of the probability grid, whereas the notion named output hints the outcome of
the probability grid by processing the input information, which means the
kinetic energies allocated to the generated muons by means of the probability
grid.

Figure 3: Contrast between the input dataset before processing with the
probability grid and the output values through the activation of our
probability grid (a) 80-bin D$\neg$CRY and (b) 36-bin EXP-BESS.

##  4 Characteristic parameters and simulation properties

For the purpose of appraisal, we express a number of characteristic parameters
consorted with the muon tomography. By defining two vectors, we bring forth
the scattering angle denoted by $\theta$ as depicted in Fig. 4, and the
scattering angle of a muon crossing the volume-of-interest (VOI) is obtained
by using these two vectors as follows [16, 17, 18]

$\theta=\arccos\left(\frac{\vec{v}_{1}\cdot\vec{v}_{2}}{\left|v_{1}\right|\left|v_{2}\right|}\right)$
(6)

The mean scattering angle due to the VOI and its standard deviation over $N$
number of the non-absorbed/non-decayed muons is calculated as written down in

$\bar{\theta}\pm\delta\theta=\frac{1}{N}\sum_{i=1}^{N}\theta_{i}\pm\sqrt{\frac{1}{N}\sum_{j=1}^{N}(\theta_{j}-\bar{\theta})^{2}}$
(7)

Furthermore, the root-mean-square (RMS) of the scattering angle over $N$
number of the non-absorbed/non-decayed muons is deternined by using the
following equation:

$\theta_{\rm RMS}=\sqrt{\frac{1}{N}\sum_{i=1}^{N}\theta_{i}^{2}}$ (8)

On top of the scattering angle, we precisely register the number of the
absorbed muons within the VOI as denoted in

$\\#_{\rm Capture}^{\rm In-target}=\mbox{\\# of muMinusCaptureAtRest in VOI}$
(9)

Figure 4: Delineation of scattering angle within the present tomographic
configuration.

Lastly, our simulation features are tabulated in Table 3, and we conduct our
GEANT4 simulations over a set of slabs including aluminum, copper, iron, lead,
and uranium with the dimensions of 40$\times$10$\times$40 $\rm cm^{3}$ in
accordance with another study [19]. The particle tracking is managed by
G4Step, and the registered hit locations on the plastic scintillators are
post-processed by dint of a Python script.

Table 3: Simulation properties. Particle | $\mbox{\textmu}^{-}$
---|---
Beam direction | Vertical
Momentum direction | (0, -1, 0)
Source geometry | Planar
Initial position (cm) | ([-0.5, 0.5], 85, [-0.5, 0.5])
Particle injector | G4ParticleGun
Number of particles | $10^{5}$
Energy distribution | Non-linear discrete
Target geometry | Rectangular prism
Target volume (cm3) | 40$\times$10$\times$40
Material database | G4/NIST
Reference physics list | FTFP$\\_$BERT

##  5 Simulation results

In the long run, we inaugurate our GEANT4 simulations founded on our discrete
approach, the initial outcomes of which are displayed in Fig. 5 in terms of
the angular counts. From Fig. 5(a)-(b), we explicitly see that the angular
populations characteristically vary depending on the atomic number as well as
the material density for the same thickness. Moreover, 80-bin D$\neg$CRY
yields a typical exponential decrease towards the higher angles, whereas this
trend appears to be a parabolic reduction in the case of 36-bin EXP-BESS.

Figure 5: Difference between the angular counts by using a step length of 1
mrad (a) 80-bin D$\neg$CRY and (b) 36-bin EXP-BESS.

Then, we also determine the average scattering angle as well as the RMS of the
scattering angle in addition to the muon absorption and we compare both the
discrete muon sources, i.e. 80-bin D$\neg$CRY and 36-bin EXP-BESS, as listed
in Table 4. From the contrast shown in Table 4, we initially observe that the
average scattering angle acquired via these discrete distributions is alike
except in the case of lead and uranium. Secondly, 80-bin D$\neg$CRY constantly
results in a strikingly higher standard deviation in comparison with 36-bin
EXP-BESS since the first five energy bins of 80-bin D$\neg$CRY constitute a
low-energy population excluded in 36-bin EXP-BESS that drastically amplifies
the propagated uncertainty. Furthermore, this increased tendency is also
examined in the RMS values owing to the same reason. Crucially, an essential
fact is traced in the event of the muon absorption, and we demonstrate that
36-bin EXP-BESS does not yield any muon capture within the VOI since the
energy level above 0.598 GeV has no potential for the muon absorption when the
material thickness is merely 10 cm by recalling that the lower bound of 36-bin
EXP-BESS is 0.598 GeV, while the lowest bin value in 80-bin D$\neg$CRY belongs
to 0.1 GeV.

Table 4: Characteristic parameters obtained via 80-bin D$\neg$CRY as well as
36-bin EXP-BESS.

Material |  $\bar{\theta}_{\rm D\neg CRY}\pm\delta\theta$ [mrad] |  $\theta^{\rm RMS}_{\rm D\neg CRY}$ [mrad] | $\\#_{\rm D\neg CRY}^{\rm Capture}$ |  $\bar{\theta}_{\rm EXP\mbox{-}BESS}\pm\delta\theta$ [mrad] |  $\theta^{\rm RMS}_{\rm EXP\mbox{-}BESS}$ [mrad] | $\\#_{\rm EXP\mbox{-}BESS}^{\rm Capture}$
---|---|---|---|---|---|---
Aluminum | 15.980$\pm$27.004 | 31.378 | - | 15.036$\pm$12.413 | 19.499 | -
Copper | 40.638$\pm$60.312 | 72.725 | 1230 | 40.759$\pm$31.898 | 51.757 | -
Iron | 35.824$\pm$51.548 | 62.773 | 1222 | 36.200$\pm$28.518 | 46.084 | -
Lead | 65.325$\pm$88.323 | 109.856 | 1272 | 68.138$\pm$52.450 | 85.987 | -
Uranium | 81.822$\pm$100.321 | 129.457 | 3836 | 95.721$\pm$75.376 | 121.836 | -

##  6 Conclusion

To sum up, we state a procedure based on the multi-group energy approximation
where we favor the utilization of the binned energy values connected with the
discrete probabilities by means of a probability grid. Consequently, we gain
the capability to control as well as to adjust our energy spectra according to
our computational goals apart from the noteworthy computation times.

## References

## References

  * [1] S. Pesente, S. Vanini, M. Benettoni, G. Bonomi, P. Calvini, P. Checchia, E. Conti, F. Gonella, G. Nebbia, S. Squarcia, et al., First results on material identification and imaging with a large-volume muon tomography prototype, Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 604 (3) (2009) 738–746.
  * [2] S. Procureur, Muon imaging: Principles, technologies and applications, Nucl. Instr. Meth. A 878 (2018) 169.
  * [3] L. Bonechi, R. D’Alessandro, A. Giammanco, Atmospheric muons as an imaging tool, Reviews in Physics 5 (2020) 100038.
  * [4] C. Hagmann, D. Lange, D. Wright, Cosmic-ray shower generator (CRY) for Monte Carlo transport codes, in: 2007 IEEE Nuclear Science Symposium Conference Record, Vol. 2, IEEE, pp. 1143–1146.
  * [5] D. Heck, J. Knapp, J. Capdevielle, G. Schatz, T. Thouw, et al., CORSIKA: A Monte Carlo code to simulate extensive air showers, Report fzka 6019 (11) (1998).
  * [6] P. Biallass, T. Hebbeker, Parametrization of the cosmic muon flux for the generator cmscgen, arXiv preprint arXiv:0907.5514 (2009).
  * [7] D. E. Cullen, Application of the probability table method to multigroup calculations of neutron transport, Nuclear Science and Engineering 55 (4) (1974) 387–400.
  * [8] M. Nakagawa, T. Mori, MORSE-DD, A Monte Carlo code using multi-group double differential form cross sections, Tech. Rep. JAERI-M 84-126, Japan Atomic Energy Research Inst. (1984).
  * [9] M. Rivard, A discretized approach to determining TG-43 brachytherapy dosimetry parameters: case study using Monte Carlo calculations for the MED3633 103Pd source, Applied Radiation and Isotopes 55 (6) (2001) 775–782.
  * [10] E. Boman, J. Tervo, M. Vauhkonen, Modelling the transport of ionizing radiation using the finite element method, Physics in Medicine & Biology 50 (2) (2004) 265\.
  * [11] J. K. Shultis, R. E. Faw, An MCNP primer, Tech. rep., Manhattan, Kansas State University (2011).
  * [12] S. Haino, T. Sanuki, K. Abe, K. Anraku, Y. Asaoka, H. Fuke, M. Imori, A. Itasaki, T. Maeno, Y. Makida, et al., Measurements of primary and atmospheric cosmic-ray spectra with the bess-tev spectrometer, Physics Letters B 594 (1-2) (2004) 35–46.
  * [13] T. Goorley, M. James, T. Booth, F. Brown, J. Bull, L. Cox, J. Durkee, J. Elson, M. Fensin, R. Forster, et al., Initial MCNP6 release overview, Nuclear technology 180 (3) (2012) 298–315.
  * [14] S. Agostinelli, J. Allison, K. Amako, J. Apostolakis, H. Araujo, P. Arce, M. Asai, D. Axen, S. Banerjee, G. Barrand, et al., GEANT4—a simulation toolkit, Nuclear instruments and methods in physics research section A: Accelerators, Spectrometers, Detectors and Associated Equipment 506 (3) (2003) 250–303.
  * [15] W. Adam, T. Bergauer, M. Dragicevic, M. Friedl, R. Frühwirth, S. Hänsel, J. Hrubec, M. Krammer, M. Oberegger, M. Pernicka, et al., Performance studies of the cms strip tracker before installation, Journal of Instrumentation 4 (06) (2009) P06009.
  * [16] T. Carlisle, J. Cobb, D. Neuffer, Multiple Scattering Measurements in the MICE Experiment, Tech. Rep. FERMILAB-CONF-12-171-APC, Fermi National Accelerator Lab. (FNAL), Batavia, IL (United States) (2012).
  * [17] J. C. Nugent, Multiple Coulomb scattering in the MICE experiment, Ph.D. thesis, University of Glasgow (2017).
  * [18] D. Poulson, J. Bacon, M. Durham, E. Guardincerri, C. Morris, H. R. Trellue, Application of muon tomography to fuel cask monitoring, Philosophical Transactions of the Royal Society A 377 (2137) (2019) 20180052.
  * [19] M. Hohlmann, P. Ford, K. Gnanvo, J. Helsby, D. Pena, R. Hoch, D. Mitra, GEANT4 Simulation of a Cosmic Ray Muon Tomography System With Micro-Pattern Gas Detectors for the Detection of High-Z Materials, IEEE Transactions on Nuclear Science 56 (3) (2009) 1356–1363.