ProtoGate: Prototype-based Neural Networks with Global-to-local Feature
  Selection for Tabular Biomedical Data
Tabular biomedical data poses challenges in machine learning because it is
often high-dimensional and typically low-sample-size (HDLSS). Previous research
has attempted to address these challenges via local feature selection, but
existing approaches often fail to achieve optimal performance due to their
limitation in identifying globally important features and their susceptibility
to the co-adaptation problem. In this paper, we propose ProtoGate, a
prototype-based neural model for feature selection on HDLSS data. ProtoGate
first selects instance-wise features via adaptively balancing global and local
feature selection. Furthermore, ProtoGate employs a non-parametric
prototype-based prediction mechanism to tackle the co-adaptation problem,
ensuring the feature selection results and predictions are consistent with
underlying data clusters. We conduct comprehensive experiments to evaluate the
performance and interpretability of ProtoGate on synthetic and real-world
datasets. The results show that ProtoGate generally outperforms
state-of-the-art methods in prediction accuracy by a clear margin while
providing high-fidelity feature selection and explainable predictions. Code is
available at https://github.com/SilenceX12138/ProtoGate.
Xiangjian Jiang, Andrei Margeloiu, Nikola Simidjievski, Mateja Jamnik
# ProtoGate: Prototype-based Neural Networks with Local Feature Selection for Tabular Biomedical Data

###### Abstract

Tabular biomedical data poses challenges in machine learning because it is often high-dimensional and typically low-sample-size. Previous research has attempted to address these challenges via feature selection approaches, which can lead to unstable performance on real-world data. This suggests that current methods lack appropriate inductive biases that capture patterns common to different samples. In this paper, we propose ProtoGate, a prototype-based neural model that introduces an inductive bias by attending to both homogeneity and heterogeneity across samples. ProtoGate selects features in a global-to-local manner and leverages them to produce explainable predictions via an interpretable prototype-based model. We conduct comprehensive experiments to evaluate the performance of ProtoGate on synthetic and real-world datasets. Our results show that exploiting the homogeneous and heterogeneous patterns in the data can improve prediction accuracy while prototypes imbue interpretability.

## 1 Introduction

In biomedical research, tabular data is frequently collected [1; 2; 3] for a wide range of applications such as detecting marker genes [4], identifying cancer sub-types [4], and performing survival analysis [5; 6]. Clinical trials, whilst collecting large amounts of high-dimensional data using modern high-throughput sequencing technologies, often consider a small number of patients due to practical reasons [7]. The resulting tabular datasets are thus often high-dimensional and typically low-sample-size (HDLSS). Moreover, given the inherent heterogeneity of biomedical data, important features often vary from sample to sample - even in the same dataset [5; 8]. Such scenarios have proven challenging for current machine learning approaches, including deep tabular models [9; 10; 11; 5; 12; 13].

Previous methods [14; 15; 16; 8; 5] have attempted to address such challenges by performing local feature selection: rather than selecting a general set of important features across all samples, local feature selection methods select specific subsets of features for each sample and these subsets may vary from sample to sample. However, existing methods have three limitations: (i) In many real-world tasks, even simple models - such as an MLP or Lasso - can outperform many existing methods [5]. One reason is the accuracy of current methods can be substantially lower for some classes than other classes, and we illustrate this in Figure 1. (ii) These methods commonly comprise a trainable feature selector to select features and a trainable predictor to make predictions with these features, which can be susceptible to the co-adaptation problem [17; 18; 19]. Because the two components are jointly trained, the predictor can fit the selected features to achieve high accuracy even when these features do not reflect the real data distribution [17]. Consequently, the prediction accuracy is
inconsistent with the quality of selected features. For instance, L2X [16] achieves 96% accuracy in digit classification on MNIST by using only one pixel as input [17]. (iii) Current methods [8; 17; 5; 16] are not explainable because they mainly use an MLP-based predictor. This lack of explainability is a major concern in high-stake applications such as medicine [20; 21; 22; 23; 24].

We hypothesise that existing local feature selection methods exhibit subpar performance on biomedical data for two reasons: (i) They lack appropriate inductive biases. These methods mainly make predictions using MLPs, although, in the biomedical domain, the clustering assumption (which states similar samples should belong to the same class [25]) has been shown effective [26; 27; 28; 29; 30]. Based on the clustering assumption, the prototype-based models can perform well on tabular data by classifying the new instances according to their similarities to the existing prototypes. For instance, a simple prototype-based model, such as \(k\)-means, can outperform complex neural networks with an accurate pre-trained feature selector [5]. (ii) Existing local feature selection models tend to overly emphasize heterogeneity, often neglecting that different samples might share some informative features. The high accuracy of global feature selection models on real-world datasets [13] suggests that informative features can indeed be shared across samples. We believe that effective local feature selection methods should be able to identify both homogeneous and heterogeneous feature patterns across samples, provided the data supports their existence.

We aim to address the challenges of suboptimal performance and opaqueness of local feature selection methods applied to tabular biomedical data. We propose ProtoGate, a novel method which performs local feature selection and makes accurate and explainable predictions in the HDLSS regime.

Firstly, ProtoGate uses a prototype-based predictor without learnable parameters - namely Differentiable K-Nearest Neighbors (DKNN) [31] - which enables explainable predictions. The prototype-based predictor confers ProtoGate two important features: (i) an inductive bias aligned with the clustering assumption in biomedical data; and (ii) consistent evaluations of the quality of selected features throughout the training process, eliminating the possibility of co-adaptation from joint training. Secondly, ProtoGate performs feature selection in a global-to-local manner with an \(\ell_{1}\)-regularised gating network. The global-to-local design helps ProtoGate consider the homogeneous and heterogeneous patterns across multiple samples.

