#  $\boldsymbol{\mu}$-Net: ConvNext-Based U-Nets for Cosmic Muon Tomography

Lim Li Xin Jed  Qiu Ziming

###### Abstract

Muon scattering tomography utilises muons, typically originating from cosmic
rays to image the interiors of dense objects. However, due to the low flux of
cosmic ray muons at sea-level and the highly complex interactions that muons
display when travelling through matter, existing reconstruction algorithms
often suffer from low resolution and high noise. In this work, we develop a
novel two-stage deep learning algorithm, $\mu$-Net, consisting of an MLP to
predict the muon trajectory and a ConvNeXt-based U-Net to convert the
scattering points into voxels. $\mu$-Net achieves a state-of-the-art
performance of 17.14 PSNR at the dosage of 1024 muons, outperforming
traditional reconstruction algorithms such as the point of closest approach
algorithm and maximum likelihood and expectation maximisation algorithm.
Furthermore, we find that our method is robust to various corruptions such as
inaccuracies in the muon momentum or a limited detector resolution. We also
generate and publicly release the first large-scale dataset that maps muon
detections to voxels. We hope that our research will spark further
investigations into the potential of deep learning to revolutionise this
field.

Machine Learning, Deep Learning, U-Net, ConvNeXt, Muon Tomography

##  1 Introduction

Muon tomography is an imaging technique that utilises muons, typically
originating from cosmic rays, to image the interiors of objects. By leveraging
the fact that muons are highly penetrating particles, muon tomography offers a
non-invasive and non-destructive means of investigating the internal
composition of dense materials. Through tracking and analysis of muon
trajectories and energies, this enables the accurate reconstruction of
internal structures and features. Consequently, muon tomography has emerged as
a tool with wide-ranging applications in fields such as geophysics (Tioukov et
al., 2019), civil engineering (Saracino et al., 2018), and archaeology
(Borselli et al., 2022; Procureur et al., 2023a).

Due to cosmic rays entering the Earth’s atmosphere, there is no need for a
specialised muon source unlike other types of tomography. However, there are
many other challenges in this task. First, the flux of cosmic muons at sea-
level is very low (Beringer et al., 2012). In order to produce a decent
reconstruction, data has to be collected for a long period of time.
Furthermore, unlike other types of tomography such as x-ray computed
tomography, muons will scatter off atomic nuclei. This makes the forward
operator highly nonlinear. In contrast, in computed tomography, x-rays only
attenuate and hence, the forward operator is the linear radon transform. In
addition, due to limitations of current day muon detectors, the momentum of
muons which significantly affects the muon scattering angle is not known. At
best, we will have an estimate of the momentum $\hat{p}$ with a significant
amount of uncertainty.

Several methods have been developed to tackle this problem. The first
algorithm developed was the Point of Closest Approach (PoCA) algorithm
(Schultz, 2003) which assumes that the muons only scatter once at the point of
closest approach of its inward and outward trajectory. The Maximum Likelihood
and Expectation Maximisation (MLEM) algorithm (Schultz et al., 2007) improves
on PoCA and iteratively optimises the reconstruction to maximise the
likelihood of producing a given scattering outcome. Other algorithms have also
been developed such as maximum a posterori (MAP) (Wang et al., 2009), most
likely path (Schulte et al., 2008; Yi et al., 2014; Chatzidakis et al., 2018),
scattering density estimation (SDE) (Jonkmans et al., 2013) and angle
statistics reconstruction (ASR) (Stapleton et al., 2014). Other methods have
also been developed for low dosages such as the binned clustering algorithm
(Thomay et al., 2012) and a method based on density clustering (Hou et al.,
2021).

However, no one has attempted to make use of developments in deep learning to
directly solve this ill-posed problem and go directly from muon detections to
a 3D reconstruction. Due to the abundance of data which can be obtained from
simulations from software such as Geant4 (Agostinelli et al., 2003) and the
non-linearity of the problem, deep learning methods are well-suited for this
task. They have been previously applied to other inverse problems such as
computed tomography (Szczykutowicz et al., 2022).

In this work, we develop a novel two-stage deep learning algorithm, $\mu$-Net,
for cosmic muon scattering tomography, based on using the point of closest
approach (PoCA) algorithm and the U-Net architecture proposed by Ronneberger
et al. (2015). To our knowledge, this is the first application of deep
learning in muon scattering tomography to directly perform the 3D
reconstruction. Instead of using the more traditional Residual Blocks, we make
use of ConvNeXt Blocks (Liu et al., 2022). We find that our model
significantly outperforms traditional reconstruction algorithms (in speed and
accuracy) and achieves state-of-the-art performance. Our method achieves a
PSNR of 17.14 with only 1024 muons while PoCA achieves a PSNR of only 13.66
while taking 22.5s to run while $\mu$-Net runs in 126 ms.

We have also generated the first large-scale dataset mapping muon detections
to voxels. In prior literature, there has been no standard benchmark to
evaluate various methods for 3D muon tomographic reconstructions and no
quantitative metrics used to compare different methods. As such, we release
our dataset and data generation code publicly. We hope this will spark more
systematic investigations into reconstruction algorithms on this task, using
deep learning or using traditional methods.

##  2 Preliminaries

###  2.1 Physical Background

Muon scattering relies primarily on modelling the scattering interaction
between muons and matter. Rossi & Greisen (1941) developed a scattering theory
for charged particles and found that charged particles travelling through a
plate of thickness $x$ that undergo Coulomb scattering have scattering angles
and lateral displacements that follow a Gaussian distribution with mean,
$\mu=0$ and variance,

$\sigma^{2}=\frac{E_{s}^{2}}{2p^{2}v^{2}}\frac{x}{L_{rad}}$ (1)

where $E_{s}=21$ MeV, $\lambda$ is the radiation length of the material, $x$
is the thickness of the plate and $p$ and $v$ are the momentum and velocity of
the muon respectively.

From this formula, we can define the parameter of interest, the scattering
density $\lambda$.

$\lambda=\frac{\sigma^{2}}{x}=\frac{15MeV}{p^{2}v^{2}}\frac{1}{L_{rad}}$ (2)

With this, the task of muon tomography is to find the distribution of
$\lambda$ within the object through looking at positions and directions of the
incoming and outgoing muons.

###  2.2 Problem Statement

Before proceeding with a literature review and the description of the method,
let us formalize the muon reconstruction problem. We shall follow a similar
notation to Schultz et al. (2007).

Let the object of interest be defined by its scattering density,

$\lambda(x,y,z)=\left(\frac{15}{p_{0}}\right)^{2}\frac{1}{L_{rad}(x,y,z)}$ (3)

where $L_{rad}(x,y,z)$ is the radiation length at each point within the
object.

We can represent the scattering density in terms of some basis functions
$\phi_{j}(x,y,z)$ such that

$\lambda(x,y,z)=\sum_{j}\alpha_{j}\phi(x,y,z)$ (4)

where $\boldsymbol{\alpha}=\left[\alpha_{j}\right]$ are the coefficients for
the basis functions.

Suppose the muon detections, $\mathbf{y}$ follow a distribution $D$
parameterized by the scattering density of the object, i.e.

$\mathbf{Y}\sim D(\boldsymbol{\alpha})$ (5)

Therefore, given a sample of $n$ muon detections
$Y_{n}=\left[\mathbf{y_{1}},\mathbf{y_{2}},\dots,\mathbf{y_{n}}\right]$ we
wish to construct a point estimate $\mathbf{a}(Y_{n})$ of
$\boldsymbol{\alpha}$, which is approximated by the deep neural network
$\boldsymbol{f_{\theta}}(Y_{n})$ parameterized by $\boldsymbol{\theta}$.

In this paper, we will take $p_{0}$ to be 15 MeV, so that we directly regress
the reciprocal of the radiation length $\frac{1}{L_{rad}(x,y,z)}$ in units of
cm-1.

###  2.3 Motivation

Some key features of the muon reconstruction problem are immediately apparent
from the problem statement, which help us design our model architecture are
listed below:

  * •

The model should be permutation-invariant, i.e. the order in which muons are
inputted into the model should not change the output.

  * •

The model should accept any number of input muons.

  * •

The model should be able to make use of paired input and output muon
detections i.e. the model needs to know which input muon corresponds to which
output.

  * •

The model should be able to take advantage of the 3D spatial structure of the
target output.

This seems to suggest making use of Transformers (Vaswani et al., 2017) with
the positional encoding removed. However, the dot product attention module
used in Transformers has a large time-complexity of $O(N^{2})$ where $N$ is
the dosage. This is not acceptable given dosages can go up to 10000 muons or
more. The transformer will also be unable to take advantage of the spatial
structure of the output. Hence, we will make use of the PoCA algorithm and the
U-Net architecture (Ronneberger et al., 2015) to create a two-stage algorithm.

