FAIR AI models in high energy physics

The project template is designed to be used with the cookiecutter [39] program, a command-line utility that creates projects from project templates using the Jinja2 [40] templating engine, and that can be installed via pip. A new FAIR AI project can be made with the command cookiecutter https://github.com/FAIR4HEP/cookiecutter4fair. The first argument corresponds to the project template that is hosted on GitHub. After asking the user for the project name, repository name, author name, author ORCID, description of the project, chosen license, DOI for the input data, DOI for the code (if available), and whether to include a template Dockerfile, cookiecutter will create the template structure as shown in figure 1.

Zoom In Zoom Out Reset image size

The questions that the repository asks the user upon project creation can be found and modified in the file cookiecutter.json. The Makefile contains commands that allow the user to do various things with their project, such as downloading the data, setting up the test environment, converting the dataset, and training and evaluating the model. It also contains global variables obtained from cookiecutter.json. This procedure makes it explicit that the analysis operations are a directed acyclic graph (DAG).

If the data is hosted on Zenodo [41], the user can download the data from the DOI link by invoking make sync_data_zenodo, which uses the zenodo_get command line utility [42] to download the data. The Dockerfile can be built and run to provide a Python environment for the project to work, which installs the dependencies specified in requirements.txt. When the Docker image is built, it can be run interactively with the command docker run -d -t <image name>. The pre-project and post-project scripts are automatically run before and after the project directory is generated and provide additional flexibility. After the project template has been generated, the user can organize their source code and documentation in order to follow the FAIR principles.

2.2.2.1. Findable

2.2.2.2. Accessible

2.2.2.3. Interoperable

2.2.2.4. Reusable

2.2.2.5. Other considerations

The IN model was first proposed [60] in order to explore the evolution of physical dynamics and was later adapted for the task of jet classification; in this case differentiating ${\mathrm{H}\to{{\mathrm{b}}\overline{\mathrm{b}}}}$ jets from QCD jets [18]. The dataset for training, validation, and testing is derived from the CMS open simulated dataset with 2016 conditions that is available from the CERN Open Data Portal [15]. It consists of jets, decomposed into constituent charged particle tracks, and SVs, labeled as either ${\mathrm{H}\to{{\mathrm{b}}\overline{\mathrm{b}}}}$ signal or QCD background. More information on the dataset can be found in Chen et al [16]. Figure 3 shows the IN model architecture and table 3 provides the values of the model hyperparameters as well as input data dimensions for the baseline model. For a detailed description of the model and chosen hyperparameters, see Moreno et al [18].

Zoom In Zoom Out Reset image size

Table 3. The choice of IN model hyperparameters and input data dimensions for the baseline model.

Hyperparameter	Value
$(P, N_\mathrm{p}, S, N_\mathrm{v})$	$(30, 60, 14, 5)$
No. of hidden layers	3
Hidden layer dimension	60
$(D_E, D_O)$	(20, 24)
Activation	ReLU

As discussed in Moreno et al [18], graphs are natural data structures to describe jets because they are permutation invariant (i.e. there is no preferred order to the constituents of the jet), they can accommodate variable-sized objects (i.e. jets may be composed of a few or many constituents), and they can describe entities as nodes (i.e. constituents) and their relations as edges. This network was trained on graph data structures based on up to $N_\mathrm{p} = 30$ particle tracks, each with P = 60 features, and up to $N_\mathrm{v} = 5$ SVs, each with S = 14 features, associated with the jet. The physical description of each feature is given in appendix C of Moreno et al [18].

