# Graph-RISE: Graph-Regularized Image Semantic Embedding

Da-Cheng Juan, Chun-Ta Lu, Zhen Li, Futang Peng, Aleksei Timofeev, Yi-Ting
Chen, Yaxi Gao, Tom Duerig, Andrew Tomkins, Sujith Ravi  Google AIMountain
ViewCA
<EMAIL_ADDRESS>
(2018)

###### Abstract.

Learning image representations to capture fine-grained semantics has been a
challenging and important task enabling many applications such as image search
and clustering. In this paper, we present Graph-Regularized Image Semantic
Embedding (Graph-RISE), a large-scale neural graph learning framework that
allows us to train embeddings to discriminate an unprecedented
$\mathcal{O}(40M)$ ultra-fine-grained semantic labels. Graph-RISE outperforms
state-of-the-art image embedding algorithms on several evaluation tasks,
including image classification and triplet ranking. We provide case studies to
demonstrate that, qualitatively, image retrieval based on Graph-RISE
effectively captures semantics and, compared to the state-of-the-art,
differentiates nuances at levels that are closer to human-perception.

Image embeddings, semantic understanding, graph regularization

††copyright: rightsretained††journalyear: 2018††conference: Woodstock ’18: ACM
Symposium on Neural Gaze Detection; June 03–05, 2018; Woodstock,
NY††booktitle: Woodstock ’18: ACM Symposium on Neural Gaze Detection, June
03–05, 2018, Woodstock, NY††doi: 10.1145/1122445.1122456††isbn:
978-1-4503-9999-9/18/06††article: 4††price: 15.00††ccs: Computing
methodologies Image representations††ccs: Information systems Web searching
and information discovery

##  1\. Introduction

“Is it possible to learn image content descriptors (a.k.a., embeddings) that
capture image semantics and similarity close to human perception?” Learning
image embeddings that capture fine-grained semantics is the core of many
modern image-related applications such as image search, either querying by
traditional keywords or by an example query image (Murphy-Chutorian et al.,
2015). Learning such embeddings is a challenging task, partly due to the large
variations seen among images that belong to the same category or class.

Several previous works consider category-level image semantics (Hadsell et
al., 2006; Taylor et al., 2011), in which two images are considered
semantically similar if they belong to the same category. As illustrated in
Figure 1, category-level similarity may not be sufficient for modern vision-
based applications such as query-by-image, which often require the distinction
of nuances among images within the same category.

Figure 1. Spectrum of image semantic similarity. We provide six image examples
(two for each granularity) to illustrate the difference from coarser (left) to
ultra-fine granularity (right). We refer to ultra fine-grained as “instance-
level” to contrast with category-level and fine-grained semantics.

Recently, deep ranking models (Wang et al., 2014; Wu et al., 2017) have been
proposed to learn fine-grained image similarity. These ranking models are
trained with image triplets, where each entry contains {query image, positive
image, negative image}. The goal is to rank (query, positive image) as more
similar than (query, negative image); this training formulation can encode
distinctions that are as fine-grained as the construction of the triplets
allows. In practice, however, it becomes increasingly difficult to generate a
large corpus of triplets that encode sufficiently fine-grained distinctions by
ensuring the negative image is similar enough to the query image, but not too
similar. Furthermore, since human raters need to be involved to provide the
triplet ranking ground truth, collecting high quality image triplets for
training is costly and labor-intensive. We instead propose moving from triplet
learning to a classification framework that learns embeddings capable of
associating an image to one of a large number of possible query strings.

Such an approach produces image embeddings that are predictive of queries that
might lead to the image. In addition, we also obtain similarity data between
the images themselves, encoding for example the fact that two images were both
clicked in a particular setting. This relational data encodes important
aspects of human image perception but is not easily encapsulated by labels
(i.e., image-label pairs for training). To incorporate image-image similarity
into the training we employ neural graph learning, in which the model is
trained by minimizing the supervised loss combined with a graph regularizer
that drives the model to reduce embedding distance between similar image
pairs.

To the best of our knowledge, this work brings the following contributions:

  * •

Effective embedding learning via large scale classification. We formulate the
problem of image embedding learning as an image classification task at an
unprecedented scale, with label space (i.e., total number of classes) in
${\mathcal{O}}(40M)$ and the number of images in ${\mathcal{O}}(260M)$. This
is the largest scale in terms of number of classes, and one of the largest in
terms of images used for learning image embeddings. Furthermore, the proposed
model is one of the largest vision models in terms of the number of parameters
(see Section 5.1 and Section 5.2). No previous literature has demonstrated the
effectiveness of such a large-scale image classification for learning image
representation.

  * •

Neural graph learning on image representation. We propose a neural graph
learning framework that leverages graph structure to regularize the training
of deep neural networks. This is the first work deploying large-scale neural
graph learning for image representation. We will describe below two techniques
to construct image-image graphs based on “co-click” rate and “similar-image
click” rate, designed to capture ultra-fine-grained notions of similarity that
emerge from human perception of result sets.

  * •

Graph-RISE for instance-level semantics. We present Graph-RISE, an image
embedding that captures ultra-fine-grained, instance-level semantics. Graph-
RISE outperforms the state-of-the-art algorithms for learning image embeddings
on several evaluations based on k-Nearest-Neighbor (kNN) search and triplet
ranking. Experimental results show that Graph-RISE improves the Top-1 accuracy
of the kNN evaluation by approximately 2X on the ImageNet dataset and by more
than 5X on the iNaturalist dataset. Case studies also show that,
qualitatively, Graph-RISE outperforms the state of the art and captures
instance-level semantics.

The remainder of this paper is organized as follows. Section 2 provides
related work on learning image embeddings. Section 3 formulates the problem
and provides the details of training datasets. Section 4 explains the proposed
learning algorithms, followed by Section 5 with the details of network
architecture and training infrastructure. Section 6 shows the experimental
results and Section 7 concludes this paper.

##  2\. Related Work

There are several prior works on learning image similarity (Wang et al., 2009;
Guillaumin et al., 2009; Taylor et al., 2011). Most of them focus on category-
level image similarity, in which two images are considered to be similar if
they belong to the same category. In general, visual and semantic similarities
tend to be consistent with each other across category boundaries (Deselaers
and Ferrari, 2011). Visual variability within a semantically-defined category
still exists, especially for broadly defined categories such as “animal,” or
“plant,” as a result of the broad semantic distinctions within such classes.
As classes become finer grained, however, the visual distinctions within a
class due to natural variations in image capture (angle, lighting, background,
etc) become larger relative to the fine distinctions between classes that are
semantically closer; hence, new techniques are required.

For learning fine-grained image similarity, local distance learning (Frome et
al., 2007) and OASIS (Chechik et al., 2010) developed ranking models based on
hand-crafted features, trained with triplets wherein each entry contains
{query image, positive image, negative image} that characterizes the ranking
orders based on relative similarity. In (Wang et al., 2014), a DeepRanking
model that integrates the deep learning and ranking model is proposed to learn
a fine-grained image similarity ranking model directly from images, rather
than from hand-crafted features. As discussed above, while these ranking
models have been widely used for learning image embeddings, the model
performance relies heavily on the quality of triplet samples, which involves
pair-wise comparisons that can be costly and labor-intensive to collect. As we
will show later in Section 3 and Section 6, Graph-RISE does not require models
to be trained by triplets and outperforms the state-of-the-art on capturing
image semantics for several evaluation tasks.

There has also been a significant amount of work on improving image
classification to near-human levels (Russakovsky et al., 2015) by increasing
the representational capacity and the depth of network architectures. See,
e.g., VGG-19 (Simonyan and Zisserman, 2014), Inception (Szegedy et al., 2016),
and ResNet (He et al., 2016). To support learning such deep networks with
millions of parameters, large-scale datasets such as ImageNet (Krizhevsky et
al., 2012), iNaturalist (Van Horn et al., 2018) and YouTube-8M (Abu-El-Haija
et al., 2016) have played a crucial role. For example, the authors of (Oquab
et al., 2014) demonstrate that the rich mid-level image features learned by
Convolutional Neural Networks (CNNs) on ImageNet can be efficiently
transferred to other visual recognition tasks with limited amount of training
data. Their study suggests that the number of images and the coverage of
classes for training in the source task are important for the performance in
the target task. In (Sun et al., 2017), the authors reveal a logarithmic
relationship between the performance on vision tasks and the amount of
training data used for representation learning. In this paper, we share the
same observation and further show that when increasing the number of classes
to ${\mathcal{O}}(40M)$ with sufficient amount of training data, the purposed
Graph-RISE is able to capture instance-level, ultra-fine-grained semantics.

##  3\. Problem Formulation

In this section, we formulate the task of learning an image embedding model,
and then provide the details of training dataset used for this task.

#### Problem Formulation

Given the following inputs:

  * •

A labeled set ${\mathcal{D}}_{L}$ that contains image-label pairs $(x,y)$,
where label $y$ provides ultra-fine-grained semantics to the corresponding
image $x$.

  * •

An unlabeled set ${\mathcal{D}}_{U}$ that contains images without labels.

The objective is to find an image embedding model that achieves instance-level
semantic understanding. Specifically, let $\phi(\cdot)$ represent a function
that projects an image to a dense vector representing an embedding; given two
images $x_{1}$ and $x_{2}$, the similarity in the image embedding space is
defined according to a distance metric: $d(\phi(x_{1}),\phi(x_{2}))$, where
$d(\cdot)$ is a distance function (e.g., Euclidean distance or cosine
distance). If $x_{1}$ and $x_{2}$ belong to the same class, an ideal
$\phi(\cdot)$ minimizes the distance between $\phi(x_{1})$ and $\phi(x_{2})$,
indicating these two images to be semantically similar. Table 1 provides the
symbols and the corresponding definitions used throughout this paper.

Table 1. List of symbols. Symbol | Definition and description
---|---
$(x,y)$ | Image-label pair
${\mathcal{D}}_{L}$, ${\mathcal{D}}_{U}$ | Labeled dataset, unlabeled dataset
${\mathcal{K}}$ | Total number of classes
${\mathcal{E}}$ | Set of edges
$w_{u,v}$ | Edge weight
$\eta$ | Margin for triplet distance
$\phi(\cdot)$ | Embedding mapping function
$d(\cdot)$ | Distance function
$p(\cdot),q(\cdot)$ | Probability and ground-truth distribution
${\mathcal{L}}(\cdot)$ | Loss function
$\Omega(\cdot)$ | Graph regularizer
${\bm{\theta}}$ | Model parameters

