# Interpretable Self-supervised Multi-task Learning
for COVID-19 Information Retrieval and Extraction

Nima Ebadi
Department of Electrical and Computer Engineering
University of Texas at San Antonio
San Antonio, TX 78249
<EMAIL_ADDRESS>
& Peyman Najafirad
Department of Information Systems and Security
University of Texas at San Antonio
San Antonio, TX 78249
<EMAIL_ADDRESS>

###### Abstract

The rapidly evolving literature of COVID-19 related articles makes it
challenging for NLP models to be effectively trained for information retrieval
and extraction with the corresponding labeled data that follows the current
distribution of the pandemic. On the other hand, due to the uncertainty of the
situation, human experts’ supervision would always be required to double check
the decision making of these models highlighting the importance of
interpretability. In the light of these challenges, this study proposes an
interpretable self-supervised multi-task learning model to jointly and
effectively tackle the tasks of information retrieval (IR) and extraction (IE)
during the current emergency health crisis situation. Our results show that
our model effectively leverage the multi-task and self-supervised learning to
improve generalization, data efficiency and robustness to the ongoing dataset
shift problem. Our model outperforms baselines in IE and IR tasks,
respectively by micro-f score of 0.08 (LCA-F score of 0.05), and MAP of 0.05
on average. In IE the zero- and few-shot learning performances are on average
0.32 and 0.19 micro-f score higher than those of the baselines.

##  1 Introduction and Background

The pandemic has resulted in a rapidly evolving literature of scientific
publications regarding the novel Coronavirus which has caused an information
crisis [1]. Therefore, researchers, healthcare practitioners, policy makers,
and other individuals fighting against COVID-19 require specialized
information retrieval and extraction systems to keep up with the literature
[2, 3]. In this regard, many models are proposed based on state-of-the-art
natural language processing (NLP) techniques, including neural ranking models
for document retrieval [4, 2, 5], and pre-trained transformers for automated
annotation and information extraction [6, 7, 8].

However, the limited amount of labeled data and drastic changes of COVID-19
landscape [9] has made it challenging for such models to effectively scale up
to the domain-specific environment of the current pandemic. Moreover, due to
these challenges, such informatic tools cannot manage the pandemic situation
on their own, and human experts’ supervision is required to double check their
decision making. Therefore, it is significantly important for these models to
provide textual evidence from the raw input such that human experts can
interpret causes of a certain decision making [10, 11].

Although these challenges are inevitable in every health crisis situation and
may be confronted again in near future, very few studies have mainly focused
on addressing these specific challenges of the pandemic. In the light of this,
we propose a specilized information retrieval and extraction approach which
provides interpretability, data efficiency and robustness to the dataset shift
problem thereby suitable for the context of the pandemic. Our approach is
based on transformers encoders with global-local attention mechanism [12],
thus providing interpretability both at document level and sub-word level. We
train the encoder in a multi-task fashion alongside a self-supervised learning
task which have shown promising results in improving data efficiency,
generalization and robustness to dataset shift problems [13, 14].

Inspired by [15] and to exploit the landscape of transfer learning in an
optimum way, we design a unified retriever-ranker framework to simultaneously
and effectively train our model on a masked language modeling (MLM) self-
supervised learning (SSL) task 111We introduce a self-supervised task of
masked language modeling which is consistent with our retriever-ranker general
framework and enables the model to acquire knowledge about the current state
of the pandemic. along with two supervised tasks: i) information retrieval:
retrieval of related articles to a given question, and ii) information
extraction: extraction of semantic indexes of PubMed articles which are
manually indexed by human experts. Using an input/output transformation
module, we translate the tasks of masked language modeling and semantic
indexing to the domain of retriever-ranker. For every task, our model gets an
input of a given query, a COVID-19 article or question, and initially,
retrieves a set of candidate related articles. Next, a transformer-based
ranker assign scores to the candidate articles using the attention mechanism
proposed by [12], both within and across articles. The final output is
computed by transformation of these scores to their corresponding domain. The
attention mechanisms within/across the input query and candidate articles
enables the model to associate textual evidence from the input to the output
decision, thus providing interpretability both on document and sub-word level.

  * •

We propose an interpretable, self-supervised, multi-task learning methodology
to effectively tackle the tasks of information retrieval and extraction in the
context of pandemic.

  * •

Our study sheds light on the importance of interpretability and transfer
learning–specifically multi-task and self-supervised learning–in addressing
the intrinsic challenges of a health crisis situation like the ongoing
pandemic; i.e. dataset/covariate shift, lack of unlabeled data and
uncertainty.

  * •

We devise a novel mechanism to simultaneously train a unified retriever-ranker
on a self-supervised task masked language modeling (MLM) and an information
extraction (IE) task semantic indexing along with an information retrieval
(IR) task of retrieving relevant articles to a given question. This enables
inter-document representation learning and interpretability which we show is
necessary for the context of the pandemic.

##  2 Methodology

In our methodology, we attempt to address the three tasks of self-supervised
learning (henceforward we refer to as SSL, or task_0), information extraction
(henceforward we refer to as IE, or task_1) and information retrieval
(henceforward we refer to as IR, or task_2) using a unified retreival-ranker
framework. Our unified framework includes the following three steps which are
identical for each task: i) an initial retrieval step which works as a weak
classifier[10] to retrieve candidate articles, ii) an I/O transformation
technique to map different formats of data into a unified I/O format suitable
for next step, and iii) a ranker which works based on transformer encodings of
the candidate retrieved articles.

In the context of our retriever-ranker framework, the three aforementioned
tasks are defined as follows:

Task_0: (SSL) A self-supervised pre-text task similar to Masked Language
Modeling in [16] is performed to introduce knowledge about the context of
pandemic. However, in this task, the masked article is treated as a query and
masked tokens are selected from a list of covid-19 related terms222We have
used the list of related terms published by NLM
https://www.nlm.nih.gov/pubs/techbull/nd20/nd20_mesh_covid_terms.html. The
model attempts to detect similar articles which include the masked term/s
rather than the term itself (order is not important). As such, not only the
pre-text task is consistent with the downstream tasks, i.e. the general
architecture of our multi-task learning model, but also it effectively learns
context matching using both intra- and inter-document information [17].

Task_1: (IE) Extraction of semantic indexes assigned to a given article
333Semantic indexing is the information extraction task we address in this
study. We use the terms ”semantic extraction,” ”extraction of semantic
indexes” and ”semantic indexing” interchangeably. In this project, we deal
with PubMed articles whose indexes are called Medical Subject Headings (MeSH).
The given article is treated as a query and outputs are the semantic indexes.

Task_2: (IR) For a given question as a query a list of related articles are
retrieved which include the target answer.

###  2.1 Initial Retrieval

To reduce the number of negative examples, we use an initial retrieval system
to retrieve a subset of related articles along with their task specific
annotations (for example for extraction of semantic indexes, those that have
manually been annotated are considered, and their annotation as candidates for
the given task). Our initial retrieval system includes a document-level
embedding model of SPECTER [17], and Bag-of-Words representation fused with
TF-IDF/BM25 following the schema of [10] and [2]. SPECTER is initialized with
SciBERT [18] and trained on a bipartite graph of citations to capture
document-level relatedness and minimize a triplet loss between related
articles and maximize over unrelated ones. We pre-train SPECTER on PubMed
articles and fine-tune it on COVID-19 dataset exclusively.

In addition, we use TF-IDF/BM25 weighted sum of article tokens to compute a
keyword-based representation as well. As such every article is represented as
follows:

$d=\frac{\sum_{i=1}^{n}{\mathrm{TF-IDF/BM25}}({w}_{i},d)\times
v_{w_{i}}}{\sum_{i=1}^{n}{\mathrm{TF-IDF/BM25}}({w}_{i},d)}$ (1)

where, $w_{i}$ is the $i^{th}$ word in article $d$, and $v_{w_{i}}$ is the
token embeddings from the pre-trained model.

We concatenate both representations for every article and query in our
database. Next using cosine similarity scores between the input query and
other ones, we find the $K$ relevant articles. As for the extraction of
semantic indexes (IE task), and masked language modeling, we further retrieve
the candidate indexes and mask terms respectively. In this regard, we use the
summation of their IDF weights as the scoring scheme to rank and retrieve the
candidate indexes and terms. The top $M$ are considered as candidates and
passed to the next stage.

###  2.2 I/O Transformation

In our multi-task learning methodology, the data for every task is fed to a
transformer-based ranking system, i.e. our ranker, which is identical for
different tasks. Therefore, we require an I/O transformation module to map
every task’s data format to the unique format of the ranker. The input and
output to our ranker is as follows:

Input: given question or article as a query and a list of candidate articles

Output: likelihood scores corresponding to every candidate article

As for the IR task, the input and outputs are consistent with our ranking
algorithm. The input is a question along with a list of candidate related
articles, and the output for every candidate article is the likelihood score
of including the target answer.

In our self-supervised task, the input is the masked articles and a list of
candidate articles which include the masked terms withing their contents.
Similarly for the information extraction task, the input is a given article
and a list of candidate articles which are manually annotated in PubMed.

In our SSL task, the goal is to predict the masked tokens based on likelihood
scores of the candidate terms. In our IE task, also, the goal is to extract
the semantic indexes of a given article based on likelihood scores of the
candidate indexes. However, the output of our ranker assign likelihood scores
to the candidate articles instead. Therefore, to train our model for SSL and
IE tasks, we need to transform the likelihood scores of candidate terms and
indexes, respectively, to those which correspond to candidate documents:

$Y_{D_{c}}=T\times Y_{M_{c}},\hskip 5.69054ptT\in\mathrm{R}^{K\times M}$ (2)

where $Y_{M_{c}}$ are the likelihood scores of the candidate terms or indexes
$M_{c}$. $T$ is the transformation matrix whose every row is bag-of-word
representation of the corresponding document based on the candidate terms.
$Y_{D_{c}}$ are the scores assigned to candidate documents. While choosing
candidate documents, we make sure the matrix $T$ is reversible so that the
likelihood scores of the terms can be extracted and fed to the loss function.
In the task of extraction of semantic indexes task, during inference, the
inverse of transformation matrix $T^{-1}$ is used to compute the indexes’
likelihood scores which are then passed through thresholds to output the final
indexes. Thresholds are optimized to maximize micro-f1 score [19].

###  2.3 Transformer-based Ranking

Similar to other studies that apply BERT to biomedical [20] or scientific
articles [18], query and the candidate articles are tokenized into frequent
sub-words using WordPiece unsupervised algorithm [21] and the original
vocabulary of BERT. The sub-word tokenization can also mitigate the problem of
out-of-vocabulary words which happens more often during the pandemic.

Query and the paired articles are separated by [SEP] token and [CLS] token is
added to the beginning of every article which is leveraged as the vector
representing the corresponding article specified by the query and other
relative articles. The same embedding scheme of BERT is used to compute the
representation of every input token–summing the token, segment444an additional
embedding is used to differentiate tokens of the query with those of the
candidate articles and position embedding.

Every candidate document ${D_{i}}^{n-1}$ is encoded along with the paired
query $Q^{n-1}$ through multiple layers of a pre-trained transformer [22].

Every candidate article’s tokens are encoded, ${D_{i}}^{n}$, through a self-
attention between its own tokens and a cross-attention between them and the
tokens of the query. The [CLS] token of other articles also attend the [CLS]
token to enable capturing the inter-document information. Since the length of
the input sequence is high in the setting of our methodology, we use global-
local attention mechanism [23] to mitigate the computational limitation of
transformers in dealing with long sequences. As for the article tokens, we use
windowed dilated attention mechanism proposed by [12] as a local attention.
The attention between [CLS] tokens are also windowed dilated since it is not
necessary that all candidate articles attend to every article in order to
capture the inter-document relationship. As for the query tokens, however, we
use global attention mechanism introduced by [24]–in which a global memory is
utilized that enables every token in the query to attend to every token of
candidate articles–to bring in task-specific flexibility.

