Learning Local Feature Masks with Variational Information Bottleneck

1 Introduction

Achieving both strong predictive performance and interpretability is a central goal in modern machine learning, particularly in high-stakes domains. Identifying informative features improves generalization while enabling practitioners to understand how inputs influence predictions [Molnar, 2020; Sokar et al., 2022; Gorishniy et al., 2021]. Instance-wise feature selection (IWFS) [Yoon et al., 2018; Chen et al., 2018; Jiang et al., 2024; Oosterhuis et al., 2025] addresses this challenge by selecting features adaptively for each instance, yielding fine-grained, context-aware explanations. IWFS is especially effective on heterogeneous datasets [Ucar et al., 2021; Si et al., 2024], where feature relevance varies across samples. For example, in electronic health records, the importance of lab tests or medications depends on patient age and medical history [Choi et al., 2016; Ma et al., 2018], while in e-commerce, fraud, and anomaly detection, feature informativeness depends on user behavior or transaction patterns [Yang et al., 2022; Liu et al., 2023].

∗Corresponding author.

Most IWFS methods adopt a data-dependent selectorpredictor architecture [Chen et al., 2018; Yoon et al., 2018], where instance-specific masks guide prediction using only locally relevant features. This design enables individualized, interpretable explanations, and recent studies [Yang et al., 2022; Oosterhuis et al., 2025; Jiang et al., 2024] show that such architectures can improve both interpretability and accuracy, challenging the presumed trade-off between the two. By adapting feature relevance per instance, IWFS offers a principled framework for transparent and reliable modeling.

Despite their empirical success, existing IWFS methods exhibit fundamental limitations. First, co-adaptation between the selector and predictor can undermine interpretability: the selector may encode label information directly into the instance-wise mask, allowing the predictor to decode the mask rather than infer the label from the selected features. This leakage under joint end-to-end training induces shortcut learning [Geirhos et al., 2020], enabling high accuracy with semantically uninformative features (e.g., 96% MNIST accuracy from a single pixel [Chen et al., 2018; Jethani et al., 2021]), and resulting in unfaithful explanations that fail to generalize to independently trained predictors. Second, reliance on a single selector limits the ability to capture heterogeneous sparsity patterns, reducing diversity in local feature selection. Third, most methods are driven primarily by architectural design within the selector-predictor framework, lacking an information-theoretic perspective and a unified formulation that could strengthen both performance and interpretability. Recent advances partially address these issues: REAL-X [Jethani et al., 2021] decouples training but sacrifices end-to-end explainability; ProtoGate [Jiang et al., 2024] enforces global and local relevance but is constrained by a non-parametric predictor; and SUWR [Oosterhuis et al., 2025] reduces shortcuts via sequential selection but incurs high cost and error propagation. These limitations motivate the need for a more principled and unified approach.

To address IWFS challenges, we propose VIBMask, which formulates instance-specific mask learning within the Variational Information Bottleneck (VIB) framework (Fig. 1). A

core difficulty in joint training is co-adaptation: the selection mask can become a shortcut communication channel where the predictor decodes labels from the mask structure rather than the feature values, undermining interpretability. VIBMask mitigates this by penalizing mutual information between unselected features and the label. This enforces a necessary condition for faithfulness that masked features retain no predictive information, thereby collapsing mask-based shortcut solutions. Furthermore, VIBMask employs an ensemble of selectors to capture heterogeneous patterns and a unified VIB formulation for discrete masks, ensuring stable training via continuous reparameterization. Extensive experiments show consistent improvements in predictive performance and explanation quality over state-of-the-art IWFS methods.

Contributions. (1) We propose VIBMask, a unified IWFS framework grounded in the variational information bottleneck, which mitigates co-adaptation through information-theoretic regularization and captures diverse local features via a multiselector mechanism. (2) We derive a novel variational lower bound for discrete masks, enabling efficient and stable optimization through reparameterization. (3) We achieve state-ofthe-art performance with more informative feature selection on both synthetic and real datasets.

2 RelatedWork