Two input graphs are used: a fully-connected directed graph with $N_\mathrm{pp} = N_\mathrm{p}(N_\mathrm{p} - 1)$ edges between the particle tracks and a separate graph with $N_\mathrm{vp} = N_\mathrm{v}N_\mathrm{p}$ connections between the particle tracks and the SVs. The node level feature space of the fully connected track graph is transformed to edge level features via two interaction matrices, identified as $R_{R[N_\mathrm{p} \times N_\mathrm{pp}]}$ and $R_{S[N_\mathrm{p} \times N_\mathrm{pp}]}$, where the former accounts for how each node receives information from other nodes and the latter encodes the information about each node sending information to other nodes. The track–vertex graph is transformed by similarly defined interaction matrices: $R_{K[N_\mathrm{p} \times N_\mathrm{vp}]}$ and $R_{V[N_\mathrm{v} \times N_\mathrm{vp}]}$. The feature spaces of these graphs are transformed via nonlinear functions, respectively called $f_R^\mathrm{pp}$ and $f_R^\mathrm{vp}$, to obtain two D_E dimensional internal state representations of these graphs. These nonlinear functions are approximated by fully connected multilayer perceptrons (MLPs).

These internal state representations, respectively given by $E_{\mathrm{pp}[D_E \times N_\mathrm{pp}]}$ and ${E}_{\mathrm{vp}[D_E \times N_\mathrm{vp}]}$ matrices, are transferred back to the particle tracks by transforming them with $R_R^\mathrm{T}$ and $R_K^\mathrm{T}$ matrices. These transformed particle level representations are given by matrices $\bar{E}_{\mathrm{pp}[D_E \times N_\mathrm{p}]}$ and $\bar{E}_{\mathrm{vp}[D_E \times N_\mathrm{p}]}$ respectively. Concatenating these particle-level internal state representations with the original track features creates a feature space with a dimension of $(P+2D_E)$ for each of the $N_\mathrm{p}$ tracks. The function f_O , represented by a trainable dense MLP, creates the post-interaction D_O dimensional internal representation that is stored in the matrix $O_{[D_O \times N_\mathrm{p}]}$. Finally, these track-level internal representations are summed to obtain a D_O dimensional state vector $\bar{O}$ and linearly combined to produce a two-dimensional output, which is transformed to individual class probabilities via a softmax function.

We created a FAIR implementation of the AI model hosted on GitHub and Zenodo [61]. The repository was initialized using the template described in section 2.2.

2.4.2.1. Features

2.4.2.2. Deployment to DLHub

In this subsection we provide details of training the benchmark experiments of the IN model with the same data input and hyperparameters setting as used by Moreno et al [18]. The training samples are saved in 57 HDF5 files, each of which contains about 100k jets. We use 52 of them for training and 5 for validation. The testing dataset is saved as a set of NumPy array files (one feature per file), where each file contains 600k jets.

There are several differences in our experiment setting compared to Moreno et al. For the training platform, we use the hardware qccelerated learning (HAL) GPU cluster at the National Center for Supercomputing Applications (NCSA) [65] as a remote GPU cluster and train on the NVIDIA V100 GPU, while Moreno et al trained their model on one NVIDIA GeForce GTX 1080 GPU. For the data splitting, we take the first five HDF5 files as validation data and the rest as training data. Moreno et al split the data into training, validation, and test samples, with 80%, 10%, and 10% of the data respectively. In our training process, each epoch takes about 450 s to finish. The training terminates following the early stopping condition when the validation loss failed to improve for eight epochs. As a first check, table 4 shows a comparison of our training results and the results from Moreno et al. We repeat the training ten times varying the random seed used for initialization and data shuffling, and report the mean and standard deviation of the validation accuracy and the area under the curve (AUC). We also report the one-sided (upper tail) p-value for the original model given the distribution of our trials. We find the reported performance of the original model is consistent (p-value $\gt5\%$) with our reproduction.

Table 4. The IN model's performance in this work and as reported in the original publication. In this work, we repeat the training ten times varying the random seed used for initialization and data shuffling, and report the mean and standard deviation of the validation accuracy and the AUC. We also report the one-sided (upper tail) p-value for the original model given the distribution of our trials. We find the reported performance of the original model is consistent (p-value $\gt5\%$) with our reproduction.

	Validation accuracy	AUC
