[Paper Review] Enhancing Feature Selection and Interpretability in AI Regression Tasks Through Feature Attribution

Alexander Hinterleitner, Thomas Bartz–Beielstein+4 more|arXiv (Cornell University)|Sep 25, 2024

Machine Learning and Data Classification4 citations

TL;DR

This paper proposes a novel, data-driven feature selection pipeline for regression tasks that leverages Integrated Gradients for feature attribution and k-means clustering to identify and retain only the most informative features, significantly improving deep neural network performance and interpretability in predicting blade vibrations in turbo machinery. The method enhances model accuracy and stability by filtering noisy inputs, with key features like oscillation mode and drivetrain temperature emerging as top predictors.

ABSTRACT

Research in Explainable Artificial Intelligence (XAI) is increasing, aiming to make deep learning models more transparent. Most XAI methods focus on justifying the decisions made by Artificial Intelligence (AI) systems in security-relevant applications. However, relatively little attention has been given to using these methods to improve the performance and robustness of deep learning algorithms. Additionally, much of the existing XAI work primarily addresses classification problems. In this study, we investigate the potential of feature attribution methods to filter out uninformative features in input data for regression problems, thereby improving the accuracy and stability of predictions. We introduce a feature selection pipeline that combines Integrated Gradients with k-means clustering to select an optimal set of variables from the initial data space. To validate the effectiveness of this approach, we apply it to a real-world industrial problem - blade vibration analysis in the development process of turbo machinery.

Motivation & Objective

• To address the lack of effective feature selection methods tailored for regression tasks in deep learning, especially in industrial applications.
• To enhance model performance and robustness by filtering out uninformative or noisy features using explainable AI (XAI) techniques.
• To develop a data-driven, model-agnostic feature selection pipeline that integrates feature attribution with clustering to eliminate arbitrary thresholding.
• To validate the approach on a real-world industrial regression problem—blade vibration prediction in turbo machinery—where interpretability and accuracy are critical.
• To demonstrate that feature attribution can serve not only for explanation but also as a tool for improving model generalization and stability.

Proposed method

• Employ Integrated Gradients (IG), a gradient-based, model-agnostic feature attribution method, to compute local importance scores for each input feature across the entire dataset.
• Aggregate IG attribution values across all samples to obtain global feature importance rankings, enabling a holistic view of feature relevance.
• Apply k-means clustering to the aggregated attribution values to group features by importance, with distinct clusters representing high-, medium-, and low-importance features.
• Use the lowest-importance cluster (identified as the violet cluster in visualizations) as the threshold for feature removal, eliminating uninformative inputs without relying on arbitrary cutoffs.
• Train and tune a deep neural network using surrogate-based hyperparameter optimization to ensure high prediction quality, which is essential for reliable attribution.
• Compare the proposed method against baseline feature selection techniques and KernelShap, using both synthetic transparent data and real industrial data from turbo machinery.

Experimental results

Research questions

• RQ1 Can feature attribution methods like Integrated Gradients be effectively repurposed for feature selection in regression tasks rather than just for model explanation?
• RQ2 How does combining feature attribution with clustering improve the robustness and performance of deep neural networks in high-dimensional regression problems?
• RQ3 To what extent does the proposed method outperform classical filter and wrapper methods in terms of predictive accuracy and error stability?
• RQ4 Which input features are most influential in predicting blade vibrations in turbo machinery, and do they align with domain knowledge?
• RQ5 Does the quality of the final model’s predictions depend on the quality of the feature attribution, and how does this affect the reliability of the selection process?

Key findings

• The proposed method significantly improves model performance and reduces prediction error standard deviation by filtering out noisy and uninformative features.
• The feature 'oscillation mode' (from the 'Oscillation Characteristics' group) was identified as the most important predictor, aligning with engineering intuition about natural frequencies.
• Drivetrain temperature measurements (indices 80 and 83) were found to be more influential than rotation speed (index 18), suggesting a delayed or indirect relationship with vibration amplitude.
• The mass flow (index 30), volume flow (index 29), and lubricating oil pressure (index 72) from thermodynamic variables also showed high importance, indicating their role in system dynamics.
• The strain gauge identifier (index 2) from the 'Experimental Setup' group was crucial, highlighting the importance of measurement configuration in prediction accuracy.
• The 'Quality Criteria' feature (index 12) was also found to be highly relevant, suggesting data quality metrics influence model predictions.