#### Dataset

Figure 2. Six potential samples of image-query pairs. Each image is labeled
with the corresponding textual search query.

In order to achieve instance-level semantic understanding, the classes should
be ultra-fine-grained. Thus, we created a training dataset ${\mathcal{D}}_{L}$
derived from Google Image Search. It contains approximately 260 million
images; selection data are used to characterize how frequently anonymous users
selected an image when that image appeared in the results of a particular
textual search query. After the characterization, a search query is then
treated as the “label” annotated to the corresponding image to provide
semantics as labeled samples. Each image is labeled with one or more queries
(2.55 queries per image on average), and the total number of unique queries
(used as classes) is around forty million. Figure 2 illustrates six potential
image-query pairs111In this paper, “image-query pairs” and “image-label pairs”
are used interchangeably since queries are used as labels.. To the best of our
knowledge, this is the largest scale of training data for learning image
embedding in terms of the number of classes222In prior arts (such as (Wang et
al., 2014)), the training dataset contains up to ${\mathcal{O}}(15M)$ samples
with ${\mathcal{O}}(100K)$ classes., and one of the largest in terms of the
number of training images (Sun et al., 2017; Wang et al., 2014; Schroff et
al., 2015).

Unlabeled dataset ${\mathcal{D}}_{U}$ contains approximately 20 million
images, without any annotation or labeling. Unlabeled dataset is mainly for
constructing similarity graphs (see Section 4.2) and for evaluation purposes
(see Section 6.1).

##  4\. Learning Algorithms

In this section, we first introduce the algorithm that takes image-query pairs
as training data in a supervised learning manner. Then we elaborate on neural
graph learning, a methodology incorporating graph signals into the learning
algorithm.

###  4.1. Discriminative Embedding Learning

#### From triplet loss to softmax loss.

In order to train image embedding models, metric learning objectives (such as
contrastive loss (Hadsell et al., 2006; Sun et al., 2014) and triplet loss
(Wang et al., 2014)), and classification objectives (such as logistic loss and
softmax loss) have been widely explored. When using metric learning
objectives, collecting high quality samples (e.g., triplets) is often
challenging. Furthermore, optimizing metric learning objectives suffers from
slow convergence or poor local optima if sampling techniques are
inappropriately applied (Wu et al., 2017; Sohn, 2016). In order to achieve
instance-level semantic understanding, we employ softmax loss for training the
image embedding model. When the size of classes is sufficiently large,
${\mathcal{O}}(40M)$ in our case, classification training (with softmax loss)
works better than triplet loss (Wang et al., 2014).

For each training example $x$, the probability of each label
$k\in\\{1,\dots,{\mathcal{K}}\\}$ in our model is computed via softmax:

(1) $p(k|x)=\frac{exp\\{z_{k}\\}}{\sum_{i=1}^{{\mathcal{K}}}exp\\{z_{i}\\}}$

where $z_{i}$ are the logits or unnormalized log probabilities. Here, the
$z_{i}$ are computed by adding a fully connected layer on top of the image
embeddings, i.e., $z_{i}={\bm{W}}_{i}^{T}\phi(x)+b_{i}$, where ${\bm{W}}_{i}$
and $b_{i}$ are weights and bias for target label, respectively. Let $q(k|x)$
denote the ground-truth distribution over classes for this training example
such that $\sum_{i=k}^{{\mathcal{K}}}q(k|x)=1$. As one image may have multiple
ground-truth labels, $q(k|x)$ is uniformly distributed to the ground-truth
labels. The cross-entropy loss for the example is computed as:

(2) $\ell=-\sum_{k=1}^{{\mathcal{K}}}log(p(k|x))q(k|x)$

While softmax loss works well when the number of classes is not large (say 10K
or 100K), several challenges arise if the number of classes is increased to
millions or even billions. First, the computational complexity involved in
computing the normalization constant of the target class probability $p(k|x)$
is prohibitively expensive (Bengio and Senécal, 2008; Jean et al., 2014).
Second, as the training objective encourages the logits corresponding to the
ground-truth labels to be larger than all other logits, the model may learn to
assign full probability to the ground-truth labels for each training example.
This would result in over-fitting and make the model fail to generalize
(Szegedy et al., 2016).

Instead of computing the normalization constant for all the classes, we sample
a small subset of the target vocabulary $L\in\\{1,\dots,{\mathcal{K}}\\}$ to
compute the normalization constant in $p(k|x)$ for each parameter update.
Then, the target label probability can be computed as:

(3) $p^{\prime}(k|x)=\frac{exp\\{z_{k}\\}}{\sum_{i\in L}exp\\{z_{i}\\}}$

which leads to much lower computational complexity and allows us to
efficiently use Tensor Processing Units (TPUs) (Jouppi et al., 2017) to train
this deep image embedding model with sampled softmax.

Furthermore, to discourage the model from assigning full probability to the
ground-truth labels (and therefore becoming prone to over-fitting), we follow
(Szegedy et al., 2016) to “smooth” the label distribution by replacing the
label distribution with a mixture of the original ground-truth distribution
$q(k|x)$ and the fixed distribution $u(k)$:

(4) $q^{\prime}(k|x)=(1-{\epsilon})q(k|x)+{\epsilon}u(k)$

where $u(k)$ is a uniform distribution over the sampled vocabulary
$u(k)=\frac{1}{|L|}$, and ${\epsilon}$ is a smoothing parameter.

Finally, the discriminative objective for training the neural network can be
defined as the cross-entropy of the target label probability on the sampled
subset and the smoothed ground-truth distribution:

(5) ${\mathcal{L}}({\bm{\theta}})=-\sum_{x_{i}\in{\mathcal{D}}_{L}}\sum_{k\in
L_{i}}log(p^{\prime}(k|x_{i}))q^{\prime}(k|x_{i})$

where ${\bm{\theta}}$ denotes the neural network parameters. The ground-truth
labels of $x_{i}$ are always selected within the sampled labels $L_{i}$. In
our experiments, we randomly sample 100K classes for each training instance
and ${\epsilon}$ is selected to be $0.1$.

###  4.2. Neural Graph Learning

While the discriminative objective indeed paves the way to learning an image
representation that captures fine-grained semantics, there is more information
available in human interactions with images. Many of these additional data
sources can be represented as graphs (such as image-image co-occurrence), and
yet current vision models (e.g., ResNet) cannot consume such graphs as inputs.
Thus, we propose to train the network using graph structure about the
relationships among images. In particular, images that are more strongly
connected in the graph should reflect stronger semantic similarity based on
user feedback (see Section 4.3 for the details of graph construction), and
should be closer in the embedding space. To achieve this goal, we deploy graph
regularization (Bui et al., 2018) to train the neural network for encouraging
neighboring images (from the graph) to lie closer in the embedding space. The
final graph-regularized objective is the sum of the discriminative loss and
the graph regularization:

(6)
${\mathcal{R}}({\bm{\theta}})={\mathcal{L}}({\bm{\theta}})+\alpha\underbrace{\sum_{(u,v)\in{\mathcal{E}}}w_{u,v}d\Big{(}\phi(x_{u}),\phi(x_{v})\Big{)}}_{\Omega({\bm{\theta}})}$

where $\Omega({\bm{\theta}})$ denotes the graph regularizer, ${\mathcal{E}}$
represents a set of edges between images, $w_{u,v}$ represents the edge weight
between image $u$ and $v$, $\phi(\cdot)$ is the representation extracted from
the embedding layer333In general, the embedding layer refers to the layer
right before the softmax layer., $d(\cdot)$ is the distance metric function,
and $\alpha\geq 0$ is the multiplier (applied on regularization) that controls
the trade-off between the discriminative loss and the regularization induced
from the graph structure. An illustration of the graph-regularized neural
network is given in Figure 3.

Figure 3. An illustration of a graph-regularized neural network. The image
similarity subgraph of a training image $x_{u}$ (with the ground-truth labels
$y_{u}$) is factorized into image-image pairs, where the neighbor image
$x_{v}$ is semantically similar to $x_{u}$. The training objective consists of
both the supervised loss $\mathcal{L}$ and the graph regularization $\Omega$;
minimizing $\Omega$ drives the distance between the embeddings of similar
images—$\phi(x_{u})$ and $\phi(x_{v})$—to be minimized, which means the neural
network is trained to encode the local structure of a graph.

The multiplier $\alpha$ (applied on regularization) controls the balance
between the discriminative information (i.e., predictive power) and the
contributions of the graph structure (i.e., encoding power). In other words,
the neural network is trained to both (a) make accurate classification (or
prediction), and (b) encode the graph structure. When $\alpha=0$, the proposed
objective ignores the graph regularization and degenerates to a neural network
with only supervised objective in Eq. (5). On the other hand, when
$p^{\prime}(x)=\phi(x)$, where $p^{\prime}(x)$ is the predicted label
distribution, we have a label propagation objective as in (Ravi and Diao,
2016) by training with the objective using $p^{\prime}(x)$ directly without
parameters ${\bm{\theta}}$ (i.e., no neural network involved), and letting the
distance function $d(\cdot)$ and the loss function $\ell(\cdot)$ to be mean
squared errors (MSE). The label propagation objective encourages the learned
label distribution $p^{\prime}(x)$ of similar images to be close. Furthermore,
the label distribution of a sample is aggregated from its neighbors to adjust
its own distribution. Thus, the proposed objective in Eq. (6) could be viewed
as a “graph-regularized version of neural network objective” or as a “non-
linear version of label propagation.”

###  4.3. Graph Construction

In this section, we provide the details for constructing graphs used by the
regularizer $\Omega({\bm{\theta}})$ in Eq. (6). In addition to image-query
pairs (described in Section 3) where images are annotated by queries for
obtaining semantics, we propose to find “image-image” pairs where the
semantics shared by these two images are closer to human perception and beyond
textual search queries. Each image is treated as a “vertex” and an image-image
pair is treated as an “edge,” together forming a graph that can be included as
ancillary training data.