Instance-Wise Feature Selection. IWFS aims to identify instance-specific relevant features, improving interpretability and generalization through context-aware explanations. Early models such as TabNet [Arik and Pfister, 2021] integrated selection and prediction end-to-end, while recent methods adopt a selector-predictor framework [Chen et al., 2018; Yang et al., 2022; Oosterhuis et al., 2025], where a selector produces instance-dependent masks and a predictor operates on selected local features. Learning differentiable sparse masks remains challenging and is addressed via Gumbelsoftmax [Chen et al., 2018; Masoomi et al., 2020], reinforcement learning [Yoon et al., 2018], stochastic gates [Yamada et al., 2020; Yang et al., 2022], influence-aware attention [Liu et al., 2023], and Gaussian copulas [Peng et al., 2023]. However, many approaches treat feature selection independently, making them vulnerable to spurious correlations from degenerate selector-predictor interactions and limited expressivity under heterogeneous sparsity. In contrast, VIBMask reformulates IWFS within the VIB framework, mitigating spurious correlations by constraining co-adaptation and addressing heterogeneous sparsity through an ensemble of diverse selectors, resulting in robust and faithful instance-wise explanations.

Mitigating Co-Adaptation. Co-adaptation in IWFS arises when the selector and predictor jointly encode and decode label information through the selection mask, yielding shortcut solutions with high accuracy but unfaithful explanations. In this regime, the selected features are predictive only for the coupled predictor and fail to transfer to external classifiers, violating a core requirement of feature selection. Existing methods mitigate co-adaptation primarily through architectural or training decoupling: REAL-X [Jethani et al., 2021] employs an amortized explanation model, ProtoGate [Jiang et al., 2024] relies on a separately trained prototype-based predictor, and SUWR [Oosterhuis et al., 2025] limits shortcuts via sequential greedy selection. In contrast, VIBMask addresses co-adaptation at the objective level by enforcing a necessary condition for faithful selection: unselected features must not retain predictive information about the target. This is implemented by penalizing the mutual information between those features and the label. This principled constraint suppresses label leakage without modifying the predictor, promoting transferable and semantically meaningful explanations.

VIB-Based Feature Selection. The VIB principle [Barber and Agakov, 2004; Mohamed and Jimenez Rezende, 2015] promotes compact yet predictive representations by minimizing input-latent mutual information while preserving target relevance, making it well suited for feature selection. Existing VIB-inspired methods, including Concrete Autoencoder [Balın et al., 2019], Gated IB [Alesiani et al., 2023], and Drop-Bottleneck [Kim et al., 2021], mainly focus on global feature selection with a shared subset across samples. IWFS introduces additional challenges, including severe selectorpredictor co-adaptation and limited expressivity under perinstance sparsity constraints. InterpreTabNet [Si et al., 2024] extends TabNet [Arik and Pfister, 2021] with variational sparsity but neither explicitly mitigates co-adaptation nor avoids high-variance relaxations. To address IWFS, VIBMask introduces a variational objective that enforces feature faithfulness by ensuring unselected features carry no predictive information. By incorporating a decorrelation regularizer, it simultaneously prevents selector-predictor co-adaptation and enables stable, reparameterized optimization.

3 Methodology

3.1 An Information-Theoretic View of IWFS

The Information Bottleneck (IB) principle [Tishby et al., 2000] seeks a latent representation Z of input X ∈ Rd that is informative about the target Y ∈ {1, . . . , c} while discarding irrelevant details. It is formalized as maxZ I(Y ; Z) - βI(X; Z), where the trade-off parameter β balances relevance and compression, and mutual information I(·; ·) quantifies uncertainty reduction. The variational IB (VIB) [Alemi et al., 2017]

approximates this objective using variational bounds, reparameterization [Kingma et al., 2014], and Monte Carlo sampling [Kalos and Whitlock, 2009]. IB has been applied to improve robustness and generalization, including adversarial learning [Kuang et al., 2023] and out-of-distribution generalization [Ahuja et al., 2021].

We reinterpret IWFS through the IB lens to ground our method, VIBMask. Unlike standard VIB with continuous latents, IWFS selects discrete feature subsets per instance. We model this using a binary mask M ∈ {0, 1}d , where Mi = 1 indicates that feature Xi is selected. The masked input XM is drawn from a distribution S(M | X), and the objective balances task relevance and redundancy with the full input:

where ⊙ denotes the element-wise product. The mutual information is taken in expectation over masks sampled from S(M | X), making Eq. (1) a discrete, conditional variant of the IB, with XM representing a hard-selected feature subset. However, the discreteness of M hinders gradient-based optimization. To overcome this, we derive a variational lower bound for discrete latents and introduce a differentiable relaxation of the selection process, enabling efficient end-to-end training while preserving sparsity and interpretability.

3.2 Mitigating Co-Adaptation in IWFS

Recent IWFS methods [Jiang et al., 2024; Jethani et al., 2021] show that directly optimizing Eq. (1) can induce co-adaptation, a degenerate regime where the selector encodes label information into the mask M. The predictor "decodes" the label from the mask rather than the feature values, creating shortcut solutions that undermine interpretability. Crucially, this occurs because task-relevant information is shifted to the mask, leaving the discarded features to retain a residual predictive signal. Consequently, co-adaptation violates a core requirement of faithful selection: once features are masked, the unselected subset must have zero mutual information with the label.

To mitigate this effect, we explicitly penalize the mutual information between the unselected features and the label. Thus, VIBMask eliminates shortcut solutions in which the mask acts as a label-encoding channel and forces the selector to surface features that are intrinsically informative for prediction. Concretely, we augment Eq. (1) with a decorrelation penalty:

where M = 1 - M denotes the complement mask and XM the unselected features. The hyperparameter γ ≥ 0 controls the strength of the decorrelation loss. When informative features are omitted or label information is encoded in the mask, I(XM ; Y ) increases, penalizing such degenerate solutions and encouraging faithful, feature-based prediction.

Many IWFS methods, such as L2X [Chen et al., 2018], STG [Yamada et al., 2020], LSPIN [Yang et al., 2022], and c-STG [Sristi et al., 2024], can be interpreted as special cases of our discrete IB framework. Further details are provided in Sec. 4. However, unlike VIBMask, these methods neither target maximal compression of informative features nor explicitly mitigate co-adaptation, limiting their ability to provide compact and faithful explanations.

3.3 IntegratingMultiple Selectors

While VIBMask mitigates co-adaptation via variational masking, it remains limited in capturing diverse sparsity patterns. To address this, we introduce multi-selector integration, which enhances both expressiveness and the ability to model heterogeneous feature relevance.

Specifically, we extend Eq. (2) by incorporating K selectors {S1, . . . , SK}, each producing a local mask Mk ∼ Sk(Mk | X) for k ∈ [K]. The final mask is formed by averaging, M = 1K ∑k Mk, and applied to the input to obtain X ⊙M for prediction. This design offers three advantages: (i) diverse selectors capture complementary data aspects, (ii) averaging improves robustness when individual selectors underperform, and (iii) architectural flexibility allows selectors to vary in width, depth, or network type. Accordingly, combining multiple selectors with (2) yields the final VIBMask objective:

To encourage selector diversity and reduce redundancy, several strategies are possible: (i) heterogeneous selector architectures with different capacities; (ii) explicit regularization that penalizes mutual information or enforces orthogonality among {Mk}; and (iii) bootstrap aggregating (bagging) [Breiman, 1996] to train selectors on different resampled datasets. For simplicity and to avoid extra hyperparameters, we adopt bagging in practice. Each selector is trained on a distinct bootstrap sample containing approximately 63.2% unique data points, naturally inducing diversity. By leveraging this ensemble of input-conditioned selectors, VIBMask enables more effective, scalable, and flexible instance-wise feature masking.

3.4 Variational Lower Bound of VIBMask

The objective in (3) is intractable due to: (i) the inaccessibility of the posterior p(Y |XM), and (ii) the exponential search space (over 2d subsets) of binary masks M ∈ {0, 1}d . To address this, we derive a variational lower bound (this subsection) and apply a continuous relaxation via reparameterization in Sec. 4. Variational bounds for each term are given in Lemmas 1-3, with the full bound in Theorem 1. The proofs are provided in Sec. A2 of the technical appendix.

Lemma 1. For a variational approximation q(Y |XM) of the true posterior p(Y |XM), the mutual information I(Y ;XM) in (2) satisfies:

The last term is constant w.r.t. the mask M and can be ignored during optimization.