##  3 Related Work

###  3.1 Deep Learning Methods

#### Point Clouds.

From the key features of the muon tomography problem, we see that it is highly
similar to point cloud problems, which also take in permutation invariant data
(a set of points), but the points themselves also exhibit some 3D structure.
Deep learning methods on point cloud data can generally be classified into 2
main categories, neural networks which operates directly on the points, and
neural networks which operate on a voxelized representation of the data (Bello
et al., 2020). In our work, we chose the latter approach since the former
tends to be slightly worse at capturing spatial structure, and our desired
output is also in the form of voxels.

#### U-Net.

The U-Net was first proposed as an medical image segmentation model
(Ronneberger et al., 2015). It consists of a downward branch, where the image
is downsampled and an upward branch, where the image is upsampled. Skip
connections are also used between the downward and upward branches to allow
high resolution features, which may be removed during downsampling, to be
retained. U-Nets have also been extended to process 3D data (Çiçek et al.,
2016; Ho et al., 2021), by replacing the standard 2D convolution operations
with 3D convolutions.

#### ConvNext.

ConvNext is a recent iteration of the family of convolutional neural networks
(Liu et al., 2022). It was proposed as an improvement of the ResNet (He et
al., 2016) by incorporating methods found in Vision Transformer architectures,
most notably the SWIN transformer (Liu et al., 2021), and has been shown to
obtain competitive performance against Transformer-based methods but with a
lower parameter count. In our work, we will use a modified 3D ConvNext block
to form our U-Net.

#### IEEE BigData 2023 Cup.

Concurrent to this work, is the ”IEEE BigData 2023 Cup: Object Recognition
with Muon Tomography using Cosmic Rays” organised by Wnuk et al. (2023).
Unlike our work which addresses full three-dimensional reconstruction, they
focus on reconstruction of 2-dimensional objects placed in the central plane.
The metric used is the mean average IoU. However, this metric equally
penalises models for misidentifying materials with a small difference in
radiation length and a large difference in radiation length. Hence, in this
paper, we adopt the mean squared error and the peak-signal to noise ratio
(PSNR) as our primary metrics. We are unfortunately unable to compare our
results to the results from this competition as at the time of writing, the
competition report and the papers of the winning entries are not accessible
online.

###  3.2 Traditional Algorithms

#### Point of Closest Approach.

The Point of Closest Approach (PoCA) method is a commonly used method for muon
scattering tomography, first proposed by Schultz (2003). It assumes that the
muon is scattered once by an object at the point of closest approach between
the input and output trajectories. However, this fails to take into account
the fact that in many cases, muons may scatter multiple times. As a result,
there will be predictions that muons scatter at points where there is actually
no object as illustrated in Figure 1. In addition, some information is lost as
some muons may have a point of closest approach outside of the object space.

Figure 1: A weakness of PoCA. Since the muon scatters more than once, the
computed PoCA is not within any object.

#### Maximum Likelihood Expectation Maximization.

The maximum likelihood expectation maximisation (MLEM) algorithm is a
statistical reconstruction method proposed by Schultz et al. (2007). It makes
use of the statistical distribution of muon scattering to frame muon
reconstruction as a maximum likelihood problem. An iterative expectation
maximisation algorithm is then used to find the scattering densities within
the object that is most likely to result in the observed data collected by the
muon detectors.

Existing literature (Zeng et al., 2020) have found that MLEM has better
qualitative performance than direct allocation to PoCA. However, in our
experiments, we observed that MLEM has significantly lower performance (see
Figure 10). We hypothesize that this is due to the lower muon dosages we used
in our experiments, which makes it difficult to apply statistical methods to
the reconstruction. Due to the lower performance and higher computational cost
of MLEM, we will be comparing our model against PoCA for most of our
experiments.

##  4 Methods

Figure 2: First stage of $\boldsymbol{\mu}$-Net. The muon features (initial
position, initial momentum, etc.) are passed through an MLP and placed at PoCA
scattering points within a 3D volume. If they overlap, the average is taken.

#### Scatter Operation.

Our goal is to convert a set of muon detections into a reconstruction output.
One simple way to do this is to first convert the muon detections into some 3D
representation, which can then be fed into a U-Net.

To achieve this, we first apply an MLP on the muon’s input parameters. The
output features are reshaped into a $d\times d\times d\times c$ block. We call
$d$ the point size and $c$ is the number of channels. Now, using the PoCA
algorithm, we find the muon’s point of closest approach and scatter the output
features into the voxels nearest to the PoCA point. We choose the PoCA point
as the information about the muons that scatter in a given area is the primary
information that is needed to help decide what the scattering density of that
area is. For muons that do not scatter, we place this block at a random point
along their trajectory. This provides improved performance since the vast
majority of muons will not scatter. Placing them randomly along their
trajectory provides additional information to the model that the scattering
density in that region is low.

After this, if there is overlap in the scattering voxels, the sum is taken and
a counter keeps track of how muons scatter in a given voxel. An MLP is used to
combine this counter with the other projected features.

This approach can be implemented in a memory-efficient manner using
TensorFlow’s $\mathtt{tf.scatter\\_nd}$ to avoid creating a large number of 3D
tensors in parallel. In contrast, other possibilities like converting each
muon detection into its own 3D tensor and then summing them will take up alot
of memory and is in practise, infeasible.

Nevertheless, this approach does come with some limitations. The most
prominent being that the scattering operation is fundamentally non-
differentiable. This means that we are unable to parameterise the position of
the scattering points using a neural network directly. As such, we resort to
using PoCA to get an approximate scattering point for the muons.

#### Feature Engineering.

Each muon detection is a 14-length vector consisting of the input position,
$\mathbf{x_{0}}$, output position, $\mathbf{x_{f}}$, input momentum,
$\mathbf{\hat{p}_{0}}$, output momentum, $\mathbf{\hat{p}_{f}}$, an estimate
of the momentum magnitude, $|\mathbf{p}|$ and an estimate of the scattering
angle, $|\mathbf{\hat{p}_{f}}-\mathbf{\hat{p}_{i}}|$. Although this estimate
of the scattering angle can be easily computed from the other features, its
inclusion improves performance as the model can easily know which muons should
have a larger activation because they scattered more.

Algorithm 1 Scattering Operation

Input: muon detections
$\\{\boldsymbol{\mu}_{1},\dots,\boldsymbol{\mu}_{n}\\}$, resolution $R$

Initialize $\mathbf{X}\in\mathbb{R}^{r\times r\times r\times(c+1)}$

for $i=1$ to $i=n$ do

$\mathbf{y}_{i}\leftarrow\mathsf{MLP}_{1}(\boldsymbol{\mu}_{i}),\mathbf{y}_{i}\in\mathbb{R}^{d\times
d\times d\times c}$

if $\boldsymbol{\mu}_{i}$ scatters then

Place $\mathbf{y}_{i}$ in X at $\mathsf{PoCA}(\boldsymbol{\mu}_{i})$

else

Place $\mathbf{y}_{i}$ at a random point along the muon’s trajectory

end if

Increment the corresponding last channel of X by 1 in the corresponding place

end for

$X^{\prime}\leftarrow\mathsf{MLP}_{2}(X),X^{\prime}\in\mathbb{R}^{r\times
r\times r\times c}$

return $X^{\prime}$

#### U-Net.

Now, we make use of the U-Net (Ronneberger et al., 2015) to process the
voxelised volume matrix from the first stage. In our U-Net, instead of using
the more traditional Residual Blocks (He et al., 2016), we make use of
ConvNeXt Blocks (Liu et al., 2022) which result in better performance than
Residual Blocks at a lower computational cost. Each layer of the U-Net
contains multiple such ConvNeXt blocks. Downsampling is done using Layer
Normalisation followed by a convolutional layer of strides = 2. Upsampling is
done by first applying a pointwise convolution and layer normalisation before
using nearest neighbour upsampling.

#### Model Sizes.

In our experiments, we tested out 3 models - $\mu$-Net-T, $\mu$-Net-B and
$\mu$-Net-L. These variants only differ in the number of blocks $B$ and the
number of channels $C$ in each stage. The number of parameters of each of
these models is listed as $P$. We do not scale up the model further due to
limited computational resources. We will leave further exploration of the
scaling of $\mu$-Net to future work.

  * •

$\boldsymbol{\mu}$-Net-T: $P=1.1M$
$B=(1,2,3,4,5),\;C=(8,16,32,64,128)$

  * •

$\boldsymbol{\mu}$-Net-B: $P=5.0M$
$B=(1,2,4,4,6),\;C=(16,32,64,128,256)$

  * •

$\boldsymbol{\mu}$-Net-L: $P=14.8M$
$B=(1,2,4,6,8),\;C=(24,48,96,192,384)$

#### Training Techniques.

