Extrapolation and explainability in hadron mass predictions with machine learning

Abstract

In this work, we study the predictive performance of four machine learning (ML) models for hadron masses, with particular attention to extrapolation beyond the conventional hadrons represented in the training data. Calculations are performed using a multilayer perceptron (MLP), eXtreme Gradient Boosting (XGBoost), Light Gradient Boosting Machine (LightGBM), and Gaussian process regression (GPR). ML models are trained on conventional hadrons, using features that describe valence quark content and quantum numbers. Then, we evaluate their performance on test sets and on a separate extrapolation set of 39 states, including exotic candidates and heavy baryons, none of which appear in the training or test sets. We find that the MLP achieves the lowest extrapolation Root Mean Square Error (RMSE) on the $30$ states with measured masses at $294.2 \pm 57.9$~MeV, roughly one-third that of the other models. We also show that the MLP reproduces the overall mass scale of nine unobserved doubly heavy baryons predicted by theoretical calculations, although it underestimates the doubly bottom masses. We use SHapley Additive exPlanations (SHAP) analysis and show that heavy quark content is the most significant factor in the MLP's mass predictions, with the remaining features contributing differently across hadron types. Our results quantify the extrapolation accuracy of ML models for hadron masses and identify the features that drive the MLP predictions.