IN: this work	$0.9545\pm0.0005$	$0.9898\pm0.0002$
IN: original model [18]	0.9550	0.9900
p-value (consistency)	12.69%	9.27%

Here we explore the portability of the AI model across software frameworks that are extensively used for AI research, optimal assembly of software and hardware solutions, and containers.

3.3.1.1. ONNX and TensorRT conversion

3.3.1.2. GPU utilization and throughput

Since NVIDIA TensorRT was developed to optimize AI models for accelerated inference, we have quantified the interplay between batch size, GPU utilization, throughput, and inference accuracy. In this context, throughput corresponds to the number of inferred events per second. In practice, throughput is calculated by computing the total number of inferences divided by total time, or batch size divided by the average running time per batch. Here, running time corresponds to the time taken to complete the analysis of one batch, including data transfer between device–host and the inference part at the GPU. In our experiments, we increased the batch size from 100 to 2400 with a step size equal of 200, while from 2400 to 4200, we used a step size of 400. For each batch size, we run ten times and draw a boxplot of the throughput. Our findings are summarized in figure 4. At a glance, we see that GPU utilization saturates at 100% for a batch size of 1000, while throughput peaks (35k inferred events per second) at a batch size of 1200. These findings exhibit the realm of applicability of TensorRT, i.e. for large scale ML inference workflows.

Zoom In Zoom Out Reset image size

Identifying feature importance has been a significant component of XAI methods and has been thoroughly studied in the context of classification models [85]. In standard feature selection tasks, a reasonable subset of the features that excels in some model performance metric is chosen. Although it is conceptually different from feature ranking in post-hoc model interpretation, the latter usually also relies on minimizing a model's performance loss [86]. One of the most useful model analysis tool of a binary classification is the ROC curve, and the corresponding AUC serves as a scalar metric for evaluating model performance. AUC-based feature ranking has been widely used in the AI literature [87–89]. We adapt those same principles for our model interpretation studies. One strategy for evaluating a feature's contribution in making predictions is to investigate the model's performance when that feature is masked, e.g. by replacing it with a population-wide average value or a zero value, whichever is contextually relevant to the model's relationship with the training dataset.

In order to identify the features that play the most important role in the IN model's decision-making process, we first train the model with its default settings, which we call the baseline model. During the training, for any event where certain input tracks or SVs are absent for a given jet, its corresponding entries are marked with zeros. Hence, we mask one feature at a time for all input tracks or SVs by replacing the corresponding entries by zero values. We obtain predictions from the trained model and evaluate the AUC score. The change observed in the AUC score when masking each of the features is presented in figure 5. It shows that while the model has been trained to take into account the entire feature space, there are 14 track features and 4 SVs features that, if removed one at a time, reduce the model's AUC score by less than 0.05%. Inspired by computer vision studies, we propose that the input features that cause the largest change in the AUC score may be regarded as the features that play the most important role in the model's decision-making process.

Zoom In Zoom Out Reset image size

The weak dependence of the model on many of its input features indicates that the model can learn the jet classification task from a subset of input features. To further investigate the cumulative impact of removing these unimportant features, we mask multiple features at the same time based on a few arbitrary thresholds for the change in AUC score compared to the baseline. The set of masked features includes every track and SV feature that causes a change in AUC score below that threshold when independently masked. To compare how individual predictions vary on average, we compute the model fidelity score [71, 90], defined as

$$\begin{equation} F\left(\mathcal{M}_1, \mathcal{M}_2\right) = 1 - \frac{1}{N}\sum_{i = 0}^{N-1} \left|\hat{y}_i^1 - \hat{y}_i^2\right|. \end{equation} \tag{ 1 }$$

Here, $\mathcal{M}_1$ and $\mathcal{M}_2$ are two different models and the corresponding classifier scores for the i-th data sample are respectively given by $\hat{y}_i^1$ and $\hat{y}_i^2$. The results are summarized in table 9. The model's performance, both in terms of AUC and fidelity scores, remains very close to the baseline even when masking up to 14 particle track and 4 SV features.