Specifically, each image-image pair contains one source vertex
$x_{u}\in{\mathcal{D}}_{L}$ and one target vertex
$x_{u}\in\\{{\mathcal{D}}_{L}\bigcup{\mathcal{D}}_{U}\\}$. In this work, we
introduce two methods to construct edges: (a) based on co-click rate of the
image pair, and (b) based on similar-image click rate of the image pair. The
co-click rate of the image pair characterizes how often users select both the
source image $x_{u}$ and the target image $x_{v}$ in response to both $x_{u}$
and $x_{v}$ being concurrently identified by search results from a textual
search query. This type of image-image relationship sheds light on the
question: “Given that one image is selected from the resulting images, what
other images that are sufficiently similar will also be selected?” If the co-
click rate between $x_{u}$ and $x_{v}$ is higher than a pre-defined threshold,
$x_{u}$ and $x_{v}$ are considered to be sufficiently similar and an edge
between them is constructed; the edge weight $w_{u,v}$ is calculated based on
the co-click rate. Then $x_{u}$ and $x_{v}$ will be used for calculating the
graph regularization $\Omega({\bm{\theta}})$ in Eq. (6).

Different from the co-click rate, the similar-image click rate of the image
pair characterizes how often users select the source image $x_{u}$ in response
to $x_{u}$ being identified by a search result for a search query using the
target image $x_{v}$ (instead of a textual query). This type of image-image
relationship sheds the light upon the question: “Given an image issued as the
query, what other images that are sufficiently similar will be selected in
response to the query image?” Similar with how edges are constructed based on
the co-click rate, if the similar-image click between $x_{u}$ and $x_{v}$ is
higher than a pre-defined threshold, $x_{u}$ and $x_{v}$ are considered to be
sufficiently similar and an edge between them is constructed. Edge weight
$w_{u,v}$ is calculated based on the similar-image click rate.

##  5\. Training Framework

Figure 4. An illustration of the Graph-RISE framework. Flow in red is added to
enable graph regularization and required only during training. In the input
layer, a labeled image is associated with one of its neighbor images, which
can be either labeled or unlabeled, and then fed into the ResNet together with
its neighbor image. Then, the image embeddings generated from ResNet are used
to both (a) compute the cross-entropy loss and (b) graph regularization.

In this section, we provide the details of network architecture and training
infrastructure along with the configurations we used in this work.

###  5.1. Network Architecture

Figure 4 illustrates the proposed network architecture. The main model is the
101-layer ResNet (referred as ResNet-101) (He et al., 2016). Compared to
Inception (Szegedy et al., 2015), ResNet-101 has larger model capacity, which
yields more than 2% of performance improvement on our internal metric for
embedding evaluation. While the major architecture of ResNet-101 remains
unchanged, several detailed configurations have been modified. The input layer
is modified to take enlarged input images from 224$\times$224 to
289$\times$289 pixels. The output 10$\times$10$\times$2K feature map is first
avg pooled to 4$\times$4$\times$2K using a 4$\times$4 kernel of stride 2, and
then flattened and projected to 64-dimensional layer representing image
embeddings. The activation function is ReLU-6444ReLU-6 computes Rectified
Linear Unit as: $\min(\max(x,0),6)$.. Finally, a softmax layer is added to
produce a multinomial distribution across 40 million classes (e.g., queries).

During training, both training samples and their neighbors provided from
graphs (described in Section 4.3) are fed into the model for enabling graph
regularization; the 64-dimensional embedding layer is selected as the target
layer to be regularized as described in Eq. (6). Furthermore, 100K out of 40
millions labels are sampled via an important sampling technique (Bengio and
Senécal, 2008) for each parameter update Eq. (3). Finally, batch normalization
is applied during training.

During the inference phase, a 64-dimensional L2 normalized embedding 555L2
normalization is not applied on embedding during training as it makes the
training hard to converge. is generated for each input image as a new,
semantically-meaningful representation. Note that neighbors and graph
regularization is not required when making inference (see the flow in red in
Figure 4). In addition to the embedding, the queries with the top-$k$
predicted probabilities are also outputted. Since the focus of this paper is
image embedding, the output queries will largely be ignored in the rest of
this paper.

###  5.2. Training Infrastructure

We implement the network architecture described in Section 5.1 using
TensorFlow(Abadi et al., 2016). The details of training configurations are as
follows. We select the batch size to be 24, and the momentum (Qian, 1999) as
the optimizer; the initial learning rate is 0.001 and will be decayed with an
exponential rate of 0.9 every 100,000 steps. The label smoothing $\epsilon$ is
0.1. The multiplier for applying L2 regularization (a.k.a. “weight decay”) is
0.00004. For configurations related to graph regularization, the multiplier
for applying graph regularization $\alpha$ is 1.0, and the distance function
$d(\cdot)$ in Eq. (6) is selected to be cosine distance666We have also
experimented with using the Euclidean (L2) distance as $d(\cdot)$; when using
Euclidean distance with $\alpha=0.01$, it achieves almost the same performance
as using cosine distance with $\alpha=1$. For constructing graphs, the
threshold is set to $0.1$, and we combine the edges (approximately 50 million
edges, built from co-clicks and similar-image clicks) into one graph for
calculating regularizer $\Omega({\bm{\theta}})$ in Eq. 6.

Since our model is one of the largest vision models in terms of number of
parameters (40M $\times 64$ plus the parameters of ResNet-101 architecture),
together with the scale of the training dataset, using TPUs (Jouppi et al.,
2017) to train the model is the only feasible solution. The training is
distributed to 8$\times$8 TPU cores, and takes two weeks to converge from
scratch after 5M steps. Training with graph regularization costs additional
computation that grows with the number of neighbors used (in this work,
approximately 40% training time increase due to the use of neighbors). To
reduce the computation, we compare the performance of two models: one is
trained without graph regularization from scratch; the other is trained with
graph regularization with $\alpha$ set to zero until it is converged, and then
set $\alpha=1$ to fine-tune the model. We find that the performance of the
first model (w/o graph regularization) is slightly better in the beginning of
the training, and two models achieves almost the same performance after 4M
steps. With these setting, the fine-tuning model (the one with graph
regularization) takes approximately 2 additional days to train for 500K extra
steps until convergence.

##  6\. Experiments

In this section, we explain the details of evaluation setup, and then show the
experimental results for performance evaluation. We also provide the case
studies for the qualitative analysis.

###  6.1. Evaluation Setup

We are interested in providing image embeddings with instance-level semantics,
such that the similarity between embeddings approximates the relationships (of
images) in terms of semantics. To evaluate the performance of the proposed
image embedding model, we conduct both k-Nearest-Neighbor (kNN) search and
triplet evaluation as metrics; these two are the most popular methods used for
evaluating embedding models.

For kNN evaluation, we conduct experiments on ImageNet (Krizhevsky et al.,
2012) and iNaturalist (Van Horn et al., 2018) datasets, and then report two
metrics in the experiments: Top-1 and Top-5 accuracy, where Top-k accuracy
calculates the percentage of query images that find at least 1 image—from the
top k searched image results—carrying the exact same labels as the query
images. For the ImageNet dataset, the images in the validation set are used as
the query images, and the training set is used as the index set to search for
top-k results. For the iNaturalist dataset, for each class two images are
randomly sampled to construct the query set, and the remainder of the images
in that class are used as the index set.

For triplet evaluation, we follow the evaluation strategy in (Wang et al.,
2014) to sample triplets $(A,P,N)$—representing Anchor, Positive, Negative
images—from Google Image Search and ask human raters to verify if P is more
semantically closer to A than N. We sample the triplets in a way such that A
and P have the same or a very similar instance-level concept, and N is
slightly less relevant (“hard-negative”). Each triplet is independently ranked
by three raters and only the triplets receiving unanimous scores are used for
evaluation. Assume positive image P is rated to be more similar to anchor
image A than negative image N, the prediction of a model is considered to be
accurate if the following condition holds:

(7) $\eta+d(\phi(A),\phi(P))-d(\phi(A),\phi(N))<0$

where $\eta$ is the hyper-parameter that controls the margin between the
distance of two image projections.

Two triplet datasets are created for calculating triplet evaluation metrics:
(a) Product-Instance-Triplets (PIT) is a dataset designed to focus on
evaluating the semantic concepts of images in the commercial product vertical,
which consists of 10,000 triplets; (b) Generic-Instance-Triplets (GIT) is a
dataset focusing on evaluating the semantic concepts of general images,
including all possible image verticals from Google Image Search, consisting of
14,000 triplets.

###  6.2. Model Comparisons

We compare the proposed method with the following state-of-the-art models:

  * •

DeepRanking model (Wang et al., 2014) that employs triplet loss on multi-scale
Inception network architecture (Szegedy et al., 2016) with an online triplet
sampling algorithm.

  * •

Inception network architecture (Szegedy et al., 2016) that employs sampled
softmax loss over 8 millions labels.

  * •

ResNet-101 network architecture (He et al., 2016) that employs sampled softmax
loss over 8 millions and 40 millions labels (referred as ResNet (8M) and
ResNet (40M) in Table 2, respectively).

  * •

Graph-RISE model based on ResNet-101 network architecture proposed in Section
4.2.

The input layers in DeepRanking model and Inception model both use
224$\times$224 image pixels, while the ResNet-101 and Graph-RISE use
289$\times$289\. Label smoothing is applied to all the classification-based
models. When the graph regularization multiplier $\alpha=0$, Graph-RISE is
equivalent to the ResNet-101 model. In all the experiments, the Euclidean (L2)
distance of the embedding vectors extracted from the penultimate layer—the
layer before the final softmax or ranking layer—is used as similarity measure.
To evaluate the effectiveness of image embeddings, no individual fine-tuning
is performed for each dataset, and all the experiments are conducted directly
based on the learned embeddings of the input images.

###  6.3. Performance Evaluation

| ImageNet (Krizhevsky et al., 2012) | iNaturalist (Van Horn et al., 2018)
---|---|---
| Top-1 | Top-5 | Top-1 | Top-5
DeepRanking | 35.20 | 60.93 | 6.03 | 13.71
Inception (8M) | 61.92 | 84.77 | 17.30 | 34.58
ResNet (8M) | 62.49 | 84.65 | 17.36 | 34.15
ResNet (40M) | 66.20 | 86.41 | 27.05 | 47.44
Graph-RISE (40M) | 68.29 | 87.75 | 31.12 | 52.76
Table 2. Performance comparisons (in %) via kNN search accuracy on publicly available datasets. | PIT | GIT
---|---|---
DeepRanking | 74.95 | 77.25
Inception (8M) | 82.81 | 86.54
ResNet (8M) | 85.46 | 87.72
ResNet (40M) | 86.70 | 88.90
Graph-RISE (40M) | 87.16 | 89.53
Table 3. Performance comparisons (in %) via triplet accuracy ($\eta=0$) on the
internal evaluation datasets. Figure 5. PIT triplet evaluation on Recall v.s.
Margin. Figure 6. GIT triplet evaluation on Recall v.s. Margin. Figure 7.
Retrieval results for 6 randomly-chosen query images. For each query image, we
provide the Top-4 images retrieved (from 1 billion images) by DeepRanking,
ResNet (40M), and Graph-RISE. Each retrieval result is rated and color-coded
by human raters: retrieved images colored by green are rated to be strongly
similar with the query image, whereas images colored by red are rated to be
not (or somewhat) similar. Notice that images retrieved by Graph-RISE
generally conform to experts’ ratings.

