Introduction

Machine learning models for tabular datasets typically expect a complete feature set during both training and inference. However, in practice, features are often missing during inference due to the high cost and difficulty of obtaining the complete feature set for some samples (e.g., gene expression counts) [1,2]. Such test-time feature unavailability necessitates models that can make accurate predictions with incomplete feature sets. Additionally, since acquiring new features may be prohibitively expensive, it is crucial for these models to offer recommendations on which missing feature values to collect to maximise their impact on the model's accuracy.

Current strategies addressing limited feature availability typically involve either: (i) imputing, or predicting, missing features at test-time [1], or (ii) selecting a minimal feature subset on which the model is retrained [3,4]. Although these methods are practical, they have clear limitations in scenarios with variable feature availability: Feature selection identifies critical features but cannot adapt to changes in feature availability, while imputation provides flexibility but lacks guidance for users on prioritising features.

In this paper, we address this gap by introducing F-Act (Feature-wise Active adaptation), a method that combines feature selection with imputation to enable adaptation to variable feature availability without retraining, all while maintaining high predictive accuracy. F-Act achieves this by, first, imputing missing features at test-time and, second, enabling new features to be incorporated through test-time interventions, where F-Act's intermediate predictions space is modified to incorporate the presence of a new feature. Technically, F-Act employs differentiable mask sampling and feature reconstruction to learn to optimally operate from an incomplete set of features. This design enables F-Act to advise on the order features should be collected to maximise their impact, permitting deployment in resource-constrained settings. Using real and synthetic datasets, we evaluate F-Act and find that it matches benchmarks' performance in imputation, feature selection, and prediction, providing recommendations that enhance model performance through adaptive feature incorporation.

Background and Related Work

Imputation and Feature Selection Our work incorporates both feature imputation and feature selection to address limited test-time feature availability. As such, our work is placed at the intersection of these two research subfields. Previous works in feature imputation can be divided into auxiliary model-based approaches and joint learning approaches. In this context, auxiliary-model-based approaches pair prediction models with separate imputation methods [5,6,7] while joint learning approaches integrate both prediction and feature selection in an end-to-end model [8,9]. Nevertheless, we emphasise that previously proposed imputation techniques lack feature prioritisation for collection, leading to uncertainty over what features one should prioritise when deploying the model in a setup with varying feature availability. This is a key gap we aim to address with this paper.

Feature selection techniques [10,11,12,13], in contrast, deal with potential feature unavailability (or redundancy) by learning to select a subset of features from which a task can be accurately solved. These approaches commonly achieve this by learning a feature importance ranking that can then inform which subset of features one should select. A shortcoming of these approaches, however, is their inflexibility to a varying set of input features, as, once features have been selected, they require a fixed subset of features to train the downstream model [14]. As such, in this work we combine feature imputation with feature selection to enable easy adaptability from a core set of initially selected features. We note that combining feature selection and imputation has been previously explored [15,16,17,3]. However, performing feature selection with joint learning for inference and imputation in a single end-to-end architecture is novel. This is worth exploring, as Bertsimas et al. [8] and Le Morvan et al. [9] note that joint learning for imputation and inference can yield improved results.

Relation to Active Feature Acquisition Active Feature Acquisition (AFA) involves learning a policy for collecting new features at test-time such that a model's accuracy is maximised after observing a small set of features [18,19,20,21]. As we are interested in providing feature collection recommendations, our work is highly related to AFA. Nevertheless, we highlight that we distinguish ourselves from traditional AFA approaches in two key ways. First, we provide feature collection recommendations at a global level rather than at a local, per-sample level. Second, we learn to both select a subset of features and impute missing features in an end-to-end

Mask Module Prediction Module Reconstruction Module ( , )( , ) ( ) ( ) ( , )

Test-time interventions

