Learning Local Feature Masks with Variational Information Bottleneck

Abstract

Instance-wise feature selection (IWFS) identifies informative features for each instance, improving generalization by discarding irrelevant information and enhancing interpretability through personalized explanations. Most IWFS methods adopt a selector--predictor architecture, where a selector generates instance-specific masks to guide prediction. This often leads to co-adaptation, in which the selector encodes label information into the mask, resulting in spurious correlations and unfaithful explanations. Existing methods also struggle to capture diverse local patterns, which is critical for IWFS under heterogeneous sparsity. We propose VIBMask, a unified IWFS framework by the variational information bottleneck. VIBMask mitigates co-adaptation by penalizing mutual information between unselected features and the label, and improves expressivity via an ensemble of diverse selectors that capture heterogeneous sparse patterns. We further derive a novel variational lower bound for discrete masks, enabling efficient end-to-end training through reparameterization. Experiments on synthetic and real datasets show that VIBMask consistently outperforms state-of-the-art IWFS methods in both predictive accuracy and informative feature discovery.