Table 2 provides the performance comparisons (in terms of percentage) on kNN
evaluations, and Table 3 shows the triplet evaluations (also in terms of
percentage). From these results, we have several observations. First, Graph-
RISE significantly outperforms the previous state-of-the-art (Wang et al.,
2014) and other models without graph regularization in all the evaluation
criteria. We attribute this to the fact that Graph-RISE leverages the graph
structure via neural graph learning to drive the embeddings of similar images
to be as close as possible. Notice that, compared to the previous state-of-
the-art (Wang et al., 2014), Graph-RISE improves the Top-1 accuracy by almost
2X (from 35.2% to 68.29%) on ImageNet dataset, and by more than 5X (from 6.03%
to 31.12%) on the iNaturalist dataset.

Second, compared to the Inception network architecture, training image
embeddings with ResNet-101 improves the performance for most datasets. This
confirms the observation from (Sun et al., 2017) that to fully exploit
${\mathcal{O}}(300M)$ images, a higher capacity model is required to achieve
better performance. Furthermore, we confirm that sampled softmax is an
effective technique to train image embedding model with datasets that have
extremely-large label space; by comparing Inception (8M) and DeepRanking (Wang
et al., 2014) (both based on Inception network), we observe that sampled
softmax helps achieve better performance in triplet loss, even if DeepRanking
directly aims at optimizing the triplet accuracy.

Moreover, increasing the number of labels (from 8M to 40M) significantly
improves the kNN accuracy. We conjecture that the labels in 40M are more fine-
grained than 8M, and therefore the learned image embeddings also need to
capture fine-grained semantics in order to distinguish these 40M labels. In
addition, we find that training ResNet-101 using larger input size
(289$\times$289) instead of smaller input size (224$\times$224) also helps
improve the model performance (from 85.13% to 86.7%, in terms of accuracy on
PIT triplet evaluation), since larger input size encapsulates more detailed
information from training images.

Figure 5 and Figure 6 depict the comparisons of triplet evaluation among four
models: DeepRanking, Inception (8M), ResNet (8M) and Graph-RISE, on PIT
dataset and GIT dataset respectively. Note that ResNet (40M) is ignored in the
figures since the curves of ResNet (40M) and Graph-RISE are visually difficult
to distinguish. In these two figures, x-axis is “Margin” and y-axis is
“Racall” rate (the higher the better). “Margin” is the $\eta$ in Eq. 7
representing the margin between the distance of “Anchor-Negative pair” and the
distance of “Anchor-Positive pair.” A large margin means “the Negative image
is further away from the Anchor image than the Positive image.” “Recall” rate
represents the percentage of triplets that satisfy
$\eta+d(\phi(A),\phi(P))<d(\phi(A),\phi(N))$. From these two figures, the
performance of Graph-RISE is consistently better the other models.

###  6.4. Qualitative Analysis

Next, we evaluate the quality of images retrieved by DeepRanking (Wang et al.,
2014), ResNet (40M), and Graph-RISE models. Given a randomly-selected query
image, each method retrieves the most semantically similar images from an
index containing one billion images. The top-ranked results are sent out to be
rated by human experts. Figure 7 illustrates the retrieval results for 6
randomly-selected query images; for each query image we provide the Top-4
images retrieved. The images colored by green are rated to be strongly similar
with the query image, whereas the images colored by red are rated to be not
(or somewhat) similar. Compared to other models, images retrieved by Graph-
RISE generally conform to experts’ ratings, meaning that Graph-RISE captures
the semantic meaning of images more effectively as judged by human raters. For
example, given a query image of “scroll with ribbon” (top-left in Figure 7),
the top three images retrieved by Graph-RISE are also “scroll with ribbon”
(with similar colors, textures and shapes), and are rated as strongly similar
by human experts. Another example is a query image of a landmark (top-right in
Figure 7); Graph-RISE is able to retrieve images of the exact same landmark,
while the other methods are only able to retrieve images of somewhat similar
buildings.

In addition, we observe that the images retrieved by DeepRanking tend to be
only visually similar to the query images, rather than semantically similar.
This is probably because generating triplets that reflect the semantic
concepts is very difficult, especially when the classes are ultra-fine-
grained.

##  7\. Conclusion

In this work, we present Graph-RISE to answer the motivational question: “Is
it possible to learn image content descriptors (a.k.a., embeddings) that
capture image semantics and similarity close to human perception?” Graph-RISE
confirms that ultra-fine-grained, instance-level semantics can be captured by
image embeddings extracted from training a sophisticated image classification
model with large-scale data: ${\mathcal{O}}(40M)$ classes and
${\mathcal{O}}(260M)$ images. Graph-RISE is also the first image embedding
model based on neural graph learning that leverages graph structures of
similar images to capture semantics close to human image perception. We
conduct extensive experiments on several evaluation tasks based on both kNN
search and triplet ranking, and experimental results confirm that Graph-RISE
consistently and significantly outperforms the state-of-the-art methods.
Qualitative analysis of image retrieval tasks also demonstrates that Graph-
RISE effectively captures instance-level semantics.

## Acknowledgement

We would like to thank Dr. Sergey Ioffe and Dr. Thomas Leung for the reviews
and suggestions. We also thank Expander, Image Understanding and several
related teams for the technical support.

## References

  * (1)
  * Abadi et al. (2016) Martín Abadi, Paul Barham, Jianmin Chen, Zhifeng Chen, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Geoffrey Irving, Michael Isard, et al. 2016\.  Tensorflow: a system for large-scale machine learning.. In _USENIX Symposium on Operating Systems Design and Implementation (OSDI)_ , Vol. 16. 265–283.
  * Abu-El-Haija et al. (2016) Sami Abu-El-Haija, Nisarg Kothari, Joonseok Lee, Paul Natsev, George Toderici, Balakrishnan Varadarajan, and Sudheendra Vijayanarasimhan. 2016\.  Youtube-8m: A large-scale video classification benchmark.  _arXiv preprint arXiv:1609.08675_ (2016).
  * Bengio and Senécal (2008) Yoshua Bengio and Jean-Sébastien Senécal. 2008\.  Adaptive importance sampling to accelerate training of a neural probabilistic language model.  _IEEE Transactions on Neural Networks_ 19, 4 (2008), 713–722.
  * Bui et al. (2018) Thang D Bui, Sujith Ravi, and Vivek Ramavajjala. 2018.  Neural Graph Learning: Training Neural Networks Using Graphs. In _ACM International Conference on Web Search and Data Mining_.
  * Chechik et al. (2010) Gal Chechik, Varun Sharma, Uri Shalit, and Samy Bengio. 2010\.  Large scale online learning of image similarity through ranking.  _Journal of Machine Learning Research_ 11, Mar (2010), 1109–1135.
  * Deselaers and Ferrari (2011) Thomas Deselaers and Vittorio Ferrari. 2011.  Visual and semantic similarity in imagenet. In _Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)_. IEEE, 1777–1784.
  * Frome et al. (2007) Andrea Frome, Yoram Singer, and Jitendra Malik. 2007\.  Image retrieval and classification using local distance functions. In _Advances in neural information processing systems_. 417–424.
  * Guillaumin et al. (2009) Matthieu Guillaumin, Thomas Mensink, Jakob Verbeek, and Cordelia Schmid. 2009.  Tagprop: Discriminative metric learning in nearest neighbor models for image auto-annotation. In _Proceedings of the IEEE 12th International Conference on Computer Vision (ICCV)_. IEEE, 309–316.
  * Hadsell et al. (2006) Raia Hadsell, Sumit Chopra, and Yann LeCun. 2006\.  Dimensionality Reduction by Learning an Invariant Mapping. In _Proceedings of the 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition - Volume 2_ _(CVPR ’06)_. IEEE Computer Society, Washington, DC, USA, 1735–1742.  https://doi.org/10.1109/CVPR.2006.100
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016\.  Deep residual learning for image recognition. In _Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)_. 770–778.
  * Jean et al. (2014) Sébastien Jean, Kyunghyun Cho, Roland Memisevic, and Yoshua Bengio. 2014.  On using very large target vocabulary for neural machine translation.  _arXiv preprint arXiv:1412.2007_ (2014).
  * Jouppi et al. (2017) Norman P Jouppi, Cliff Young, Nishant Patil, David Patterson, Gaurav Agrawal, Raminder Bajwa, Sarah Bates, Suresh Bhatia, Nan Boden, Al Borchers, et al. 2017\.  In-datacenter performance analysis of a tensor processing unit. In _Computer Architecture (ISCA), 2017 ACM/IEEE 44th Annual International Symposium on_. IEEE, 1–12.
  * Krizhevsky et al. (2012) Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. 2012\.  Imagenet classification with deep convolutional neural networks. In _Advances in neural information processing systems_. 1097–1105.
  * Murphy-Chutorian et al. (2015) Erik Murphy-Chutorian, Charles J Rosenberg, Nemanja Petrovic, Sergey Ioffe, and Sean O’malley. 2015\.  Visual content retrieval.  US Patent 8,983,941.
  * Oquab et al. (2014) Maxime Oquab, Leon Bottou, Ivan Laptev, and Josef Sivic. 2014\.  Learning and transferring mid-level image representations using convolutional neural networks. In _Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)_. IEEE, 1717–1724.
  * Qian (1999) Ning Qian. 1999\.  On the momentum term in gradient descent learning algorithms.  _Neural networks_ 12, 1 (1999), 145–151.
  * Ravi and Diao (2016) Sujith Ravi and Qiming Diao. 2016.  Large scale distributed semi-supervised learning using streaming approximation. In _Artificial Intelligence and Statistics_. 519–528.
  * Russakovsky et al. (2015) Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. 2015\.  Imagenet large scale visual recognition challenge.  _International Journal of Computer Vision_ 115, 3 (2015), 211–252.
  * Schroff et al. (2015) Florian Schroff, Dmitry Kalenichenko, and James Philbin. 2015\.  Facenet: A unified embedding for face recognition and clustering. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 815–823.
  * Simonyan and Zisserman (2014) Karen Simonyan and Andrew Zisserman. 2014.  Very deep convolutional networks for large-scale image recognition.  _arXiv preprint arXiv:1409.1556_ (2014).
  * Sohn (2016) Kihyuk Sohn. 2016\.  Improved deep metric learning with multi-class n-pair loss objective. In _Advances in Neural Information Processing Systems_. 1857–1865.
  * Sun et al. (2017) Chen Sun, Abhinav Shrivastava, Saurabh Singh, and Abhinav Gupta. 2017.  Revisiting unreasonable effectiveness of data in deep learning era. In _Proceedings of the IEEE International Conference on Computer Vision (ICCV)_. IEEE, 843–852.
  * Sun et al. (2014) Yi Sun, Yuheng Chen, Xiaogang Wang, and Xiaoou Tang. 2014\.  Deep learning face representation by joint identification-verification. In _Advances in neural information processing systems_. 1988–1996.
  * Szegedy et al. (2015) Christian Szegedy, Wei Liu, Yangqing Jia, Pierre Sermanet, Scott Reed, Dragomir Anguelov, Dumitru Erhan, Vincent Vanhoucke, and Andrew Rabinovich. 2015.  Going deeper with convolutions. In _Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)_. 1–9.
  * Szegedy et al. (2016) Christian Szegedy, Vincent Vanhoucke, Sergey Ioffe, Jon Shlens, and Zbigniew Wojna. 2016.  Rethinking the inception architecture for computer vision. In _Proceedings of the IEEE conference on computer vision and pattern recognition (CVPR)_. 2818–2826.
  * Taylor et al. (2011) Graham W Taylor, Ian Spiro, Christoph Bregler, and Rob Fergus. 2011.  Learning invariance through imitation. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_. IEEE, 2729–2736.
  * Van Horn et al. (2018) Grant Van Horn, Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. 2018.  The inaturalist species classification and detection dataset.  (2018).
  * Wang et al. (2009) Gang Wang, Derek Hoiem, and David Forsyth. 2009.  Learning image similarity from flickr groups using stochastic intersection kernel machines. In _Proceedings of the IEEE 12th International Conference on Computer Vision (ICCV)_. IEEE, 428–435.
  * Wang et al. (2014) Jiang Wang, Yang Song, Thomas Leung, Chuck Rosenberg, Jingbin Wang, James Philbin, Bo Chen, and Ying Wu. 2014\.  Learning fine-grained image similarity with deep ranking. In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_. 1386–1393.
  * Wu et al. (2017) Chao-Yuan Wu, R Manmatha, Alexander J Smola, and Philipp Krähenbühl. 2017.  Sampling matters in deep embedding learning. In _Procceedings of the IEEE International Conference on Computer Vision (ICCV)_.


