Table of Contents
Fetching ...

Evidential Deep Learning for Uncertainty Quantification and Out-of-Distribution Detection in Jet Identification using Deep Neural Networks

Ayush Khot, Xiwei Wang, Avik Roy, Volodymyr Kindratenko, Mark S. Neubauer

TL;DR

This work investigates evidential deep learning (EDL) as an efficient method for uncertainty quantification (UQ) and anomaly detection in jet tagging at the LHC. By modeling outputs with a Dirichlet distribution parameterized by evidence $e_k$, it derives class beliefs and uncertainty, and optimizes a loss combining $\mathcal{L}_{MSE}$ and a KL regularizer $\mathcal{L}_{KL}$ with an annealing schedule $\lambda_t$. Across TopData, JetNet, and JetClass, EDL is shown to achieve competitive UQ with single-pass inference, often improving AUROC relative to Bayesian baselines while retaining high ID accuracy; however, OOD detection remains challenging, and the results depend on hyperparameters like $\zeta$ and training strategy (standard EDL vs. EDL-CT). Latent-space analyses reveal that high-uncertainty misclassified jets occupy overlapping regions in the PFIN latent space, linking uncertainty to physical jet features and interactions. The findings offer practical guidelines for applying EDL in jet tagging, with potential extensions to edge-triggering and anomaly-aware analyses, while highlighting areas for future improvement in OOD discrimination and ensemble integration.

Abstract

Current methods commonly used for uncertainty quantification (UQ) in deep learning (DL) models utilize Bayesian methods which are computationally expensive and time-consuming. In this paper, we provide a detailed study of UQ based on evidential deep learning (EDL) for deep neural network models designed to identify jets in high energy proton-proton collisions at the Large Hadron Collider and explore its utility in anomaly detection. EDL is a DL approach that treats learning as an evidence acquisition process designed to provide confidence (or epistemic uncertainty) about test data. Using publicly available datasets for jet classification benchmarking, we explore hyperparameter optimizations for EDL applied to the challenge of UQ for jet identification. We also investigate how the uncertainty is distributed for each jet class, how this method can be implemented for the detection of anomalies, how the uncertainty compares with Bayesian ensemble methods, and how the uncertainty maps onto latent spaces for the models. Our studies uncover some pitfalls of EDL applied to anomaly detection and a more effective way to quantify uncertainty from EDL as compared with the foundational EDL setup. These studies illustrate a methodological approach to interpreting EDL in jet classification models, providing new insights on how EDL quantifies uncertainty and detects out-of-distribution data which may lead to improved EDL methods for DL models applied to classification tasks.

Evidential Deep Learning for Uncertainty Quantification and Out-of-Distribution Detection in Jet Identification using Deep Neural Networks

TL;DR

This work investigates evidential deep learning (EDL) as an efficient method for uncertainty quantification (UQ) and anomaly detection in jet tagging at the LHC. By modeling outputs with a Dirichlet distribution parameterized by evidence , it derives class beliefs and uncertainty, and optimizes a loss combining and a KL regularizer with an annealing schedule . Across TopData, JetNet, and JetClass, EDL is shown to achieve competitive UQ with single-pass inference, often improving AUROC relative to Bayesian baselines while retaining high ID accuracy; however, OOD detection remains challenging, and the results depend on hyperparameters like and training strategy (standard EDL vs. EDL-CT). Latent-space analyses reveal that high-uncertainty misclassified jets occupy overlapping regions in the PFIN latent space, linking uncertainty to physical jet features and interactions. The findings offer practical guidelines for applying EDL in jet tagging, with potential extensions to edge-triggering and anomaly-aware analyses, while highlighting areas for future improvement in OOD discrimination and ensemble integration.

Abstract

