[Paper Review] Feature Selection with Distance Correlation

Ranit Das, Gregor Kasieczka+1 more|arXiv (Cornell University)|Nov 30, 2022

Machine Learning in Materials Science4 citations

TL;DR

This paper introduces a novel forward feature selection method based on Distance Correlation (DisCo) to identify the most informative energy flow polynomials (EFPs) for jet tagging in high-energy physics. By selecting only 9 EFPs from over 7,000 candidates, the method achieves near state-of-the-art performance in boosted top and W tagging with a compact DNN using significantly fewer parameters, outperforming prior methods like ADO-based selection.

ABSTRACT

Choosing which properties of the data to use as input to multivariate decision algorithms -- a.k.a. feature selection -- is an important step in solving any problem with machine learning. While there is a clear trend towards training sophisticated deep networks on large numbers of relatively unprocessed inputs (so-called automated feature engineering), for many tasks in physics, sets of theoretically well-motivated and well-understood features already exist. Working with such features can bring many benefits, including greater interpretability, reduced training and run time, and enhanced stability and robustness. We develop a new feature selection method based on Distance Correlation (DisCo), and demonstrate its effectiveness on the tasks of boosted top- and $W$-tagging. Using our method to select features from a set of over 7,000 energy flow polynomials, we show that we can match the performance of much deeper architectures, by using only ten features and two orders-of-magnitude fewer model parameters.

Motivation & Objective

• To develop a more efficient and effective feature selection method for multivariate jet tagging in high-energy physics.
• To reduce the number of input features while preserving or improving classifier performance, especially for lightweight inference in trigger systems.
• To enable both ab initio training and explanation of black-box classifiers without requiring a pre-trained reference model.
• To improve upon existing forward feature selection techniques like ADO by using a more robust statistical dependence measure.

Proposed method

• Uses Distance Correlation (DisCo) as a scoring function to measure statistical dependence between candidate features and the target label (signal vs. background).
• Applies a forward selection strategy: iteratively selects the EFP with the highest DisCo score relative to the target, then re-scores remaining features conditioned on the selected set.
• Employs a non-parametric, affine-invariant measure of dependence that captures both linear and nonlinear relationships, unlike Pearson correlation.
• Trains a compact deep neural network (DNN) on only the top-ranked EFPs, enabling fast inference and reduced model complexity.
• Supports two operational modes: (1) ab initio training using truth labels, and (2) explanation of a pre-trained black-box classifier.
• Validates performance using standard jet tagging benchmarks: boosted top quark and W boson jets, with EFPs derived from simulated jet events.

Experimental results

Research questions

• RQ1 Can Distance Correlation-based feature selection outperform existing methods like ADO in identifying the most informative EFPs for jet tagging?
• RQ2 To what extent can a small set of EFPs selected via DisCo match the performance of deep learning models trained on full feature sets?
• RQ3 Does the DisCo-based method maintain high performance when trained directly from truth labels, without relying on a pre-trained reference classifier?
• RQ4 How does the selected feature set compare in terms of model efficiency, parameter count, and inference speed to state-of-the-art architectures?

Key findings

• The DisCo-based method selects only 9 EFPs from over 7,000 candidates to achieve near state-of-the-art performance in boosted top tagging, matching the rejection power of ParticleNet-Lite.
• The resulting compact DNN trained on the selected 9 EFPs uses two orders of magnitude fewer parameters than ParticleNet-Lite, enabling faster inference and reduced training cost.
• The method outperforms the ADO-based feature selection method from [30] in both top and W tagging tasks, achieving better performance with fewer features.
• The DisCo method achieves strong performance even when trained ab initio from truth labels, unlike the ADO method, which degrades significantly under such conditions.
• The selected EFPs are highly informative and capture complex, nonlinear correlations in jet substructure that are critical for discrimination.
• The approach demonstrates that a small, carefully selected set of theoretically motivated features can rival deep learning models trained on vast, unprocessed inputs.