SHAP-Guided Regularization in Machine Learning Models

TL;DR

This paper introduces a SHAP-guided regularization method that improves model accuracy and interpretability by incorporating feature importance constraints during training, applicable to regression and classification tasks.

Contribution

It presents a novel regularization framework that uses SHAP-based feature importance constraints to enhance model performance and interpretability, especially in tree-based models.

Findings

Improves generalization performance on benchmark datasets.

Ensures robust and interpretable feature attributions.

Applicable to both regression and classification models.

Abstract

Feature attribution methods such as SHapley Additive exPlanations (SHAP) have become instrumental in understanding machine learning models, but their role in guiding model optimization remains underexplored. In this paper, we propose a SHAP-guided regularization framework that incorporates feature importance constraints into model training to enhance both predictive performance and interpretability. Our approach applies entropy-based penalties to encourage sparse, concentrated feature attributions while promoting stability across samples. The framework is applicable to both regression and classification tasks. Our first exploration started with investigating a tree-based model regularization using TreeSHAP. Through extensive experiments on benchmark regression and classification datasets, we demonstrate that our method improves generalization performance while ensuring robust and interpretable feature attributions. The proposed technique offers a novel, explainability-driven regularization approach, making machine learning models both more accurate and more reliable.