# A Massively Parallel Performance Portable Free-space Spectral Poisson Solver

Sonali Mayani  Veronica Montanaro  Antoine Cerfon  Matthias Frey
Sriramkrishnan Muralikrishnan  Andreas Adelmann

###### Abstract

Vico et al. (2016) suggest a fast algorithm for computing volume potentials,
beneficial to fields with problems requiring the solution of Poisson’s
equation with free-space boundary conditions, such as the beam and plasma
physics communities. Currently, the standard method for solving the free-space
Poisson equation is the algorithm of Hockney and Eastwood (1988), which is
second order in convergence at best. The algorithm proposed by Vico et al.
converges spectrally for sufficiently smooth functions i.e. faster than any
fixed order in the number of grid points. In this paper, we implement a
performance portable version of the traditional Hockney-Eastwood and the novel
Vico-Greengard Poisson solver as part of the IPPL (Independent Parallel
Particle Layer) library. For sufficiently smooth source functions, the Vico-
Greengard algorithm achieves higher accuracy than the Hockney-Eastwood method
with the same grid size, reducing the computational demands of high resolution
simulations since one could use coarser grids to achieve them. More
concretely, to get a relative error of $10^{-4}$ between the numerical and
analytical solution, one requires only $16^{3}$ grid points in the former, but
$128^{3}$ in the latter, more than a 99% memory footprint reduction.
Additionally, we propose an algorithmic improvement to the Vico-Greengard
method which further reduces its memory footprint. This is particularly
important for GPUs which have limited memory resources, and should be taken
into account when selecting numerical algorithms for performance portable
codes. Finally, we showcase performance through GPU and CPU scaling studies on
the Perlmutter (NERSC) supercomputer, with efficiencies staying above $50\%$
in the strong scaling case.

###### keywords:

Fast free-space Poisson solvers , Exascale , Performance portability , PIC

[inst1]organization=Paul Scherrer Institute, Villigen, country=Switzerland
[inst2]organization=ETH Zürich, country=Switzerland
[inst3]organization=Courant Institute of Mathematical Sciences, New York
University, country=USA [inst4]organization=Type One Energy Group - Canada
Inc, addressline=Vancouver, BC, country=Canada
[inst5]organization=Mathematical Institute, addressline=University of St
Andrews, postcode=KY16 9SS, country=UK [inst6]organization=Jülich
Supercomputing Centre, Forschungszentrum Jülich GmbH, country=Germany

##  1 Introduction

Simulation codes in many areas of physics need to solve elliptic partial
differential equations (PDEs), such as the Poisson equation or the Helmholtz
equation. Many computational methods have been developed for this purpose. In
particular, we look at integral-based approaches for open boundary conditions,
as is customary in previous work [1, 2, 3]. This approach requires dealing
with free-space Green’s functions, which are singular at the origin and long-
range in nature, requiring special attention during computational
implementations.

We take the example of the Poisson equation in the context of electrostatics.
The PDE governing our system in this case is given by:

$\Delta\phi(\vec{x})=-\frac{\rho(\vec{x})}{\epsilon_{0}},\\\ $ (1)

where $\rho(\vec{x})$ is the charge density and $\phi(\vec{x})$ the scalar
potential, $\epsilon_{0}$ the vacuum permittivity, and
$\vec{x}\in\mathbb{R}^{3}$. For free-space boundaries
($\phi(\vec{x})\rightarrow 0\text{ as }|\vec{x}|\rightarrow\infty$), the
solution $\phi$ can be written as a convolution:

$\phi(\vec{x})=\frac{1}{\epsilon_{0}}\int
G(\vec{x}-\vec{x}^{\prime})\rho(\vec{x}^{\prime})d\vec{x}^{\prime},$ (2)

where the Green’s function $G(\vec{x})=1/(4\pi|\vec{x}|)$ is the solution of
$\Delta\phi=-\delta(\vec{x})$. A computational brute-force approach would
consist of discretizing $G$ and $\rho$ on the computational mesh, and
approximating the integral as a sum. However, the computational cost of such
an operation is $\mathcal{O}(N^{2})$ [4], where $N=N_{x}N_{y}N_{z}$, with
$N_{x}$, $N_{y}$, and $N_{z}$ being the number of grid points in each
direction of the 3D grid.

For uniform meshes, the computational cost can be reduced by making use of
fast Fourier transforms (FFTs). The convolution theorem states that in Fourier
space a convolution can be written as a multiplication [5]. Equation 2 can
therefore be rewritten as
$\phi=\mathcal{F}^{-1}\\{\mathcal{F}\\{G\\}\mathcal{F}\\{\rho\\}\\}$ where
$\mathcal{F}$ denotes a Fourier transform and $\mathcal{F}^{-1}$ its inverse.
This aproach drastically reduces the cost to $\mathcal{O}(N\text{log}N)$.
However, FFTs work for periodic signals, and in this case $\rho$ and $G$ are
not periodic, since we are interested in free space boundary conditions.
Hockney and Eastwood introduced an algorithm which converts the convolution to
a cyclic one, making all involved signals periodic, while still obtaining the
correct free-space solution at the end [1].

Vico, Greengard, and Ferrando have proposed a novel, fast and spectrally
accurate method which can be applied to many scientific computing problems
requiring a convolution with a free-space Green’s function [6]. For smooth
distributions, the accuracy of this solver converges faster than any fixed
polynomial order of the gridsize as the resolution is increased, whereas the
Hockney-Eastwood algorithm is only second-order accurate. As suggested in [7],
this method could replace the Hockney-Eastwood algorithm for beam and plasma
physics problems, and provide better accuracy with lower resolution and
therefore a smaller memory footprint. Using a faster, more accurate, and less
memory intensive code is suitable for memory-bound architectures such as GPUs,
which are becoming increasingly common in scientific computing and super-
computing clusters.

However, the new Vico-Greengard method adds another memory bottleneck, even if
there is a gain in accuracy: while in the standard Hockney-Eastwood method one
needs to double the domain $(2N)^{3}$, in the new method a pre-computation of
the Green’s function on a $(4N)^{3}$ grid is required to capture the full
frequency content [6]. Storing this $(4N)^{3}$ grid causes the memory
footprint of the Poisson solver to double for each dimension, resulting in an
eight-fold memory increase in total.

We propose an algorithmic improvement by changing from a discrete Fourier
transform to the discrete cosine transform to constrain the pre-computation to
a $(2N+1)^{3}$ grid, reducing the memory footprint to be similar to the
Hockney-Eastwood algorithm. With this improvement, the novel Vico-Greengard
method becomes attractive both from the memory requirement as well as its
ability to deliver high accuracy solutions.

In this paper, we present a parallel, performant, and portable version of the
algorithm originally proposed by Vico, Greengard, and Ferrando [6]. We outline
our implementation in Section 2. A comprehensive evaluation of our
implementation through convergence studies, scaling studies, and memory
benchmarks, demonstrating its superior performance compared to the state-of-
the-art Hockney-Eastwood algorithm is given in Section 3.

##  2 Methodology

###  2.1 IPPL: Independent Parallel Particle Layer

IPPL [8] is an open-source C++ framework that provides the tools to develop
particle-mesh methods, such as particle-in-cell, explained in A.1. The
framework is built upon Kokkos [9] which enables cross-platform performance
portability for shared memory parallelism. Inter-process communication takes
place via the message passing inteface (MPI), the de facto standard for
distributed memory parallelism. IPPL further depends on the heFFTe (highly
efficient FFT for exascale) library to carry out fast Fourier transforms
(FFTs). A general overview of the dependencies and features of IPPL is
depicted in Figure 1. IPPL has a set of benchmark problems in plasma physics
applications, namely ALPINE [10], which helps in developing novel algorithms,
optimizing HPC implementations, and benchmarking individual components.
Further details about the structure of IPPL can be found in [8, 10].