Lemma 2. Let πk,i = P(Mk,i = 1 | X) denote the conditional probability that the k-th selector (k ∈ [K]) chooses feature Xi given input X. Then, for every input X, we have

where H(Xi) is the entropy of the input feature Xi.

Lemma 3. For every input X, we have

where πk,i = P(Mk,i = 1 | X) is the conditional selection probability of feature Xi by selector Sk, and I(Xi; Y ) denotes the mutual information between feature Xi and the label Y .

The lower bound in Lemma 1 follows from the nonnegativity of the KL divergence. The upper bounds in Lemmas 2 and 3 are derived using the chain rule of mutual information under mild conditional independence assumptions shown in Sec. A2 of the technical appendix.

Theorem 1. By combining the bounds from Lemmas 1-3 and discarding terms that are independent of the selection probabilities, we obtain the following variational lower bound for VIBMask in (3):

whereM = 1K ∑k Mk and ωi = βH(Xi) - γI(Xi; Y ). For the k-th selector (k ∈ [K]), each feature i is independently retained with probability πk,i ∈ [0, 1], which depends on the input X.

Building on Theorem 1, our method is inspired by the IB principle [Balın et al., 2019; Kim et al., 2021; Alesiani et al., 2023] but differs in three key aspects: (i) it enables instance-wise feature selection for personalized, sample-level interpretability; (ii) it explicitly penalizes mutual information between unselected features and the target, mitigating co-adaptation; and (iii) it integrates multiple selectors to capture diverse local feature relevance.

4 Optimization

Theorem 1 replaces the intractable posterior p(Y | XM) with a variational approximation q(Y | XM) and relaxes the binary mask M ∈ {0, 1}d to a continuous vector π, enabling gradient-based optimization while preserving a probabilistic notion of feature relevance. We parameterize π with a selector gψ(X) and q with a predictor fθ(XM), and optimize the bound L(q, π) over ψ and θ.

Continuous Relaxation. Directly optimizing the objective in Theorem 1 is challenging due to the discreteness of the binary mask M, which leads to high-variance gradients [Mnih and Rezende, 2016]. Although estimators such as REINFORCE [Williams, 1992], soft attention [Bahdanau, 2014], and hard concrete [Louizos et al., 2018] alleviate this issue, they remain unstable for feature selection [Yamada et al., 2020]. We adopt a continuous relaxation inspired by stochastic gates [Yamada et al., 2020], approximating Mk with a relaxed Bernoulli vector ˜ Mk ∈ [0, 1]d :

˜ Mk = clip(µk + ϵ, 0, 1), µk = gψk (X), ϵ ∼N(0, σ2 I), where µk is generated by the k-th selector network gψk and clip(x, 0, 1) = max(0,min(1, x)). This enables the reparameterization trick [Figurnov et al., 2018], reducing gradient variance. The regularizer of L(q, π) in Theorem 1 becomes:

where Φ is the standard Gaussian CDF. In practice, we estimate of H(Xi) and I(Xi; Y ) through discretizations, as detailed in the supplementary materials. A detailed derivation and choice of σ are discussed in Sec. A4 and Fig. A1 of the technical appendix.

Optimization Algorithm. Based on the continuous relaxation and the lower bound in Theorem 1, the differentiable objective of VIBMask becomes:

where L is the loss function and fθ denotes the predictor network parameterized by θ. Gradients w.r.t. θ are computed directly, while those w.r.t. ψk are estimated via Monte Carlo [Kalos and Whitlock, 2009]:

where N denotes the number of instances. Therefore, the entire VIBMask model is trained end-to-end using gradient descent with the continuous reparameterization trick. A detailed algorithmic procedure and complexity analysis are provided in Sec. A5 of the technical appendix.

Inference. At inference, we remove stochasticity by setting ϵ = 0, which yields a deterministic mask ˜ Mk = clip(µk, 0, 1). Most elements of ˜ Mk converge to binary values. We avoid rounding intermediate values-often arising under weak training signals (e.g., overlapping classes)-since rounding may distort feature magnitudes and degrade performance. For interpretability, we threshold each Mk at 0.5 and apply majority voting across the K local masks {Mk} to identify the selected instance-wise features.