Table 9. The table shows the performance of a baseline model when multiple features are simultaneously masked based on AUC score drop threshold. $\Delta P$ ($\Delta S$) represents the number of particle (secondary vertices) features that have been masked. The fidelity score, see equation (1), is measured with respect to the baseline model.

Threshold (%)	$\Delta P$	$\Delta S$	AUC (%)	Fidelity (%)
0	0	0	99.02	100
0.001	8	2	99.02	99.69
0.005	9	2	99.05	99.46
0.01	11	3	98.99	98.85
0.05	14	4	98.81	97.12
1.00	25	8	80.49	70.83

While the AUC and fidelity scores allow determining which features play important roles in the IN's decision making process, we can inspect the importance of these features for individual tracks and vertices by the layerwise relevance propagation (LRP) technique [91, 92]. The LRP technique propagates the classification score predicted by the network backwards through the layers of the network and attributes a partial relevance score to each input. The original LRP method has been developed for simple MLP networks. Variants of this method have been explored to propagate relevance across convolutional neural networks [82, 93] and GNNs [84, 94].

Since some of the input features show a high degree of correlation with each other, we use the LRP-γ method described by Montavon et al [92], which is designed to skew the LRP score distributions to nodes with positive weights in the network and thus, avoiding propagation of large but mutually canceling relevance scores. In order to apply the LRP method for the IN model, we propagate scores across (i) the aggregation of internal representation of track features obtained from the aggregator network

$$\begin{eqnarray} O_{\left[D_O \times N_\mathrm{p}\right]} & \to &\bar{O}_{\left[D_O\right]}, \end{eqnarray} \tag{ 2 }$$

and (ii) the interaction matrices that send edge-level representations to the individual particle tracks

$$\begin{eqnarray} {E}_{\mathrm{pp}\left[D_E \times N_\mathrm{pp}\right]} & \to &\bar{E}_{\mathrm{pp}\left[D_E \times N_\mathrm{p}\right]} \end{eqnarray} \tag{ 3 }$$

$$\begin{eqnarray} {E}_{\mathrm{vp}\left[D_E \times N_\mathrm{vp}\right]} &\to &\bar{E}_{\mathrm{vp}\left[D_E \times N_\mathrm{p}\right]}. \end{eqnarray} \tag{ 4 }$$

The relevance scores for the output, $O_{[D_O \times N_\mathrm{p}]}$, of the f_O function can be obtained as

$$\begin{equation} r_{kn} = \bar{r}_k\left(\frac{o_{kn}}{\displaystyle\sum_m o_{km}}\right), \end{equation} \tag{ 5 }$$

where $\bar{r}$ represents the LRP scores for the summed internal representation. On the other hand, the relevance scores $\bar{R}_{kn}$ for the track level internal representations in $\bar{E}_{\mathrm{pp}[D_E \times N_\mathrm{p}]}$ can be propagated to edge level representations, $E_{\mathrm{pp}[D_E \times N_\mathrm{pp}]}$, using the relation

$$\begin{equation} r_{km} = e_{km}\sum_n \left(\frac{\bar{r}_{kn}}{\bar{e}_{kn}}\right) \left(R_{R}\right)_{nm}, \end{equation} \tag{ 6 }$$

where R_R is the receiver matrix for particle–particle interactions. A similar expression allows translating the relevance scores of the track level representation in $\bar{E}_{\mathrm{pv}[D_E \times N_\mathrm{p}]}$ to track-vertex edge representations in ${E}_{\mathrm{vp}[D_E \times N_\mathrm{vp}]}$ using the the receiver matrix R_K for vertex–particle interactions.

