A new audit for jet taggers: KAN surrogates map Particle Transformer attention back to energy correlators
Synopsis
Jet taggers in high-energy physics are judged mainly by accuracy, yet it has been hard to say whether their decisions track the radiation structure inside a jet. Researchers fit a Kolmogorov-Arnold Network (KAN) locally to a trained model's perturbation response, then read two-body and three-body interaction scores from the surrogate's mixed partial derivatives and rank-correlate them with the dominant contributions to two-point and three-point energy correlation functions. On JetClass, the Particle Transformer shows clear positive two-body alignment while a standard multi-head attention baseline does not; three-body alignment is weaker for both, with ParT still stronger and its error bars excluding null correlation.
Interpretation
A jet tagger's decision can now be checked pair by pair and triplet by triplet against the energy-correlation structure inside the jet, rather than only through overall accuracy. Prior interpretability work in this area is limited, and common component-attribution methods assume additive linear effects between components, so they capture independent contributions but cannot tell one irreducible three-body fluctuation from three unrelated pairwise effects. For jet taggers trained on the JetClass dataset for multiclassification, the top 256 particle pairs or triplets by contribution to the energy correlators are compared with the surrogate's second- and third-order gradient tensors using Spearman rank correlation; in the two-prong case the Particle Transformer shows a clear positive correlation while the standard multi-head attention baseline does not exhibit similar alignment.
The approach learns a local surrogate near a trained model's decision boundary and then reads multi-particle interactions off the surrogate's analytic derivatives. A Kolmogorov-Arnold Network (KAN) places learnable spline functions on network edges, and each edge acts on a single input variable, so mixed partial derivatives over distinct particles involve only first derivatives of the input-layer splines. Local Gaussian perturbations are applied to input particle kinematic vectors, and a KAN with cubic b-spline activations is fit to map those perturbations to shifts in the target model's output logits; each jet's surrogate has shape two layers by the number of constituents, uses cubic splines on 5 grid intervals plus a linear residual branch on every edge, has at most 5832 parameters, is fit with Adam on 50000 perturbations, selected among three configurations on 5000 validation perturbations, and evaluated on 10000 held-out perturbations.
Perspective
The demonstrated scope is transformer-based jet taggers trained on the JetClass dataset, covering a standard multi-head attention baseline and the Particle Transformer, with comparisons against the top 256 particle pairs or triplets by contribution to the two-point and three-point energy correlators. The authors propose extending the method to larger datasets, a wider class of signal models with more exotic substructure, and interactions beyond third order.
For three-prong jets the rank correlation between two-point and three-point energy correlators is weaker for both the baseline and ParT; ParT is stronger but its error bars only exclude null correlation. The authors note that ParT receives no physics-informed inductive bias at third order, and whether structure and alignment persist beyond third order remains to be studied. They also report that adding attention blocks has a negative effect on MHA alignment while ParT stays consistent, and that the two-prong decay of the boson shows anomalous behavior in both two-point and three-point correlations.