Our contributions can be summarised as follows:

1. We propose ProtoGate, a novel method which addresses the challenge of high-dimensional and low-sample-size (HDLSS) biomedical data by achieving local feature selection and explainable predictions with a global-to-local feature selector and a prototype-based classifier.
2. We show that ProtoGate generally outperforms 12 benchmark methods on seven real-world biomedical datasets (Section 4.1) while selecting fewer features (Section 4.2), paving the path to more robust and interpretable local feature selection models.
3. We demonstrate that ProtoGate effectively handles the co-adaptation problem with a prototype-based predictor by comparing its performance against nine feature selection benchmark methods on three synthetic datasets (Section 4.4).

Figure 1: Illustration of the unstable performance of LSPIN [5] on the lung dataset. (a) The class-wise classification accuracy, with class distribution in the parentheses. (b) The mean number of selected features, displayed on a logarithmic scale. MLP and ProtoGate achieve stable performance, while LSPIN has a large variance in the accuracy and number of selected features across different classes.
Related Work

**Feature Selection Methods** Feature selection is a common technique for improving the accuracy and interpretability of machine learning models on HDLSS datasets. An extensive line of work selects features globally with Lasso-based regularisation [32; 33; 34; 35; 36] or specialised layers in neural networks [13; 37; 38; 39; 40]. However, the global feature selection ignores the heterogeneous nature of biomedical data, leading to insufficient interpretability [5; 8].

Prior studies attend to the heterogeneity between samples by designing local feature selection models that select instance-wise features for explaining a pre-trained predictor [17; 41; 42; 43; 44; 45; 46]. These methods are limited because the post-hoc analysis on feature importance does not improve the performance of pre-trained predictors.

Recent work proposes to select instance-wise features for making predictions [8; 5; 16; 15; 47]. L2X uses mutual information for instance-wise feature selection with Concrete distribution, but it requires specifying the number of selected features [16]. INVASE addresses such limitation by modelling each feature's mask/gate value with independent Bernoulli distributions [8]. However, both methods utilise computationally expensive gradient estimators: REINFORCE [48] or REBAR [49]. Similar to STG [40], LSPIN/LLSPIN re-formalises the mask/gate value with injected Gaussian noise and extends Localized Lasso [50] with a gating network that can select similar features for similar samples [5]. However, the poor performance of a vanilla KNN on real-world datasets (Table 1) demonstrates that the similarity in the initial high-dimensional feature space is inaccurate because a large proportion of features can be noise for the prediction. In contrast, ProtoGate measures the similarity across samples within an intrinsically interpretable DKNN predictor. The predictor takes the samples after feature selection as input, and thus the similarity is measured in a feature space with fewer dimensions than LSPIN/LLSPIN and Localized Lasso.

**Co-adaptation Problem** In feature selection, co-adaptation refers to the situation where the model encodes predictions into the feature selection, leading to high accuracy with features that do not reflect the real data distributions [17; 18; 19; 51]. REAL-X proves that co-adaptation can happen in models with jointly trained feature selectors and predictors [17], and addresses this problem by decoupling the training objectives of the feature selector and predictor. But it only provides post-hoc explanations of the feature importance for a predictor trained with all features. In ProtoGate, we propose to address the co-adaptation problem with DKNN, a prototype-based predictor without learnable parameters. The DKNN predictor can consistently evaluate the selected features throughout the training process, eliminating the possibility of co-adaptation from joint training.

## 3 Method

### Problem Setup

We consider the classification task on tabular biomedical data with \(\mathcal{Y}\) classes. Let \(X\coloneqq[\mathbf{x}^{(1)},\ldots,\mathbf{x}^{(N)}]^{\top}\in\mathbb{R}^{N\times D}\) be the data matrix consisting of \(N\) samples \(\mathbf{x}^{(i)}\in\mathbb{R}^{D}\) with \(D\) features, and let \(Y\coloneqq[y^{(1)},\ldots,y^{(N)}]\in\mathbb{R}^{N}\) be the corresponding labels. We denote \(x^{(i)}_{d}\) as the \(d\)-th feature of the \(i\)-th sample. To simplify the notation, we assume all samples in \(X\) are used for training.

A common local feature selection model contains two components: (i) an instance-wise feature selector \(S_{\mathbf{W}}:\mathbb{R}^{D}\rightarrow\{0,1\}^{D}\) that takes as input a sample \(\mathbf{x}^{(i)}\) and generates a mask \(\mathbf{s}^{(i)}\in[0,1]^{D}\) for its features, and (ii) a predictor model \(F_{\theta}:\mathbb{R}^{D}\rightarrow\mathcal{Y}\) which takes as input both the sample \(\mathbf{x}^{(i)}\) and the mask \(\mathbf{s}^{(i)}\) for prediction:

\[y^{(i)}=F_{\theta}\left(S_{\mathbf{W}}(\mathbf{x}^{(i)}),\mathbf{x}^{(i)}\right)=F_{ \theta}(\mathbf{x}^{(i)}\odot\mathbf{s}^{(i)}), \tag{1}\]

where \(\hat{y^{(i)}}\) is the predicted label and \(\odot\) is the element-wise multiplication. Here, we define the \(d\)-th feature is selected if and only if the mask value is positive (\(s^{(i)}_{d}>0\)).

### Rationale for Model Architecture