The query also gets encoded, $Q^{n}$, using a full self-attention mechanism
within its own tokens and a cross-attention between them and the [CLS] tokens
of the candidate articles.

##  3 Experiment

Model | MiF
---|---
Medical Text Indexer (MTI) (default) | 0.6578
MTI (first line indes) | 0.6491
average top score | 0.7143
R+TR (full attention) | 0.5532
R+TR (increasing $w$) (from 32-512) | 0.6277
R+TR (fixed $w$) (=230) | 0.6142
R+TR (decreasing $w$) (from 512-32) | 0.6002
R+TR (increasing $w$) (dilation on 2 heads) | 0.6334
R+TR (global + dilated sliding window*) | 0.6602
R+TR (base) (w/o multi-task) | 0.6602
R+TR (base) (w/ multi-task) | 0.6671
R+TR (large) (w/o multi-task) | 0.6980
R+TR (large) (w/ multi-task) | 0.7054
(a)

Model | MAP
---|---
average top score | 0.4638
R+TR (full attention) | 0.1911
R+TR (increasing $w$) (from 32-512) | 0.2933
R+TR (fixed $w$) (=230) | 0.2805
R+TR (decreasing $w$) (from 512-32) | 0.2581
R+TR (increasing $w$) (dilation on 2 heads) | 0.3032
R+TR (global + dilated sliding window*) | 0.3276
R+TR (base) (w/o multi-task) | 0.3276
R+TR (base) (w/ multi-task) | 0.3441
R+TR (large) (w/o multi-task) | 0.3550
R+TR (large) (w/ multi-task) | 0.4097
(b)

Table 1: Information extraction (a) and retrieval (b) performances of our models along with the baselines (best performing models of BioASQ Task 8a for IE, and Task 8b Phase A for IR). The baseline scores are the average of their provided Micro-F1 and MAP (respectively for IE and IR) scores across all test batches. We also evaluate our models on all the provided test batches. * dilated sliding window uses increasing window from 32 to 512 and dilation on 2 heads. Hyperparameter | Value(s)
---|---
$|V|$ | 20M, 30M
$K$ | 128, 256 , 512, 1024
$M$ | 128, 256, 512, 1024
$w$ | 32,…, 512, $\mathbf{inc}{[32:512]}$, $\mathrm{dec}{[32:512]}$
dilation | 0, 1, 2, 3, $\mathbf{inc}{[0:3]}$
dilation heads | 1, 2, 3
dorpout | 0.1, 0.2, 0.3, 0.4∗
batch size | 8, 16, 32, 64 (gpu memory limit)
output vector size | 512, 1024, 2048
w.e. size | 128, 256, 512∗
hidden size | 128, 256, 512∗
#layers |  4, 5, 6∗, 7, 8
learning rate | 0.001, 0.0005, 0.00025, 0.0001
Table 2:  Hyperparameters value. w.e.: embedding size for initial retrieval
step. We use bold text for the optimal ones among all tried values. ∗ refer to
those for large ranker. Best dilation size is achieved by increasing it by 1
from first layer to the last.

###  3.1 Dataset

Prior to training on COVID-19 datasets, the models are initially train on
general IE and IR datasets. For IE, we use BioASQ’s [25] Task 8a dataset which
includes almost 15 million article abstracts and titles manually annotated in
PubMed. We select 8M recent articles published from 2007-2019. For IR, we use
BioASQ’s Task 8b dataset which includes 3,243 question paired with related
article abstracts. Test sets of each task is used for hyperparameter tuning.
We also make sure there is no overlap between these general sets and their
corresponding COVID-19 ones.

After the initial training, The models are trained and evaluated on the
following three COVID-19 datasets corresponding to the three tasks in hand. As
for our self-supervised task, we use CORD-19 dataset [26] which includes 200K
research articles about Coronavirus published in peer-reviewed venues and
archival services such as bioRxiv555https://www.biorxiv.org and
medRxiv666https://www.medrxiv.org.

We select CORD-19 articles whose MeSH indexes are manually annotated in PubMed
for our semantic indexing task. We crawl and append their MeSH indexes.
Consequently, our semantic indexing dataset contains 17K articles which we
chronologically sort and split to 13.6K for training (the oldest 80%) and 3.4K
for testing (the latest 20%).

Note: for self-supervised and semantic indexing, we only use the title and
abstracts of the articles.

As for IR in the context of COVID-19, we choose TREC-COVID[1] dataset which is
an information retrieval dataset for question answering similar to BioASQ QA
task phase b. It includes 50 topics, as queries, represented by tuples of
(concept, question, narrative). In also includes a dataset of 191K candidate
documents from CORD-19. The relevance of 69,317 topic-document pairs are
manually evaluated by experts and annotated with three labels of unrelated,
partially related, related.

Model | LCA-F | MiF | MaF | Accu. | MiF | MiF | MiF | MiF
---|---|---|---|---|---|---|---|---
(0%) | (5%) | (10%) | (20%)
MTI (default) | 0.5630 | 0.7299 | 0.5058 | 0.4913 | 0.2216 | 0.3322 | 0.4592 | 0.5637
MTI (first line indes) | 0.5534 | 0.7220 | 0.5011 | 0.5066 | 0.2185 | 0.3092 | 0.4617 | 0.5774
attention mesh | 0.5788 | 0.7645 | 0.5298 | 0.5576 | 0.2710 | 0.3965 | 0.5041 | 0.6186
R+TR (base) (w/o ssl & mt) | 0.5396 | 0.7002 | 0.4924 | 0.4855 | 0.3066 | 0.4334 | 0.5038 | 0.5910
R+TR (base) (w/ ssl) | 0.5521 | 0.7282 | 0.5060 | 0.5099 | 0.3804 | 0.4855 | 0.6160 | 0.6634
R+TR (base) (w/ ssl & mt) | 0.5632 | 0.7555 | 0.5111 | 0.5234 | 0.4852 | 0.5920 | 0.6762 | 0.7241
R+TR (large) (w/o ssl & mt) | 0.5616 | 0.7419 | 0.5023 | 0.5225 | 0.3631 | 0.4743 | 0.5594 | 0.5952
R+TR (large) (w/ ssl) | 0.5974 | 0.7767 | 0.5319 | 0.5688 | 0.4902 | 0.6191 | 0.6978 | 0.7330
R+TR (large) (w/ ssl & mt) | 0.6115 | 0.8102 | 0.5587 | 0.5864 | 0.5644 | 0.6759 | 0.7411 | 0.7888
Table 3: Information extraction performance of our proposed models in
comparison with baselines. ssl: self-supervised learing; mt: multi-task
learning. The second half of the table shows the Micro F1 score based on the
size of the COVID-19 training dataset ranging from 0% (zero-shot) to 20% (few-
shot).

###  3.2 Evaluation Metrics

To evaluate the performance of IE we use two sets of evaluation measures: i)
flat measures such as micro- and macro-f1 scores, and ii) hierarchical: for
which we leverage BioASQ suggested
algorithm777https://github.com/BioASQ/Evaluation-
Measures/tree/master/hierarchical Lowest Common Ancestor F-measure (LCA-F)
[27].

For evaluation of IR we leverage trec_eval, the evaluation metrics and
algorithms provided by TREC-COVID 888https://trec.nist.gov/trec_eval/. The
evaluation metrics includes normalized discounted cumulative gain (nDCG@N),
P@N, Mean Average Precision (MAP), and Binary preference (Bpref) (see [2] for
more details on these metrics).

###  3.3 Analysis

In this section, we analyze various elements of our proposed methodology
architecture which is essentially retriever + transformer-based ranker
(henceforward we refer to as R+TR). In this regard, we only use the general
datasets, and do not use any sample from COVID-19 datasets for our design
decisions or hyperparametes tuning. Table 2 shows the hyperparameters we have
tried and selected for our methodology.

We start with the retriever step which has two hyperparameters for IE task;
the number of retrieved candidates $K$, and $M$ the number of candidate MeSH
terms; and one for IR task; the number of retrieved candidates
$K^{{}^{\prime}}$. We use different values and choose those that yields the
highest recalls. Since this step is not trainable we can choose different
values for $K$ and $K^{{}^{\prime}}$. It would only change the size of input
and output, not their formats. As shown in Table 2, for IE we choose $K=M=512$
because not only it gives the highest recall but also we need $K=M$ so that
the $T$ matrix in Equation 2 is reversible and likelihood scores of documents
can be mapped to those of MeSH indexes during inference. For IR we choose
$K^{{}^{\prime}}=1024$.

Ablation Study Moreover, we analyze several design decision for our
transformer-based ranking system, the effect of multi-task learning on the
general datasets, and experimentally compare our use of different attention
mechanisms. The ablation test results for IE and IR are reported in Table 1(a)
and Table 1(b) respectively.

We use two versions of our ranker, a base version (4 layers, 256 hidden size,
8 heads) and a large version (6 layers, 512 hidden size, 8 heads). We evaluate
different attention mechanism on the base model. The implementation of full
attention mechanism requires truncating the input documents resulting in poor
performance. However, the combination of global and dilated sliding window
with increasing window size shows better performance than other combinations
in both IE and IR. The multi-task learning improves the performance of IR
without affecting the IE’s performance. Such improvement is expected not only
because of the effect of transfer learning but also because IE task is design
to improve retrieval and reinforce the latent space to be closer to those of
the semantic indexes which human experts believed to be a better
representation. Such improvement is relatively higher for larger version of
our ranker, showing that the knowledge transfer capability increases with the
size of the transformer model. Note: the proposed self-supervised task is
exclusively for COVID-19 datasets and disregarded in our ablation analysis.
However, its effect is analysed in the following sections.

model | nDCG@20 | P@20 | Bpref | MAP
---|---|---|---|---
top score | 0.8496 | 0.8760 | 0.6378 | 0.4731
ranke#1 in nCDG@20 | 0.8496 | 0.8760 | 0.6372 | 0.4718
ranke#1 in P@20 | 0.8496 | 0.8760 | 0.6372 | 0.4718
ranke#1 in Bpref | 0.8490 | 0.8690 | 0.6378 | 0.4731
ranke#1 in MAP | 0.8490 | 0.8690 | 0.6378 | 0.4731
R+TR (base) (w/o ssl & mt) | 0.7915 | 0.8383 | 0.6021 | 0.4498
R+TR (base) (w/ ssl) | 0.8198 | 0.8556 | 0.6259 | 0.4685
R+TR (base) (w/ ssl & mt) | 0.8571 | 0.8707 | 0.6420 | 0.4713
R+TR (large) (w/o ssl & mt) | 0.8047 | 0.8489 | 0.6204 | 0.4566
R+TR (large) (w/ ssl) | 0.8302 | 0.8607 | 0.6331 | 0.4752
R+TR (large) (w/ ssl & mt) | 0.893 | 0.8911 | 0.6573 | 0.4911
R+TR (base) (w/ ssl & mt) (w/ f.t.) | 0.8986 | 0.9152 | 0.6635 | 0.5056
R+TR (large) (w/ ssl & mt) (w/ f.t.) | 0.9237 | 0.9459 | 0.6904 | 0.5230
Table 4: Information retrieval performance of our model with and without pre-
training on self-supervised and semantic extraction tasks. The performance of
SSL models are improved which shows the models learn the current time period
context of medicine literature. During the IE task training the model’s
attention gets well trained to extract the relative semantic indexes from
articles which is basically the main part of the IR procedure in bio-medicine
articles.

###  3.4 Information Extraction on COVID-19

In this section, we discuss about the information extraction performance of
our methodology exclusively on COVID-19. Table 3 shows the IE performance of
our models and baselines on COVID-19 IE testing set. Results on the left side
shows the performance of the models once trained on COVID-19 IE training set.
The baselines are similarly fine-tuned with the training data. Our proposed
R+TR perform similar to the state-of-the-art baselines without but the
inferior without leveraging the proposed self-supervised task and multi-task
learning with IR. However, each of these transfer learning techniques
significantly improves the IE performance. Leveraging the self-supervised
learning task contributes more since R+TR models get acquainted with the
context of novel pandemic and it’s distributions. Our large model with both
self-supervised and multi-task learning achieves higher scores than the
baselines by 0.05 LCA-F score and 0.08 Micro F1 score.