Figure 1: A general overview of the IPPL framework and its external library
dependencies. The framework provides the essential components to develop
particle-mesh methods. A set of mini-applications is gathered in ALPINE [10]
which is currently distributed as part of IPPL.

It is in the context of these particle-mesh methods in IPPL that we seek to
solve the Poisson equation (Eq. 1). As mentioned previously, for free-space
boundary conditions, this is usually done by writing the solution in terms of
a convolution, which is efficiently computed in Fourier space and then Fourier
transformed back to real space. Until recently, the standard method for doing
so was the Hockney-Eastwood method, explained in the following section.

###  2.2 Hockney-Eastwood Method

In order to make the calculation of the convolution in Eq. 2 efficient, the
Hockney-Eastwood method makes the convolution cyclic, such that FFTs can be
used. Assuming a 3-dimensional physical domain of size
$[0,L_{x}]\times[0,L_{y}]\times[0,L_{z}]$ and a mesh of $N_{x}\times
N_{y}\times N_{z}$ grid-points, such that the mesh spacing is
$h_{x}=L_{x}/N_{x}$, $h_{y}=L_{y}/N_{y}$, $h_{z}=L_{z}/N_{z}$ in the
corresponding directions, the steps of the algorithm are as follows [1]:

  1. 1.

Double the computational grid in the non-periodic directions (those which have
open boundary conditions). If all boundary conditions are open, this means we
will have a grid of size $2N_{x}\times 2N_{y}\times 2N_{z}$.

  2. 2.

The source term $\rho(\vec{x})$, which is defined on
$\vec{x}\in[0,L_{x})\times[0,L_{y})\times[0,L_{z})$, is then zero-padded to be
defined on the full doubled domain. So for
$\vec{x}\in[0,2L_{x})\times[0,2L_{y})\times[0,2L_{z})$, we define
$\rho_{2}(\vec{x})$ such that:

$\rho_{2}(\vec{x})=\begin{cases}\rho(\vec{x}),&\text{if
$\vec{x}\in[0,L_{x})\times[0,L_{y})\times[0,L_{z})$}\\\
0,&\text{otherwise}.\end{cases}$

  3. 3.

The Green’s function $G(\vec{x})$ also needs to be extended to the doubled
domain. This is accomplished by circular shifting and making $G(\vec{x})$
periodic, such that the new function $G_{2}(\vec{x})$ satisfies:

$G_{2}(\vec{x})=\begin{cases}G(\vec{x}),&\text{if
$\vec{x}\in[0,L_{x})\times[0,L_{y})\times[0,L_{z})$}\\\
G(2L_{x}-x,y,z),&\text{if
$\vec{x}\in[L_{x},2L_{x})\times[0,L_{y})\times[0,L_{z})$}\\\
G(x,2L_{y}-y,z),&\text{if
$\vec{x}\in[0,L_{x})\times[L_{y},2L_{y})\times[0,L_{z})$}\\\
G(x,y,2L_{z}-z),&\text{if
$\vec{x}\in[0,L_{x})\times[0,L_{y})\times[L_{z},2L_{z})$}\\\
G(2L_{x}-x,2L_{y}-y,z),&\text{if
$\vec{x}\in[L_{x},2L_{x})\times[L_{y},2L_{y})\times[0,L_{z})$}\\\
G(2L_{x}-x,y,2L_{z}-z),&\text{if
$\vec{x}\in[L_{x},2L_{x})\times[0,L_{y})\times[L_{z},2L_{z})$}\\\
G(x,2L_{y}-y,2L_{z}-z),&\text{if
$\vec{x}\in[0,L_{x})\times[L_{y},2L_{y})\times[L_{z},2L_{z})$}\\\
G(2L_{x}-x,2L_{y}-y,2L_{z}-z),&\text{otherwise}.\end{cases}$

The Green’s function is singular at the origin and must therefore be
regularized. It is customary to replace the singularity by
$G(\vec{0})=-1/4\pi$.

  4. 4.

Since $\rho_{2}$ and $G_{2}$ are now periodic with period
$(2L_{x},2L_{y},2L_{z})$, we can Fourier transform them (after discretization)
using FFTs. In the Fourier domain, the convolution becomes a simple
multiplication of the Fourier transforms of $\rho_{2}$ and $G_{2}$. In this
case, the potential on the doubled domain is given by $\phi_{2}$:

$\phi_{2}=h_{x}h_{y}h_{z}\mathcal{F}^{-1}\\{\mathcal{F}\\{\rho_{2}\\}\times\mathcal{F}\\{G_{2}\\}\\}.$

  5. 5.

The physical solution for the potential is obtained by restricting the double-
grid solution to the physical domain, i.e. $\phi(\vec{x})=\phi_{2}(\vec{x})$
for $\vec{x}\in[0,L_{x})\times[0,L_{y})\times[0,L_{z})$. The proof of the
equivalence in the physical domain between the cyclic convolution and the
original convolution can be found in [11, p. 12].

The computational cost to solve the Poisson equation with this method is
$\mathcal{O}((8N)\text{log}(8N))$, where $N=N_{x}N_{y}N_{z}$. The trade-off is
the increased memory requirement to store the fields $\rho_{2}$ and $G_{2}$ on
the doubled domain. The Hockney-Eastwood method is second-order accurate.
Thus, if high accuracy is desired, the computation quickly becomes expensive.
It is therefore advantageous to explore more accurate methods, such as the
spectrally accurate Vico-Greengard method [7], which is briefly summarised
below.

###  2.3 Vico-Greengard Method

Vico, Greengard, and Ferrando describe a new method which differs from Section
2.2 by the choice of Green’s function. However, this change enables spectral
accuracy for smooth functions [6]. For the Poisson equation, where the
differential operator is the Laplace operator, the choice of the Green’s
function in the Fourier space as per [6] is

$G_{L}(\vec{s})=2\left(\frac{\sin(L\absolutevalue{\vec{s}}/2)}{\absolutevalue{\vec{s}}}\right)^{2},$
(3)

where $\vec{s}\in\mathbb{R}^{3}$ is in frequency domain and $L$ is the
truncation window size, which must be chosen larger than the maximum distance
between any two points in the computational domain [6], i.e.
$L=\alpha\sqrt{L_{x}^{2}+L_{y}^{2}+L_{z}^{2}}$ with $\alpha>1$ for a domain
$[0,L_{x}]\times[0,L_{y}]\times[0,L_{z}]$. Here, we have chosen $\alpha=1.1$.

To avoid any loss of information when calculating the inverse Fourier
transform of Eq. 3, the transform must be performed on a $4N_{x}\times
4N_{y}\times 4N_{z}$ grid due to the oscillatory nature of $G_{L}$. First, as
explained by Hockney and Eastwood [1, p. 213]), we need to have a grid of
double the original size to perform the aperiodic convolution. The additional
factor of two in the mesh size compared to the Hockney-Eastwood method comes
from the fact that $G_{L}$ is an oscillatory signal. As stated by the Nyquist
sampling theorem, an oscillatory signal must be sampled at a sampling rate
$f_{s}$ two times larger than the maximum frequency $f_{\max}$ of that signal,
i.e. $f_{s}>2f_{\max}$, to not lose any information.

However, if all the convolution computation is performed on the quadruple
grid, the source term must also be zero-padded on the quadruple domain, which
makes the method prohibitively memory-intensive for computational simulations.
Fortunately, this can be avoided by calculating the inverse transform of
$G_{L}$ on the fourfold grid in only one pre-computation step, which is then
restricted to the doubled grid of size $2N_{x}\times 2N_{y}\times 2N_{z}$ to
obtain $G_{2}$. The restricted Green’s function $G_{2}$ is afterwards reused
every time step in which the Poisson solve is called during the simulation.
The quadruple grid is therefore never needed again during the course of the
simulation unless the grid spacing changes in which case the Green’s function
again needs to be updated. Thus, the algorithm follows the same steps as in
Section 2.2 except for the Green’s function computation, where we use the pre-
computed $G_{2}$. In this case, the complexity is similar to the Hockney-
Eastwood method, but the method has a higher accuracy for sufficiently smooth
functions. The main drawback is the need for a $4N_{x}\times 4N_{y}\times
4N_{z}$ grid in the pre-computation step, which increases memory demands by a
factor of eight compared to the Hockney-Eastwood method. This is especially a
problem with GPU architectures, which generally have less memory available
than their CPU counterparts. In order to circumvent this, we present an
algorithmic change to reduce the memory required by the pre-computation step.

###  2.4 A Modified Vico-Greengard Method

A simple modification in the above algorithm allows us to reduce the grid size
of the pre-computation step from $(4N_{x}\times 4N_{y}\times 4N_{z})$ to
$((2N_{x}+1)\times(2N_{y}+1)\times(2N_{z}+1))$. We make use of the fact that
the Green’s function $G_{L}$ used in the Vico-Greengard method, given by Eq.
3, is purely real and has even symmetry i.e. $G_{L}(\vec{s})=G_{L}(-\vec{s})$.
In view of the even symmetry, the Green’s function on the Fourier domain
$G_{L}$, which is defined on a $4N_{x}\times 4N_{y}\times 4N_{z}$ grid to
contain the full frequency content, actually only has
$((2N_{x}+1)\times(2N_{y}+1)\times(2N_{z}+1))$ unique values. Furthermore, the
discrete Fourier transform reduces to a discrete cosine transform (DCT) thanks
to the same properties of the signal, namely that $G_{L}$ is purely real and
even. In this way, we can perform the pre-computation step of the Vico-
Greengard method on the $(2N+1)^{3}$ domain instead of the fourfold domain,
and use the DCT to compute the inverse transform of $G_{L}$. This way, the
memory footprint of the Vico-Greengard method is the same order-of-magnitude
to the Hockney-Eastwood method.

##  3 Results and Discussion

###  3.1 Code verification and accuracy analysis

To test the correctness and accuracy of the numerical methods used in our
implementation, we use the analytical solution of the Poisson’s equation for a
Gaussian source

$\rho=-\frac{1}{\sqrt{(2\pi)^{3}}}\exp\left(-\frac{r^{2}}{2\sigma^{2}}\right),$
(4)

which is given by [7]