Remarks. Notably, our VIBMask formulation in (4) generalizes several prior IWFS methods as special cases when restricted to a single selector (K = 1). Specifically, L2X [Chen et al., 2018] corresponds to γ, β = 0; STG [Yamada et al., 2020] to constant γ, λ = 0 with π independent of X; and both LSPIN [Yang et al., 2022] and c-STG [Sristi et al., 2024] to constant γ, λ = 0. This highlights that VIBMask offers strong generalization capability and flexibility, making it applicable across a wide range of IWFS settings. A detailed comparison with existing IWFS approaches is presented in Table A1 in the technical appendix.

5 Experiments

We evaluate VIBMask to address three key questions: Q1 Interpretability-Can it identify informative features, mitigate co-adaptation, and enhance prediction? Q2 Performance-Does it achieve efficient, competitive results in prediction and feature selection vs. SOTA IWFS methods? Q3 Ablation-What is the contribution of each component (coadaptation avoidance, multi-selector integration) to overall performance? To this end, we first assess VIBMask's interpretability on six synthetic datasets (Sec. 5.1), then compare it with sixteen methods on eight real-world and two MNIST datasets (Sec. 5.2). Finally, we analyze the impact of each component of VIBMask (Sec. 5.3).

Datasets. We use six benchmark synthetic datasets [Chen et al., 2018; Yoon et al., 2018; Yang et al., 2022; Jiang et al., 2024] with 11-dimensional Gaussian features (X ∼ N(0, I)) and Bernoulli labels generated by varying logit(X). Each has 5,000 training and 5,000 testing samples. For real-world evaluation, we use ten popular datasets [Yang et al., 2022; Jiang et al., 2024]: six from scikit-feature (Relathe, PCMAC, Colon, Tox-171, Lung, Prostate), two from TabZilla (Cnae9, Mfeat-Fourier), and MNIST [Deng, 2012] plus FashionMNIST [Xiao et al., 2017]. MNIST and Fashion-MNIST each contain 70,000 grayscale 28 × 28 images in 10 classes; we use 6,000 training and 1,000 testing samples.

Comparison methods. We compare VIBMask 1 with sixteen methods: four baselines (Ridge regression [Cox, 1958], KNN [Fix, 1985], SVM [Cortes, 1995], MLP [Gorishniy et al., 2021]), four global Feature Selection (FS) methods (Lasso [Tibshirani, 1996], Random Forest [Breiman, 2001], XGBoost [Chen and Guestrin, 2016], STG [Yamada et al., 2020]), and eight local (instance-wise) FS methods (INVASE [Yoon et al., 2018], L2X [Chen et al., 2018], LSPIN, LLSPIN [Yang et al., 2022], ProtoGate [Jiang et al., 2024], SUWR [Oosterhuis et al., 2025], InterpreTabNet [Si et al., 2024], REAL-X [Jethani et al., 2021]).

Configuration. We perform 5-fold cross-validation and report the average results over 5 independent runs. Within each training fold, 10% of the samples are held out for validation. Hyperparameters are selected based on aggregated validation performance across all folds. Additional details on datasets, experimental setup, and complementary results are provided in Sec. A6 of the technical appendix.

5.1 Q1: Experiments on Synthetic Datasets

We evaluate interpretability on six synthetic datasets (Syn1Syn6) with known ground-truth informative features, which enable a precise assessment of faithfulness-whether the model consistently selects the features that are causally responsible for the label. Each dataset has 11 standard-Gaussian features and a Bernoulli label whose logit is a temperature-scaled nonlinear function of a designated feature subset; Syn4-Syn6 are conditional, with X11 as a control-flow switch. Full setup and per-dataset best configurations are in Secs. A6.1.1 and A6.3.6 of the technical appendix. Table 1 reports TPR↑, FDR↓, and Accuracy↑; for conditional datasets we additionally report the Control-Flow Selection Rate (CFSR↑) [Oosterhuis et al., 2025], the fraction of test samples whose mask includes X11. A low CFSR with high TPR is the diagnostic signature of co-adaptation.

1 Code and technical appendix: https://github.com/futuresun912/ VIBMask.