The right side shows the performances based on the size of the training data.
We chronologically sort the data and train the IE models with the a proportion
of the from the beginning. As shown in Table 3 the partitions include: 0%
which represents the zero shot learning ability, 5%, 10%, and 20% denoting the
few-shot learning. Our large R+TR’s zero-shor micro-f1 score is significantly
higher than the baselines, by 0.32 on average. It achieves 97% of its optimum
performance by using only 20% of the training data.

###  3.5 Information Retrieval on COVID-19

Table 4 shows the information retrieval performance of our models evaluated on
TREC-COVID round 5 dataset. The result for other benchmarks are announced
scores for round 5 999https://ir.nist.gov/covidSubmit/archive.html. Our models
which are only trained on BioASQ QA task shows inferior performance which is
due to the inconsistencies between two tasks, e.g. different distributions of
related/unrelated documents for every query. However, leveraging SSL and
multi-task learning, our base model beats the top nDCG@20 and Bpref scores.
This shows how the proposed transfer learning framework improves model’s
ability to scale up to a new domain. Our large R+RT achieves significantly
superior performance in every metric score.

Furthermore, to analyze the zero- and few-shot learning ability of our model,
we fine-tune our SSL multi-task learning models with TREC-COVID dataset. In
this regard, we choose round 3 dataset for training which has 40 topics
identical to the first 40 topics in round 5. This is because the competition
stated from 30 topics in round 1 and every time added 5 topics for the next
round. We leave the last 10 topics of round 5 for evaluation.

We also expand the 40 samples of (topic , set of candidate documents) to 1,530
samples by randomly selecting a subset of 128 candidate documents for every
given topic rather than 1024. As shown in the bottom of Table 4, our base and
large model are able to effectively leverage such fine-tuning and achieve
significantly better scores than the top ones, by 0.05 MAP score. Note that in
TREC-COVID challenge also participants could use results from previous rounds.

Figure 1: Illustration of attention weights between the input query and
candidate articles along with the extracted outputs. The intensity and the
color of the highlights denotes attention weights’ values which is averaged
and set to three scalar between the highly correlated terms.

###  3.6 Interpretability

As mentioned in Section 1, the COVID-19 infodemic has entangled automatic
information retrieval and extraction models, so much so that human expers’
supervision would always be required. The interpretability can help human
experts comprehend the decision making of a model and what has caused a
mistaken output. As shown in Figure 1, the local-global attention of our model
can assist human experts even when it makes an error by providing evidence for
the mistaken output and suggesting other alternatives. The model extract the
semantic index of SARS-CoV-2 while the manual annotator believes the article
is about the general SARS viruses rather than a specific variant. The
highlights shows the global attention between the related articles and the
query article, and the local attention withing the query article. The weights
are averaged and set to three scalar values, following [28], to make the
visualization simple [29]. As depicted by Figure 1, the extraction of SARS-
CoV-2 is because of the highly matched context about COVID-19 (the top related
article) and the last sentence. However, the global attention provide another
related article along with suggestions for the correct index. Knowing these,
one can quickly identify and fix the error.

Figure 2: Attention weights of terms attending to COVID-19 and SARS-CoV-2 over
different time frames. These weights are normalized for visualization purpose,
following [30].

The interpretability can also help to understand the performance of the model
in mitigating the challenges of COVID-19 infodemic. As a case study, to
analyze the performance of our model in handling the shift in the topics and
terminologies of COVID-19 related literature, we take a look at attention
weights between the various stigmatized and standard terms for the novel
Coronavirus over the time. The stigmatized terms include those which have been
used prior to the provisional standard term "2019-nCoV", such as "Wuhan
Coronavirus," "Chinese Virus," "Wuhan Novel Pneumonia" to name a few [31]. We
use the aggregated weights101010Summed and averaged over all sample queries
and candidate articles, using both local and global attentions. when these
terms attend to or get attended by the standard ones (i.e. COVID-19 and SARS-
CoV-2). We use the chronologically sorted dataset and looked at the weights as
the model gets trained over the different time frames.

As shown in Figure 2, as the distribution of terminologies changes over time,
the attention mechanism learns to relate to the well-established terms
mitigating the effect of dataset shift. In the beginning, the model shows high
attention weights toward SARS-CoV as it is another variant of Coronavirus
which has also originated from Chine; thereby substantially matches with the
new context. Then it quickly starts to related stigmatized terms even prior to
the introduction of their standard terms. With the advent of the standard
terms, the model pays less attention to stigmatized and provisional terms. The
attention over SARS-CoV-1 and other related variants decreases as the model
dissolves the confusion between them.

##  4 Conclusion

In this study, we have unified the tasks of information retrieval for question
answering and extraction of semantic indexes along with a self-supervised pre-
text task to leverage the advantages of transfer learning in addressing the
data efficiency, generalization and dataset shift issues confronting these
application in the current pandemic. In comparison to benchmarks, our model
learns with fewer number of labeled data and shows substantially higher zero-
shot performance. Furthermore, the performance of the tasks which have scarce
amount of labeled data improves dramatically once simultaneously trained with
other tasks for which relatively decent amount of labeled data is available.

In addition, the global-local attention mechanism, which is the perfect fit
for retrieval and ranking models, provided our model with not only robustness
to the drastic changes of the literature, but also interpretibility both on
document and token level. We have shown the potentials of such
interpretability in understanding the decision making of the model as well as
lightening the load on human experts.

Fututre Works: Our study brings focus towards state-of-the-art remedies to the
current challenges of the pandemic which opens up new doors to a more
systematic analysis of each of these challenges and more sophisticated
algorithms in that regard. In future, we would like to combine more
information retrieval and extraction tasks as more data is being annotated and
prepared for domain-specific environment of the pandemic. To better evaluate
the performance of the global-local interpretability, we plan to perform
qualitative analysis by providing this tool to a human expert and analyze
their time efficiency and performance. Furthermore, we believe that such
unifying mechanism would contribute to more applications and situations other
than the current pandemic. We aim to identify other situations which cause
drastic changes and growth in the literature and examine how such models would
perform in mitigating those challenges.

## FUNDING

The authors gratefully acknowledge the use of the services of Jetstream cloud,
funded by National Science Foundation (NSF) awards 1445604, and the Cloud
Technology Endowed Professorship.

## References

  * [1] Kirk Roberts, Tasmeer Alam, Steven Bedrick, Dina Demner-Fushman, Kyle Lo, Ian Soboroff, Ellen Voorhees, Lucy Lu Wang, and William R Hersh.  Trec-covid: Rationale and structure of an information retrieval shared task for covid-19.  Journal of the American Medical Informatics Association, 2020.
  * [2] Andre Esteva, Anuprit Kale, Romain Paulus, Kazuma Hashimoto, Wenpeng Yin, Dragomir Radev, and Richard Socher.  Co-search: Covid-19 information retrieval with semantic search, question answering, and abstractive summarization.  arXiv preprint arXiv:2006.09595, 2020.
  * [3] Jingqi Wang, Huy Anh, Frank Manion, Masoud Rouhizadeh, and Yaoyun Zhang.  Covid-19 signsym–a fast adaptation of general clinical nlp tools to identify and normalize covid-19 signs and symptoms to omop common data model.  ArXiv.
  * [4] Edwin Zhang, Nikhil Gupta, Raphael Tang, Xiao Han, Ronak Pradeep, Kuang Lu, Yue Zhang, Rodrigo Nogueira, Kyunghyun Cho, Hui Fang, et al.  Covidex: Neural ranking models and keyword search infrastructure for the covid-19 open research dataset.  arXiv preprint arXiv:2007.07846, 2020.
  * [5] Abdullatif Köksal, Hilal Dönmez, Rıza Özçelik, Elif Ozkirimli, and Arzucan Özgür.  Vapur: A search engine to find related protein–compound pairs in covid-19 literature.  arXiv preprint arXiv:2009.02526, 2020.
  * [6] Nico Colic, Lenz Furrer, and Fabio Rinaldi.  Annotating the pandemic: Named entity recognition and normalisation in covid-19 literature.  In Proceedings of the 1st Workshop on NLP for COVID-19 (Part 2) at EMNLP 2020, 2020.
  * [7] Panagiotis Lymperopoulos, Haoling Qiu, and Bonan Min.  Concept wikification for covid-19.  In Proceedings of the 1st Workshop on NLP for COVID-19 (Part 2) at EMNLP 2020, 2020.
  * [8] Alexander Spangher, Nanyun Peng, Jonathan May, and Emilio Ferrara.  Enabling low-resource transfer learning across covid-19 corpora by combining event-extraction and co-training.  In Proceedings of the 1st Workshop on NLP for COVID-19 at ACL 2020, 2020.
  * [9] Farhad Shokraneh and Tony Russell-Rose.  Lessons from covid-19 to future evidence synthesis efforts: first living search strategy and out of date scientific publishing and indexing industry (submitted).  Journal of Clinical Epidemiology, 2020.
  * [10] Qiao Jin, Bhuwan Dhingra, William Cohen, and Xinghua Lu.  Attentionmesh: Simple, effective and interpretable automatic mesh indexer.  In Proceedings of the 6th BioASQ Workshop A challenge on large-scale biomedical semantic indexing and question answering, pages 47–56, 2018.
  * [11] Guangxu Xun, Kishlay Jha, Ye Yuan, Yaqing Wang, and Aidong Zhang.  Meshprobenet: a self-attentive probe net for mesh indexing.  Bioinformatics, 35(19):3794–3802, 2019.
  * [12] Iz Beltagy, Matthew E Peters, and Arman Cohan.  Longformer: The long-document transformer.  arXiv preprint arXiv:2004.05150, 2020.
  * [13] Dan Hendrycks, Mantas Mazeika, Saurav Kadavath, and Dawn Song.  Using self-supervised learning can improve model robustness and uncertainty.  arXiv preprint arXiv:1906.12340, 2019.
  * [14] Yu Sun, Xiaolong Wang, Zhuang Liu, John Miller, Alexei A Efros, and Moritz Hardt.  Test-time training for out-of-distribution generalization.  2019\.
  * [15] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  arXiv preprint arXiv:1910.10683, 2019.
  * [16] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  BERT: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.
  * [17] Arman Cohan, Sergey Feldman, Iz Beltagy, Doug Downey, and Daniel S Weld.  Specter: Document-level representation learning using citation-informed transformers.  In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2270–2282, 2020.
  * [18] Iz Beltagy, Kyle Lo, and Arman Cohan.  Scibert: A pretrained language model for scientific text.  arXiv preprint arXiv:1903.10676, 2019.
  * [19] Ignazio Pillai, Giorgio Fumera, and Fabio Roli.  Threshold optimisation for multi-label classifiers.  Pattern Recognition, 46(7):2055–2065, 2013.
  * [20] Jinhyuk Lee, Wonjin Yoon, Sungdong Kim, Donghyeon Kim, Sunkyu Kim, Chan Ho So, and Jaewoo Kang.  Biobert: a pre-trained biomedical language representation model for biomedical text mining.  Bioinformatics, 36(4):1234–1240, 2020.
  * [21] Yonghui Wu, Mike Schuster, Zhifeng Chen, Quoc V Le, Mohammad Norouzi, Wolfgang Macherey, Maxim Krikun, Yuan Cao, Qin Gao, Klaus Macherey, et al.  Google’s neural machine translation system: Bridging the gap between human and machine translation.  arXiv preprint arXiv:1609.08144, 2016.
  * [22] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In Advances in neural information processing systems, pages 5998–6008, 2017.
  * [23] Anirudh Ravula, Chris Alberti, Joshua Ainslie, Li Yang, Philip Minh Pham, Qifan Wang, Santiago Ontanon, Sumit Kumar Sanghai, Vaclav Cvicek, and Zach Fisher.  Etc: Encoding long and structured inputs in transformers.  2020\.
  * [24] Ankit Gupta and Jonathan Berant.  Gmat: Global memory augmentation for transformers.  arXiv preprint arXiv:2006.03274, 2020.
  * [25] George Tsatsaronis, Georgios Balikas, Prodromos Malakasiotis, Ioannis Partalas, Matthias Zschunke, Michael R Alvers, Dirk Weissenborn, Anastasia Krithara, Sergios Petridis, Dimitris Polychronopoulos, Yannis Almirantis, John Pavlopoulos, Nicolas Baskiotis, Patrick Gallinari, Thierry Artieres, Axel Ngonga, Norman Heino, Eric Gaussier, Liliana Barrio-Alvers, Michael Schroeder, Ion Androutsopoulos, and Georgios Paliouras.  An overview of the bioasq large-scale biomedical semantic indexing and question answering competition.  BMC Bioinformatics, 16:138, 2015.
  * [26] Lucy Lu Wang, Kyle Lo, Yoganand Chandrasekhar, Russell Reas, Jiangjiang Yang, Darrin Eide, Kathryn Funk, Rodney Michael Kinney, Ziyang Liu, William. Merrill, Paul Mooney, Dewey A. Murdick, Devvret Rishi, Jerry Sheehan, Zhihong Shen, Brandon Stilson, Alex D. Wade, Kuansan Wang, Christopher Wilhelm, Boya Xie, Douglas M. Raymond, Daniel S. Weld, Oren Etzioni, and Sebastian Kohlmeier.  Cord-19: The covid-19 open research dataset.  ArXiv, 2020.
  * [27] Aris Kosmopoulos, Ioannis Partalas, Eric Gaussier, Georgios Paliouras, and Ion Androutsopoulos.  Evaluation measures for hierarchical classification: a unified view and novel approaches.  Data Mining and Knowledge Discovery, 29(3):820–865, 2015.
  * [28] Abeed Sarker, Ari Z Klein, Janet Mee, Polina Harik, and Graciela Gonzalez-Hernandez.  An interpretable natural language processing system for written medical examination assessment.  Journal of biomedical informatics, 98:103268, 2019.
  * [29] Tao Lei et al.  Interpretable neural models for natural language processing.  PhD thesis, Massachusetts Institute of Technology, 2017.
  * [30] Toan Q Nguyen and Julian Salazar.  Transformers without tears: Improving the normalization of self-attention.  arXiv preprint arXiv:1910.05895, 2019.
  * [31] Zhiwen Hu, Zhongliang Yang, Qi Li, and An Zhang.  The covid-19 infodemic: Infodemiology study analyzing stigmatizing search terms.  Journal of medical Internet research, 22(11):e22639, 2020.