We show the average scores attributed to the different features for QCD and ${\mathrm{H}\to{{\mathrm{b}}\overline{\mathrm{b}}}}$ jets in figure 6. When compared with the change in AUC score by individual features in figure 5, the track and SV features with largest relevance scores are also the features that individually cause the largest drop in AUC score. We additionally observe that the track features are generally assigned larger relevance scores for QCD jets and SV features play a more important role in identifying the ${\mathrm{H}\to{{\mathrm{b}}\overline{\mathrm{b}}}}$ jets. This behavior is also justified from a physics standpoint, since the presence of high energy SVs is an important signature for jets from ${\mathrm{b}}$ quarks because of its relatively longer lifetime. This is also illustrated in figure 6, where the cumulative relevance score for each track and vertex is shown. The tracks and vertices are ordered according to their relative energy and our results show that the higher energy tracks and vertices are generally attributed with higher relevance scores for both jet classes. However, feature representing relative track energy, track_erel, itself does not carry notable relevance weight. On the other hand, the relevance attributed to sv_pt, which is strongly correlated with sv_erel, is very large.

Zoom In Zoom Out Reset image size

We also note that while the SV features sv_ptrel and sv_erel are assigned relatively low relevance scores, masking them independently leads to very large drops in the AUC score. This apparent discrepancy can be explained by the very high correlation between these variables, each of which also displays a very large correlation (correlation coefficient of 0.85) with sv_pt, as shown in figure 7. Because the LRP-γ method skews the relevance distribution between highly correlated features, it suppresses the LRP scores for those two variables while assigning a large relevance score to the variable sv_pt.

Zoom In Zoom Out Reset image size

We make an additional observation regarding the importance attributed to the feature called track_quality. This feature is a qualitative tag denoting the track reconstruction status, and has an almost identical, doubly peaked distribution for both jet categories. In figure 7, the peak at 0 represents absent tracks. With such an underlying distribution, this variable does not contribute to the classifier's ability to distinguish the jet categories. However, the large relevance score associated with it, along with the large drop in AUC score upon masking this feature, indicates that the classifier's class-predictive output for each class somehow receives a large contribution from the numerical embedding used to represent this feature and eventually gets canceled by the softmax operation.

We have found that the two previously mentioned SV features, along with track_quality, have no discernible impact on the IN model's ability to tell the jet categories apart by retraining the model without these variables. The model that was trained without these variables, along with the 11 (3) track (SV) features that report a change in AUC of less than 0.01%, converged with an AUC score of 99.00%. In the absence of these redundant features, we observed some differences in the relative distribution of the relevance scores. Thus, we are better able to understand which features play a more important role in the identification of ${\mathrm{H}\to{{\mathrm{b}}\overline{\mathrm{b}}}}$ or QCD jets, respectively. These physics-informed validation of model explanation pinpoints two major drawbacks of the existing XAI methods. First, explanations for models trained with highly correlated input features can be inconsistent across approaches and second, treating categorical and continuous variables on equal footing in XAI methods might lead to misleading attribution of feature importance.

3.4.1.1. Inspecting the activation layers

Here we aim to gain new insights on the IN model's decision-making process at the layer level. As the IN processes the input, it is passed through three different MLPs that approximate arbitrary nonlinear functions identified as $f_R^\mathrm{pp}, f_R^\mathrm{vp},$ and f_O . In order to explore the activity of each neuron and compare it with the activity of neurons in the same layer, we define relative neural activity (RNA) [83] as

$$\begin{equation} \mathrm{RNA}\left(j,k;\mathcal{S}\right) = \frac{\displaystyle\sum_{i = 1}^{N} a_{j,k} \left(s_i\right)}{\displaystyle\max_j\sum_{i = 1}^{N} a_{j,k}\left(s_i\right)} \end{equation} \tag{ 7 }$$