$\phi_{exact}=\frac{1}{4\pi
r}\text{erf}\left(\frac{r}{\sqrt{2}\sigma}\right),$

with $r=\sqrt{(x-\mu)^{2}+(y-\mu)^{2}+(z-\mu)^{2}}$. We position the Gaussian
in the center of a unit box, i.e. $\mu=0.5$ and $L_{x}=L_{y}=L_{z}=1$, with
standard deviation $\sigma=0.05$.

In a second test case, we show the behaviour of the numerical schemes when the
source is non-smooth. The test case described in [12] has a source term given
by

$\rho=\begin{cases}4\pi G,&\text{if $r\leq$ 1}\\\ 0,&\text{if $r>$
1}\end{cases},$ (5)

and the exact solution

$\phi_{exact}=\begin{cases}-\frac{2}{3}\pi G(3-r^{2}),&\text{if $r\leq$ 1}\\\
-\frac{4}{3}\pi\frac{G}{r},&\text{if $r>$ 1}\end{cases},$

where $G=6.67408\cdot 10^{-11}$ m3kg-1s-2 is the gravitational constant and
$r$ is as above. We choose $\mu=1.2$ to center the distribution in the 3D
cubic box of extents $L_{x}=L_{y}=L_{z}=2.4$. The solution $\phi$ is the
gravitational potential of a sphere.

Figure 2: Relative error between analytical and computed solutions for all
three solvers (Hockney-Eastwood, Vico-Greengard, and its modified version).
The 2nd-order convergence theoretical line is plotted for reference. Left:
Source term given by Equation 4. Right: Source term given by Equation 5.

In order to compare the numerical and analytical solutions, we use the
relative error L2 norm. By means of a convergence study we compare the error
convergence rates with their theoretical values to verify the correctness of
the implementation. In Figure 2, we obtain the expected second-order
convergence with the Hockney-Eastwood solver, whereas with the Vico-Greengard
solver we get spectral convergence for the smooth source and second-order
convergence for the non-smooth source terms. The modified Vico-Greengard
solver follows the same trend as the original, therefore proving the
correctness of our algorithmic modifications.

These results also shed light on the limitations of the new solver: the
convergence of Vico-Greengard solver depends on the smoothness of the source
term. For the second test case with a non-smooth source term, i.e. Equation 5,
the Vico-Greengard solver is only as good as the Hockney-Eastwood solver.
However, since the modified Vico-Greengard solver does not cause any
additional computational costs in terms of memory and complexity compared to
the Hockney-Eastwood solver, it can replace it as the default method. The
accuracy is therefore limited by the properties of the source term instead of
the solver itself.

###  3.2 Scaling studies

We run strong and weak scaling studies of both the Hockney-Eastwood and Vico-
Greengard solvers implemented in IPPL. These studies are conducted on the
Perlmutter system, a HPE Cray machine at NERSC (USA). To showcase portability
and scaling across architectures, the studies are done both on the GPU and CPU
partitions. The GPU partition contains 1536 nodes where each of them is
equipped with 4 NVIDIA A100 GPUs. The CPU partition consists of 3072 nodes,
with 2 AMD EPYC 7763 CPUs of 64 cores each. We run the simulations with 1 MPI
rank per GPU for the GPU scaling studies and 128 MPI ranks per CPU node (1 MPI
task per physical core) for the CPU scaling studies.

The strong scaling study is done by keeping the problem size fixed and
measuring the timing as we increase the number of nodes for the simulation. We
run this for two different problem sizes: $512^{3}$ and $1024^{3}$. The weak
scaling study consists in increasing the number of nodes used for the
simulation while maintaining a constant workload, i.e. we increase the problem
size with the number of nodes. Starting with $512^{3}$ as our problem size, we
increase to $512^{2}\times 1024$ when doubling the number of nodes, then
$512\times 1024^{2}$ for four times the initial node count and finally ending
at $1024^{3}$.

IPPL was compiled with Cuda 12.0 (GPU-specific), FFTW 3.3.10.5 (CPU-specific),
Kokkos 4.1.0, Heffte 2.4.0, and manual installations of MPICH (due to some
technical issues, a manual installation of MPICH was needed in order to be
able to compile with Cuda 12 on the cluster). One small caveat is that the
MPICH version used to compile IPPL on CPU is 4.1.2 whereas for GPU it is
4.1.1, due to some technical issues on the cluster. Since this is a minor
version change with bug fixes related to “user-reported crashes and build
issues” [13], we do not expect differences in our scaling studies.

Figure 3: Strong scaling results on GPU for both the Hockney-Eastwood (left)
and the modified Vico-Greengard solver (right) with a problem size of
$N^{3}=512^{3}$, using the heFFTe parameters of pencil decomposition, a2av
communication, no reordering, and with GPU-aware enabled. The efficiency stays
above 65%.

In Fig. 3 we show the GPU strong scaling results of the $512^{3}$ grid
simulation for both the Hockney-Eastwood algorithm (left panel) and the
modified Vico algorithm (right panel). The blue lines denote the overall
algorithm execution times. The timings of the forward and backward FFT
transforms are highlighted in green. The costs for the data transfer between
the grids of size $N^{3}$ and $(2N)^{3}$ and vice versa are shown in red. The
algorithms have nearly identical execution times and they show very good
scaling behaviour up to 256 nodes (i.e. 1024 GPUs). The parallel efficiency
does not drop below 65%. After 32 nodes, the data copy kernel timings (red
lines) start to flatten since the communication costs associated with the data
transfer break even with the actual data copy.

Figure 4: Strong scaling results on CPU for both the Hockney-Eastwood (left)
and the modified Vico-Greengard solver (right) with a problem size of
$N^{3}=512^{3}$, using the heFFTe parameters of pencil decomposition, a2av
communication, no reordering. The efficiency stays above 50%.

The CPU strong scaling of the same setup is given in Fig. 4. We again observe
very good scaling behaviour, i.e. the parallel efficiency remains above 50%.
Note that here we start with 1 node rather than 4 nodes as for the GPU scaling
due to the smaller amount of memory on the GPUs. Also, the data transfer
timings do not flatten as we have seen for the GPU runs because the actual
data copy on the CPU is slower than on the GPU and therefore hides the
associated communication costs. The dominating cost on both hardware systems
(GPU and CPU) are the FFT transforms as expected. Note that the FFT timers
also include communication costs.

The GPU and CPU strong scalings for a problem size of $1024^{3}$ grid points
are provided in Fig. 5 and Fig. 6, respectively. We obtain similar scaling
behaviour on both hardware systems as for $512^{3}$ grid points and the
parallel efficiency remains above 65%.

Figure 5: Strong scaling results on GPU for both the Hockney (left) and the
modified Vico solver (right) with a problem size of $N^{3}=1024^{3}$, using
the heFFTe parameters of pencil decomposition, a2av communication, no
reordering, and with GPU-aware enabled. The efficiency stays above 80%.

Figure 6: Strong scaling results on CPU for both the Hockney-Eastwood (left)
and the modified Vico-Greengard solver (right) with a problem size of
$N^{3}=1024^{3}$, using the heFFTe parameters of pencil decomposition, a2av
communication, no reordering. The efficiency stays above 75%.

The weak scaling results on both hardware systems are given in Fig. 7 and Fig.
8. We start at a problem size of $512^{3}$ grid points and increase the
problem size up to $1024^{3}$ grid point while keeping the workload per node
approximately constant. On both hardware systems we experience a loss in
parallel efficiency because the cost of the all-to-all communication in the
FFTs increases with the node count while the workload is kept constant.

Figure 7: Weak scaling results on GPU for both the Hockney-Eastwood (left) and
the modified Vico-Greengard solver (right), starting from a problem size of
$N^{3}=512^{3}$ and increasing the workload proportionally to the node
increase until reaching a problem size of $N^{3}=1024^{3}$. The heFFTe
parameters used are pencil decomposition, a2av communication, no reordering,
and GPU-aware enabled. The efficiency stays above 75%.

Figure 8: Weak scaling results on CPU for both the Hockney-Eastwood (left) and
the modified Vico-Greengard solver (right), starting from a problem size of
$N^{3}=512^{3}$ and increasing the workload proportionally to the increase in
node counts until reaching a problem size of $N^{3}=1024^{3}$. The heFFTe
parameters used are pencil decomposition, a2av communication, no reordering.
The efficiency goes down to $40\%$.

###  3.3 Memory footprint

We can quantify the memory gain obtained from using the discrete cosine
transform by measuring the memory usage on one processor for all three
solvers: the Hockney-Eastwood solver, the Vico-Greengard solver, and the
modified version of Vico-Greengard. The CPU results are shown in Figure 9, and
are similar for the GPU case. As we can see, without the improvement, the
memory footprint of Vico-Greengard is $\sim$5-8 times larger than the Hockney-
Eastwood method. With the algorithmic modification of using the discrete
cosine transform, the memory footprint is brought down to be similar to the
Hockney-Eastwood method. The Vico-Greengard algorithm therefore becomes a
viable solver for the free-space Poisson equation in the context of mesh-based
methods.

Figure 9: Memory footprints for various problem sizes obtained using Kokkos
Highwater on CPU. Left: Memory comparison of Hockney-Eastwood, original Vico-
Greengard, and modified Vico-Greengard solvers. Right: Comparison of only
Hockney-Eastwood and the modified Vico-Greengard solver, which has an
algorithmic optimization.

##  4 Conclusion

In this work, we introduce an enhancement to the fast free space Poisson
solver initially proposed by Vico et al. in [6]. Our improvement reduces the
memory usage of the algorithm to just one-eighth of its original amount,
bringing its memory footprint into the same order-of-magnitude range as the
Hockney-Eastwood solver [1], the current state-of-the-art method for
addressing the free space Poisson problem in particle-mesh simulations. The
new solver offers a significant advantage by providing spectral accuracy for
smooth source terms. As a result, it needs fewer grid points compared to the
state-of-the-art second-order Hockney-Eastwood solver to achieve the same
level of accuracy for sufficiently smooth distributions. This makes it an
excellent choice for overcoming memory and runtime constraints in large-scale,
high-resolution simulations.

We have implemented this newly improved solver in a massively parallel and
portable framework targeted to exascale machines. We quantify the correctness
and accuracy improvements with respect to the Hockney-Eastwood method, as well
as the memory improvements with respect to the original algorithm in [6].
Finally, we demonstrate the efficiency of the solver by running strong and
weak scaling analyses on Perlmutter, for both CPU and GPU, using up to $\sim
70\%$ of the full machine in the latter case. The results of the strong
scaling study show that the efficiency stays above $50\%$ for the problem size
of $512^{3}$, and above $65\%$ in the case of the larger problem size
$1024^{3}$.