We train our model using the AdamW optimizer with a learning rate of $2\times
10^{-3}$ and a weight decay of $4\times 10^{-}3$. The model is trained for 15
epochs.

#### Universal Approximation.

We have shown that $\mu$-Net is an universal function approximator for
continuous set functions given the model is large enough and either of the
following conditions are met:

  * •

The resolution of the reconstructed volume is sufficiently large such that
there is no overlap in the reconstructed points or,

  * •

the resolution of the reconstructed volume is finite but the point size is
large enough to cover the entire reconstructed volume.

Formally, we can define
$\chi=\\{S:S\subseteq\mathbb{R}^{m}\;\mathit{and}\;|S|=n\\}$ and
$f:\chi\rightarrow\mathbb{R}$ is a continuous set function w.r.t the Hausdorff
distance $d_{H}(\cdot,\cdot)$. Our theorem proves that $f$ can be arbitrarily
approximated by our model if the resolution is sufficiently high, or if the
resolution is fixed but the point size $d$ is the same as the resolution.

######  Theorem 4.1.

Suppose $f:\chi\rightarrow\mathbb{R}^{p}$ is a continuous set function w.r.t
$d_{H}(\cdot,\cdot)$, such that for all $\epsilon>0$, there exists some
configuration of the model parameters $\theta$ for sufficiently large $p$ or
$\phi(\eta(x_{i}))=J_{p\times d}$ (i.e. the indicator function maps to every
point), such that for any $S\in\chi$,

$\left|f(S)-\gamma_{\theta}\left(\left[\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot h_{\theta}(x_{i})\\},\right.\right.\right.\\\
\left.\left.\left.\sum_{x_{i}\in S}\\{\phi(\eta(x_{i}))\cdot J_{d\times
c}\\}\right]\right)\right|<\epsilon$ (6)

where $\gamma_{\theta}:\mathbb{R}^{p\times c}\rightarrow\mathbb{R}^{p}$ is any
continuous function, $h_{\theta}:\mathbb{R}^{m}\rightarrow\mathbb{R}^{d\times
c}$ is any continuous function, $\eta:\mathbb{R}^{m}\rightarrow\mathbb{R}$,
$\phi:\mathbb{R}^{m}\rightarrow\mathbb{R}^{p\times d}$ and $J_{d\times c}$ is
the ones matrix of shape $(d,c)$. $\eta$ represents the PoCA function that
generates a scattering point from the muon detection. $\phi$ is an indicator
function for a set of intervals of length $d$ derived from its input. The
indicator function for each of these intervals is placed along one row in the
last dimension. $\gamma_{\theta}$ and $h_{\theta}$ can be taken to be any
continuous function due to the universal approximation theorem for CNNs and
MLPs. $[\textbf{A},\textbf{B}]$ represents the concatenation of 2 matrices
along their last axes.

A brief proof is provided in the Appendix. We also note this theorem
generalises to all dimensions easily by replacing $p$ with $r\times r$,
$r\times r\times r$, etc.