We propose ProtoGate as a method for selecting instance-wise features with inductive bias from the prototype-based model, as shown in Figure 2. Instead of predicting the local mask from all features,
ProtoGate selects instance-wise features in a global-to-local manner with an \(\ell_{1}\)-regularised gating network. This allows the feature selector to attend to both homogeneity and heterogeneity across samples. Additionally, ProtoGate leverages the selected features with a prototype-based predictor of DKNN. The DKNN predictor makes explainable predictions and encodes the clustering assumption into feature selection. Without learnable parameters, DKNN can further avoid the co-adaptation problem by providing consistent evaluations for the selected features while training the feature selector. The pseudocode for model training is summarised in Algorithm 1.

```
0: training samples \(X\in\mathbb{R}^{N\times D}\), ground truth labels \(Y\in\mathbb{R}^{N}\), global-to-local feature selector \(S_{\mathbf{W}}\), prototype-based classifier \(F\), sparsity hyper-parameters (\(\lambda_{g}\), \(\lambda_{l}\)), number of nearest neighbours \(K\), total training epochs \(E\), learning rate \(\alpha\) trained model \(S_{\mathbf{W}}\), prototype base \(\mathcal{B}\)\(\mathbf{W}\leftarrow\) GaussianInitialisation()\(\triangleright\) Initialise the weights of feature selector for\(e\gets 1\) to \(E\)do\(\mathcal{B}\leftarrow\{\}\)\(\triangleright\) Initialise the prototype base as an empty set for\(i\gets 1\) to \(N\)do\(\triangleright\) Construct the prototype base \(\mathbf{x}^{(i)}_{\text{masked}}\leftarrow\mathbf{x}^{(i)}\odot S_{\mathbf{W}}(\mathbf{x} ^{(i)})\)\(\triangleright\) Select instance-wise features for training samples \(\mathcal{B}\leftarrow\mathcal{B}\cup\{(\mathbf{x}^{(i)}_{\text{masked}},y^{(i)})\}\)\(\triangleright\) Add masked sample and its label to the prototype base endforfor\(i\leftarrow\) to \(N\)do\(\mathbf{x}^{(i)}_{\text{masked}}\leftarrow\mathbf{x}^{(i)}\odot S_{\mathbf{W}}(\mathbf{x} ^{(i)})\)\(\triangleright\) Select instance-wise features for training samples \(P^{(i)}_{\mathcal{B}}\leftarrow\) NeuralSort(\(\mathcal{B},\mathbf{x}^{(i)}_{\text{masked}}\))\(\triangleright\) Compute the permutation matrix for \(\mathbf{x}^{(i)}_{\text{masked}}\)\(y^{(i)}\gets F(\mathcal{B},\mathbf{x}^{(i)}_{\text{masked}},K)\)\(\triangleright\) Classify the query sample with \(K\) nearest prototypes endfor\(L=\frac{1}{N}\sum_{i=1}^{N}\left(\ell_{\text{pred}}(P^{(i)}_{\mathcal{B}},\mathbf{x}^{(i)},y^{(i)})+R( \mathbf{W}^{[1]},\mathbf{s}^{(i)},\lambda_{g},\lambda_{l})\right)\)\(\triangleright\) Compute the training loss \(\mathbf{W}\leftarrow\mathbf{W}-\alpha\nabla_{\mathbf{W}}L\)\(\triangleright\) Update the weights of feature selector endfor return \(S_{\mathbf{W}}\), \(\mathcal{B}\)
```

**Algorithm 1** Training Procedure of ProtoGate

### Global-to-local Feature Selection

The global-to-local feature selector \(S_{\mathbf{W}}:\mathbb{R}^{D}\rightarrow[0,1]^{D}\) (Figure 2 (A)) is a neural network that maps feature values \(\mathbf{x}^{(i)}\) into mask values \(\mathbf{s}^{(i)}\). The feature selector attends to the homogeneity

Figure 2: The architecture of ProtoGate. **(A)** Given a sample \(\mathbf{x}\in\mathbb{R}^{D}\), the global-to-local feature selector performs global feature selection in the first layer. The orange dashed lines denote sparsified weights in \(\mathbf{W}^{[1]}\) under \(\ell_{1}\) regularisation. The neural network then computes the instance-wise mask values \(\{s_{d}\}_{d=1}^{D}\in[0,1]^{D}\) with a thresholding function. **(B)** The mask is applied to the sample for local feature selection by element-wise multiplication. **(C)** The prototype-based predictor classifies \(\mathbf{x}\) by retrieving the \(K\) nearest neighbours to the masked sample in base \(\mathcal{B}\). The majority class of neighbours is used as the predicted label \(\hat{y}\).
between samples via applying \(\ell_{1}\)-regularisation on \(\mathbf{W}^{[1]}\), the weights of the first layer. Intuitively, the regularisation can lead to sparse weights in the first layer, which implicitly selects features globally for all samples. The output \(\boldsymbol{\mu}^{(i)}\) from the last layer is thresholded to obtain instance-wise mask values by

\[s_{d}^{(i)}=\max(0,\min(1,\mu_{d}^{(i)}+\epsilon_{d}^{(i)})) \tag{2}\]

where \(\epsilon_{d}^{(i)}\) is the injected noise sampled from Gaussion distribution \(\mathcal{N}(0,\sigma^{2})\). The standard deviation \(\sigma\) is fixed during training, and it is removed during the inference time for deterministic mask values. With the injected noise, \(\boldsymbol{s}^{(i)}\) can be re-formalised as random vectors with parameters \(\boldsymbol{\mu}^{(i)}\) predicted by a neural network. Therefore, the sparsity regularisation on mask values can be computed by