Currently, the new solver has been verified in the context of a plasma physics
application, a charge-neutral Penning trap, the results of which can be found
in A.

In future work, we will apply this free space solver in the context of beam
dynamics simulations using OPAL [2], a particle accelerator modelling code,
which is currently being interfaced with the IPPL framework.

## Availability

IPPL is an open source project. The source code can be found here:
https://github.com/IPPL-framework/ippl. The version of IPPL used for this
study is tagged scaling_study_vico_paper. The versions of the build
dependencies used on the Perlmutter system at NERSC are listed in the Results
section of the paper. The convergence study can be found in
test/solver/TestGaussian_convergence.cpp, and the scaling studies were done
with test/solver/TestGaussian.cpp. The solver type can be chosen by passing
either HOCKNEY, VICO, or VICO_2, which correspond to the Hockney-Eastwood,
Vico-Greengard, and modified Vico-Greengard methods respectively.

## Acknowledgements

This research used resources of the National Energy Research Scientific
Computing Center (NERSC), a Department of Energy Office of Science User
Facility using NERSC award ERCAP-M888. The Gwendolen cluster at PSI was
essential to the successful development of the presented method. We
particularly acknowledge the invaluable technical advice provided by the PSI
High-Performance Computing and Engineering (HPCE) group, especially Marc
Caubet Serrabou. We would also like to acknowledge Stanimire Tomov and Hartwig
Anzt from the Innovative Computing Laboratory - University of Tennessee and
Miroslav Stoyanov from Oak Ridge National Laboratory for their help with
introducing the Discrete Cosine Transform of Type 1 in heFFTe.

##  Appendix A Penning trap simulation using the Particle-In-Cell scheme

As an application of the Vico-Greengard solver, we use it in combination with
the Particle-In-Cell (PIC) scheme for an electrostatic plasma physics
simulation. More concretely, we simulate the dynamics of an electron bunch in
a Penning trap with external electric and magnetic fields to confine them.
First, we introduce the theory behind the PIC method, and then we follow-up
with the results, where we compare the simulation using the Vico-Greengard
solver with the Hockney-Eastwood solver.

###  A.1 The Particle-In-Cell (PIC) scheme

The particle-in-cell method combines particle tracking in phase space with a
mesh-based approach to speedup force calculations. To reduce the computational
cost needed to model physical phenomena, a cloud of plasma particles is
replaced by a macro-particle with the same charge-to-mass ratio. In the case
considered here, these macro-particles evolve in phase space over time
according to the Vlasov-Poisson equation and Newton’s second law of motion. In
order to compute the force fields, the charge density of the particles is
computed on the grid points of a fixed mesh (scatter), on which we solve for
the electrostatic fields (field solve). These are then interpolated to the
macro-particle locations (gather) in order to compute the force which will
drive the particles to their new position at the next time-step (particle
push) [14]. The PIC loop consists of the scatter, solve, gather, and push, and
is repeated at each time-step until the end of the simulation, schematically
shown in Figure 10.

Figure 10: The PIC loop, which is repeated every time-step after the
initialization.

Let $\vec{r}_{i}$ be the position of macro-particle $i$ with mass $m_{i}$ and
charge $q_{i}$, $\vec{p}_{i}=m\gamma\vec{v}_{i}$ its relativistic momentum,
$\vec{v}_{i}$ its velocity, and
$\gamma_{i}=1/\sqrt{1-|\vec{v}_{i}|^{2}/c^{2}}$ the Lorentz factor, where $c$
is the speed of light. The equations of motion are given by the Newton-Lorentz
equations:

$\displaystyle\frac{d\vec{r}_{i}}{dt}$
$\displaystyle=\frac{\vec{p}_{i}}{m_{i}\gamma_{i}}$
$\displaystyle\frac{d\vec{p}_{i}}{dt}$
$\displaystyle=q_{i}(\vec{E}_{i}+\frac{\vec{p}_{i}}{m_{i}\gamma_{i}}\times\vec{B}_{i}).$

The above are integrated using a Leapfrog algorithm [15] in the particle push
phase to evolve all of the macro-particles ($i=1,...,N_{p}$ where $N_{p}$ is
the number of macro-particles) to the next timestep. The fields $\vec{E}_{i}$
and $\vec{B}_{i}$ are the sum of the self-fields of the particles and any
external fields that may be applied, $\vec{E}_{ext}$ and $\vec{B}_{ext}$. In
the electrostatic case, the magnetic self-field is zero in the non-
relativistic regime. In the case of relativistic particles, a magnetic self-
field appears due to the moving charges.

Let $\vec{x}\in\mathbb{R}$ represent a position in the computational mesh, for
fields (not the same as the macro-particle positions). At each time-step, the
electric field $\vec{E}(\vec{x})$ is obtained from

$\vec{E}(\vec{x})=-\vec{\nabla}\phi(\vec{x}),$ (6)

where the potential $\phi(\vec{x})$ is computed from the charge density
$\rho(\vec{x})$ using Eq. 1. After calculating the electric field
$\vec{E}(\vec{x})$ on the mesh, it is interpolated as $\vec{E}_{i}$ to each
macro-particle located at $\vec{r}_{i}$ (gather phase).

As mentioned in the previous section, the Poisson equation in the solve phase
of the PIC loop is written as a convolution and solved in Fourier space. In
the case of free-space boundaries, one needs to to use the Hockney-Eastwood or
the Vico-Greengard method as explained before.

###  A.2 Penning trap as a physics example

For the Penning trap, we follow the same set-up as in [10], with open boundary
conditions. We do a qualitative comparison between using the Hockney-Eastwood
solver and the Vico-Greengard solver by looking at the evolution of the
electron charge density, as shown in Figure 11. Both solvers produce the same
expected behaviour with similar runtimes.

(a) Time = 0

(b) Time = 1.95

(c) Time = 3.91

(d) Time = 0

(e) Time = 1.95

(f) Time = 3.91

Figure 11: Evolution of electron bunch dynamics in a 3D charge-neutral Penning
trap, shown as snapshots in time during the simulation. The first and second
rows correspond to using the Hockney-Eastwood and Vico-Greengard solvers,
respectively.

Identifying the regimes or test cases where the high accuracy of the Vico-
Greengard solver is beneficial is beyond the scope of the current work and
will be carried as a part of future work.

## References

  * [1] R. Hockney and J. W. Eastwood, Computer Simulation Using Particles.  CRC Press, 1988.
  * [2] A. Adelmann, P. Calvo, M. Frey, A. Gsell, U. Locans, C. Metzger-Kraus, N. Neveu, C. Rogers, S. Russell, S. Sheehy, J. Snuverink, and D. Winklehner, “OPAL a Versatile Tool for Charged Particle Accelerator Simulations,” arXiv:1905.06654 [physics], May 2019.
  * [3] R. A. James, “The solution of Poisson’s equation for isolated source distributions,” Journal of Computational Physics, vol. 25, pp. 71–93, Oct. 1977.
  * [4] J. Qiang, S. Lidia, R. D. Ryne, and C. Limborg-Deprey, “Three-dimensional quasistatic model for high brightness beam dynamics simulation,” Physical Review Accelerators and Beams, vol. 9, p. 044204, Apr. 2006.
  * [5] E. W. Weisstein, “Convolution Theorem.”  Publisher: Wolfram Research, Inc.
  * [6] F. Vico, L. Greengard, and M. Ferrando, “Fast convolution with free-space Green’s functions,” Journal of Computational Physics, vol. 323, pp. 191–203, Oct. 2016.
  * [7] J. Zou, E. Kim, and A. J. Cerfon, “FFT-based free space Poisson solvers: why Vico-Greengard-Ferrando should replace Hockney-Eastwood,” arXiv:2103.08531 [physics], Mar. 2021.
  * [8] M. Frey, A. Vinciguerra, S. Muralikrishnan, S. Mayani, V. Montanaro, and A. Adelmann, “IPPL-framework/ippl: IPPL-3.1.0,” Sept. 2023.  URL: https://github.com/IPPL-framework/ippl/tree/IPPL-3.1.0.
  * [9] H. Carter Edwards, C. R. Trott, and D. Sunderland, “Kokkos: Enabling manycore performance portability through polymorphic memory access patterns,” Journal of Parallel and Distributed Computing, vol. 74, July 2014.  Institution: Sandia National Lab. (SNL-NM), Albuquerque, NM (United States) Number: SAND-2013-5603J Publisher: Elsevier.
  * [10] S. Muralikrishnan, M. Frey, A. Vinciguerra, M. Ligotino, A. J. Cerfon, M. Stoyanov, R. Gayatri, and A. Adelmann, “Scaling and performance portability of the particle-in-cell scheme for plasma physics applications through mini-apps targeting exascale architectures,” in Proceedings of the 2024 SIAM Conference on Parallel Processing for Scientific Computing (PP), Proceedings, pp. 26–38, Society for Industrial and Applied Mathematics, Jan. 2024.
  * [11] R. D. Ryne, “On FFT-based convolutions and correlations, with application to solving Poisson’s equation in an open rectangular pipe,” arXiv:1111.4971 [physics], Nov. 2011.  arXiv: 1111.4971.
  * [12] R. D. Budiardja and C. Y. Cardall, “Parallel FFT-based Poisson solver for isolated three-dimensional systems,” Computer Physics Communications, vol. 182, pp. 2265–2275, Oct. 2011.
  * [13] K. J. Raffenetti, “MPICH 4.1.2 released | MPICH.”  URL: https://www.mpich.org/2023/06/08/mpich-4-1-2-released/, Accessed 09/01/2024.
  * [14] X. Sáez, A. Soba, J. M. Cela, E. Sánchez, and F. Castejón, “Particle-in-Cell Algorithms for Plasma Simulations on Heterogeneous Architectures,” in 2011 19th International Euromicro Conference on Parallel, Distributed and Network-Based Processing, pp. 385–389, Feb. 2011.  ISSN: 2377-5750.
  * [15] J. P. Verboncoeur, “Particle simulation of plasmas: review and advances,” Plasma Physics and Controlled Fusion, vol. 47, pp. A231–A260, Apr. 2005\.  Publisher: IOP Publishing.