Feature selection method improves neural network interpretability in multiclass tasks
A Feature Selection Layer Approach for Multi-Class Classification on Neural Networks
Machine Learning
Summary
Understanding which features in data influence predictions is important but hard for complex neural networks, especially when there are many categories to choose from. The authors improved a built-in method that highlights important features during training, making it better at showing which features matter for each specific category. Their approach works as well as existing methods overall and is better at revealing class-specific insights without extra steps after training. This helps users see how the model makes decisions in a clearer way.
What this means in practice
- •For business data teams: Identify which data features drive predictions for different categories to improve trust and decision-making in business AI models.
- •For medical informatics teams: Use feature relevance insights per diagnosis category to help explain AI-based clinical decision support models to healthcare providers.
Authors
João Pedro Silveira e Silva, Márcio Dorn, Bruno Iochins Grisci
Abstract
In the age of data-driven decision-making, deep learning models have become powerful tools for uncovering patterns and making predictions that influence many aspects of society. However, their broader adoption in scientific and business applications is often limited by a lack of interpretability. Feature weighting methods, which assess the relevance of individual features, are among the most widely studied approaches for improving interpretability. While post hoc techniques have advanced significantly, especially in modeling nuanced relationships between features and labels in multiclass problems, embedded methods typically fall short in this regard. Embedded techniques integrate feature relevance learning directly into the model training process, but they are often unable to capture label-specific relationships. To address the limitation, we propose multiple extensions to the Feature Selection Layer, an embedded feature weighting mechanism for neural networks. Through extensive experimentation, we identify the most effective extension and benchmark it against established feature weighting methods. Our results show that the proposed method matches existing techniques in identifying globally relevant features and outperforms them in capturing label-specific relevance. By leveraging the internal dynamics of the trained neural network, it reveals how individual features contribute to the prediction of specific labels, requiring no post-processing beyond the training process.