The limitations of the theorem to arbitrarily large resolutions or point sizes
comes from our lack of assumptions about the PoCA and scatter operations (i.e.
$\phi(\eta(x_{i}))$. The nature of these functions depends on the interactions
between the muons and the object and is very difficult to analyse due to the
complex interactions that muons exhibit with matter.

However, despite these limitations, we hypothesise that our model performs
well with a fixed resolution and point size since most of the information
about the muon scattering is contained near the the scattering point.

##  5 Experiments

###  5.1 Experimental Setup

Our data is generated using CERN’s Geant4 simulation software (Agostinelli et
al., 2003). To generate 3D objects to be analysed by our system, we make use
of fractal noise to generate objects of various shapes. The materials of the
object is randomly chosen from a set list of materials of different radiation
lengths, taken from the IEEE BigData Competition on Muon Tomography (Wnuk et
al., 2023).

The geometry of the system can be found in Figure 3. The target object is
contained within a cube of side length 1 m. The input and output detectors are
squares of side length 2 m. They are separated from the object by a distance
of 0.5 m.

Figure 3: Simulation setup. A scale diagram (except the voxels) illustrating
the simulation setup within Geant4.

For the muon beam, we use a beam with a $\cos^{2}$ angular distribution and a
power law distribution, in accordance with characterised values of the cosmic
muon flux (Shukla & Sankrith, 2018). Muons that are calculated not to hit the
detector are killed at the start of the simulation to ensure the simulation
runs at a reasonable speed for data generation.

We use 20000 samples for the training set, 1600 samples for the validation set
and 1600 samples for the test set. Our model is implemented using TensorFlow
and trained using 2 T4 GPUs.

###  5.2 Ablations

#### Point Size.

First, we vary the point size at various dosages to see its impact on the
model’s performance. Interestingly, we find that a smaller point size of 1
results in the best performance.

Table 1: Smaller is better. The results of the model at different dosages for
various point sizes. The inference times are evaluated on 2 T4 GPUs with a
batch size of 8. The best results are bolded. The smaller point size of 1
outperforms the larger point size of 3 on almost all dosages.
Point Size | Dosage | Time$\downarrow$ | MSE$\downarrow$ | MAE$\downarrow$ | PSNR$\uparrow$
---|---|---|---|---|---
1 | 1024 | 126 ms | 0.2276 | 0.2204 | 17.1426
3 | 1024 | 134 ms | 0.2289 | 0.2366 | 17.0971
1 | 2048 | 135 ms | 0.1965 | 0.1989 | 17.7786
3 | 2048 | 143 ms | 0.1947 | 0.1949 | 17.8280
1 | 4096 | 141 ms | 0.1653 | 0.1725 | 18.5347
3 | 4096 | 178 ms | 0.1685 | 0.1741 | 18.4482
1 | 8192 | 169 ms | 0.1388 | 0.1438 | 19.2979
3 | 8192 | 219 ms | 0.1403 | 0.1465 | 19.2434
1 | 16384 | 246 ms | 0.1169 | 0.1207 | 20.0433
3 | 16384 | 318 ms | 0.1205 | 0.1340 | 19.9047
1 | 32768 | 347 ms | 0.0993 | 0.1156 | 20.7906
3 | 32768 | 574 ms | 0.1048 | 0.1224 | 20.4946

#### Estimate of Scattering Angle.

We then test out the impact of not providing an estimate of the scattering
angle. For a dosage of 1024 muons using $\mu$-Net-T, this results in a PSNR of
16.9142, in contrast with a PSNR of 17.1426 when it is provided. This
demonstrates that including the scattering angle, although it can be computed
from the other features in the muon detection, helps to improve the model’s
accuracy.

#### Random Placement of Muons.

We also test out the impact of placing the non-scattered muons at the centre
of the their trajectory, rather than at a random point along their trajectory.
For a dosage of 1024 muons using $\mu$-Net-T, this results in a PSNR of
16.8650, in contrast with a PSNR of 17.1426 when the muons are placed randomly
along their trajectory. This similarly demonstrates the effectiveness of this
method in spreading out the information about unscattered muons across the 3D
voxels, enabling a more accurate reconstruction.

###  5.3 Model Scaling

Figure 4: Model scaling has little impact. The PSNR of the different model
sizes plotted against dosage levels. We see that for low dosages, the model
size has little impact. However, at larger dosages, the impact of model size
is greater.

Now, we explore the impact of the scaling of the model on the performance. The
results are plotted in Figure 4. We notice that the improvements in
performance are actually quite small, especially for small dosages. We
hypothesise that this is because we are near the ”limit” of how good the
reconstruction can be given the U-Net’s input of the PoCA points. It is likely
that attaining a better estimate of the scattering points would lead to better
performance.

Furthermore, we notice that at larger dosages, the improvements in performance
increase, likely due to the problem being more complex, hence, providing more
room for improvement with a larger model.

###  5.4 Comparison with Traditional Algorithms

Figure 5: $\boldsymbol{\mu}$-Net outperforms PoCA. The PSNR of $\mu$-Net and
PoCA plotted against dosage levels. $\mu$-Net consistently outperforms PoCA
over all dosages and all model sizes.

#### Dosage.

We vary the dosage of muons from 1024 to 32768. The results are shown in
Figure 4. We see that as the dosage increases, the performance of the models
increase as well. It is also clear that our model, $\mu$-Net, significantly
outperforms the the traditional PoCA algorithm in Figure 5. However, we notice
that at the gradient of the graph for PoCA appearing to be increasing. Future
study is needed to ascertain if this is just a statistical anomaly or if the
PoCA indeed curves upwards.

Figure 6: $\boldsymbol{\mu}$-Net is robust to momentum error. The PSNR of
various methods plotted against percentage error in the momentum estimate.
$\mu$-Net∗ represents the model’s performance when it is finetuned on the new
data. Our model is consistently shown to be robust to perturbations and to
significantly outperform the PoCA algorithm.

#### Momentum Estimate.

We also look at how varying levels of error in the momentum will affect
predictions in Figure 6. We again find that our model significantly
outperforms the traditional PoCA algorithm. After fine-tuning, we also see
that the model’s performance stays relatively constant as the error in the
momentum increases, indicating our model is robust.

Figure 7: $\boldsymbol{\mu}$-Net is robust to detector resolution. The PSNR of
various methods plotted against detector resolution. $\mu$-Net∗ represents the
model’s performance when it is finetuned on the new data. Our model is
consistently shown to be robust to perturbations and to significantly
outperform the PoCA algorithm.

#### Detector Resolution.

Finally, we look at how the model’s performance changes with the detector
resolution in Figure 7. Again, we find that our model significantly
outperforms the PoCA algorithm. In addition, we find that our model performs
well at a variety of resolutions, showing that it is very robust. We also
notice that the performance of PoCA stays constant across all resolutions and
this is because the reconstruction volume has a resolution of $64\times
64\times 64$.

Figure 8: 2D Reconstruction Results. This is a cross-section of an object made
of 3 different materials. The ability of the model to reconstruct the
approximate shapes and differentiate materials is shown clearly. However,
there is still a significant amount of blurring. (left) ground truth, (middle-
left) $\mu$-Net, (middle-right) PoCA (right) MLEM. MLEM was performed at a
much lower resolution because it requires muon tracks to intersect with every
voxel.

Figure 9: 3D Reconstruction Results. This is a 3D plot and a cross-section of
one of the imaging targets from the testing set. The ability of the model to
reconstruct the approximate shapes and differentiate materials is shown
clearly. However, there is still a significant amount of blurring. (left)
ground truth, (middle) $\mu$-Net, (right) PoCA

#### Visual Comparison.

Now, we visually compare the reconstructions of PoCA and MLEM. These
reconstructions are shown in Figure 8 and 9. More reconstructions from
different dosages can be found in Appendix B. A dosage of 32768 muons is used.
The MLEM reconstruction is significantly worse since it had to be done using a
lower resolution, because the algorithm requires the muon tracks to pass
through every voxel, which is not possible at higher resolutions.

We see that $\mu$-Net provides superior reconstruction quality compared to
PoCA. Furthermore, $\mu$-Net is also able to distinguish different materials
by their radiation length and accurately reconstruct the approximate shape of
objects. Nevertheless, we still observe significant levels of blurring.
Further analysis of the artifacts found in the reconstruction can be found in
Appendix C.

##  6 Discussion

###  6.1 Applications

#### Nuclear Non-proliferation.

Muon tomography has been applied in detecting the presence of high-Z materials
such as radioactive materials since these materials result in a large amount
of scattering. It can also be applied for general screening of cargo for
high-Z materials (Barnes et al., 2023). Our model can be used to help improve
the quality of these reconstructions, potentially reducing the dosage required
to screen the cargo, enabling lower screening times.

#### Archaeology.

In addition, muon tomography also has numerous applications in archaeology. In
particular, the ability of muons to easily penetrate thick layers of material
such as rock, enables their use to image the interiors of large structures.
For instance, researchers have used muon tomography to discover secret
chambers within Khufu’s Pyramid (Procureur et al., 2023a, b), image
underground cavities in Mount Echia (Cimmino et al., 2019; Saracino et al.,
2017) and discover a secret ancient Greek burial chamber in the centre of
Naples (Tioukov et al., 2023). In cases where the input detections are
available, our model can be used to significantly improve the accuracy of the
reconstructions.

###  6.2 Future Work

#### No Input Muon Information.

One limitation of the model is that it depends on the presence of information
about the input muons because of its use of the PoCA algorithm. However, in
some cases, such as in some archaeological applications, the direction of the
input muons is unavailable.

#### Trajectory Prediction.

As shown in Figure 1, when muons scatter more than once, their PoCA will be
outside of the boundaries of any object. This is a significant limitation of
PoCA and sometimes results in false hotspots of scattering density in the
final prediction. One solution to this could be to attempt to predict the
muon’s full trajectory using a method like (Benton et al., 2012) given an
initial guess of the scattering density, which can be obtained from our
current model. Using this information, more accurate coordinates of scattering
points can be obtained, allowing a more accurate scattering density to be
obtained. This process can then be iterated until convergence.

#### Out-of-Distribution Generalisation.

It is not clear to what extent the model’s performance is independent of
factors such as the angular and spacial distribution of cosmic ray muons. In
addition, in real life, the spacial distribution of the scattering densities
will not be that of fractal noise, like what the model was trained on. Before
such a model can be deployed in the real-world, it would be important to check
its ability to generalise.

##  7 Conclusion

In conclusion, we have constructed a state-of-the-art model for muon
scattering tomography which outperforms traditional methods such as PoCA and
MLEM. Furthermore, we find that our model is robust to various corruptions,
with its performance barely changing when they are applied. We hope that our
research will spark further investigation into the usage of deep learning in
this field. Improvements in imaging techniques for muon scattering tomography
will have wide-ranging applications from ensuring nuclear non-proliferation to
the discovery of secret chambers in ancient structures.

## Acknowledgments

The authors would like to thank their friend Kannan Vishal for helping with
the writing of an initial simulation in Geant4, their friends Prannaya Gupta,
Kabir Jain, Mahir Hitesh for providing computing power for this project and
their teachers Mr Silas Yeem Kai Ean and Mr Ng Chee Loong for providing useful
advice.

## References

  * Agostinelli et al. (2003) Agostinelli, S., Allison, J., Amako, K., Apostolakis, J., Araujo, H., Arce, P., Asai, M., Axen, D., Banerjee, S., Barrand, G., Behner, F., Bellagamba, L., Boudreau, J., Broglia, L., Brunengo, A., Burkhardt, H., Chauvie, S., Chuma, J., Chytracek, R., Cooperman, G., Cosmo, G., Degtyarenko, P., Dell'Acqua, A., Depaola, G., Dietrich, D., Enami, R., Feliciello, A., Ferguson, C., Fesefeldt, H., Folger, G., Foppiano, F., Forti, A., Garelli, S., Giani, S., Giannitrapani, R., Gibin, D., Cadenas, J. G., González, I., Abril, G. G., Greeniaus, G., Greiner, W., Grichine, V., Grossheim, A., Guatelli, S., Gumplinger, P., Hamatsu, R., Hashimoto, K., Hasui, H., Heikkinen, A., Howard, A., Ivanchenko, V., Johnson, A., Jones, F., Kallenbach, J., Kanaya, N., Kawabata, M., Kawabata, Y., Kawaguti, M., Kelner, S., Kent, P., Kimura, A., Kodama, T., Kokoulin, R., Kossov, M., Kurashige, H., Lamanna, E., Lampén, T., Lara, V., Lefebure, V., Lei, F., Liendl, M., Lockman, W., Longo, F., Magni, S., Maire, M., Medernach, E., Minamimoto, K., de Freitas, P. M., Morita, Y., Murakami, K., Nagamatu, M., Nartallo, R., Nieminen, P., Nishimura, T., Ohtsubo, K., Okamura, M., O'Neale, S., Oohata, Y., Paech, K., Perl, J., Pfeiffer, A., Pia, M., Ranjard, F., Rybin, A., Sadilov, S., Salvo, E. D., Santin, G., Sasaki, T., Savvas, N., Sawada, Y., Scherer, S., Sei, S., Sirotenko, V., Smith, D., Starkov, N., Stoecker, H., Sulkimo, J., Takahata, M., Tanaka, S., Tcherniaev, E., Tehrani, E. S., Tropeano, M., Truscott, P., Uno, H., Urban, L., Urban, P., Verderi, M., Walkden, A., Wander, W., Weber, H., Wellisch, J., Wenaus, T., Williams, D., Wright, D., Yamada, T., Yoshida, H., and Zschiesche, D.  Geant4—a simulation toolkit.  _Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment_ , 506(3):250–303, July 2003.  doi: 10.1016/s0168-9002(03)01368-8.  URL https://doi.org/10.1016/s0168-9002(03)01368-8.
  * Barnes et al. (2023) Barnes, S., Georgadze, A., Giammanco, A., Kiisk, M., Kudryavtsev, V. A., Lagrange, M., and Pinto, O. L.  Cosmic-ray tomography for border security.  _Instruments_ , 7(1), 2023.  ISSN 2410-390X.  doi: 10.3390/instruments7010013.  URL https://www.mdpi.com/2410-390X/7/1/13.
  * Bello et al. (2020) Bello, S. A., Yu, S., Wang, C., Adam, J. M., and Li, J.  Deep learning on 3d point clouds.  _Remote Sensing_ , 12(11):1729, 2020.
  * Benton et al. (2012) Benton, C. J., Smith, N. D., Quillin, S. J., and Steer, C. A.  Most probable trajectory of a muon in a scattering medium, when input and output trajectories are known.  _Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment_ , 693:154–159, November 2012.  doi: 10.1016/j.nima.2012.07.008.  URL https://doi.org/10.1016/j.nima.2012.07.008.
  * Beringer et al. (2012) Beringer, J., Arguin, J. F., Barnett, R. M., Copic, K., Dahl, O., Groom, D. E., Lin, C. J., Lys, J., Murayama, H., Wohl, C. G., Yao, W. M., Zyla, P. A., Amsler, C., Antonelli, M., Asner, D. M., Baer, H., Band, H. R., Basaglia, T., Bauer, C. W., Beatty, J. J., Belousov, V. I., Bergren, E., Bernardi, G., Bertl, W., Bethke, S., Bichsel, H., Biebel, O., Blucher, E., Blusk, S., Brooijmans, G., Buchmueller, O., Cahn, R. N., Carena, M., Ceccucci, A., Chakraborty, D., Chen, M. C., Chivukula, R. S., Cowan, G., D'Ambrosio, G., Damour, T., de Florian, D., de Gouvêa, A., DeGrand, T., de Jong, P., Dissertori, G., Dobrescu, B., Doser, M., Drees, M., Edwards, D. A., Eidelman, S., Erler, J., Ezhela, V. V., Fetscher, W., Fields, B. D., Foster, B., Gaisser, T. K., Garren, L., Gerber, H. J., Gerbier, G., Gherghetta, T., Golwala, S., Goodman, M., Grab, C., Gritsan, A. V., Grivaz, J. F., Grünewald, M., Gurtu, A., Gutsche, T., Haber, H. E., Hagiwara, K., Hagmann, C., Hanhart, C., Hashimoto, S., Hayes, K. G., Heffner, M., Heltsley, B., Hernández-Rey, J. J., Hikasa, K., Höcker, A., Holder, J., Holtkamp, A., Huston, J., Jackson, J. D., Johnson, K. F., Junk, T., Karlen, D., Kirkby, D., Klein, S. R., Klempt, E., Kowalewski, R. V., Krauss, F., Kreps, M., Krusche, B., Kuyanov, Y. V., Kwon, Y., Lahav, O., Laiho, J., Langacker, P., Liddle, A., Ligeti, Z., Liss, T. M., Littenberg, L., Lugovsky, K. S., Lugovsky, S. B., Mannel, T., Manohar, A. V., Marciano, W. J., Martin, A. D., Masoni, A., Matthews, J., Milstead, D., Miquel, R., Mönig, K., Moortgat, F., Nakamura, K., Narain, M., Nason, P., Navas, S., Neubert, M., Nevski, P., Nir, Y., Olive, K. A., Pape, L., Parsons, J., Patrignani, C., Peacock, J. A., Petcov, S. T., Piepke, A., Pomarol, A., Punzi, G., Quadt, A., Raby, S., Raffelt, G., Ratcliff, B. N., Richardson, P., Roesler, S., Rolli, S., Romaniouk, A., Rosenberg, L. J., Rosner, J. L., Sachrajda, C. T., Sakai, Y., Salam, G. P., Sarkar, S., Sauli, F., Schneider, O., Scholberg, K., Scott, D., Seligman, W. G., Shaevitz, M. H., Sharpe, S. R., Silari, M., Sjöstrand, T., Skands, P., Smith, J. G., Smoot, G. F., Spanier, S., Spieler, H., Stahl, A., Stanev, T., Stone, S. L., Sumiyoshi, T., Syphers, M. J., Takahashi, F., Tanabashi, M., Terning, J., Titov, M., Tkachenko, N. P., Törnqvist, N. A., Tovey, D., Valencia, G., van Bibber, K., Venanzoni, G., Vincter, M. G., Vogel, P., Vogt, A., Walkowiak, W., Walter, C. W., Ward, D. R., Watari, T., Weiglein, G., Weinberg, E. J., Wiencke, L. R., Wolfenstein, L., Womersley, J., Woody, C. L., Workman, R. L., Yamamoto, A., Zeller, G. P., Zenin, O. V., Zhang, J., Zhu, R. Y., Harper, G., Lugovsky, V. S., and Schaffner, P.  Review of particle physics.  _Physical Review D_ , 86(1), July 2012.  doi: 10.1103/physrevd.86.010001.  URL https://doi.org/10.1103/physrevd.86.010001.
  * Borselli et al. (2022) Borselli, D., Beni, T., Bonechi, L., Bongi, M., Brocchini, D., Casagli, N., Ciaranfi, R., Cimmino, L., Ciulli, V., D’Alessandro, R., Dini, A., Frosin, C., Gigli, G., Gonzi, S., Guideri, S., Lombardi, L., Nocentini, M., and Saracino, G.  Three-dimensional muon imaging of cavities inside the temperino mine (italy).  _Scientific Reports_ , 12(1), December 2022.  doi: 10.1038/s41598-022-26393-7.  URL https://doi.org/10.1038/s41598-022-26393-7.
  * Chatzidakis et al. (2018) Chatzidakis, S., Liu, Z., Hayward, J. P., and Scaglione, J. M.  A generalized muon trajectory estimation algorithm with energy loss for application to muon tomography.  _Journal of Applied Physics_ , 123(12), mar 2018.  doi: 10.1063/1.5024671.  URL https://doi.org/10.1063%2F1.5024671.
  * Çiçek et al. (2016) Çiçek, Ö., Abdulkadir, A., Lienkamp, S. S., Brox, T., and Ronneberger, O.  3d u-net: learning dense volumetric segmentation from sparse annotation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2016: 19th International Conference, Athens, Greece, October 17-21, 2016, Proceedings, Part II 19_ , pp. 424–432. Springer, 2016\.
  * Cimmino et al. (2019) Cimmino, L., Baccani, G., Noli, P., Amato, L., Ambrosino, F., Bonechi, L., Bongi, M., Ciulli, V., D’Alessandro, R., D’Errico, M., Gonzi, S., Melon, B., Minin, G., Saracino, G., Scognamiglio, L., Strolin, P., and Viliani, L.  3d muography for the search of hidden cavities.  _Scientific Reports_ , 9(1), February 2019.  doi: 10.1038/s41598-019-39682-5.  URL https://doi.org/10.1038/s41598-019-39682-5.
  * He et al. (2016) He, K., Zhang, X., Ren, S., and Sun, J.  Deep residual learning for image recognition.  In _2016 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2016, Las Vegas, NV, USA, June 27-30, 2016_ , pp. 770–778. IEEE Computer Society, 2016.  doi: 10.1109/CVPR.2016.90.  URL https://doi.org/10.1109/CVPR.2016.90.
  * Ho et al. (2021) Ho, N.-V., Nguyen, T., Diep, G.-H., Le, N., and Hua, B.-S.  Point-unet: A context-aware point-based neural network for volumetric segmentation.  In _Medical Image Computing and Computer Assisted Intervention–MICCAI 2021: 24th International Conference, Strasbourg, France, September 27–October 1, 2021, Proceedings, Part I 24_ , pp. 644–655. Springer, 2021.
  * Hou et al. (2021) Hou, L., Zhang, Q., Yang, J., Cai, X., Yao, Q., Huo, Y., and Chen, Q.  A novel reconstruction algorithm based on density clustering for cosmic-ray muon scattering inspection.  _Nuclear Engineering and Technology_ , 53(7):2348–2356, July 2021.  doi: 10.1016/j.net.2021.01.014.  URL https://doi.org/10.1016/j.net.2021.01.014.
  * Jonkmans et al. (2013) Jonkmans, G., Anghel, V., Jewett, C., and Thompson, M.  Nuclear waste imaging and spent fuel verification by muon tomography.  _Annals of Nuclear Energy_ , 53:267–273, March 2013.  doi: 10.1016/j.anucene.2012.09.011.  URL https://doi.org/10.1016/j.anucene.2012.09.011.
  * Liu et al. (2021) Liu, Z., Lin, Y., Cao, Y., Hu, H., Wei, Y., Zhang, Z., Lin, S., and Guo, B.  Swin transformer: Hierarchical vision transformer using shifted windows.  In _Proceedings of the IEEE/CVF international conference on computer vision_ , pp. 10012–10022, 2021.
  * Liu et al. (2022) Liu, Z., Mao, H., Wu, C., Feichtenhofer, C., Darrell, T., and Xie, S.  A convnet for the 2020s.  In _IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022_ , pp. 11966–11976. IEEE, 2022.  doi: 10.1109/CVPR52688.2022.01167.  URL https://doi.org/10.1109/CVPR52688.2022.01167.
  * Procureur et al. (2023a) Procureur, S., Morishima, K., Kuno, M., Manabe, Y., Kitagawa, N., Nishio, A., Gomez, H., Attié, D., Sakakibara, A., Hikata, K., Moto, M., Mandjavidze, I., Magnier, P., Lehuraux, M., Benoit, T., Calvet, D., Coppolani, X., Kebbiri, M., Mas, P., Helal, H., Tayoubi, M., Marini, B., Serikoff, N., Anwar, H., Steiger, V., Takasaki, F., Fujii, H., Satoh, K., Kodama, H., Hayashi, K., Gable, P., Guerriero, E., Mouret, J.-B., Elnady, T., Elshayeb, Y., and Elkarmoty, M.  Precise characterization of a corridor-shaped structure in khufu’s pyramid by observation of cosmic-ray muons.  _Nature Communications_ , 14(1), March 2023a.  doi: 10.1038/s41467-023-36351-0.  URL https://doi.org/10.1038/s41467-023-36351-0.
  * Procureur et al. (2023b) Procureur, S., Morishima, K., Kuno, M., Manabe, Y., Kitagawa, N., Nishio, A., Gomez, H., Attié, D., Sakakibara, A., Hikata, K., Moto, M., Mandjavidze, I., Magnier, P., Lehuraux, M., Benoit, T., Calvet, D., Coppolani, X., Kebbiri, M., Mas, P., Helal, H., Tayoubi, M., Marini, B., Serikoff, N., Anwar, H., Steiger, V., Takasaki, F., Fujii, H., Satoh, K., Kodama, H., Hayashi, K., Gable, P., Guerriero, E., Mouret, J.-B., Elnady, T., Elshayeb, Y., and Elkarmoty, M.  Precise characterization of a corridor-shaped structure in khufu’s pyramid by observation of cosmic-ray muons.  _Nature Communications_ , 14(1), March 2023b.  doi: 10.1038/s41467-023-36351-0.  URL https://doi.org/10.1038/s41467-023-36351-0.
  * Ronneberger et al. (2015) Ronneberger, O., Fischer, P., and Brox, T.  U-net: Convolutional networks for biomedical image segmentation.  In Navab, N., Hornegger, J., III, W. M. W., and Frangi, A. F. (eds.), _Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015 - 18th International Conference Munich, Germany, October 5 - 9, 2015, Proceedings, Part III_ , volume 9351 of _Lecture Notes in Computer Science_ , pp. 234–241. Springer, 2015.  doi: 10.1007/978-3-319-24574-4“˙28.  URL https://doi.org/10.1007/978-3-319-24574-4_28.
  * Rossi & Greisen (1941) Rossi, B. and Greisen, K.  Cosmic-ray theory.  _Reviews of Modern Physics_ , 13(4):240, 1941\.
  * Saracino et al. (2017) Saracino, G., Amato, L., Ambrosino, F., Antonucci, G., Bonechi, L., Cimmino, L., Consiglio, L., Alessandro, R. D., Luzio, E. D., Minin, G., Noli, P., Scognamiglio, L., Strolin, P., and Varriale, A.  Imaging of underground cavities with cosmic-ray muons from observations at mt. echia (naples).  _Scientific Reports_ , 7(1), April 2017.  doi: 10.1038/s41598-017-01277-3.  URL https://doi.org/10.1038/s41598-017-01277-3.
  * Saracino et al. (2018) Saracino, G., Ambrosino, F., Bonechi, L., Cimmino, L., D'Alessandro, R., D'Errico, M., Noli, P., Scognamiglio, L., and Strolin, P.  Applications of muon absorption radiography to the fields of archaeology and civil engineering.  _Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences_ , 377(2137):20180057, December 2018.  doi: 10.1098/rsta.2018.0057.  URL https://doi.org/10.1098/rsta.2018.0057.
  * Schulte et al. (2008) Schulte, R. W., Penfold, S. N., Tafas, J. T., and Schubert, K. E.  A maximum likelihood proton path formalism for application in proton computed tomography.  _Medical Physics_ , 35(11):4849–4856, oct 2008\.  doi: 10.1118/1.2986139.  URL https://doi.org/10.1118%2F1.2986139.
  * Schultz (2003) Schultz, L. J.  _Cosmic Ray Muon Radiography_.  PhD thesis, Portland State University, 2003.
  * Schultz et al. (2007) Schultz, L. J., Blanpied, G. S., Borozdin, K. N., Fraser, A. M., Hengartner, N. W., Klimenko, A. V., Morris, C. L., Orum, C., and Sossong, M. J.  Statistical reconstruction for cosmic ray muon tomography.  _IEEE transactions on Image Processing_ , 16(8):1985–1993, 2007.
  * Shukla & Sankrith (2018) Shukla, P. and Sankrith, S.  Energy and angular distributions of atmospheric muons at the earth.  _International Journal of Modern Physics A_ , 33(30):1850175, October 2018.  doi: 10.1142/s0217751x18501750.  URL https://doi.org/10.1142/s0217751x18501750.
  * Stapleton et al. (2014) Stapleton, M., Burns, J., Quillin, S., and Steer, C.  Angle statistics reconstruction: a robust reconstruction algorithm for muon scattering tomography.  _Journal of Instrumentation_ , 9(11):P11019–P11019, November 2014.  doi: 10.1088/1748-0221/9/11/p11019.  URL https://doi.org/10.1088/1748-0221/9/11/p11019.
  * Szczykutowicz et al. (2022) Szczykutowicz, T. P., Toia, G. V., Dhanantwari, A., and Nett, B.  A review of deep learning CT reconstruction: Concepts, limitations, and promise in clinical practice.  _Current Radiology Reports_ , 10(9):101–115, July 2022.  doi: 10.1007/s40134-022-00399-5.  URL https://doi.org/10.1007/s40134-022-00399-5.
  * Thomay et al. (2012) Thomay, C., Baesso, P., Cussans, D., Davies, J., Glaysher, P., Quillin, S., Robertson, S., Steer, C., Vassallo, C., and Velthuis, J.  A novel technique to detect special nuclear material using cosmic rays.  _Geoscientific Instrumentation, Methods and Data Systems_ , 1(2):235–238, December 2012.  doi: 10.5194/gi-1-235-2012.  URL https://doi.org/10.5194/gi-1-235-2012.
  * Tioukov et al. (2019) Tioukov, V., Alexandrov, A., Bozza, C., Consiglio, L., D’Ambrosio, N., Lellis, G. D., Sio, C. D., Giudicepietro, F., Macedonio, G., Miyamoto, S., Nishiyama, R., Orazi, M., Peluso, R., Sheshukov, A., Sirignano, C., Stellacci, S. M., Strolin, P., and Tanaka, H. K. M.  First muography of stromboli volcano.  _Scientific Reports_ , 9(1), April 2019.  doi: 10.1038/s41598-019-43131-8.  URL https://doi.org/10.1038/s41598-019-43131-8.
  * Tioukov et al. (2023) Tioukov, V., Morishima, K., Leggieri, C., Capriuoli, F., Kitagawa, N., Kuno, M., Manabe, Y., Nishio, A., Alexandrov, A., Gentile, V., Iuliano, A., and Lellis, G. D.  Hidden chamber discovery in the underground hellenistic necropolis of neapolis by muography.  _Scientific Reports_ , 13(1), April 2023.  doi: 10.1038/s41598-023-32626-0.  URL https://doi.org/10.1038/s41598-023-32626-0.
  * Vaswani et al. (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L., and Polosukhin, I.  Attention is all you need.  In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pp. 5998–6008, 2017.  URL https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html.
  * Wang et al. (2009) Wang, G., Schultz, L., and Qi, J.  Bayesian image reconstruction for improving detection performance of muon tomography.  _IEEE Transactions on Image Processing_ , 18(5):1080–1089, May 2009.  doi: 10.1109/tip.2009.2014423.  URL https://doi.org/10.1109/tip.2009.2014423.
  * Wnuk et al. (2023) Wnuk, M., Dziuba, J., Janusz, A., and Slezak, D.  Ieee bigdata 2023 cup: Object recognition with muon tomography using cosmic rays, 2023.  URL https://knowledgepit.ai/object-recognition-with-muon-tomography/.
  * Yi et al. (2014) Yi, H., Zeng, Z., Yu, B., Cheng, J., Zhao, Z., Wang, X., Zeng, M., and Wang, Y.  Bayesian-theory-based most probable trajectory reconstruction algorithm in cosmic ray muon tomography.  In _2014 IEEE Nuclear Science Symposium and Medical Imaging Conference (NSS/MIC)_ , pp. 1–4, 2014.  doi: 10.1109/NSSMIC.2014.7431084.
  * Zaheer et al. (2017) Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R., and Smola, A. J.  Deep sets.  In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/f22e4747da1aa27e363d86d40ff442fe-Paper.pdf.
  * Zeng et al. (2020) Zeng, W., Zeng, M., Pan, X., Zeng, Z., Ma, H., and Cheng, J.  Principle study of image reconstruction algorithms in muon tomography.  _Journal of Instrumentation_ , 15(02):T02005, 2020\.

##  Appendix A Proofs

To prove Theorem 4.1, we shall split the theorem into the arbitrarily large
resolution case and the arbitrary large point size case.

###  A.1 Arbitrarily Large Resolution

######  Theorem A.1.

Suppose $f:\chi\rightarrow\mathbb{R}^{p}$ is a continuous set function w.r.t
$d_{H}(\cdot,\cdot)$, such that for all $\epsilon>0$, there exists some
configuration of the model parameters $\theta$ for sufficiently large $p$ or
$\phi(\eta(x_{i}))=J_{p\times d}$ (i.e. the indicator function maps to every
point), such that for any $S\in\chi$,

$\left|f(S)-\gamma_{\theta}\left(\left[\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot h_{\theta}(x_{i})\\},\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot J_{d\times c}\\}\right]\right)\right|<\epsilon$