Across all six datasets, VIBMask attains the highest accuracy and is at least tied for the lowest FDR. On the unconditional datasets (Syn1-Syn3) it achieves perfect feature recovery (TPR = 100.0, FDR = 0.0), with the largest accuracy gains on the harder distributions: +5.0 pp on Syn3 (0.950 vs. SUWR 0.903) and Syn6 (0.951 vs. SUWR 0.896), and +1.1 and +2.7 pp on Syn4 and Syn5. On the conditional datasets Syn4 and Syn6, VIBMask selects the control-flow feature on every (Syn4) or nearly every (Syn6, 99.9%) instance, confirming that the decorrelation objective surfaces the gating variable rather than routing label information through the mask. The remaining trade-offs are local: on Syn5 the CFSR and TPR trail REAL-X/SUWR by ∼8 pp while accuracy still leads, and on Syn6 a higher FDR (15.0) accompanies the largest accuracy jump in the table. Importantly, all top-3 CFSR slots on every conditional dataset are occupied by methods with explicit co-adaptation mitigation (REAL-X, ProtoGate, SUWR, VIBMask), validating that closing the mask shortcut is the operative mechanism behind faithful conditional-dependency capture across all six benchmarks.

5.2 Q2: Experiments on Real-World Datasets

Classification Accuracy. Table 2 reports balanced accuracy on the eight real-world datasets. VIBMask ranks first on Lung (95.34%, +1.90 pp over the strongest IWFS baseline), Relathe (89.54%, +4.87 pp over SUWR), and PCMAC (89.43%, +1.54 pp over ProtoGate), second on Cnae-9 (96.83%), and third on Colon and Prostate; on TOX-171 it lands within 0.5 pp of the best IWFS baseline. The only column with a meaningful gap is Mfeat-Fourier (87.73%), where the strongest method, SUWR, retains a 4.8 pp lead. Aggregated across the eight datasets, VIBMask attains the highest mean balanced accuracy (90.41% vs. 84.32% for the strongest published IWFS), confirming the benefit of unifying co-adaptation avoidance and multi-selector integration within a single informationbottleneck objective.

Two patterns are worth highlighting. First, VIBMask dominates on the two text-classification datasets (Relathe +8.93 pp, PCMAC +4.32 pp over the previous IWFS best), where the relevant features differ markedly across documents and an instance-adaptive mask is essential. Second, on the HDLSS biomedical datasets (Colon, Lung, Prostate, TOX-171) the spread among the top IWFS methods is small (≤ 2 pp), reflecting both the high variance of small test folds and the fact that several methods already operate near the Bayes-optimal regime. In this regime VIBMask remains competitive while still delivering compact and instance-adaptive masks. Competing IWFS approaches share related goals but lack a unified information-theoretic design: ProtoGate extends global to local masking without explicit information-bottleneck regularization; SUWR and REAL-X encourage diversity or conditional importance but do not directly mitigate selectorpredictor co-adaptation; L2X and LSPIN can be viewed as restricted variants of VIBMask's objective.

Feature Selection Sparsity and Computational Efficiency.

We evaluate selector-predictor IWFS methods in terms of sparsity and inference latency, reporting results averaged over eight datasets (Fig. 2). INVASE and L2X select relatively large feature subsets (75.43% and 51.24%, respectively) yet achieve lower predictive performance (81.73% and 62.58%). In contrast, VIBMask attains the highest mean accuracy (90.41%) while selecting only 38.67% of features on average, demonstrating its ability to identify compact and informative subsets that improve both interpretability and reliability. The sparsityaccuracy frontier in Fig. 2 (left) shows that VIBMask sits at the top of the accuracy axis while keeping a mask fraction comparable to LSPIN and clearly below those of LLSPIN, L2X, and INVASE. From a computational perspective, L2X and SUWR exhibit the highest inference latency (2.24 ms and 2.54 ms), with L2X also yielding the lowest accuracy. ProtoGate and LLSPIN achieve a favorable trade-off, maintaining accuracy above 84% with inference times below 0.7 ms. Although VIBMask employs multiple selectors and thus has more trainable parameters than ProtoGate, LLSPIN, and L2X, its parameter count remains comparable to INVASE. Importantly, VIBMask achieves state-of-the-art accuracy with the lowest inference time among the compared IWFS methods (0.51 ms), about 20% faster than ProtoGate, underscoring its practical scalability for downstream deployment.