Figure 1: F-Act. Given a sample x, we apply an element-wise mask m soft ∈ [0, 1] 𝑛 , which acts as a sparsity-inducing global feature selection mask, eliminating noisy features and resulting in x ˜. We then apply an element-wise mask m hard ∈ {0, 1} 𝑛 , sampled from a Gumbel-Softmax distribution with learnable probabilities 𝜋 hard . The masked features are reconstructed by the Reconstruction Module, 𝑓 𝑅 , to approximate x ˜. Subsequently, we perform test-time interventions on this reconstructed sample by re-incorporating 𝑘 previously masked features according to a greedy policy, 𝜇 derived from the rankings of 𝜋 hard , resulting in x ˜𝐼 . Finally, the Prediction Module, 𝑓 𝑃 predicts the sample's label.

fashion, enabling missing features to be predicted at test-time.

Relation to Human Interpretable Artificial Intelligence Human Interpretable Artificial

Intelligence (HI-AI) refers to AI systems designed to ensure their decisions and workings are understandable and transparent to humans. Our work is related to HI-AI methods as it enables (1) reconstruction of missing features through test-time imputations, providing insights into a model's understanding of how missing features relate to provided ones, (2) construction of feature importance rankings through its feature collection recommendations, and (3) test-time interventions, where users can provide previously missing features by intervening on F-Act's intermediate predictions using these features' values. As such, our work is related to previous interpretable imputation techniques [22,7] and methods in the concept-based explainable AI literature [23,24,25,26] that provide test-time feedback to models via human-aligned concepts.

Feature-wise Active adaptation

We present a joint learning framework for inference and missing value imputation that offers users insights into which missing feature values should be prioritised for collection and intervention. Formally, our goal is to learn a predictor 𝑓 𝜃 , parameterised by 𝜃, which can operate on any subset of features 𝑆 ⊆ ℱ. Concurrently, the predictor 𝑓 𝜃 should also suggest which missing feature values to collect for intervention at test-time, prioritised by their importance to improve the predictor's performance. We achieve this by introducing F-Act (Figure 1), a method for the joint learning of feature selection, missing value imputation, and prediction. Our architecture comprises three modules: (i) a Mask module that facilitates feature selection, (ii) a Reconstruction module for feature imputation, and (iii) a Prediction module for making predictions.

The Mask module serves three objectives: i) global feature selection to eliminate irrelevant features, ii) learning feature importance rankings to provide recommendations for feature collection, and iii) simulating a missing feature scenario to train the Reconstruction module. We achieve all this functionality through hierarchical masking, first employing a soft mask m soft ∈ [0, 1] 𝑑 for feature selection and then a hard mask m hard ∈ {0, 1} 𝑑 to simulate missing features. The hard mask is sampled from a Gumbel-Softmax distribution [27] with a learnable probability 𝜋 ℎ𝑎𝑟𝑑 .

To enable predictions from any truncated feature space and facilitate test-time interventions, we use the Reconstruction Module to reconstruct features from the truncated feature space 𝒳 ˜𝑆 generated by the hard mask. The Reconstruction Module outputs "reconstructed" samples, containing feature values for all values, even though they are missing from the original input. The reconstructed samples are then processed by the Prediction Module, which maps from the complete feature space 𝒳 ˜to make the predictions 𝒴. This setup ensures the model can make predictions even when in the presence of missing features.

To provide feature collection recommendations, we define a greedy intervention policy that corrects reconstructed data based on feature selection probabilities from the mask module. This is the same approach used by feature importance-based selection methods [11,4,13,12].

In order to jointly learn to perform feature selection, missing feature imputation and, downstream task prediction, we train our model using a composite loss function ℒ = ℒ P + 𝛼 𝑆 ℒ S + 𝛼 𝑅 ℒ R where 𝛼 𝑆 and 𝛼 𝑅 are hyperparameters controlling how much we value feature selection (i.e., ℒ S ) and feature reconstruction (i.e., ℒ R ) over task accuracy (i.e., ℒ P ).

To encourage our model to perform a sparse feature selection, we follow previous works [28,4] and let ℒ S be the ℓ 1 norm of the soft and hard learnable mask probabilities:

ℒ S := 𝑑 ∑︁ 𝑖=1 (︁ 𝜋 soft 𝑖 + 𝜋 hard 𝑖 )︁

In contrast, to encourage accurate imputation of masked features, we include a reconstruction loss term ℒ R that minimises the ℓ 2 norm of the difference between reconstructed/imputed feature values and ground truth feature values:

ℒ R := 1 |ℱ ∖ 𝑆| ∑︁ 𝑖∈ℱ ∖𝑆 (x ˜𝑖 − 𝑓 𝑅 𝑖 (m hard ⊙ x ˜; 𝜃 𝑅 )) 2

where 𝑆 is the set of features selected by the Mask module. This loss encourages our model to learn to select a set of core features from which other, dependent features, may be easily imputed.

Finally, to enable our model to predict downstream tasks both in and outside the presence of potential feature interventions, we follow the work in [25] and define our prediction loss, ℒ P , as

ℒ P := 𝐿 pred (𝑓 𝜃 (x; 𝜃, 0), X, Y) + 𝜔 𝑘max 𝐿 pred (𝑓 𝜃 (x; 𝜃, 𝑘 max ), X, Y)

Here, 𝑓 𝜃 (x; 𝜃, 𝑖) represents the output of the task predictor for sample x when the top-𝑖 dependent features are intervened on and 𝑘 max is the maximum number of interventions one may perform (i.e., the number of dependent features |ℱ ∖ 𝑆|). We clarify that, in this context, an intervention for a feature x 𝑗 involves setting the 𝑗-th feature of the reconstructed features x 𝑅 to x ˜𝑗. This loss, therefore, encourages the model to minimise a task-specific loss (e.g., cross-entropy) before and after interventions, with higher penalties incurred when a mistake is made after a higher number of features have been intervened on at train time (controlled by a hyperparameter 𝜔 > 1). Inference By thresholding mask probabilities, we can identify core necessary features and recommend which features to prioritize for collection. During inference, F-Act imputes missing features dynamically and allows the re-incorporation of non-selected features in the form of test-time interventions. In practice, given an incomplete sample at test-time, we replace the missing values with 0. To perform feature selection, imputation and prediction, we apply the mask, reconstruction and prediction modules in order. A change from the training procedure is that, at test-time, the Gumbel Softmax function's temperature is set to 0, making its performance deterministic. This is equivalent to thresholding the hard mask probabilities at 0.5. Test-time interventions are performed by replacing the reconstructions of the hard-masked features with their true values, as shown in Figure 1.

Experiments

Datasets and benchmark methods We consider various real-world datasets commonly referenced in feature selection literature. These include image datasets (COIL20 and USPS), a voice audio dataset (Isolet) sourced from [29], a synthetic dataset (Madelon) from [30], genomic datasets (PBMC [31] and Mice Protein [32]), and a financial dataset (Finance) from [33].

Beyond predictive accuracy, we assess F-Act's capabilities in selecting important features and recommending which features to collect at test-time to enhance performance. To this end, we consider several feature selection methods, including LASSO [11], Random Forest [13], Concrete Autoencoders (CAE) [34], XGBoost [12], and SEFS [4]. All methods except CAE rank the features by their importance, which allows us to evaluate F-Act's ability to recommend features for collection. We train each feature selection method on the prediction task, using a simple MLP as a baseline for comparison.

For missing data imputation, we evaluate four methods: Mean, Iterative Chained Equations (ICE) [6], and MissForest [7]. Additionally, we assess the performance of all combinations of downstream models and imputation methods.

We train F-Act by minimising the loss 𝐿. We pre-train the reconstruction module following [4], with tasks that include reconstructing input vectors and estimating gate vectors. Following this, we tune the intervention number 𝑘 to minimise the prediction loss on the validation dataset. F-Act outperforms other methods when the features selected for collection at test-time are based on the predictor's learned feature ranking. F-Act outperforms all baseline methods when many features are missing. However, with fewer missing features, F-Act can be outperformed by XGBoost combined with ICE. This indicates that F-Act is especially effective when feature collection can be prioritised.