where $\gamma_{\theta}:\mathbb{R}^{p\times c}\rightarrow\mathbb{R}^{p}$ is any
continuous function, $h_{\theta}:\mathbb{R}^{m}\rightarrow\mathbb{R}^{d\times
c}$ is any continuous function, $\eta:\mathbb{R}^{m}\rightarrow\mathbb{R}$,
$\phi:\mathbb{R}^{m}\rightarrow\mathbb{R}^{p\times d}$ and $J_{d\times c}$ is
the ones matrix of shape $(d,c)$. $\eta$ represents the PoCA function that
generates a scattering point from the muon detection. $\phi$ is an indicator
function for a set of intervals of length $d$ derived from its input. The
indicator function for each of these intervals is placed along one row in the
last dimension. $\gamma_{\theta}$ and $h_{\theta}$ can be taken to be any
continuous function due to the universal approximation theorem for CNNs and
MLPs. $[\textbf{A},\textbf{B}]$ represents the concatenation of 2 matrices
along their last axes.

###### Proof.

The idea is that for a sufficiently large $p$, the indicator functions of
$\phi$ will not overlap, allowing $S$ to be recovered exactly using an inverse
function $T$.

Let $h_{\theta}$ simply be the identity function. 111Since $h_{\theta}$ is the
identity function, $c=m$ Now, consider a function
$\mathcal{T}:\mathbb{R}^{p\times m}\rightarrow\chi$,
$\mathcal{T}(X)=\\{x:|x|>0,x\in\mathbb{R}^{m},\textit{x is an entry in the
last dimension of X}\\}$. Now, we define $\gamma_{\theta}$ as
$f\circ\mathcal{T}$. Clearly,