\[R(\mathbf{W}^{[1]},\boldsymbol{s}^{(i)},\lambda_{g},\lambda_{l})=\lambda_{g} ||\mathbf{W}^{[1]}||_{1}+\mathbb{E}\left[\lambda_{l}||\boldsymbol{s}^{(i)}|| _{0}\right]\!=\!\lambda_{g}||\mathbf{W}^{[1]}||_{1}+\lambda_{l}\!\sum_{d=1}^{D }\!\left(\frac{1}{2}-\frac{1}{2}\text{erf}(-\frac{\mu_{d}^{(i)}}{\sqrt{2} \sigma})\right) \tag{3}\]

where \((\lambda_{g},\lambda_{l})\) is a pair of hyper-parameters to balance the effects of global and local feature selection, and \(\text{erf}(\cdot)\) is the Gauss error function. The full derivations are available in Appendix B.

### Prototype-based Prediction

The prototype-based predictor \(F:\mathds{R}^{D}\rightarrow\mathcal{Y}\) is a DKNN model (Figure 2 (C)). The DKNN predictor first constructs a prototype base \(\mathcal{B}\) with training samples. After masking the training samples with the mask generated by \(S_{\mathbf{W}}(X)\), DKNN retains the masked samples and their labels as prototypes in the base \(\mathcal{B}\coloneqq\{(S_{\mathbf{W}}(\boldsymbol{x}^{(i)})\odot\boldsymbol{ x}^{(i)},y^{(i)})\}_{i=1}^{N}\). With the acquired prototypes, the predictor can classify a query sample \(\boldsymbol{x}_{\text{query}}\in X\) by retrieving the base \(\mathcal{B}\). The predictor sorts the prototypes by their similarities to the masked query sample with NeuralSort [31], a differentiable relaxed sorting operator. Note that ProtoGate computes the Euclidean distance between samples as the similarity evaluation metric. According to the sorting results, the predictor uses the majority class of the \(K\) closest prototypes as the predicted label \(\hat{y}_{\text{query}}\). Because the feature selector is learnable, the mask can change and thus the base \(\mathcal{B}\) is dynamic over the training time. After training, the prototype base \(\mathcal{B}\) is fixed and query samples are from unseen test data.

For each query sample, the loss of prototype-based classification is defined as:

\[\ell_{\text{pred}}(P_{\mathcal{B}}^{\text{query}},\boldsymbol{x}_{\text{query }},y_{\text{query}})=K-\frac{1}{K}\sum_{j=1}^{K}\sum_{i=1}^{N}\mathbbm{1}\left( y^{(i)}=y_{\text{query}}\right)P_{\mathcal{B}}^{\text{query}}[i,j] \tag{4}\]

where \(P_{\mathcal{B}}^{\text{query}}\in\mathds{R}^{N\times N}\) denotes the relaxed permutation matrix and \(\mathbbm{1}(\cdot)\) denotes the indicator function. In the permutation matrix, \(P_{\mathcal{B}}^{\text{query}}[i,j]\) denotes the possibility that the \(i\)-th prototype is the \(j\)-th closest to query sample \(\boldsymbol{x}_{\text{query}}\) under NeuralSort. Among the \(K\) nearest prototypes, Equation 4 estimates the number of prototypes that have different labels to \(\boldsymbol{x}_{\text{query}}\). DKNN encodes the clustering assumption into feature selection by encouraging samples of the same class to have similar representations. Additionally, DKNN measures the similarity between masked samples and mitigates the effects of the noisy features in the initial high-dimensional feature space.

### Training Loss

The training loss is comprised of the average classification loss in Equation 4 and the sparsity regularisation in Equation 3:

\[L=\frac{1}{N}\sum_{i=1}^{N}\left(\ell_{\text{pred}}(P_{\mathcal{B}}^{(i)}, \boldsymbol{x}^{(i)},y^{(i)})+R(\mathbf{W}^{[1]},\boldsymbol{s}^{(i)},\lambda _{g},\lambda_{l})\right) \tag{5}\]

Because the loss function is fully differentiable, the global-to-local feature selector and the prototype-based predictor can be trained in tandem. The whole model can be optimised with standard gradient-based approaches, such as stochastic gradient descent. We did not observe optimisation issues when training over 3,000 models (Appendix A.4).

## 4 Experiments

We now evaluate ProtoGate on both synthetic and real-world datasets to substantiate the model design choices. Firstly, we compare ProtoGate against 12 benchmark methods on real-world classification
tasks (Section 4.1 and Section 4.2). Secondly, we investigate the impact of the prototype-based predictor by replacing it with a linear or MLP-based prediction head (Section 4.3) and adjusting the number of nearest neighbours (Appendix D). Thirdly, we investigate the impact of the global-to-local feature selector by considering the interplay between global and local feature selection (Section 4.3). Finally, we analyse the co-adaptation problem by considering the performance misalignment between feature selection and classification on the synthetic datasets (Section 4.4). We also provide the comparison of training time in Appendix E.

**Real-world datasets.** Following [13], we utilise seven HDLSS tabular biomedical datasets. The datasets contain \(2000-5966\) features with \(62-197\) samples of \(2-4\) different classes. We are interested in datasets with much fewer samples than LSPIN [5], which uses \(\sim 1,500\) samples. Full descriptions of the real-world datasets are available in Appendix A.1.

