English

SLM: End-to-end Feature Selection via Sparse Learnable Masks

Machine Learning 2023-04-07 v1

Abstract

Feature selection has been widely used to alleviate compute requirements during training, elucidate model interpretability, and improve model generalizability. We propose SLM -- Sparse Learnable Masks -- a canonical approach for end-to-end feature selection that scales well with respect to both the feature dimension and the number of samples. At the heart of SLM lies a simple but effective learnable sparse mask, which learns which features to select, and gives rise to a novel objective that provably maximizes the mutual information (MI) between the selected features and the labels, which can be derived from a quadratic relaxation of mutual information from first principles. In addition, we derive a scaling mechanism that allows SLM to precisely control the number of features selected, through a novel use of sparsemax. This allows for more effective learning as demonstrated in ablation studies. Empirically, SLM achieves state-of-the-art results against a variety of competitive baselines on eight benchmark datasets, often by a significant margin, especially on those with real-world challenges such as class imbalance.

Keywords

Cite

@article{arxiv.2304.03202,
  title  = {SLM: End-to-end Feature Selection via Sparse Learnable Masks},
  author = {Yihe Dong and Sercan O. Arik},
  journal= {arXiv preprint arXiv:2304.03202},
  year   = {2023}
}
R2 v1 2026-06-28T09:53:13.915Z