CaloFlow for CaloChallenge Dataset 1

Claudius Krause1, 2, Ian Pang1, and David Shih1

1 NHETC, Dept. of Physics and Astronomy, Rutgers University, Piscataway, NJ,
USA

2 Institut für Theoretische Physik, Universität Heidelberg, Germany

## Abstract

CaloFlow is a new and promising approach to fast calorimeter simulation based
on normalizing flows. Applying CaloFlow to the photon and charged pion Geant4
showers of Dataset 1 of the Fast Calorimeter Simulation Challenge 2022, we
show how it can produce high-fidelity samples with a sampling time that is
several orders of magnitude faster than Geant4\. We demonstrate the fidelity
of the samples using calorimeter shower images, histograms of high level
features, and aggregate metrics such as a classifier trained to distinguish
CaloFlow from Geant4 samples.

###### Contents

  1. 1 Introduction
  2. 2 Dataset
  3. 3 Method
    1. 3.1 Normalizing Flows
    2. 3.2 CaloFlow
  4. 4 Results
    1. 4.1 Average shower images
    2. 4.2 Histograms
      1. 4.2.1 Energy histograms
      2. 4.2.2 Shower shape histograms
    3. 4.3 Classifier Scores
    4. 4.4 Generation timing
  5. 5 Conclusions/Outlook
  6. A Classifier Architecture

##  1 Introduction

The LHC has just restarted for its third run in the spring of 2022, and the
need for accurate fast simulation at the LHC is ever more urgent. Simulation
of calorimeter showers with Geant4[1, 2, 3] is already a major computational
bottleneck at the LHC, and this is expected to further intensify as the
detectors are upgraded and the luminosity is increased[4].

Recently, there have been significant advancements in fast calorimeter
simulation through the application of deep generative models such as
Generative Adversarial Networks (GANs), normalizing flows, Variational
Autoencoders (VAEs) and other approaches[5, 6, 7, 8, 9, 10, 11, 12, 13, 14,
15, 16, 17, 18, 19, 20] In CaloFlow-I[9] and CaloFlow-II[10], two of us have
demonstrated the effectiveness of normalizing flows in emulating Geant4 events
with high-fidelity at generation speeds that are $10^{4}$ times faster than
Geant4\. Also, we found that CaloFlow is robust against mode collapse that is
a common problem in GAN-based generative models. In fact, CaloFlow proved to
be the first ever generative model in HEP that could pass the following
stringent test: a binary classifier trained on the raw voxels of (CaloFlow)
generated vs. (Geant4) reference showers could not distinguish the two with
100% accuracy.

In this paper, we adapt CaloFlow to Dataset 1 from the Fast Calorimeter
Simulation Challenge 2022[21] (hereafter referred to as the CaloChallenge).
The CaloChallenge is aimed at encouraging the development of fast and high-
fidelity calorimeter shower generation through the application of deep
generative models. It includes three datasets with increasing dimensionality.
Dataset 1[22] has the lowest dimensionality and consists of photon $\gamma$
(368 voxels) and charged pion $\pi^{+}$ (533 voxels) showers. This dataset is
actually the official ATLAS simulation dataset used to develop the FastCaloGAN
model [8] used in a portion of AtlFast3[7], which is currently the official
fast calorimeter simulation framework of the ATLAS collaboration. It is more
realistic than the Geant4 dataset[23] used in [5, 6, 9, 10] because it
includes a realistic sampling fraction and calorimeter geometry.

We will demonstrate the robustness of CaloFlow even when applied to this new
dataset. We will show that CaloFlow can achieve comparably fast generation
times as FastCaloGAN, with a significant improvement in generation quality. We
will also show that CaloFlow compares favorably (on Dataset 1), in both
generation time and quality, to CaloScore [12]. This is a recent approach that
uses score-based diffusion models and is currently the only approach that has
been trained successfully on all three datasets of the CaloChallenge.

The outline of the paper is as follows. In Section 2, we provide a brief
description of CaloChallenge Dataset 1 (for more details, see [21]) and
outline the differences compared to the CaloGAN dataset. In Section 3, we
elaborate on the method we used to learn the distribution of showers from the
Geant4 reference datasets and the modifications we have made to adapt the
previous version of CaloFlow to CaloChallenge Dataset 1. We include the main
results of the study in Section 4. Here we show detailed comparison between
the CaloFlow generated samples and the reference Geant4 samples. Finally, we
summarize our findings in Section 5 and include possibilities for future
research.

##  2 Dataset

As stated in Section 1, there are three datasets in the CaloChallenge and our
study focuses only on Dataset 1[22]. This dataset comprises calorimeter shower
events for $\gamma$ and $\pi^{+}$ particle types and is a subset of the ATLAS
Geant4 datasets that are published at [24]. The datasets described here have
been used to train the ATLAS GAN-based model FastCaloGAN [8] used in
AtlFast3[7].

The calorimeter setup is based on a voxelized version [7] of the current ATLAS
detector configuration in the $\eta$ range $[0.2,0.25]$. As shown in left
diagram in Figure 1, the calorimeter layers are simulated as concentric
cylinders aligned along the $z$-axis, defined as the direction of travel of
the particle. Working in cylindrical coordinates, $\Delta\phi$ and
$\Delta\eta$ are defined as the $x$ and $y$-axes respectively. Based on the
voxelization, each calorimeter is further divided into radial and angular bins
as shown in the right diagram in Figure 1. Here
$r\equiv\sqrt{(\Delta\phi)^{2}+(\Delta\eta)^{2}}$ and
$\alpha\equiv\arctan\left(\frac{\Delta\eta}{\Delta\phi}\right)$. The total
number of voxels, and the arrangement of voxels in $z$, $\alpha$ and $r$ for
the $\gamma$/$\pi^{+}$ datasets are shown in Table 1. For comparison, the
CaloGAN calorimeter setup consists of 3 calorimeter layers and a total of 504
voxels regardless of the particle type.

The incident energies are discrete and range from $2^{8}$ MeV to $2^{22}$ MeV,
increasing in powers of 2. Hence, there are 15 different incident energies in
total. (In contrast, the incident energies in the CaloGAN dataset were
uniformly sampled.) The photon and pion datasets each contain 10000 events for
every low incident energy. Fewer events were generated for each high incident
energy due to the longer generation time with Geant4\. Figure 2 shows the
distribution of incident energies in the photon and pion datasets.

In Dataset 1, there are two independent photon datasets with 121000 events
each, and two independent pion datasets with 120800 events each. For both
particle types, the first dataset is used for training the generative models
(with a 70/30 train/val split); the second is used for training the classifier
metric (with a 60/20/20 train/val/test split) and for producing all the
evaluation plots. This is ideal, since then any overfitting of the density
estimator to the training set could in principle be diagnosed by the
independent dataset used for the classifier metric and evaluation plots.

| Number of voxels | $N_{z}$ | $N_{\alpha}$ | $N_{r}$
---|---|---|---|---
$\gamma$ | 368 | 5 | [1, 10, 10, 1, 1] | [8, 16, 19, 5, 5]
$\pi^{+}$ | 533 | 7 | [1, 10, 10, 1, 10, 10, 1] | [8, 10, 10, 5, 15, 16, 10]
Table 1: Details of $\gamma$ and $\pi^{+}$ shower datasets: $N_{z}$ is the
total number of calorimeter layers, $N_{\alpha}$ is the number of $\alpha$
bins in a given layer, and $N_{r}$ is the number of radial bins in a given
layer. $N_{\alpha}$ and $N_{r}$ are listed according to increasing layer
numbers (i.e. the number of $\alpha$/$r$ bins in layer 0 are leftmost on the
list)  Figure 1: Diagram of coordinate system used in the CaloChallenge
dataset[21].

Figure 2: Breakdown of incident energies in CaloChallenge Dataset 1.

##  3 Method

Our approach follows the algorithm of CaloFlow [9, 10] to a large extent. Here
we briefly describe the approach, focusing primarily on the differences with
[9, 10].

###  3.1 Normalizing Flows

Normalizing Flows (NFs) (see [25, 26, 27] for reviews on NFs) are a class of
machine learning models for density estimation and generative modeling that
aim to learn a bijective transformation (with tractable Jacobian) between data
and a latent space following a simple base distribution (e.g. the normal
distribution). Since the probability of a point of the base distribution, as
well as the Jacobian of the transformation is known, NFs can give the
probability density $p(x)$ of each data point $x$ in data space, i.e. NFs are
density estimators. When run in the other direction, starting from samples $z$
in the latent space generated from the base distribution, NFs can be used as
generative models. Since by construction, the log-likelihood (LL) of
datapoints under the flow are known, a NF can be trained by minimizing the
negative LL.

Following [9, 10], the NFs used here are based on the Masked Autoregressive
Flow (MAF)[28] and Inverse Autoregressive Flow (IAF)[29] architectures with
rational quadratic spline (RQS) transformations [30, 31]. The autoregressive
parameters of the RQS transformations are determined using neural networks
known as MADE[32] blocks. This maximizes the expressivity of the NF, but comes
at the expense of unbalanced evaluation times in the forward and inverse
direction. MAFs are fast for density estimation but slow for sampling, while
IAFs are fast for sampling but slow for density estimation. As in [9, 10], we
will refer to the MAF as the “teacher” model and the IAF as the “student”
model, as the MAF was used as an intermediate step to obtain the trained IAF
model, which is final product of the CaloFlow method.