$\left|f(S)-f\left(\mathcal{T}\left(\left[\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot h_{\theta}(x_{i})\\},\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot J_{d\times
c}\\}\right]\right)\right)\right|=\left|f(S)-f(S)\right|=0<\epsilon$

∎

###  A.2 Arbitrarily Large Point Size

######  Theorem A.2.

Suppose $f:\chi\rightarrow\mathbb{R}^{p}$ is a continuous set function w.r.t
$d_{H}(\cdot,\cdot)$, such that for all $\epsilon>0$, there exists some
configuration of the model parameters $\theta$ for sufficiently large $p$ or
$\phi(\eta(x_{i}))=J_{p\times d}$ (i.e. the indicator function maps to every
point), such that for any $S\in\chi$,

$\left|f(S)-\gamma_{\theta}\left(\left[\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot h_{\theta}(x_{i})\\},\sum_{x_{i}\in
S}\\{\phi(\eta(x_{i}))\cdot J_{d\times c}\\}\right]\right)\right|<\epsilon$

where $\gamma_{\theta}:\mathbb{R}^{p\times c}\rightarrow\mathbb{R}^{p}$ is any
continuous function, $h_{\theta}:\mathbb{R}^{m}\rightarrow\mathbb{R}^{d\times
c}$ is any continuous function, $\eta:\mathbb{R}^{m}\rightarrow\mathbb{R}$,
$\phi:\mathbb{R}^{m}\rightarrow\mathbb{R}^{p\times d}$ and $J_{d\times c}$ is
the ones matrix of shape $(d,c)$. $\eta$ represents the PoCA function that
generates a scattering point from the muon detection. $\phi$ is an indicator
function for a set of intervals of length $d$ derived from its input. The
indicator function for each of these intervals is placed along one row in the
last dimension. $\gamma_{\theta}$ and $h_{\theta}$ can be taken to be any
continuous function due to the universal approximation theorem for CNNs and
MLPs. $[\textbf{A},\textbf{B}]$ represents the concatenation of 2 matrices
along their last axis.