where $\mathcal{S} = \{s_1, s_2, \dots, s_N\}$ represents a set of samples over which the RNA score is evaluated. The quantity $a_{j,k} (s_i)$ is the activation of jth neuron in the kth layer when the input to the network is s_i . When summed over all the samples in the evaluation set $\mathcal{S}$, this represents the cumulative neural response of a node, which is normalized with respect to the largest cumulative neural response in the same layer to obtain the RNA score. Hence, in each layer, there will be at least one node with an RNA score of 1. Since the neurons are activated with ReLU activation in the IN model, the RNA score will be strictly between 0 and 1.

At a qualitative level, this study aims to identify which neurons are most actively engaged when the IN model produces an output. Since the MLPs in the IN model consist of only fully-connected layers, each layer takes all the activations from the previous layer as inputs. As all nodes within a given layer are subject to the same set of inputs, we can reliably estimate how strongly they perceive and transfer that information to the next layer by looking at their activation values. For the same reason, we normalize the cumulative activation of a node with respect to the largest aggregate in the same layer.

Figure 8 (left) shows the (NAP) diagram for the baseline model, showing the RNA scores for the different activation layers. The scores are separately evaluated for QCD and ${\mathrm{H}\to{{\mathrm{b}}\overline{\mathrm{b}}}}$. To simultaneously visualize these scores, we project the RNA scores of the former as negative values. The NAP diagram clearly shows that the network's activity level is quite sparse. In some layers, more than half of the nodes show RNA scores less than 0.2. This implies that while some nodes are playing very important roles in propagating the necessary information, other nodes do not participate as much. We additionally observe that right until the very last layer of the aggregator network f_O , the same nodes show the largest activity level for both jet categories. This is better illustrated in figure 8 (right), where the absolute difference in RNA scores for the two jet categories are mapped. For most nodes in every layer but the very last one, the difference in RNA scores is very close to zero. However, different nodes are activated in the last layer for the two jet categories, indicating an effective disentanglement of the jet category information in this layer. However, even in this layer, the activity level appears to be sparse—only a few nodes showing large activation for each category.

Zoom In Zoom Out Reset image size

The studies presented in sections 3.4.1 and 3.4.1.1 suggest that the baseline IN model can be made simpler by reducing both the number of input features it relies on and the number of trainable parameters. To explore this observation, we trained alternate variants of the IN models where the features sv_ptrel, sv_etrel, and track_quality were dropped along with additional 11 track and 3 SV features that reduce the AUC less than 0.01%, as shown in figure 5.

The details and performance metrics of these models are given in table 10. It should be noted that the ablated models presented here represent neither an exhaustive list of such choices nor any result of some rigorous optimization. These results demonstrate that a simpler IN model may be developed without compromising the quality of its performance. As can be seen from the results in table 10, both AUC score and fidelity of the alternate models are very close to that of the baseline model, though the number of trainable parameters is significantly lower.

Table 10. The performance of a baseline and ablated models. $\Delta P$ represents the number of particle track features that have been dropped and h is the number of nodes in the hidden layers. The fidelity score is measured with respect to a baseline model. Sparsity is measured by the fraction of activation nodes with an RNA score less than 0.2.

$\Delta P$, $\Delta S$	$h, D_E, D_O$	Parameters	AUC score (%)	Fidelity (%)	Sparsity
0, 0 (baseline)	60, 20, 24	25 554	99.02	100	0.56
12, 5	32, 16, 16	8498	98.87	96.93	0.52
	32, 8, 8	7178	98.84	96.79	0.48
	16, 8, 8	2842	98.62	96.12	0.40

Figure 9 shows the NAP diagrams for the model with 15 (5) dropped track (vertex) features with 32 nodes per hidden layer where the internal representation dimensions D_E and D_O are set to 16 and 8 for the left and right figures, respectively. Sparsity of the latter, as measured by the number of activation nodes with $\mathrm{RNA} \lt0.2$, is noticeably lower than the baseline model though the former has increased sparsity. With reduced size for the post interaction internal space representation, the alternate models do not completely disentangle the jet classes at the output stage of f_O .

Zoom In Zoom Out Reset image size