MAFs can be trained efficiently with the LL objective described above, but for
IAFs this is usually not possible due to time and memory constraints. Training
the IAF efficiently requires a different approach known as Probability Density
Distillation (PDD)[33] or teacher-student training. Instead of fitting the IAF
directly to the data, the approach involves fitting the IAF to the MAF. In
practice, the fitting is implemented based on two training loss terms that we
refer to as $z$ and $x$-losses. To compute the $z$-loss, we begin with a
sample $z$ which is then passed through the student IAF to obtain a sample
$x^{\prime}$ in data space and the corresponding likelihood $s(x^{\prime})$.
The data sample $x^{\prime}$ is then mapped via the teacher MAF back to the
latent space which obtains the likelihood $t(x^{\prime})$. Similarly, to
compute the $x$-loss, one can start with a data sample $x$ which maps to
latent space $z^{\prime}$ via the teacher, and then map back to data space via
the student. In the original PDD study[33] study, the KL divergence of
$s(x^{\prime})$ and $t(x^{\prime})$ was initially used as the training loss.
However, the authors noted that it does not converge well. Hence, as in [10],
we used a training loss function that is based on a mean square error that
compares relevant values111These consist of coordinates before and after
passing them through the flows and RQS parameters from individual MADE blocks
within the bijectors. at each equivalent stage of the teacher and student
passes. The total loss function is then the sum of the $x$ and $z$-losses.
Such a loss function has proven effective in matching the student model to the
teacher model. For more details we refer to [10].

###  3.2 CaloFlow

CaloFlow [9, 10] learns the distribution of calorimeter showers in voxel space
conditioned on the incident energy, $p(\vec{\mathcal{I}}|E_{\text{inc}})$, in
a two-step setup. In the first step, a normalizing flow called flow-I learns
the distribution of energy depositions in the layers of the calorimeter
$E_{i}$ conditioned on the incident energy, $p_{1}(E_{i}|E_{\text{inc}})$. A
second normalizing flow, called flow-II, learns the normalized showers
conditioned on these energies,
$p_{2}(\vec{\mathcal{I}}|E_{i},E_{\text{inc}})$. Normalized in this context
means that the sum of energy depositions in all voxels per calorimeter layer
is 1.

When adapting CaloFlow to CaloChallenge Dataset 1, we retained most of its key
features, while making several important modifications that are elaborated on
below:

  1. 1.

Preprocessing:

    1. (a)

Unit-space definition:

In [9, 10], the energy deposition in the calorimeter layers were recursively
transformed to $\vec{u}\in[0,1]^{N_{z}}$, where $u_{0}$ (the first element in
$\vec{u}$) was defined by:

$u_{0}=\frac{\sum_{i=0}^{N_{z}-1}E_{i}}{E_{\text{inc}}},$ (1)

In this study, we had to modify $u_{0}$ to account for cases where
$\sum_{i=0}^{N_{z}-1}E_{i}>E_{\text{inc}}$:222Naively, we might think that
$E_{\text{inc}}\geq\sum_{i=0}^{N_{z}-1}E_{i}$ due to the conservation of
energy. However, this is not always true due to rescaling that is done on the
true deposited energy to account for the sampling fraction $(<1)$ of the
calorimeter when creating the training/evaluation datasets. we rescaled
$u_{0}$ by its maximum value taken over all showers in the training
set.333There was an exception when training on the $\gamma$ dataset. We found
that the event with the largest ratio of total deposited energy to incident
energy was an outlier ($u_{0}=3.04$) and hence decided to ignore that single
event when training our model. Specifically, we rescaled by $\max(u_{0})=2.42$
for $\gamma$ and $\max(u_{0})=7.0$ for $\pi^{+}$. This rescaling of $u_{0}$
ensures that $u_{0}\in[0,1]$ for the showers of Dataset 1.

    2. (b)

Flow-I noise level for $\pi^{+}$:

For the flow-I of $\pi^{+}$ dataset, we used a smaller range for the uniform
random noise $[0,0.1]$ keV that is added to the voxel energies prior to
summing them to get the layer energies. We observe that a lower noise level
improved the training of $\pi^{+}$ flow-I. We kept the original noise range
$[0,1]$ keV when training flow-I on the $\gamma$ dataset and flow-II for both
particle types as lowering the noise level worsened the classifier score.

    3. (c)

Preprocessing of incident energies:

The incident energy $E_{\text{inc}}$ is fed to flow-I and flow-II as a
conditional label in the form:

$\log_{10}\left(E_{\text{inc}}/33.3~{}\text{GeV}\right)\in[-2.5,2.5]$ (2)

(In [9, 10], 10 GeV was used as the normalization point.)

    4. (d)

Preprocessing of layer energies (Flow-II): The preprocessing of layer energies
used as conditional labels in flow-II is:

$\log_{10}\left(\left(E_{i}+1~{}\text{keV}\right)/100~{}\text{GeV}\right)-1\in[-2,4]$
(3)

which is slightly different than that of [9, 10].

  2. 2.

Number of training epochs: See Table 2 for changes in the number of training
epochs.

| Number of training epochs (Training time)
---|---
Teacher | Student
$\gamma$ | flow-I | 100 (46 min) | -
flow-II | 100 (77 min) | 100 (360 min)
$\pi^{+}$ | flow-I | 100 (52 min) | -
flow-II | 100 (98 min) | 150 (520 min)
Table 2: Number of training epochs used in CaloFlow when training on
CaloChallenge Dataset 1. The corresponding training times, obtained on our
TITAN V GPU, are included in parentheses. Note that flow-I is only trained
once per dataset ($\gamma$ or $\pi^{+}$), and then used for both the teacher
and the student models.

  3. 3.

Hyperparameters:

    1. (a)

Hidden layer/Batch sizes:

We experimented with different hidden layer sizes when adapting CaloFlow to
CaloChallenge Dataset 1. At times, larger hidden layer size leads to lower
losses. However, at other times, it was surprising that the performance went
down and the loss went up during training; this could be a sign of
overfitting. Memory size was also a limitation when deciding on hidden layer
size. Finally, we settled with the following hidden layer sizes: 378 for
$\gamma$ teacher; 736 for $\gamma$ student; 533 for $\pi^{+}$ teacher; 500 for
$\pi^{+}$ student. A summary of the MADE[32] block architecture features for
the various models are shown in Table 3.

When training flow-I, we used a batch size of 200 for both particle types.
When training flow-II, we used batch sizes of 500 for $\pi^{+}$ teacher and
200 for $\gamma$ teacher. For the student models, we used the batch size of
175 as stated in [9, 10]. This choice of hyperparameters enabled us to achieve
better overall classifier scores.

| input | hidden | output
---|---|---|---
$\gamma$ | Teacher | 378 | $1\times 378$ | 8464
Student | 736 | $1\times 736$ | 8464
$\pi^{+}$ | Teacher | 533 | $1\times 533$ | 12259
Student | 500 | $1\times 500$ | 12259
Table 3: Summary of MADE block architecture of the various teacher and student
models. The number nodes in the input layer, hidden layer and output layer are
shown for each of the 4 models.

    2. (b)

$E_{\text{inc}}$ distribution in $z$-loss:

Previously in [9, 10], we used the same uniform $E_{\text{inc}}$ distribution
as the training data when feeding the conditional label through the student
model to compute the $z$-loss. In general, it does not have to be the case. In
this study, we have a discrete $E_{\text{inc}}$ distribution in the training
data. We found that using this same distribution works well when computing the
$z$-loss for the $\gamma$ student model. In contrast, we observed that using a
uniform $E_{\text{inc}}$ distribution that has the same range as the training
data works better when computing the $z$-loss for the $\pi^{+}$ student model.

  4. 4.

Learning rate (LR) schedule:

Instead of the simple multi-step decreasing LR schedule used in [9, 10], here
we adopted a cyclic LR schedule [34] when training flow-I and flow-II, as this
was found to result in significantly improved performance. With a cyclic LR
schedule (see fig. 3(a)), the LR begins at a chosen base LR and then increases
up to a maximum LR. The number of batches taken to increase from the base LR
to the maximum LR is defined as the step size. After which, the LR decreases
from the maximum LR to the base LR. This increase and subsequent decrease in
LR is defined as a cycle, and this repeats for a chosen number of cycles. In
certain cases, the maximum LR can be made to decrease with increasing number
of batches. The effectiveness of cyclic LR schedule may be due the use of
higher LRs at certain points in the training which helps the model move out of
saddle points or local minima. One key difference in the implementation of
cyclic LR is that the LR is updated after each batch in contrast to updating
after each epoch milestone in [9, 10]. OneCycle LR[35] (see fig. 3(b)) is a
special case of cyclic LR which only relies on a single cycle followed by an
annihilation phase where the base LR is gradually decreased up to a factor of
10,000. As in [35], we observe that OneCycle LR policy generally allows us to
shorten the training time by obtaining a lower loss with a smaller number of
training epochs. This is especially useful when training the student model
(IAF) which generally has a longer training time per epoch compared to
training the teacher model (MAF).

When experimenting with LR schedules, we found that training with OneCycle LR
often achieved better results compared to regular cyclic LR, so we adopted
this throughout. The only exception was for $\gamma$ teacher where training
with regular cyclic LR had a better outcome.

For most of the models using OneCycle LR, the length of the annihilation phase
was chosen to be $20\%$ of the total number of training epochs shown in Table
2. The only exception was for $\pi^{+}$ teacher where the annihilation phase
was chosen to be $10\%$ of the total number of training epochs. The step size
was taken to be half of the remaining number of epochs.

For $\gamma$ teacher, which used regular cyclic LR, the length of one cycle
was chosen to be 10 epochs. After each cycle, the max LR is decreased such
that the difference between the max LR and the base LR is half that of the
previous cycle. The step size in our study was fixed to be 5 epochs (i.e. half
of a cycle).

The base and max LRs used when training the various models are shown in Table
4.

| base LR | max LR
---|---|---
$\gamma$ | Teacher | $5\times 10^{-5}$ | $2\times 10^{-3}$
Student | $4\times 10^{-5}$ | $1\times 10^{-3}$
$\pi^{+}$ | Teacher | $4\times 10^{-5}$ | $1\times 10^{-3}$
Student | $2\times 10^{-5}$ | $1\times 10^{-3}$
Table 4: Base and max LRs used when training $\gamma$ and $\pi^{+}$
teacher/student models. For $\gamma$ teacher, the max LR here refers to the
max LR in the first cycle.

(a) Regular cyclic LR

(b) OneCycle LR

Figure 3: (a) Illustration of regular cyclic LR schedule with decreasing max
LR. The length of one cycle was chosen to be 10 epochs. The step size in our
study was fixed to be half the number of batches in one cycle. After each
cycle, the max LR is decreased such that the difference between the max LR and
the base LR is half that of the previous cycle. (b) Illustration of OneCycle
LR schedule with annihilation phase.

##  4 Results

###  4.1 Average shower images

In Figures 4 and 5, we compare the average shower images between the samples
generated by CaloFlow and Geant4 for the $\gamma$ and $\pi^{+}$ datasets
respectively. Looking at the average shower images for the $\gamma$ samples,
we see that they are virtually indistinguishable by eye. In the case of
$\pi^{+}$, we note that the shower averages are also very much similar to the
eye. There is a slight noticeable difference when comparing the layer 0 image
for CaloFlow $\pi^{+}$ student and that from the Geant4 reference dataset. A
more detailed comparison between the CaloFlow and Geant4 samples can be seen
by looking at the histograms in Section 4.2.

Figure 4: Shower averages for $\gamma$ teacher (top), $\gamma$ student
(middle) and Geant4 $\gamma$ reference dataset (bottom) respectively.

Figure 5: Shower averages for $\pi^{+}$ teacher (top), $\pi^{+}$ student
(middle) and Geant4 $\pi^{+}$ reference dataset (bottom) respectively.

###  4.2 Histograms

####  4.2.1 Energy histograms