###### Proof.

In the case of $\phi(\eta(x_{i}))=J_{p\times d}$, the theorem reduces to

$\left|f(S)-\gamma_{\theta}\left(\left[\sum_{x_{i}\in
S}\\{h_{\theta}(x_{i})\\},J_{p\times 1}\right]\right)\right|<\epsilon$

which can be equivalently expressed as

$\left|f(S)-\gamma_{\theta}\left(\sum_{x_{i}\in
S}\\{h_{\theta}(x_{i})\\}\right)\right|<\epsilon$

given that $\gamma_{\theta}$ is a universal function approximator.

By invoking Theorem 7 in (Zaheer et al., 2017), any continuous permutation
invariant set function can be decomposed into the form

$f(S)=\rho(\sum_{x_{i}\in S}\\{\psi(x_{i})\\})$

where $\rho$ and $\psi$ are continuous functions. By the universal
approximation theorem for CNNs and MLPs, there exists some $\theta$ such that

$\rho(\sum_{x_{i}\in S}\psi(x_{i}))=\gamma_{\theta}\left(\sum_{x_{i}\in
S}\\{h_{\theta}(x_{i})\\}\right)$

Therefore,

$\left|f(S)-\gamma_{\theta}\left(\sum_{x_{i}\in
S}\\{h_{\theta}(x_{i})\\}\right)\right|=\left|f(S)-f(S)\right|<\epsilon$

∎

##  Appendix B Reconstruction Results

$\sf{Ground\;Truth}$ | $\sf{1024}$ | $\sf{2048}$ | $\sf{4096}$ | $\sf{8192}$ | $\sf{16384}$ | $\sf{32768}$ |
---|---|---|---|---|---|---|---
|  |  |  |  |  |  |
|  |  |  |  |  |  |
|  |  |  |  |  |  |
|  |  |  |  |  |  |
|  |  |  |  |  |  |
|  |  |  |  |  |  |
|  |  |  |  |  |  |
|  |  |  |  |  |  |
Figure 10: 2D Cross-sections. 2D cross-sections of the 3D reconstructions produced by $\mu$-Net-L at various dosages. The improvement in reconstruction quality as the dosage increases can be seen clearly. We also see that some cross-sections appear to have worse cross-sections. This is because the materials being reconstructed have a high radiation length, so the muons do not scatter very much. $\sf{Ground\;Truth}$ | $\sf{1024}$ | $\sf{2048}$ | $\sf{4096}$ |
---|---|---|---|---
|  |  |  |
| $\sf{8192}$ | $\sf{16384}$ | $\sf{32768}$ |
|  |  |  |
$\sf{Ground\;Truth}$ | $\sf{1024}$ | $\sf{2048}$ | $\sf{4096}$ |
|  |  |  |
| $\sf{8192}$ | $\sf{16384}$ | $\sf{32768}$ |
|  |  |  |
Figure 11: 3D Reconstructions. More 3D reconstructions produced by $\mu$-Net-L at various dosages. The improvement in reconstruction quality as the dosage increases can be seen clearly $\sf{Ground\;Truth}$ | $\sf{1024}$ | $\sf{2048}$ | $\sf{4096}$ |
---|---|---|---|---
|  |  |  |
| $\sf{8192}$ | $\sf{16384}$ | $\sf{32768}$ |
|  |  |  |
$\sf{Ground\;Truth}$ | $\sf{1024}$ | $\sf{2048}$ | $\sf{4096}$ |
|  |  |  |
| $\sf{8192}$ | $\sf{16384}$ | $\sf{32768}$ |
|  |  |  |
Figure 12: 3D Reconstructions (continued). Even more 3D reconstructions
produced by $\mu$-Net-L at various dosages. The improvement in reconstruction
quality as the dosage increases can be seen clearly

##  Appendix C Artifact Analysis

#### Little Squares.

We can see these little square in Figure 10 (first row). These are caused by
the muons that do not scatter and are placed randomly along their
trajectories. Since the model typically will see points placed within the
voxels corresponding to there being actual material there, there is a slightly
larger scattering density predicted at these regions where there should be
nothing. These artifacts are only visible within the cross-section when there
is nothing else inside (as is the case for the first row of Figure 10).