For further implementation details, please refer to Appendix A.

Predictive Accuracy Table 1 illustrates the predictive accuracy of F-Act compared to other benchmark methods. F-Act demonstrates competitive performance, consistently ranking in the top three across datasets and outperforming all baselines in three cases. Overall, F-Act ranks the best across datasets. However, besides making predictions, F-Act offers two additional functionalities without any re-training: imputing missing data at test-time and recommending which feature values to collect. We next explore these capabilities.

Test-time Imputation First, we evaluate F-Act's imputation capabilities when features are missing completely at random (MCAR). To simulate this, we randomly remove features at test-time (without considering the potential to collect these feature values). At low levels of missing features, Figure 2a shows that F-Act is generally outperformed by Random Forest and MLP, and at higher levels of missing values, it is outperformed by Lasso and MLP. These mixed results suggest F-Act achieves relatively average performance when one cannot utilise its learned feature ranking.

Second, we consider prioritising test-time feature collection based on the model's learned feature ranking. Figure 2b shows that when this ranking is used, F-Act outperforms all other methods with only a few features collected.

Test-Time Interventions

As a feature selection method, F-Act uses a threshold to separate core from non-core features. Unlike standard approaches, F-Act can incorporate non-core features during inference. Figure 3 illustrates that F-Act's performance improves with test-time interventions, sometimes even surpassing the best-performing method on that dataset. For more results, please see Appendix B.

Conclusion

This paper introduces F-Act, a method combining feature selection and missing data imputation to enable the model to operate when there are missing features at test-time. More importantly, F-Act provides recommendations of which features one should prioritise collecting at testtime to improve the model's performance. Our empirical analysis shows F-Act performs competitively or better than previous baselines in inference tasks with missing features when incorporating feature collection recommendations. Additionally, we show how F-Act can incorporate missing feature values at test-time through test-time interventions, improving performance without retraining and boosting F1 scores across datasets. This work highlights the benefit of designing methods that learn, in an end-to-end fashion, to adapt to different feature availability while providing feature collection recommendations. for 𝑖 = 1 to 𝑛 mb do ◁ For each point in the mini-batch 4:

(x 𝑖 , 𝑦 𝑖 ) ∼ (X, Y) ◁ Sample a data point 5:

m hard 𝑖 ← Bernoulli(𝜋) ◁ Sample hard mask with probability 𝜋 6:

x 𝑆 𝑖 ← m hard 𝑖 ⊙ x 𝑖 ◁ Apply hard mask to the data point 7:

x

𝜃 ← 𝜃 − 𝜂∇ 𝜃 ∑︀ 𝑛 mb 𝑖=1 (︀ ℒ R (x 𝑆 𝑖 , x 𝑅 𝑖 ) + 𝛼 pre • CE(m ^hard 𝑖 , m hard 𝑖 )

)︀ ◁ Update parameters using gradient descent 11: until convergence hyperparameter tuning and final evaluation phases. We pre-train out model as per Appendix A.3 and train our model as per Algorithm 2.

A.5. Hyper-parameter Tuning

Random Forest. For the Random Forest model, we conducted a hyper-parameter sweep on the max_depth parameter. The values considered for max_depth were {3, 5, 7}. This tuning was performed to control the complexity of the individual trees in the forest, with a goal of balancing the bias-variance tradeoff.

Lasso. In our implementation of the Lasso model, we performed hyper-parameter tuning on two key parameters: l1_ratio and C. The l1_ratio was varied over {0, 0.25, 0.5, 0.75, 1}, allowing us to explore the impact of the ElasticNet mixing parameter which adjusts the balance between L1 and L2 penalties. The C parameter, which controls the inverse of regularisation strength, was swept over {10, 100, 1000}, providing a wide range of regularisation effects.

