[Paper Review] Explainable AI for ML jet taggers using expert variables and layerwise relevance propagation

Garvita Agarwal, L. Hay+6 more|arXiv (Cornell University)|Nov 26, 2020

Reservoir Engineering and Simulation Methods61 references18 citations

TL;DR

This paper introduces a novel Explainable AI framework that enhances deep neural network (DNN) jet taggers by integrating expert variables (XAUG) and applying layerwise relevance propagation (LRP). By concatenating XAUG variables to intermediate network layers, the method improves classification performance by up to 30–40% and enables interpretable feature ranking. Key results show that XAUG variables alone can capture classifier behavior in simple cases, while combined use with low-level features yields optimal discrimination, especially in realistic particle-level models with AUCs exceeding 0.99.

ABSTRACT

A framework is presented to extract and understand decision-making information from a deep neural network (DNN) classifier of jet substructure tagging techniques. The general method studied is to provide expert variables that augment inputs ("eXpert AUGmented" variables, or XAUG variables), then apply layerwise relevance propagation (LRP) to networks both with and without XAUG variables. The XAUG variables are concatenated with the intermediate layers after network-specific operations (such as convolution or recurrence), and used in the final layers of the network. The results of comparing networks with and without the addition of XAUG variables show that XAUG variables can be used to interpret classifier behavior, increase discrimination ability when combined with low-level features, and in some cases capture the behavior of the classifier completely. The LRP technique can be used to find relevant information the network is using, and when combined with the XAUG variables, can be used to rank features, allowing one to find a reduced set of features that capture part of the network performance. In the studies presented, adding XAUG variables to low-level DNNs increased the efficiency of classifiers by as much as 30-40\%. In addition to performance improvements, an approach to quantify numerical uncertainties in the training of these DNNs is presented.

Motivation & Objective

• To develop an explainable AI framework for deep learning jet taggers in high-energy physics (HEP), where model decisions are often opaque.
• To improve classification performance by integrating expert-derived variables (XAUG) with low-level DNN features.
• To enable interpretability of DNN decisions using layerwise relevance propagation (LRP) to rank feature importance.
• To quantify numerical uncertainties in DNN training through ensemble averaging of multiple training runs.
• To demonstrate that XAUG variables can capture or guide classifier behavior, reducing reliance on complex, low-level features.

Proposed method

• Augmenting DNN inputs with expert variables (XAUG) by concatenating them after convolutional or recurrent network layers.
• Applying layerwise relevance propagation (LRP) to decompose final prediction scores and attribute relevance to input features and XAUG variables.
• Using normalized relevance scores to rank and identify the most influential features, enabling dimensionality reduction.
• Training multiple DNN instances with different random seeds and averaging predictions to estimate numerical uncertainties.
• Employing 2D CNNs, 1D CNNs, and 1D RNNs on simulated jet data (Z→bb vs. QCD) to evaluate performance and interpretability.
• Using particle-level simulations with 20 particles per event and 17 XAUG variables (e.g., jet mass, grooming, lifetime, impact parameters) to test real-world applicability.

Experimental results

Research questions

• RQ1 Can XAUG variables alone replicate the decision-making of a complex DNN jet tagger in simplified models?
• RQ2 How does combining XAUG variables with low-level features affect classification performance in realistic jet substructure tasks?
• RQ3 To what extent can LRP and XAUG variables identify and rank the most relevant features driving DNN decisions?
• RQ4 Can the relevance scores from LRP be used to construct a minimal, high-performing feature set for jet tagging?
• RQ5 How can numerical uncertainties in DNN training be quantified and managed using ensemble training and LRP analysis?

Key findings

• Adding XAUG variables to low-level DNNs increased classifier efficiency by 30–40% in particle-level models, significantly improving AUC to 0.994–0.995.
• In the toy model, XAUG variables alone fully captured the DNN’s decision-making, demonstrating that the network’s behavior could be replicated by a simple multilayer perceptron using only expert variables.
• For realistic particle-level models, the combined use of XAUG variables and low-level features achieved the highest performance, with AUCs exceeding 0.99 for 2D CNNs and 0.975 for RNNs.
• The top 5 XAUG variables, ranked by LRP relevance, captured nearly all the discriminatory power of the full set, enabling effective feature reduction.
• LRP revealed that the image input (low-level features) had higher relevance than any single XAUG variable in the particle-level model, but the sum of the top XAUG variables exceeded the image’s relevance, indicating their collective importance.
• Multiple training runs showed variation in low-relevance feature dependence, suggesting that averaging predictions across runs provides a robust method for estimating numerical uncertainties in DNN training.