Performance on MNIST. We evaluate VIBMask on image classification using MNIST [Deng, 2012] and FashionMNIST [Xiao et al., 2017]. Following [Oosterhuis et al., 2025], methods select 3 × 3 patches on Fashion-MNIST to enhance interpretability. Compared to recent IWFS methods (excluding L2X and INVASE due to known unfaithfulness [Jethani et al., 2021]), including w/o FS, LSPIN, LLSPIN, ProtoGate, and SUWR, VIBMask consistently ranks second in accuracy, excelling especially when selecting very few pixels or patches (Fig. 3). While w/o FS (using all features) achieves the highest accuracy, VIBMask better balances interpretability and performance. SUWR performs comparably, ProtoGate is competitive but slightly behind, and LSPIN/LLSPIN lag with less effective selection. These results demonstrate VIBMask's superior ability to select robust, informative features while maintaining strong predictive accuracy.

5.3 Q3: Ablation Study

We assess two key design choices in VIBMask: the number of selectors K and the co-adaptation regularization strength γ. As shown on the left of Fig. 4, the cross-dataset mean accuracy peaks at K = 3 (87.20%), exceeding both K=1 (86.74%) and larger ensembles (K=4: 84.97%, K=5: 85.52%); the standard deviation also grows sharply for K ≥ 4, indicating that overly large ensembles enlarge the optimization burden. Per-dataset optima vary with the underlying sparsity, but the cross-dataset average consistently favors K = 3, which we adopt as a balance between expressivity and efficiency.

The right panel of Fig. 4 shows the γ-ablation on the Colon dataset, a small-sample setting where co-adaptation is hard to diagnose by accuracy alone. We report a min-maxnormalized ˜ I(XM ; Y ) ∈ [0, 1] so that the value at γ = 0 sits at the top of the scale and full leakage suppression at 0. Activating the decorrelation term progressively drives ˜ I(XM ; Y ) down, from ≈ 0.95 at γ = 0 to 0.00 at γ = 102 ; the rebound at γ = 103 reflects over-regularization, where the mask becomes nearly all-selecting and the complement set collapses. Meanwhile, the accuracy curve actually rises above the unregularized baseline (79.40 ± 9.60%) for moderate regularization, peaking at 82.35% for γ = 102 . Suppressing the mask-encoded shortcut thus not only restores faithfulness but also marginally improves predictive accuracy, by forcing the predictor to rely on the genuine feature signal rather than on label leakage [Jiang et al., 2024; Jethani et al., 2021]. Together with the per-dataset predictorreplacement diagnostic in the appendix, these results demonstrate that both multi-selector integration and explicit coadaptation mitigation are critical for robust and interpretable feature selection, validating the two core design choices of our information-bottleneck framework.

6 Conclusion

We introduce VIBMask, a novel IWFS method grounded in a unified framework that generalizes several existing IWFS approaches. Central to VIBMask is a newly derived variational bound that is broadly applicable to the discrete VIB problem, enabling principled and efficient feature selection. VIBMask addresses co-adaptation by minimizing mutual information between unselected features and the target, thus retaining only the most relevant features. Its design also leverages multiple selectors to capture diverse sparse patterns while fostering robustness and diversity. Using continuous reparameterization, VIBMask supports efficient end-to-end optimization and consistently outperforms existing IWFS methods across diverse synthetic and real-world datasets in terms of both accuracy and feature precision.

Despite these strengths, VIBMask has several limitations. It focuses on instance-wise mask selection and may overlook global feature patterns in datasets with an underlying sparse global structure. The method also introduces multiple hyperparameters requiring careful tuning, which may limit scalability to very high-dimensional inputs. In addition, VIBMask currently lacks explicit mechanisms to model complex instance-feature interactions such as context-dependent relevance. Future work will address these limitations through global feature awareness, adaptive hyperparameter tuning, and interaction-aware components (e.g., attention mechanisms) to further enhance the expressivity and applicability of VIBMask across diverse application domains.

Acknowledgments

We would like to thank the anonymous reviewers for their valuable comments and constructive feedback, which helped improve the quality of this paper. This work was supported in part by JST under Grant Nos. JPMJNX25C2, JPMJKP24C3, and JPMJCR21D3, and by JSPS KAKENHI under Grant Nos. 23H00483, 120251002, and 26K21320.