**Experimental setup.** For each dataset, we perform 5-fold cross-validation on 5 different splits, summing up to 25 runs per model. We obtain the validation set by randomly selecting 10% of training data. For each benchmark model, the training loss is a weighted loss, and we perform a hyper-parameter search for model selection on the validation set. Full details about the reproducibility and hyper-parameter tuning are available in Appendix A.5.

**Evaluation metrics.** We report the results averaged over 25 runs on test sets. (i) For classification, we measure the performance by the mean \(\pm\) std test balanced accuracy. (ii) Note that the proportion of selected features varies across samples for local feature selection methods. Therefore, we measure the sparsity of feature selection by the mean \(\pm\) std proportion of selected features across samples. (iii) To distinguish between "similar number of selected features" and "similar selected features", we introduce a new metric: degree of local sparsity \(\mathcal{Q}\), which is computed by

\[\mathcal{Q}=\frac{1}{D\cdot N}\sum_{j=1}^{N}\text{card}\left(\bigcup_{i=1}^{N} \text{nonzero}(\mathbf{s}^{(i)})-\text{nonzero}(\mathbf{s}^{(j)})\right) \tag{6}\]

where \(\text{card}(\cdot)\) returns the cardinality of a set and \(\text{nonzero}(\cdot)\) returns the indices of non-zero elements in a vector. \(\mathcal{Q}\) measures the difference between the union set of selected features for all samples and the selected features for a specific sample. Intuitively, a non-zero \(\mathcal{Q}\) denotes selected features are different across samples, and thus the feature selection is local. For global feature selection, the degree of local sparsity is zero (\(\mathcal{Q}\equiv 0\)).

**ProtoGate implementation.** The global-to-local feature selector is flexible on the number of hidden layers, and we implement it as a three-layer feed-forward neural network. The numbers of neurons in the input and output layers are the same as the number of features of the input data and the number of neurons in the hidden layer is set to 100. The feature selector has batch normalisation and \(tanh\) activation for all layers. We train the models with a batch size of 64 and utilise an SGD optimizer with a weight decay of \(1e-4\). The number of nearest neighbours \(K\) is searched in \(\{1,2,3,4,5\}\). The global sparsity hyper-parameter \(\lambda_{g}\) is searched in \(\{1e-4,2e-4,3e-4,4e-4,6e-4\}\), and the local sparsity hyper-parameter \(\lambda_{l}\) is set as \(1e-3\).

**Benchmark methods.** We evaluate the classification accuracy of ProtoGate and compare it with several benchmark models, including global feature selection models (LightGBM [52], Random Forest (RF) [53], Lasso [32] and STG [40]) and local feature selection models (TabNet [15], L2X [16], INVASE [8], REAL-X [17] and LSPIN/LLSPIN [5]). Additionally, we also compare ProtoGate with some standard models, including KNN [54] and MLP.

### Classification Performance

Table 1 shows that ProtoGate consistently achieves better than or comparable balanced accuracy to the benchmark models. We compute the average rank across different datasets and ProtoGate ranks first, followed by Lasso. ProtoGate outperforms all other local feature selection models by a clear margin. We also find that the existing local feature selection methods cannot outperform even the simple linear Lasso or vanilla MLPs on HDLSS datasets. Note that REAL-X trains the MLP-based predictor with all features, and thus it achieves comparable performance as the MLP model.

The stable and competitive performance of ProtoGate shows the suitability of the clustering assumption in the biomedical field. Moreover, ProtoGate intrinsically provides explanations for the predictions by explicitly pointing out the \(K\) nearest prototypes, while other local feature selection methods can be unexplainable with MLP-based predictors. Poor performance of the vanilla KNN
model also demonstrates that a large proportion of features can be irrelevant to the predictions, and thus the similarity in the high-dimensional feature space can introduce noise to feature selection, which can be one reason for the failure of LSPIN and LLSPIN.

In most HDLSS cases, ProtoGate consistently outperforms both Lasso and MLP. The exceptions are the lung, prostate and toxicity datasets, where the ProtoGate accuracy is slightly lower. As mentioned in [13, 5], Lasso and MLP can outperform other feature selection models when they are well-regularised on some datasets, such as the toxicity dataset. Compared with well-regularised Lasso and MLP, the prototype-based predictor could have limited expressivity, resulting in the suboptimal performance of ProtoGate.

### Feature Selection Performance

We compare ProtoGate against both global feature selection methods (RF, Lasso and STG) and local feature selection methods (L2X, LSPIN and LLSPIN). We plot the mean \(\pm\) std of the proportion of selected features across samples in Figure 3. The numerical results and full visualisation of selected features are available in Appendix C.

Figure 3 and Figure 4 show that ProtoGate consistently selects fewer features per sample than other benchmark methods, except L2X. Because the performance of L2X is the worst among the 12 benchmark models, we argue that the L2X model does not perform better than ProtoGate on feature selection