Figures 6–9 show histograms corresponding to energy distributions produced by
CaloFlow compared to those from the Geant4 reference sample. In Figure 6, we
see that, despite having more calorimeter layers than before, flow-I is still
able to precisely learn the $\gamma$ layer energies. Also, the transformation
to unit-space in the preprocessing ensured that
$E_{\text{tot}}/E_{\text{inc}}$ is learned well and this is clearly reflected
in the $E_{\text{tot}}/E_{\text{inc}}$ histogram at the bottom right of Figure
6. We are pleasantly surprised at how well flow-I was able to learn the spikes
found at higher energies in the $E_{2}$ plot. There are some minor
discrepancies at higher energies which can be attributed to there being fewer
events with high $E_{\text{inc}}$ as shown in Figure 2.

Figure 6: Energy distributions for $\gamma$ dataset.

Similarly, we see from Figure 7 that flow-I was able to learn the $\pi^{+}$
layer energies very accurately. Initially, we found that flow-I struggled with
learning the bimodal distribution found in the first 3 layer energy histograms
(top row in Figure 7). However, training flow-I with a lower noise level
helped it to better learn complicated distributions in the $\pi^{+}$ dataset.
Like for the $\gamma$ dataset, $E_{\text{tot}}/E_{\text{inc}}$ is well-learned
by flow-I. The slight discrepancy at high energies can be attributed to there
being fewer events with high $E_{\text{inc}}$ as shown in Figure 2.

Figure 7: Energy distributions for $\pi^{+}$ dataset.

Next, we compare the distributions of $E_{\text{tot}}/E_{\text{inc}}$ from the
CaloFlow and Geant4 samples for various discrete incident energies
$E_{\text{inc}}$. The corresponding histograms are shown in Figures 8 and 9.
An equivalent comparison can be found for the AtlFast3 study in Figures 10 and
11 of [7]. Overall, we find that CaloFlow is able to properly learn the mean
and widths of the $E_{\text{tot}}/E_{\text{inc}}$ distributions. CaloFlow is
even able to reproduce most of the distributions corresponding to higher
$E_{\text{inc}}$ which have lower statistics. By using lower noise level when
training $\pi^{+}$ flow-I, we were able to properly describe the bimodal
distribution found in the $\pi^{+}$ histograms corresponding to 256 MeV, 512
MeV, and 1024 MeV. This is a significant improvement from the results found in
the AtlFast3 study.

In Figures 10 and 11, we include plots that show a detailed comparison between
the histograms generated based on CaloFlow and Geant4 shown in Figures 8 and
9. Following FastCaloGAN [36], we calculated the mean and RMS from the
histograms which have excluded outliers. In Figure 10, we find that there is
less than $0.3\%$ discrepancy in the mean value of each $\gamma$ histogram in
Figure 8. There is a larger discrepancy in the RMS values. Nevertheless, the
maximum discrepancy in RMS is still only about $5\%$ and it occurs at high
$E_{\text{inc}}=2.1$ TeV which has lower statistics. Most of the RMS
discrepancies for lower $E_{\text{inc}}$ are less than $2\%$. In Figure 11, we
find that all the $\pi^{+}$ histogram means in Figure 9 have less than $2\%$
discrepancy. We also see excellent agreement in the RMS with less than $3\%$
discrepancy for most of the histograms. Comparing with the equivalent results
from AtlFast3 (see Figures 12(a) and 12(c) in [7]), we see that CaloFlow
managed to achieve a greater degree of agreement with Geant4.

As an additional comparison with the AtlFast3 results, we computed the
$\chi^{2}$ per degree of freedom ($\chi^{2}$/NDF) for the histograms in
Figures 8 and 9. Although we used a slightly different binning compared to
AtlFast3, dividing $\chi^{2}$ by NDF should make the calculated values robust
against different binning choices. The results can be found in Table 5. We
find that the $\chi^{2}/\text{NDF}$ values of the CaloFlow results in Figures
8 and 9 are $\mathcal{O}(10)$ better than the corresponding results found in
AtlFast3.

|  $\chi^{2}$/NDF
---|---
CaloFlow (this work) |  AtlFast3 [7]
$\gamma$ | 526/450 = 1.17 | 5657/419 = 13.5
$\pi^{+}$ | 592/450 = 1.32 | 5503/435 = 12.7
Table 5: Comparison of $\chi^{2}$/NDF values between CaloFlow and AtlFast3 for
histograms of $E_{\text{tot}}/E_{\text{inc}}$ at discrete values of
$E_{\text{inc}}$ (see Figures 8 and 9). Figure 8: Histograms of
$E_{\text{tot}}/E_{\text{inc}}$ for various discrete values of
$E_{\text{inc}}$ in the $\gamma$ dataset. See Figure 10 in [7] for comparison
with AtlFast3 study.  Figure 9: Histograms of $E_{\text{tot}}/E_{\text{inc}}$
for various discrete values of $E_{\text{inc}}$ in the $\pi^{+}$ dataset. See
Figure 11 in [7] for comparison with AtlFast3 study.  Figure 10: Comparison of
mean and RMS of $\gamma$ histograms between CaloFlow and Geant4 in Figure 8.
Following FastCaloGAN [36], we calculated the mean and RMS from the histograms
which have excluded outliers. Top: Mean of $E_{\text{tot}}/E_{\text{inc}}$ vs
$E_{\text{inc}}$. Middle: Ratio of $E_{\text{tot}}$ mean between CaloFlow and
Geant4\. Bottom: Ratio of $E_{\text{tot}}$ RMS between CaloFlow and Geant4\.
The low statistical uncertainty in the data points of the lower two points
resulted in the error bars being smaller than the size of the markers. See
Figure 12(a) in [7] for comparison with AtlFast3 study. Figure 11: Same as
Fig. 10 but for $\pi^{+}$ showers. See Figure 12(c) in [7] for comparison with
AtlFast3 study.

####  4.2.2 Shower shape histograms

Figures 12–14 show histograms corresponding to shower shape distributions
produced by CaloFlow compared to those from the Geant4 reference sample. The
center of energy in the $\eta$ and $\phi$ directions are defined based on the
location of the voxel centers. The center of energy in the $\eta(\phi)$
direction is defined via the locations of the voxel centers in mm, $H(F)$, as
shown in eq. (4)

$\langle\eta_{k}\rangle=\frac{\sum\left(\vec{\mathcal{I}}_{k}\odot
H\right)}{E_{k}}\quad\text{and}\quad\langle\phi_{k}\rangle=\frac{\sum\left(\vec{\mathcal{I}}_{k}\odot
F\right)}{E_{k}}$ (4)

where $\vec{\mathcal{I}}_{k}$ contains the voxel energies in the $k$th layer,
$\odot$ denotes the Hadamard product, and $\sum$ denotes sum over all
elements.

The corresponding width of the center of energy is defined as shown in eq.
(5).

$\sigma^{\eta}_{k}=\sqrt{\frac{\sum\left(\vec{\mathcal{I}}_{k}\odot
H^{2}\right)}{E_{k}}-\left(\frac{\sum\left(\vec{\mathcal{I}}_{k}\odot
H\right)}{E_{k}}\right)^{2}}\quad\text{and}\quad\sigma^{\phi}_{k}=\sqrt{\frac{\sum\left(\vec{\mathcal{I}}_{k}\odot
F^{2}\right)}{E_{k}}-\left(\frac{\sum\left(\vec{\mathcal{I}}_{k}\odot
F\right)}{E_{k}}\right)^{2}}$ (5)

We see in Figures 12–14 that the center of energy in both the $\eta$ and
$\phi$ directions modelled by CaloFlow closely match the ones found in the
Geant4 reference sample. Furthermore, teacher and student histograms are
mostly overlapped. This implies that the CaloFlow student was very well
trained on the CaloFlow teacher. Due to the close similarity between the
teacher and student histogram results, it will not be necessary for us to
distinguish between the teacher and student results in the rest of Section
4.2. This result is remarkable considering the fact that there are spikes in
the some of the center of energy histograms (e.g. see Fig 13) that might make
them difficult to learn. As for the widths of the center of energy, we find
two peaks in the distribution for both $\gamma$ and $\pi^{+}$. For $\gamma$,
the global maximum is found away from zero, while the global maximum is found
at zero for $\pi^{+}$. There is some disagreement at larger widths found in
Figures 13-14. Still, we find excellent overall agreement in the bimodal
distribution of the center of energy widths. Even at larger widths, CaloFlow
is able to model the general shape of the distribution.

Finally, looking at the distribution of voxel energies for all layers in Fig.
15, we see that they are closely modelled by CaloFlow over 6 orders of
magnitude for both, $\gamma$ and $\pi^{+}$ showers. There is some discrepancy
in the low voxel energy region between $[10^{-2},10^{-1}]$ MeV especially for
the pion showers. There is a peak around $10^{-1}$ MeV corresponding to
minimum ionizing particles (MIPs) in the pion voxel energy distribution.
However, CaloFlow is unable to properly model this feature.

Fig. 16 includes 3 voxel energy distributions with each corresponding to a
different $E_{\text{inc}}$. Showers with $E_{\text{inc}}=512$ MeV are taken to
be representative of low energy events, while showers with
$E_{\text{inc}}=32768$ MeV are representative of mid energy events, and
showers with $E_{\text{inc}}=2097152$ MeV are representative of high energy
events. We found that the MIP peak is present for low and mid $E_{\text{inc}}$
events as shown in Fig. 16. The inability to model the MIP peak might be the
reason for the poor fit at low voxel energies for the low and mid
$E_{\text{inc}}$ plots. In Section 4.3, we show that a trained binary
classifier can be sensitive to the poor fit at low voxel energies.

Figure 12: Distributions of the center of energy in the $\eta$ and $\phi$
directions for the $\gamma$ dataset. Note that layers 0, 3 and 12 only have a
single $\alpha$ bin each. Hence, there is no positional information to extract
for these layers. Figure 13: Distributions of the center of energy in the
$\eta$ direction for the $\pi^{+}$ dataset. Note that layers 0, 3 and 14 only
have a single $\alpha$ bin each. Hence, there is no positional information to
extract for these layers. Figure 14: Distributions of the center of energy in
the $\phi$ direction for the $\pi^{+}$ dataset. Note that layers 0, 3 and 14
only have a single $\alpha$ bin each. Hence, there is no positional
information to extract for these layers.

Figure 15: Distribution of voxel energies across all layers for $\gamma$ and
$\pi^{+}$ datasets.

Figure 16: Distribution of voxel energies across all layers at
$E_{\text{inc}}=512$ MeV (low energy), $E_{\text{inc}}=32768$ MeV (mid
energy), and $E_{\text{inc}}=2097152$ MeV (high energy) for $\pi^{+}$ dataset.

###  4.3 Classifier Scores

After looking at the similarities in 1D histograms between the CaloFlow and
Geant4 samples in Section 4.2, it is important to perform a full phase space
analysis by comparing the probability density that CaloFlow learned to the one
induced by Geant4\. As in [9, 10], we train a binary, neural classifier based
on a fully-connected, deep neural network (DNN) to distinguish between
generated CaloFlow samples and Geant4 samples. This serves as an approximation
to the Neyman-Pearson classifier, which is the ultimate test of whether
$p_{\text{generated}}(x)=p_{\text{data}}(x)$. According to the Neyman-Pearson
lemma, we expect the AUC to be 0.5 if the true and generated probability
densities are equal. The AUC is 1 if the classifier is able to perfectly
distinguish between generated and true samples. The second metric, JSD
$\in[0,1]$, is the Jensen-Shannon divergence which also measures the
similarity between the two probability distributions. The JSD is 0 if the two
distributions are identical and 1 if they are disjoint.

We detail the classifier architecture and training procedure444We also trained
a classifier with the architecture and preprocessing of [12] and found
quantitatively similar results to tab. 6. in appendix A. The resulting scores
are summarized in Table 6. Low-level input refers to incident energies
$E_{\text{inc}}$ and the calorimeter samples in the format provided in
CaloChallenge datasets. Training the classifier on low-level input allows us
to measure the difference between the generated and reference datasets based
on the voxel energies and their physical location.

