Learning Local Feature Masks with Variational Information Bottleneck
Learning Local Feature Masks with Variational Information Bottleneck
Lu Sun, Jun Sakuma
Proceedings of the Thirty-Fifth International Joint Conference on Artificial Intelligence
Main Track. Pages 4867-4875.
https://doi.org/10.24963/ijcai.2026/542
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.
Keywords:
Machine Learning: Explainable/Interpretable machine learning
Machine Learning: Feature extraction, selection and dimensionality reduction
Machine Learning: Learning sparse models
Machine Learning: Variational Inference