\begin{table}
\begin{tabular}{l c c c c c c c c c} \hline \hline Methods & lung & meta-dr & meta-pm & prostate & \multicolumn{2}{c}{tega-2y} & toxicity & colon & **Avg. Rank** \\ \hline LightGBM & \(93.42\pm 5.91\) & \(58.23\pm 8.56\) & \(94.98\pm 5.19\) & \(91.38\pm 5.71\) & \(57.09\pm 7.87\) & \(81.98\pm 6.25\) & \(76.60\pm 11.67\) & \(5.71\) \\ RF & \(91.73\pm 6.61\) & \(51.48\pm 3.41\) & \(88.73\pm 6.24\) & \(90.38\pm 7.31\) & \(58.70\pm 6.84\) & \(79.78\pm 7.10\) & \(80.05\pm 10.37\) & \(7.14\) \\ KNN & \(91.06\pm 7.92\) & \(54.64\pm 7.95\) & \(82.79\pm 9.20\) & \(78.78\pm 6.71\) & \(58.83\pm 7.07\) & \(83.86\pm 1.20\) & \(77.33\pm 5.41\) & \(8.00\) \\ Lasso & \(94.47\pm 3.71\) & \(85.88\pm 9.49\) & \(91.55\pm 2.83\) & \(91.19\pm 6.39\) & \(56.99\pm 6.26\) & \(91.98\pm 5.27\) & \(94.09\pm 5.00\) & \(4.29\) \\ MLP & \(\mathbf{95.81\pm 2.69}\) & \(\mathbf{54.68\pm 9.63}\) & \(95.71\pm 2.89\) & \(87.22\pm 7.41\) & \(55.52\pm 7.24\) & \(93.84\pm 4.28\) & \(90.80\pm 8.70\) & \(5.14\) \\ STG & \(93.30\pm 6.28\) & \(81.58\pm 8.67\) & \(67.18\pm 8.91\) & \(89.38\pm 5.85\) & \(57.04\pm 5.76\) & \(87.95\pm 5.01\) & \(79.55\pm 10.53\) & \(6.29\) \\ TabNet & \(77.65\pm 11.56\) & \(49.18\pm 15.02\) & \(82.66\pm 7.81\) & \(65.66\pm 9.03\) & \(51.58\pm 8.26\) & \(40.06\pm 12.23\) & \(56.75\pm 7.31\) & \(12.00\) \\ L2X & \(50.02\pm 8.30\) & \(52.54\pm 13.75\) & \(62.64\pm 13.69\) & \(61.78\pm 6.29\) & \(52.30\pm 9.11\) & \(31.72\pm 13.48\) & \(57.60\pm 14.26\) & \(12.43\) \\ INVASE1 & \(91.22\pm 6.16\) & \(-\) & \(91.70\pm 6.84\) & \(-\) & \(-\) & \(55.98\pm 6.45\) & \(80.40\pm 6.60\) & \(-\) & \(9.00\) \\ REAL-X & \(93.27\pm 4.32\) & \(60.01\pm 7.12\) & \(95.95\pm 3.04\) & \(86.75\pm 6.68\) & \(59.30\pm 7.49\) & \(90.97\pm 4.75\) & \(76.55\pm 12.21\) & \(5.14\) \\ ILSPIN & \(70.10\pm 12.31\) & \(56.77\pm 9.65\) & \(95.50\pm 3.60\) & \(85.71\pm 5.98\) & \(57.87\pm 6.02\) & \(61.67\pm 9.01\) & \(79.35\pm 7.34\) & \(7.14\) \\ LSPIN & \(76.92\pm 9.38\) & \(53.98\pm 8.00\) & \(\mathbf{97.18\pm 3.16}\) & \(87.75\pm 6.74\) & \(55.95\pm 4.75\) & \(83.47\pm 8.59\) & \(\mathbf{81.30\pm 7.97}\) & \(6.71\) \\ \hline
**ProtoGate** & \(93.44\pm 6.37\) & \(\mathbf{60.43\pm 7.61}\) & \(95.96\pm 3.93\) & \(90.58\pm 5.64\) & \(\mathbf{61.18\pm 6.47}\) & \(92.34\pm 5.67\) & \(\mathbf{81.10\pm 12.14}\) & \(\mathbf{2.00}\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Evaluation comparison of ProtoGate with 12 benchmark methods on seven real-world tabular biomedical datasets. We report the mean \(\pm\) std balanced accuracy (averaged across 25 runs) and average accuracy rank across datasets. A higher rank implies higher accuracy. We highlight the **First**, Second and Third ranking accuracy for each dataset. ProtoGate consistently ranks Top-3 across datasets and achieves the best overall performance.

Figure 4: Heatmaps of mask values on the meta-dr dataset. ProtoGate has a sparser feature selection result, since it can better identify the homogeneity across samples.

Figure 3: Comparison of the feature selection sparsity on real-world datasets. We report the mean \(\pm\) std of the proportion of selected features on test samples, averaged over 25 runs. ProtoGate learns sparser patterns than other methods by a clear margin except for L2X.
although it has the fewest selected features. Compared with the rest local feature selection methods, ProtoGate has smaller standard deviations in the proportion of selected features across test samples. Note that this does not mean that ProtoGate selects features globally, because similar proportions of selected features only denote similar numbers of selected features, not necessarily similar selected features (Section 4.3). The sparse feature selection results from ProtoGate demonstrate the effectiveness of global information in feature selection, and the global-to-local process helps ProtoGate attend to both homogeneity and heterogeneity across samples.

### Model Design Ablations

**Impact of prototype-based predictor.** We now investigate how the prototype-based predictor impacts classification performance. For a fair comparison, we replace the DKNN predictor with a linear head network or an MLP, and then tune the hyper-parameter for global sparsity \(\lambda_{g}\) by searching within \(\{1e-4,2e-4,3e-4\}\).

As shown in Table 2, the DKNN predictor consistently outperforms other predictors. We attribute the performance improvement to the appropriate inductive bias in prototype-based classification and the reduction in learnable parameters. In ProtoGate, only the feature selector needs training, while other local feature selection methods have learnable predictors with vast amounts of parameters to optimise. We also find that simply combining a global-to-local feature selector and an MLP/linear prediction head does not outperform LSPIN/LLSPIN. This further indicates that a prototype-based predictor is the key to the high accuracy of ProtoGate.

**Impact of global-to-local feature selector.** In Section 4.2, we discussed how the global-to-local feature selector helps ProtoGate generate a sparser feature selection result. We further examine how different hyper-parameter values of global sparsity \(\lambda_{g}\) impact the feature selection behaviour.

Figure 5(a) shows that increasing \(\lambda_{g}\) can lead to a lower degree of local sparsity. We also find in Figure 5(b) that ProtoGate achieves the best test accuracy when selecting features locally (\(\mathcal{Q}>0\)), which aligns with the domain knowledge that heterogeneity across samples is important for accurate predictions on biomedical data. This also suggests that the outstanding performance of ProtoGate is due to its considering both homogeneity and heterogeneity for feature selection.

\begin{table}
\begin{tabular}{l c c c c c c c} \hline \hline Predictors & Lung & meta-dr & meta-span & prostate & tega-2y & toxicity & colon \\ \hline MLP & \(69.97\pm 9.17\) & \(560.00\pm 6.37\) & \(93.62\pm 6.04\) & \(89.13\pm 6.36\) & \(54.74\pm 8.11\) & \(90.36\pm 5.61\) & \(80.95\pm 7.77\) \\ linear Head & \(66.51\pm 12.45\) & \(56.10\pm 8.95\) & \(93.20\pm 6.18\) & \(89.87\pm 5.80\) & \(56.60\pm 8.20\) & \(90.29\pm 5.93\) & \(79.45\pm 6.23\) \\
**DKNN** & \(\mathbf{93.44\pm 6.37}\) & \(\mathbf{60.43\pm 7.61}\) & \(\mathbf{95.96\pm 3.93}\) & \(\mathbf{90.58\pm 5.64}\) & \(\mathbf{61.18\pm 6.47}\) & \(\mathbf{92.34\pm 5.67}\) & \(\mathbf{81.10\pm 12.14}\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: Balanced accuracy for different predictors on real-world datasets, averaged over 25 runs. We **bold** the highest accuracy for each dataset. The prototype-based classifier consistently outperforms linear and MLP predictors on all datasets.

Figure 5: Comparison of different values of global sparsity hyper-parameter \(\lambda_{g}\). (a) The degree of local sparsity averaged over 25 runs. Increasing \(\lambda_{g}\) reduces the diversity of selected features between samples. (b) The balanced accuracy averaged over 25 runs. Increasing \(\lambda_{g}\) does not guarantee improvement in the prediction accuracy.
**Training considerations.** ProtoGate can require larger training overhead, mostly for tuning the hyper-parameters compared to some existing models, since we need to consider the interplay between \(\lambda_{g}\), \(\lambda_{l}\) and \(K\). ProtoGate also stores all training samples in the prototype base \(\mathcal{B}\), leading to higher memory consumption on large datasets than benchmark methods. Because we mainly focus on the HDLSS datasets, memory consumption is not a major problem in this regime.

### Co-adaptation Analysis

We evaluate ProtoGate and benchmark feature selection models on the synthetic datasets to examine their correctness in feature selection and susceptibility to the co-adaptation problem. We use the same experimental settings as real-world datasets and change the range of hyper-parameter searching for each model to achieve their optimal performance. Following [5; 17], we measure the quality of selected features by computing the F1 score with predicted masks and ground truth masks, and the results are averaged over 25 runs.

**Synthetic datasets.** We generate three synthetic datasets by adapting the nonlinear datasets used in [5; 8; 17], and the exact data models are described in Appendix A.2. Each dataset has 200 samples of 100 features, which is only 10% of the samples and 10 times more features compared to [5]. All feature values are sampled independently from \(\mathcal{N}(0,\mathbf{I})\), where \(\mathbf{I}\) is an \(100\times 100\) identity matrix. Each dataset has two classes, and we make the data distribution imbalance by generating 50 and 150 samples for two classes respectively.

_We purposely design Syn3\({}_{(-)}\) to examine the inductive bias in ProtoGate._ Note that the absolute value function is an even function. Two samples with opposite values of the same feature are likely to have equal logit values, and then they belong to the same class. However, the opposite values mean a long distance between them, and they should not belong to the same class according to the clustering assumption. Therefore, prototype-based models are expected to perform poorly in this regime. We implement it by adding absolute value function \(|x_{9}|\) in the first class of Syn3\({}_{(-)}\) to observe the performance degradation in ProtoGate.

**Results.** On Syn1\({}_{(+)}\) and Syn2\({}_{(+)}\), ProtoGate achieves better or comparable performance in feature selection and classification than benchmark methods. On Syn3\({}_{(-)}\), ProtoGate performs poorly as expected. Although Syn1\({}_{(+)}\) and Syn2\({}_{(+)}\) also contain even functions like square and absolute value, they also have many other informative features that do not utilise the even functions to compute logit value. Therefore, the side effect of even functions is diluted in Syn1\({}_{(+)}\) and Syn2\({}_{(+)}\).

We also find LSPIN achieves the highest accuracy on Syn1\({}_{(+)}\), but has poor F1\({}_{\text{select}}\) in the selected features, denoting a severe problem of co-adaptation between the feature selector and predictor. In other words, LSPIN simply overfits the dataset without correctly identifying the informative features, making the feature selection results meaningless. In contrast, ProtoGate has consistently non-positive rank differences between F1\({}_{\text{select}}\) and ACC\({}_{\text{pred}}\), showing that the co-adaptation does not occur. The results demonstrate that ProtoGate can achieve a well-aligned performance of feature selection and classification, guaranteeing the quality of selected features.

\begin{table}
\begin{tabular}{l c c c c c|c c c c} \hline \hline \multirow{2}{*}{Methods} & \multicolumn{3}{c|}{Syn1\({}_{(+)}\)} & \multicolumn{3}{c|}{Syn2\({}_{(+)}\)} & \multicolumn{3}{c}{Syn3\({}_{(-)}\)} \\ \cline{2-10}  & \multicolumn{2}{c}{F1\({}_{\text{select}}\)} & \multicolumn{2}{c}{ACC\({}_{\text{pred}}\)} & \multicolumn{1}{c}{**Diff.**} & \multicolumn{2}{c}{F1\({}_{\text{select}}\)} & \multicolumn{2}{c|}{ACC\({}_{\text{pred}}\)} & \multicolumn{1}{c|}{**Diff.**} & \multicolumn{2}{c}{F1\({}_{\text{select}}\)} & \multicolumn{2}{c}{ACC\({}_{\text{pred}}\)} & \multicolumn{1}{c}{**Diff.**} \\ \hline RF & 0.1461 \(\pm\) 0.0367 & 57.08 \(\pm\) 6.48 & 3 & 0.1921 \(\pm\) 0.0230 & 59.44 \(\pm\) 5.24 & 1 & 0.2232 \(\pm\) 0.0241 & 56.33 \(\pm\) 9.08 & -1 \\ Lasso & 0.0905 \(\pm\) 0.0197 & 54.55 \(\pm\) 6.14 & 2 & 0.1130 \(\pm\) 0.0070 & 52.42 \(\pm\) 6.69 & 0 & 0.0900 \(\pm\) 0.0179 & 55.30 \(\pm\) 7.44 & 2 \\ STG & 0.2656 \(\pm\) 0.0420 & 86.85 \(\pm\) 9.3 & -1 & 0.2247 \(\pm\) 0.0994 & 58.28 \(\pm\) 8.36 & -2 & **0.2846** \(\pm\) **0.1820** & 54.00 \(\pm\) 9.09 & -7 \\ ThNet & 0.0843 \(\pm\) 0.0172 & 48.59 \(\pm\) 6.55 & 1 & 0.0642 \(\pm\) 0.0246 & 49.57 \(\pm\) 5.38 & 0 & 0.0605 \(\pm\) 0.0200 & 48.45 \(\pm\) 8.31 & 0 \\ L2X & 0.1599 \(\pm\) 0.0710 & 52.89 \(\pm\) 7.51 & -3 & 0.1873 \(\pm\) 0.0976 & 55.78 \(\pm\) 6.97 & -1 & 0.0984 \(\pm\) 0.0889 & 55.92 \(\pm\) 7.30 & 2 \\ INVAE & 0.1763 \(\pm\) 0.0456 & 53.36 \(\pm\) 9.00 & -1 & 0.1553 \(\pm\) 0.0338 & 60.28 \(\pm\) 8.61 & 6 & 0.1332 \(\pm\) 0.0256 & **50.75** \(\pm\) **8.70** & 5 \\ REAL-X & **0.1850** & 0.0438 & 47.54 \(\pm\) 9.51 & -7 & 0.2328 \(\pm\) 0.0729 & 52.50 \(\pm\) 6.38 & -6 & 0.2630 \(\pm\) 0.0567 & 56.48 \(\pm\) 9.34 & -1 \\ LLSPIN & 0.1060 \(\pm\) 0.0246 & 54.96 \(\pm\) 9.49 & 2 & 0.1692 \(\pm\) 0.0795 & 56.18 \(\pm\) 5.80 & 1 & 0.1031 \(\pm\) 0.0635 & 52.35 \(\pm\) 8.32 & -2 \\ LSPIN & 0.1466 \(\pm\) 0.0380 & **59.04 \(\pm\) 9.24** & 5 & 0.1911 \(\pm\) 0.0389 & 59.40 \(\pm\) 8.07 & 1 & 0.1927 \(\pm\) 0.0645 & 58.09 \(\pm\) 6.41 & 2 \\ \hline ProtoGate & **0.2948** \(\pm\) **0.0728** & 58.68 \(\pm\) 6.28 & **-1 & **0.2922** \(\pm\) **0.0943** & **60.67 \(\pm\) 8.21** & 0 & 0.1653 \(\pm\) 0.0554 & 56.16 \(\pm\) 6.82 & 0 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Evaluation comparison of ProtoGate and nine benchmark methods on three synthetic datasets. We report the F1 score of selected features (F1\({}_{\text{select}}\)) and the balanced accuracy for prediction (ACC\({}_{\text{pred}}\)). “Diff.” refers to the difference between the ranks of F1\({}_{\text{select}}\) and ACC\({}_{\text{pred}}\), and a positive value indicates a high possibility of co-adaptation. We highlight the First, Second and Third performance for each dataset. ProtoGate achieves well-aligned performance for feature selection and prediction.
Conclusion

We present ProtoGate, a prototype-based neural model for local feature selection on high-dimensional and low-sample-size datasets. ProtoGate selects features in a global-to-local manner and makes predictions with an interpretable prototype-based model. The experimental results on real-world datasets demonstrate that ProtoGate improves classification accuracy and interpretability by attending to both homogeneity and heterogeneity across samples. The analysis of synthetic datasets further reveals that ProtoGate can effectively avoid the co-adaptation problem by utilising a prototype-based predictor without learnable parameters. Although we evaluate ProtoGate only on classification tasks in this paper, it is readily extendable and applicable to other biomedical tasks, including regression.