Meanwhile, the “high-level” scores in Table 6 correspond to a DNN classifier
trained on a set of physically-relevant high-level features, here chosen to
be: incident energy $E_{\text{inc}}$, layer energies $E_{i}$, the center of
energies and their corresponding widths in the $\eta$ and $\phi$ directions.
The definitions of the center of energies and corresponding widths were stated
in eqs. (4) and (5).

AUC / JSD | DNN based classifier
---|---
Geant4 vs. |  Geant4 vs.
|  CaloFlow (teacher) |  CaloFlow (student)
$\gamma$ | low-level (regular) | 0.701(3) / 0.092(3) | 0.739(3) / 0.131(4)
low-level (logit) | 0.678(4)/ 0.075(4) | 0.706(3)/ 0.100(3)
high-level | 0.551(3) / 0.013(2) | 0.556(3) / 0.015(2)
$\pi^{+}$ | low-level (regular) | 0.827(3)/ 0.260(5) | 0.866(2)/ 0.341(3)
low-level (logit) | 0.911(3)/ 0.457(7) | 0.914(2)/0.464(6)
high-level | 0.692(2)/ 0.098(2) | 0.706(4)/ 0.108(4)
Table 6: Classifier results, based on 10 independent runs. The first (second)
numbers in every entry are AUCs (JSDs). See text for more detailed
explanations.

From the scores in Table 6, we see that CaloFlow is able to produce generated
samples of sufficiently high fidelity to fool the DNN-based classifier. This
is evident from most of the classifier AUC and JSD scores in Table 6 being
significantly lower than unity. Overall, the teacher models performed better
than the student models. Nevertheless, the student models were still able to
achieve impressive classifier scores which quantitatively demonstrates the
effectiveness of normalizing flows in generative modeling tasks. We note that
the high-level classifier scores are even better than the low-level classifier
scores for a given model. This is good considering that these high-level
features are probably more directly relevant for subsequent analysis steps
(such as reconstruction).

Recently, the authors of CaloScore[12] have produced results for the $\gamma$
portion of Dataset 1 (as well as Datasets 2 and 3), using a score-based
diffusion model. For the classifier test between CaloScore-generated and
Geant4 showers (i.e. low-level features), they found an AUC of 0.98. Even
accounting for the possibility of different classifier architectures and
training hyperparameters, this likely indicates considerably more separability
than our CaloFlow-generated showers (whose AUC is 0.739 for the low-level
classifier trained on $\gamma$ showers from the student flow).

Figure 17: Plots of AUC vs $E_{\text{inc}}$. The AUCs are obtained using a
classifier trained (10 independent runs) on low-level input from all showers
from the CaloFlow $\pi^{+}$ student and Geant4 samples, and then evaluated on
showers of specific $E_{\text{inc}}$. The left plot is based on voxel energy
inputs with the regular proprecessing decribed in appendix A. The right plot
is based on voxel energy inputs with logit preprocessing. The reference AUC
based on showers of all $E_{\text{inc}}$ is shown by the horizontal line.

The low-level classifiers, with regular preprocessing as in [21], obtained
from the 10 independent runs (see Table 6) were evaluated on $\pi^{+}$ showers
with specific $E_{\text{inc}}$. The corresponding plot of AUC vs
$E_{\text{inc}}$ is shown in Fig. 17 (left plot). The plot seems to imply that
CaloFlow is fitting the voxel energy distribution better at low
$E_{\text{inc}}$. However, this does not agree with our observations from the
plots in Fig. 16. We found that preprocessing voxel energy inputs in the
following way allows the classifier to become sensitive to deviations at low
voxel energies: First we normalize the voxel energies by layer energies
$E_{i}$ (such that normalized voxel energies in each layer sum to one). Then,
we transform the normalized voxel energies to logit space as done during
training. The overall classifier scores for this modified logit preprocessing
for both particle types are also shown in Table 6.

The corresponding plot of AUC vs $E_{\text{inc}}$ for this modified
preprocessing is shown in Fig. 17 (right plot). Here we see that the low
$E_{\text{inc}}$ showers generally have higher AUC scores compared to the
scores at the same $E_{\text{inc}}$ in the left plot. This is more consistent
with our observations from Fig. 16.

###  4.4 Generation timing

|  | Student generation time per event
---|---|---
Particle type | Batch size | GPU | CPU
$\gamma$ | 1 | 57.23 ms | 115.88 ms
10 | 5.81 ms | 12.68 ms
100 | 0.62 ms | 2.50 ms
1000 | 0.11 ms | -
| 10000 | 0.07 ms | -
$\pi^{+}$ | 1 | 74.09 ms | 126.05 ms
10 | 7.46 ms | 14.03 ms
100 | 0.80 ms | 2.91 ms
1000 | 0.15 ms | -
| 10000 | 0.09 ms | -
Table 7: Average time taken to generate a single shower event by CaloFlow
student for the two particle types. The timing was computed for different
generation batch sizes on our Intel i9-7900X CPU at 3.30GHz and our TITAN V
GPU. We were not able to generate the shower events on the CPU for batch sizes
of 1000 and 10000 due to memory constraints.

We found that the average time taken to generate a single shower events by
CaloFlow teacher for a generation batch size of 10000 on our TITAN V GPU is
18.91 ms for $\gamma$ and 43.91 ms for $\pi^{+}$. For the same batch size, we
were able to significantly reduce the generation time for CaloFlow student.

In Table 7, we include the average time taken to generate a single shower
events by CaloFlow student for different generation batch sizes. The timings
were separately evaluated on our Intel i9-7900X CPU at 3.30GHz and our TITAN V
GPU. We observe that increasing the batch size reduces the generation time per
shower event. However, we could not generate with large (1000 and 10000) batch
sizes on the CPU due to memory constraints. We observed that the generation
time is largely dependent on the sizes of layers in the MADE block (see Table
3 for MADE block layer sizes). This accounts for the difference (similarity)
in generation timing between $\gamma$ teacher (student) and $\pi^{+}$ teacher
(student). With CaloFlow student, we are able to achieve impressively low GPU
generation times of 0.07 ms per event for $\gamma$ and 0.09 ms per event for
$\pi^{+}$.

Figure 18 compares the average CPU generation time per event found using
CaloFlow, FastCaloGAN555Note that the timing of FastCaloGAN also includes the
time it takes to map the voxel energies back to calorimeter cells and we do
not know how the total time splits between voxel energy generation by the GAN
and cell assignment of the subsequent step. and Geant4\. The generation times
for 65 GeV and 2 TeV are shown for FastCaloGAN and Geant4 based on timings
provided in Section 6.4 of [8]. For CaloFlow, we included generation times for
these two energies and also two other intermediate energies (262 GeV and 1
TeV). As we do not know the batch size used when timing FastCaloGAN, we timed
CaloFlow separately using batch sizes of 1 (abbreviated as CF(1)) and 100
(abbreviated as CF(100)). We see from Figure 18 that CaloFlow generation times
are constant regardless of $E_{\text{inc}}$. In contrast, Geant4 generation
time increases with $E_{\text{inc}}$. At $E_{\text{inc}}=65$ GeV, Geant4 has a
generation time that is $\mathcal{O}(10^{2})$ slower compared to CF(1) and
$\mathcal{O}(10^{3})$ slower compared to CF(100). At $E_{\text{inc}}=2$ TeV,
Geant4 has a generation time that is $\mathcal{O}(10^{3})$ slower compared to
CF(1) and $\mathcal{O}(10^{4})$ slower compared to CF(100). The timings for
FastCaloGAN are on the same order as those found using CF(1).

Figure 18: Comparison of CPU generation time per event for CaloFlow
(abbreviated as CF in legend), FastCaloGAN (abbreviated as FCGAN in legend)
and Geant4 at several different $E_{\text{inc}}$. Numbers in parentheses refer
to the generation batch sizes used in CaloFlow. The CaloFlow timing was based
on our Intel i9-7900X CPU at 3.30GHz. The other timings are based on CPU used
by FastCaloGAN.

We can also compare directly against the results of CaloScore for the $\gamma$
portion of Dataset 1. Table II of [12] reports a time of 4s to generate 100
$\gamma$ showers using the score-based diffusion model approach, on a GPU with
batch size of 100. Meanwhile the CaloFlow student accomplishes the same task
on similar hardware in 0.062s, which is nearly 2 orders of magnitude faster.

##  5 Conclusions/Outlook

We have shown that the CaloFlow algorithm can be applied successfully to
realistic detector simulations, namely Dataset 1 of the Fast Calorimeter
Simulation Challenge 2022. Our results significantly outperform FastCaloGAN, a
deep generative model that was trained on the same dataset and is used in a
(very limited) portion of AtlFast3, the current fast-simulation framework of
the ATLAS experiment. Based on the results shown here, we are confident that
state-of-the-art deep generative models (and particularly those based on
normalizing flows) will be able to play a much larger role in future versions
of the fast-simulation framework.

From a machine-learning perspective, we have seen some drawbacks in the choice
to use discrete incident energies, in particular concerning the need to
interpolate between energies. We expect CaloFlow to interpolate well, but with
the datasets provided by the ATLAS collaboration, there is no way of checking
this in detail. We think training a conditional generative model using
continuous, uniformly (or log-uniformly) sampled energies, instead of discrete
energies, could be advantageous and should be considered in the future.

One of the main future directions stemming from this work is figuring out how
to extend the CaloFlow framework to larger numbers of voxels, e.g. those of
Datasets 2 [37] and 3 [38] of the CaloChallenge. Indeed, the current CaloFlow
algorithm scales badly with the number of simulated voxels and cannot be
applied as-is to Datasets 2 and 3. However, we expect straightforward
modifications of the basic CaloFlow approach to be strong contenders for fast-
and-accurate simulation of these higher-dimensional datasets. In the future,
it will be interesting to compare flow-based approaches to Datasets 2 and 3 to
other approaches, such as the score-based diffusion models of [12] (which are
currently the only generative models successfully trained on Datasets 2 and
3).

It would also be interesting to consider generalizing the CaloFlow approach
beyond the incident particles considered here, and beyond the narrow $\eta$
slice. For example, it should be very straightforward to train a single
conditional flow or pair of flows, conditioned on $\eta$ as well as the
incident energy. This could potentially simplify the AtlFast3 framework, where
300 GANs were trained on individual narrow $\eta$ slices.

## Acknowledgements

We are grateful to Michele Faucci Gianelli and Ben Nachman for helpful
discussions and comments on the draft. This work was supported by DOE grant
DOE-SC0010008. CK would like to thank the Baden-Württemberg-Stiftung for
support through the program Internationale Spitzenforschung, project
Uncertainties — Teaching AI its Limits (BWST_IF2020-010). IP would like to
thank the Rutgers Dept. of Physics and Astronomy for support through the Boyd
scholarship.

In this work, we used the NumPy 1.16.4 [39], Matplotlib 3.1.0 [40], sklearn
0.21.2 [41], h5py 2.9.0 [42], pytorch 1.11.0 [43], and nflows 0.14 [44]
software packages.

##  Appendix A Classifier Architecture

The classifier-based performance metric uses the classifier architecture that
was provided in the evaluation script of the CaloChallenge [21]. It is based
on the classifiers that were used in [9, 10]. In detail, the classifier is a
deep, fully-connected neural network with an input and two hidden layers with
512 nodes each. The output layer returns a single number which is passed
through a sigmoid activation function. All other activation functions are
leaky ReLUs, with default negative slope of $0.01$. We do not use any dropout
or batch normalization regulators.

In the classifier of the low-level features, we use as input the incident
energy (preprocessed as $\log_{10}{E_{\text{inc}}}$) and the energy deposition
in each voxel (preprocessed as $\mathcal{I}/E_{\text{inc}}$). High-level
features are the incident energy (preprocessed as
$\log_{10}{E_{\text{inc}}}$), the energy deposited in each layer (preprocessed
as $\log_{10}{(E_{i}+10^{-8})}$), the center of energy in $\eta$ (normalized
with a factor 100), the center of energy in $\phi$ (normalized with a factor
100), the width of the $\eta$ distribution (normalized with a factor 100), and
the width of the $\phi$ distribution (normalized with a factor 100).