XGBoost. For the XGBoost model, we focused our hyper-parameter sweep on the eta (learning rate) and max_depth. The eta values considered were {0.1, 0.3, 0.5}, providing a spectrum of learning rates to control the step size during optimisation. For max_depth, the values were {3, 6, 9}, allowing us to examine different depths for the trees to manage the model's complexity and prevent overfitting.

Neural Network Based Models. For the neural network-based models, which include MLP, SEFS, CAE, Supervised CAE, and F-Act, we conducted a hyper-parameter sweep. Key parameters included learning rate, lr, ({1e-3, 3e-4, 1e-4}), number of hidden layers ({1, 2, 4}), and dropout rates ({0, 0.2}). Additionally, for the Concrete Autoencoder models (CAE, Supervised CAE), we swept the neurons_ratio over {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0} to explore various proportions of neurons in the encoder and decoder layers. This extensive tuning process was aimed at optimising each of the methods architectures and regularisation techniques to x ˜𝑖 ← m soft ⊙ x 𝑖 ◁ Apply soft mask to input data 8:

x ˜𝑆𝑖 ← m hard ⊙ x ˜𝑖 ◁ Apply hard mask to soft-masked data Reconstruct the data, apply interventions, and make prediction 9:

x ˜𝑅𝑖 ← 𝑓 𝑅 (x ˜𝑆𝑖 ; 𝜃 𝑅 ) ◁ Reconstruct the hard-masked features In Figure 2b, we observe the interesting phenomenon that at low levels of missing data, ICE imputation enables the tree-based models to achieve considerably improved performance. Concurrently, ICE negatively affects the Neural Network and Lasso models. Here, we discuss that phenomenon. Note that, as per the structure of the experiment, the missing features at those levels are relatively lower-ranked features which have less impact on predictions in the case of tree-based models. One possible explanation for the increase in performance of the tree-based models, is that they were initially negatively affected by over-fitting to noise in the lower-ranked features. As a result, the models might benefit from the removal of these features in the test set. This then explains why replacing the lower-ranked features with expected values conditioned by the other observed features of that data point (as per the ICE imputation strategy [22]) would reduce their noisy impact. To explain the poor performance of ICE with the Neural Network-based and Lasso models, we note that those models are known to be more sensitive to scale and domain shift [35]. We also note that ICE is known to suffer from misspecification [36], where imputed values are "implausible", falling out of the domain.

B.2. Optimal feature availability

The adaptive nature of our model enables the ability to provide a more finely tuned optimal feature selection recommendation than the standard "selected" vs "non-selected" features dichotomy. Rather than being derived from what is typically an arbitrarily set threshold, we are able to tune the number of selected features without re-training. By contrast, to implement feature selection with methods such as Lasso it is standard to re-train the model with the reduced feature set. In this section, we hypothesise that this functionality will enable improved performance, due to the ability to more finely tune the number of selected features. The "optimal feature selection" is found through post-training hyperparameter tuning of the number of interventions, 𝑘. That is, we evaluate the model at varying degrees of test-time interventions on the validation set, and then set this as the number of features used by the model. In Table 3, we present the results of this approach. In the table, we compare F-Act, which uses post-training hyperparameter tuning of 𝑘, to F-Act "selected only", which only uses only the selected features. We find that in most datasets, this enables an improvement over the "selected only" variant of the model. Exceptions include the finance and mice protein datasets, where F-Act "selected only" outperforms F-Act ῾ 0.005, which is within the standard deviation of our model's F1 score for those datasets. However, on other datasets, such as PBMC, the improvement of F-Act over F-Act "selected only" is rather notably greater than twice the standard deviation. Overall we find that the difference in average F1 score falls within the standard deviation of the model, indicating that the potential gains for this mechanism are limited, however the substantial gains on the PBMC dataset, as well as the small computational costs associated with its implementation, indicate that the mechanism is worth exploring when deploying F-Act.