Current methods commonly used for uncertainty quantification (UQ) in deep learning (DL) models utilize Bayesian methods which are computationally expensive and time-consuming. In this paper, we provide a detailed study of UQ based on evidential deep learning (EDL) for deep neural network models designed to identify jets in high energy proton-proton collisions at the Large Hadron Collider and explore its utility in anomaly detection. EDL is a DL approach that treats learning as an evidence acquisition process designed to provide confidence (or epistemic uncertainty) about test data. Using publicly available datasets for jet classification benchmarking, we explore hyperparameter optimizations for EDL applied to the challenge of UQ for jet identification. We also investigate how the uncertainty is distributed for each jet class, how this method can be implemented for the detection of anomalies, how the uncertainty compares with Bayesian ensemble methods, and how the uncertainty maps onto latent spaces for the models. Our studies uncover some pitfalls of EDL applied to anomaly detection and a more effective way to quantify uncertainty from EDL as compared with the foundational EDL setup. These studies illustrate a methodological approach to interpreting EDL in jet classification models, providing new insights on how EDL quantifies uncertainty and detects out-of-distribution data which may lead to improved EDL methods for DL models applied to classification tasks.
Paper Structure (16 sections, 13 equations, 14 figures, 3 tables)

This paper contains 16 sections, 13 equations, 14 figures, 3 tables.

Figures (14)

  • Figure 3: Distribution of \ref{['fig:jetclass-Nconst']} number of constituent particles, \ref{['fig:jetclass-jet-pt']} jet transverse momentum ($p_{T,J}$), and \ref{['fig:jetclass-jet-m']} jet mass ($m_J$) for QCD ($q/g$), Higgs ($H$), boson ($W,Z$) and top ($t$) jets.
  • Figure 4: Model architecture and data flow for the PFIN model. $N_b$ represents the batch size. The InTra block computes the pairwise particle interaction features given in Eqn. \ref{['eqn:phys-feats']}. The Cat/Sum block creates the augmented particle embeddings by either concatenating or summing the outputs of $\Phi_{I,2}$ with $\Phi$. The $\Phi, \Phi_I, \Phi_{I,2}$, and $F$ networks are implemented as fully connected MLPs with ReLU activation. From Figure 17 in Ref. Khot_2023.
  • Figure 6: For the TopData dataset, on each row, from left to right, \ref{['fig:topdata_baseline_0.1_unc_total']},\ref{['fig:topdata_baseline_0.7_unc_total']} uncertainty distribution, \ref{['fig:topdata_baseline_0.1_unc']},\ref{['fig:topdata_baseline_0.7_unc']} logarithmic uncertainty distributions separated by correct and incorrect jets, and \ref{['fig:topdata_baseline_0.1_up']},\ref{['fig:topdata_baseline_0.7_up']} 2D histogram of maximum probability versus uncertainty, for baseline EDL $\lambda_t (0.1)$ (top row) and $\lambda_t (0.7)$ (bottom row).
  • Figure 7: For the JetNet dataset, on each row, from left to right, \ref{['fig:jetnet_0.1_baseline_us']},\ref{['fig:jetnet_0.7_baseline_us']} uncertainty distribution, separated by correctly and incorrectly classified jets, \ref{['fig:jetnet_0.1_baseline_unc_class']},\ref{['fig:jetnet_0.7_baseline_unc_class']} uncertainty distribution for correctly classified jets, separated by initiating particle jet type, and \ref{['fig:jetnet_0.1_baseline_up']},\ref{['fig:jetnet_0.7_baseline_up']} 2D histogram of maximum probability versus uncertainty for baseline EDL $\lambda_t (0.1)$ (top row) and $\lambda_t (0.7)$ (bottom row).
  • Figure 8: Uncertainty Aware Confusion Matrix for, respectively, baseline JetNet EDL \ref{['fig:jetnet_0.1_baseline_lpu']}$\lambda_t (0.1)$ and \ref{['fig:jetnet_0.7_baseline_lpu']}$\lambda_t (0.7)$.
  • ...and 9 more figures