We split all the data in train/test/validation sets of the ratio (60:20:20).
For the both the $\gamma$ and $\pi^{+}$ showers, we used the second file that
was provided at [22]. From our flow models, we used a generated sample of the
same size and $E_{\text{inc}}$ distribution as the Geant4 data.

The networks are then optimized by training 50 epochs with an Adam [45]
optimizer with initial learning rate of $2\cdot 10^{-4}$ and a batch size of
1000, minimizing the binary cross entropy. We use the model state with the
highest accuracy on the validation set for the final evaluation and we
subsequently calibrate the classifier using isotonic regression [46] of
sklearn [41] based on the validation dataset before evaluating the test set.

## References

  * [1] S. Agostinelli _et al._ ,  _GEANT4–a simulation toolkit_ ,  Nucl. Instrum. Meth. A 506, 250 (2003),  10.1016/S0168-9002(03)01368-8.
  * [2] J. Allison, K. Amako, J. Apostolakis, H. Araujo, P. Arce Dubois, M. Asai, G. Barrand, R. Capra, S. Chauvie, R. Chytracek, G. Cirrone, G. Cooperman _et al._ ,  _Geant4 developments and applications_ ,  IEEE Transactions on Nuclear Science 53(1), 270 (2006),  10.1109/TNS.2006.869826.
  * [3] J. Allison, K. Amako, J. Apostolakis, P. Arce, M. Asai, T. Aso, E. Bagli, A. Bagulya, S. Banerjee, G. Barrand, B. Beck, A. Bogdanov _et al._ ,  _Recent developments in geant4_ ,  Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 835, 186 (2016),  https://doi.org/10.1016/j.nima.2016.06.125.
  * [4] P. Calafiura, J. Catmore, D. Costanzo and A. Di Girolamo,  _ATLAS HL-LHC Computing Conceptual Design Report_ ,  Tech. rep., CERN, Geneva (2020).
  * [5] M. Paganini, L. de Oliveira and B. Nachman,  _CaloGAN: Simulating 3d high energy particle showers in multilayer electromagnetic calorimeters with generative adversarial networks_ ,  Physical Review D 97(1) (2018),  10.1103/physrevd.97.014021.
  * [6] M. Paganini, L. de Oliveira and B. Nachman,  _Accelerating Science with Generative Adversarial Networks: An Application to 3D Particle Showers in Multilayer Calorimeters_ ,  Phys. Rev. Lett. 120(4), 042003 (2018),  10.1103/PhysRevLett.120.042003,  1705.02355.
  * [7] G. Aad _et al._ ,  _AtlFast3: the next generation of fast simulation in ATLAS_ ,  Comput. Softw. Big Sci. 6, 7 (2022),  10.1007/s41781-021-00079-7,  2109.02551.
  * [8] ATLAS Collaboration,  _Fast simulation of the ATLAS calorimeter system with Generative Adversarial Networks_ ,  Tech. rep., CERN, Geneva,  All figures including auxiliary figures are available at https://atlas.web.cern.ch/Atlas/GROUPS/PHYSICS/PUBNOTES/ATL-SOFT-PUB-2020-006 (2020).
  * [9] C. Krause and D. Shih,  _CaloFlow: Fast and Accurate Generation of Calorimeter Showers with Normalizing Flows_ ,  Phys. Rev. D 107(11), 113003 (2023),  10.1103/PhysRevD.107.113003,  2106.05285.
  * [10] C. Krause and D. Shih,  _CaloFlow II: Even Faster and Still Accurate Generation of Calorimeter Showers with Normalizing Flows_ ,  Phys. Rev. D 107(11), 113004 (2023),  10.1103/PhysRevD.107.113004,  2110.11377.
  * [11] E. Buhmann, S. Diefenbacher, E. Eren, F. Gaede, D. Hundhausen, G. Kasieczka, W. Korcari, A. Korol, K. Krüger, P. McKeown _et al._ ,  _Fast and accurate electromagnetic and hadronic showers from generative models_ ,  In _EPJ Web of Conferences_ , vol. 251, p. 03049. EDP Sciences (2021).
  * [12] V. Mikuni and B. Nachman,  _Score-based generative models for calorimeter shower simulation_ ,  Phys. Rev. D 106(9), 092009 (2022),  10.1103/PhysRevD.106.092009,  2206.11898.
  * [13] L. de Oliveira, M. Paganini and B. Nachman,  _Controlling Physical Attributes in GAN-Accelerated Simulation of Electromagnetic Calorimeters_ ,  J. Phys. Conf. Ser. 1085(4), 042017 (2018),  10.1088/1742-6596/1085/4/042017,  1711.08813.
  * [14] M. Erdmann, L. Geiger, J. Glombitza and D. Schmidt,  _Generating and refining particle detector simulations using the Wasserstein distance in adversarial networks_ ,  Comput. Softw. Big Sci. 2(1), 4 (2018),  10.1007/s41781-018-0008-x,  1802.03325.
  * [15] M. Erdmann, J. Glombitza and T. Quast,  _Precise simulation of electromagnetic calorimeter showers using a Wasserstein Generative Adversarial Network_ ,  Comput. Softw. Big Sci. 3(1), 4 (2019),  10.1007/s41781-018-0019-7,  1807.01954.
  * [16] ATLAS Collaboration,  _Deep generative models for fast photon shower simulation in ATLAS_ (2022),  2210.06204.
  * [17] D. Belayneh _et al._ ,  _Calorimetry with deep learning: particle simulation and reconstruction for collider physics_ ,  Eur. Phys. J. C 80(7), 688 (2020),  10.1140/epjc/s10052-020-8251-9,  1912.06794.
  * [18] E. Buhmann, S. Diefenbacher, E. Eren, F. Gaede, G. Kasieczka, A. Korol and K. Krüger,  _Getting High: High Fidelity Simulation of High Granularity Calorimeters with High Speed_ ,  Comput. Softw. Big Sci. 5(1), 13 (2021),  10.1007/s41781-021-00056-0,  2005.05334.
  * [19] E. Buhmann, S. Diefenbacher, E. Eren, F. Gaede, G. Kasieczka, A. Korol and K. Krüger,  _Decoding Photons: Physics in the Latent Space of a BIB-AE Generative Network_ ,  EPJ Web Conf. 251, 03003 (2021),  10.1051/epjconf/202125103003,  2102.12491.
  * [20] E. Buhmann, S. Diefenbacher, D. Hundhausen, G. Kasieczka, W. Korcari, E. Eren, F. Gaede, K. Krüger, P. McKeown and L. Rustige,  _Hadrons, better, faster, stronger_ ,  Mach. Learn. Sci. Tech. 3(2), 025014 (2022),  10.1088/2632-2153/ac7848,  2112.09709.
  * [21] M. F. Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast calorimeter simulation challenge 2022_ ,  https://calochallenge.github.io/homepage/ (2022).
  * [22] M. F. Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast calorimeter simulation challenge 2022 - dataset 1_ ,  https://doi.org/10.5281/zenodo.6234054 (2022).
  * [23] B. Nachman, L. de Oliveira and M. Paganini,  _Electromagnetic calorimeter shower images_ ,  Mendeley Data (2017).
  * [24] ATLAS collaboration,  _Datasets used to train the generative adversarial networks used in ATLFast3_ ,  10.7483/OPENDATA.ATLAS.UXKX.TXBN (2021).
  * [25] I. Kobyzev, S. J. Prince and M. A. Brubaker,  _Normalizing flows: An introduction and review of current methods_ ,  IEEE transactions on pattern analysis and machine intelligence 43(11), 3964 (2020).
  * [26] G. Papamakarios, E. T. Nalisnick, D. J. Rezende, S. Mohamed and B. Lakshminarayanan,  _Normalizing flows for probabilistic modeling and inference._ ,  J. Mach. Learn. Res. 22(57), 1 (2021).
  * [27] D. Rezende and S. Mohamed,  _Variational inference with normalizing flows_ ,  In _International conference on machine learning_ , pp. 1530–1538. PMLR (2015).
  * [28] G. Papamakarios, T. Pavlakou and I. Murray,  _Masked autoregressive flow for density estimation_ ,  Advances in neural information processing systems 30 (2017).
  * [29] D. P. Kingma, T. Salimans, R. Jozefowicz, X. Chen, I. Sutskever and M. Welling,  _Improved variational inference with inverse autoregressive flow_ ,  Advances in neural information processing systems 29 (2016).
  * [30] C. Durkan, A. Bekasov, I. Murray and G. Papamakarios,  _Neural spline flows_ ,  Advances in neural information processing systems 32 (2019).
  * [31] J. Gregory and R. Delbourgo,  _Piecewise rational quadratic interpolation to monotonic data_ ,  IMA Journal of Numerical Analysis 2(2), 123 (1982).
  * [32] M. Germain, K. Gregor, I. Murray and H. Larochelle,  _Made: Masked autoencoder for distribution estimation_ ,  In _International conference on machine learning_ , pp. 881–889. PMLR (2015).
  * [33] A. van den Oord, Y. Li, I. Babuschkin, K. Simonyan, O. Vinyals, K. Kavukcuoglu, G. van den Driessche, E. Lockhart, L. Cobo, F. Stimberg, N. Casagrande, D. Grewe _et al._ ,  _Parallel WaveNet: Fast high-fidelity speech synthesis_ ,  In J. Dy and A. Krause, eds., _Proceedings of the 35th International Conference on Machine Learning_ , vol. 80 of _Proceedings of Machine Learning Research_ , pp. 3918–3926. PMLR (2018).
  * [34] L. N. Smith,  _Cyclical learning rates for training neural networks_ ,  In _2017 IEEE winter conference on applications of computer vision (WACV)_ , pp. 464–472. IEEE (2017).
  * [35] L. N. Smith and N. Topin,  _Super-convergence: Very fast training of neural networks using large learning rates_ ,  In _Artificial intelligence and machine learning for multi-domain operations applications_ , vol. 11006, pp. 369–386. SPIE (2019).
  * [36] ATLAS collaboration,  _Fastcalogan training project_ ,  10.5281/zenodo.5589623 (2021).
  * [37] M. F. Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast calorimeter simulation challenge 2022 - dataset 2_ ,  https://doi.org/10.5281/zenodo.6366271 (2022).
  * [38] M. F. Giannelli, G. Kasieczka, C. Krause, B. Nachman, D. Salamani, D. Shih and A. Zaborowska,  _Fast calorimeter simulation challenge 2022 - dataset 3_ ,  https://doi.org/10.5281/zenodo.6366324 (2022).
  * [39] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus _et al._ ,  _Array programming with NumPy_ ,  Nature 585(7825), 357 (2020),  10.1038/s41586-020-2649-2.
  * [40] J. D. Hunter,  _Matplotlib: A 2d graphics environment_ ,  Computing in Science Engineering 9(3), 90 (2007),  10.1109/MCSE.2007.55.
  * [41] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos _et al._ ,  _Scikit-learn: Machine learning in Python_ ,  Journal of Machine Learning Research 12, 2825 (2011).
  * [42] A. Collette,  _Python and HDF5_ ,  O’Reilly (2013).
  * [43] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf _et al._ ,  _Pytorch: An imperative style, high-performance deep learning library_ ,  In H. Wallach, H. Larochelle, A. Beygelzimer, F. Alché-Buc, E. Fox and R. Garnett, eds., _Advances in Neural Information Processing Systems 32_ , pp. 8024–8035. Curran Associates, Inc. (2019).
  * [44] C. Durkan, A. Bekasov, I. Murray and G. Papamakarios,  _nflows: normalizing flows in PyTorch_ ,  10.5281/zenodo.4296287 (2020).
  * [45] D. P. Kingma and J. Ba,  _Adam: A method for stochastic optimization_ (2014), 1412.6980.
  * [46] C. Guo, G. Pleiss, Y. Sun and K. Q. Weinberger,  _On Calibration of Modern Neural Networks_ ,  arXiv e-prints arXiv:1706.04599 (2017),  1706.04599.