$\sf{Ground\;Truth}$ | $\sf{1024}$ | $\sf{2048}$ | $\sf{4096}$ | $\sf{8192}$ | $\sf{16384}$ | $\sf{32768}$ |
---|---|---|---|---|---|---|---
|  l7.2mmb2.6mm  |  l7.2mmb2.6mm  |  l7.2mmb2.6mm  |  |  |  |
|  l10mmb5.5mm  |  l10mmb5.5mm  |  l10mmb3.5mm  |  l10mmb3.5mm  |  l10mmb3.5mm  |  l10mmb3.5mm  |
Figure 13: Hotspots. 2D cross-sections of the 3D reconstructions produced by
$\mu$-Net-L at various dosages. The false PoCA hotspots are circled in red. We
see that as the dosage increases, these hotposts fade in prominence.

#### False Hotspots.

We can see false hotspots in Figure 10 (2nd and 3rd law row) and highlighted
in Figure 13. These are the result of muons scattering twice when passing
through materials. PoCA assumes that muons scatter only once. This means that
if a muon scatters twice, its PoCA point will end up somewhere end the
midpoint of its actual scattering points. We also notice that as the dosage
increases, these hotspots tend to fade in prominence. This is likely because
with a larger dosages, the model is better able to distinguish between real
scattering points and these false hotspots.

$\sf{Ground\;Truth}$ | $\sf{Reconstructions}$ |  |  |  |  |
---|---|---|---|---|---|---
|  |  |  |  |  |
|  |  |  |  |  |
Figure 14: Distortions. Several sample reconstructions of the same target
object with a dosage of 8192. Each reconstruction uses a different sample of
8192 muons from the distribution of muon detections.

#### Distortions and Blurring.

In Figure 14, many of the objects in the cross-section are visibly distorted,
with these distortions shifting as the dosage increases (particularly
noticeable in last row and 2nd row). This is due to the fundamentally random
nature of this type of tomography and the low dosages of muons. Some
scattering points will inevitably fall outside of what is the actual material.
The model will then assume that these are part of the actual material, leading
to blurring and distortions due to the fundamentally random nature of where
these points will be found.

##  Appendix D Raw Results

Table 2: Model Scaling. The results of the model at different dosages for
various sizes. The inference times are evaluated on 2 T4 GPUs with a batch
size of 8. The best results are bolded. For PoCA, the inference times are
evaluated on a P100 GPU.
Model | Dosage | Time$\downarrow$ | MSE$\downarrow$ | MAE$\downarrow$ | PSNR$\uparrow$
---|---|---|---|---|---
$\mu$-Net-T | 1024 | 126 ms | 0.2276 | 0.2204 | 17.1426
$\mu$-Net-B | 1024 | 200 ms | 0.2318 | 0.2313 | 17.0255
$\mu$-Net-L | 1024 | 288 ms | 0.2295 | 0.2178 | 17.0646
PoCA | 1024 | 22.5s | 0.4595 | 0.2447 | 13.6627
$\mu$-Net-T | 2048 | 135 ms | 0.1965 | 0.1989 | 17.7786
$\mu$-Net-B | 2048 | 208 ms | 0.1936 | 0.1911 | 17.8486
$\mu$-Net-L | 2048 | 306 ms | 0.1929 | 0.1918 | 17.8633
PoCA | 2048 | 43.8s | 0.4338 | 0.2465 | 13.9112
$\mu$-Net-T | 4096 | 141 ms | 0.1653 | 0.1725 | 18.5347
$\mu$-Net-B | 4096 | 218 ms | 0.1649 | 0.1738 | 18.5504
$\mu$-Net-L | 4096 | 301 ms | 0.1644 | 0.1697 | 18.5388
PoCA | 4096 | 79.8s | 0.3950 | 0.2466 | 14.3228
$\mu$-Net-T | 8192 | 169 ms | 0.1388 | 0.1438 | 19.2979
$\mu$-Net-B | 8192 | 234 ms | 0.1350 | 0.1457 | 19.3958
$\mu$-Net-L | 8192 | 325 ms | 0.1348 | 0.1389 | 19.4232
PoCA | 8192 | 164s | 0.3660 | 0.2420 | 15.0769
$\mu$-Net-T | 16384 | 246 ms | 0.1169 | 0.1207 | 20.0433
$\mu$-Net-B | 16384 | 293 ms | 0.1118 | 0.1238 | 20.2685
$\mu$-Net-L | 16384 | 384 ms | 0.1062 | 0.1180 | 20.4322
PoCA | 16384 | 310s | 0.3285 | 0.2315 | 15.5586
$\mu$-Net-T | 32768 | 347 ms | 0.0993 | 0.1156 | 20.7906
$\mu$-Net-B | 32768 | 434 ms | 0.0919 | 0.1040 | 21.1595
$\mu$-Net-L | 32768 | 538 ms | 0.0875 | 0.0983 | 21.3530
PoCA | 32768 | 612s | 0.3092 | 0.2258 | 17.1091
Table 3: Detector Resolutions. Results of various methods for resolutions of
the detector. $\mu$-Net-T∗ indicates that the model was finetuned on the new
data for 10 epochs. The best results are bolded.
Model | Detector Resolution | MSE$\downarrow$ | MAE$\downarrow$ | PSNR$\uparrow$
---|---|---|---|---
$\mu$-Net-T | $64\times 64$ | 0.2545 | 0.2647 | 16.5660
$\mu$-Ne-Tt∗ | $64\times 64$ | 0.2317 | 0.2107 | 17.0075
PoCA | $64\times 64$ | 0.5210 | 0.2690 | 13.1784
$\mu$-Net-T | $128\times 128$ | 0.2259 | 0.2383 | 17.1177
$\mu$-Net-T∗ | $128\times 128$ | 0.2174 | 0.2349 | 17.3046
PoCA | $128\times 128$ | 0.5226 | 0.2696 | 13.1625
$\mu$-Net-T | $256\times 256$ | 0.2123 | 0.2221 | 17.4097
$\mu$-Net-T∗ | $256\times 256$ | 0.2068 | 0.2226 | 17.5387
PoCA | $256\times 256$ | 0.5210 | 0.2690 | 13.1784
$\mu$-Net-T | $1024\times 1024$ | 0.1991 | 0.2057 | 17.7123
$\mu$-Net-T∗ | $1024\times 1024$ | 0.2035 | 0.2119 | 17.6165
PoCA | $1024\times 1024$ | 0.5210 | 0.2690 | 13.1784
$\mu$-Net-T | $2048\times 2048$ | 0.1970 | 0.2023 | 17.7616
$\mu$-Net-T∗ | $2048\times 2048$ | 0.2013 | 0.1901 | 17.6695
PoCA | $2048\times 2048$ | 0.5210 | 0.2690 | 13.1784
$\mu$-Net-T | $\infty$ | 0.1936 | 0.1911 | 17.8486
PoCA | $\infty$ | 0.4338 | 0.2465 | 13.9112
Table 4: Momentum Error. Results of various models for different levels of
error in the momentum estimate. $\mu$-Net∗ indicates that the model was
finetuned on the new data for 10 epochs. The best results are bolded.
Model | $\Delta\mathbf{p}$ | MSE$\downarrow$ | MAE$\downarrow$ | PSNR$\uparrow$
---|---|---|---|---
$\mu$-Net | 0% | 0.1951 | 0.1977 | 17.8110
$\mu$-Net∗ | 0% | 0.1920 | 0.1938 | 17.8865
PoCA | 0% | 0.4224 | 0.2442 | 14.0280
$\mu$-Net | 20% | 0.1965 | 0.1989 | 17.7786
PoCA | 20% | 0.4338 | 0.2465 | 13.9112
$\mu$-Net | 40% | 0.2004 | 0.2033 | 17.6844
$\mu$-Net∗ | 40% | 0.1987 | 0.1940 | 17.7258
PoCA | 40% | 0.4577 | 0.2515 | 13.6717
$\mu$-Net | 60% | 0.2230 | 0.2207 | 17.2291
$\mu$-Net∗ | 60% | 0.1989 | 0.1920 | 17.7326
PoCA | 60% | 0.5316 | 0.2616 | 13.1310
$\mu$-Net | 80% | 0.2761 | 0.2558 | 16.2882
$\mu$-Net∗ | 80% | 0.2020 | 0.2023 | 17.6497
PoCA | 80% | 0.6228 | 0.2730 | 12.5438
$\mu$-Net | 100% | 0.3293 | 0.2866 | 15.6005
$\mu$-Net∗ | 100% | 0.2000 | 0.2014 | 17.7015
PoCA | 100% | 0.8279 | 0.2933 | 11.5817