Active feature acquisition via explainability-driven ranking
Authors: Osman Berke Guney
Ketan Suhaas Saichandran
Karim Elzokm
Ziming Zhang
Vijaya B Kolachalama
Venue: ICML 2025 poster (2025)
Keywords: Active feature acquisition
Policy network
Feature importance
Primary area: general_machine_learning->online_learning_active_learning_and_bandits
TL;DR: An explainability-driven AFA method uses local explanations and a decision-transformer policy network to dynamically select the most relevant features, achieving better accuracy and efficiency than existing methods.

Abstract: In many practical applications, including medicine, acquiring all relevant data for machine learning models is often infeasible due to constraints on time, cost, and resources. This makes it important to selectively acquire only the most informative features, yet traditional static feature selection methods fall short in scenarios where feature importance varies across instances. Here, we propose an active feature acquisition (AFA) framework, which dynamically selects features based on their importance to each individual case. Our method leverages local explanation techniques to generate instance-specific feature importance rankings. We then reframe the AFA problem as a feature prediction task, introducing a policy network grounded in a decision transformer architecture. This policy network is trained to select the next most informative feature by learning from the feature importance rankings. As a result, features are acquired sequentially, ordered by their predictive significance, leading to more efficient feature selection and acquisition. Extensive experiments on multiple datasets demonstrate that our approach outperforms current state-of-the-art AFA methods in both predictive accuracy and feature acquisition efficiency. These findings highlight the promise of an explainability-driven AFA strategy in scenarios where the cost of feature acquisition is a key concern.

Active Feature Acquisition Via Explainability-rriven Ranking
Osman Berke Guney'
Ketan Suhaas Saichandran 23 Karim Elzokm
Ziming Zhang
Vijaya B. Kolachalama" 235
Abstract
1. Introduction
In traditional machine learning, all features are typically as- sumed be available at inference. However, in real-world settings, feature acquisition is often costly, time-consuming, and sequential. Developing models that can make accurate predictions while minimizing feature acquisition is impor- tant for efficiency and practical implementation. This can be achieved by selecting a static global subset features, but it suboptimal since the important set of features may vary across different instances (Kachuee et al., 2019; Covert al., 2023b). Furthermore, the chosen subset might not provide sufficient information for some cases, which requires the acquisition of more features to ensure confident prediction. A more effective strategy is to identify important features sequentially for each individual instance, a technique known as active feature acquisition (AFA), which has gained in- creasing attention in recent years (He Chen, 2022; von Kleist al., 2023; Chattopadhyay et al., 2024).
In many practical applications, including medicine, acquiring all relevant data for machine learning models is often infeasible due to constraints on time, cost, and resources. This makes important to selectively acquire only the most informative features, yet traditional static feature selection methods fall short in scenarios where feature importance varies across instances. Here, we propose an active feature acquisition (AFA) framework, which dynamically selects features based on their importance each individual case. Our method leverages local explanation techniques to generate instance-specific feature importance rankings. We then reframe the AFA problem as feature prediction task, introducing a policy network grounded a decision transformer architecture. This policy network is trained to select the next most informative feature by learning from the feature importance rankings. As result, features are acquired sequentially, ordered by their predictive significance, leading to more efficient feature selection and acquisition. Extensive experiments on multiple datasets demonstrate that our approach outperforms current state-of-the-art AFA methods in predictive accuracy and feature acquisition efficiency. These findings highlight the promise of an explainability-diie -driven AFA strategy in scenarios where feature acquisition is concern.
The literature mainly contains two different ways of ap- proaching AFA: reinforcement learning (RL)-based and greedy-based methods. Both approaches aim to develop feature selection policy through exploration, as instance- wise feature importance rankings are typically unavailable. RL-based methods (Kachuee et al., 2019; Yin al., 2020; von Kleist al., 2023) train policy networks by maximizing different reward functions. While the RL-based approach is intuitive for this sequential task and theoretically capa- ble of finding the optimal policy, empirical evidence shows that RL-based methods often underperform compared to greedy-based methods (Gadgil et al., 2024). Greedy-based methods attempt to predict the next most important avail- able feature by calculating conditional mutual information (CMI). To compute CMI, researchers have proposed gener- ative approaches (Rangrej & Clark, 2021; He et al., 2022) and methods based on the variational perspective (Covert et al., 2023b; Gadgil et al., 2024). However, calculating CMI directly remains challenging, and methods leverag- ing the variational perspective have demonstrated superior performance compared to generative alternatives.
Department of Electrical & Computer Engineering, Boston University, MA, USA Department of Computer Seience, Boston University, MA, USA 'Department of Medicine, Boston Uni- versity Chobanian& & Avedisian School of Medicine, MA, USA "Department of Electrical & Computer Engineering, Worcester Polytechnic Institute, MA, USA FFauulyy of Computing Data Sciences, Boston University, MA, USA. Correspondence to: Vi- jaya B. Kolachalama <vkola @ bbu.euu..
In this work! we approached the problem by empirically
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
Our code is available at https://github.com "kola-lab/icmle

---

Active Feature Acquisition Via Explainabiltty-drien Ranking
observing that local explanation methods, such as SHap-
many RL-based approaches proposed (Dulac-Amold et al., 2011; Shim al., 2018; Kachuee et al., 2019; Yin et al., 2020; Li & Oliva, 2021; von Kleist et al., 2023). These methods generally train their policy networks with the ob- jective of maximizing the defined reward functions. Namely, they try to approximate the action-value function (i.e., Q- function). For example, (Dulac- et al., 2011), the Q-function is approximated linearly and later is extended in (Janisch et al., 2019) using deep Q network (Mnih et al., 2015; van Hasselt et al., 2016). A similar approach was taken by the opportunistic learning (OPL. method in (Kachuee et al., 2019). Another type of mainstream meth- ods (Rangrej & Clark, 2021; He et al., 2022; Covert et al., 2023b; Chattopadhyay et al., 2023; Gadgil et al., 2024) are the greedy-based frameworks. These methods acquire the features by estimating the conditional mutual information (CMI) between the current available subset of features and the unacquired features. For CMI estimation, there are generative approaches (Rangrej & Clark, 2021; He et al., 2022) that use variational autoencoders (Kingma & Welling, 2013), and discriminative approaches (Covert et al., 2023b; Chattopadhyay et al., 2023; Gadgil et al., 2024) directly pre- dicting the feature index with the highest CMI without ex- plicitly calculating CMI. Although, the MDP formulation is theoretically appealing. RL-baase methods often underper- form compared to discriminative greedy- based approaches (Covert et al., 2023b; Gadgil et al., 2024).
ley Additive exPlanations (SHAP) (Lundberg & Lee, 2017) and LIME (Ribeiro et al., 2016), can be utilized to iden- tify instance-wis wise feature importance rankings. With this insight, we framed the AFA problem as a feature prediction task rather than a feature exploration one, Specifically, we trained our policy network to select the next unacquired feature with the highest importance ranking based on the current observations. This approach contrasts with feature exploration, which would involve searching for relevant fea- tures without predefined rankings, and instead focuses on prioritizing and acquiring the most informative feature in a structured, instance-spectt specific manner. Our contributions are listed below:
We demonstrated that local explanation methods effec- tively determine instance-specific feature importance rankings. Additionally, we showed that an ideal (or- acle) policy network, following a precomputed fea- ture acquisition order based on these rankings, out- performed state-of-thh-arr (SOTA) AFA techniques for any fixed feature budget. These findings empir- ically underscore the potential of explainability meth- ods in guiding instance-specifi feature acquisition. Similar trends have been observed in the local expla- nation literature (Petsiuk et al., 2018; Jethani et al., 2021; 2022), where inserting or removing features based on explanation-derive derived rankings leads to im- provements or degradations in model performance, re- spectively. However, prior work has not formally com- pared explanation-based feature ranking approaches with AFA techniques, leaving a gap in understanding how these methods align or diverge in sequential fea- ture acquisition tasks. Our study addresses this gap and highlights explanation-base rankings as a principled and effective baseline for AFA.
Alternative AFA approaches also exist that avoid policy networks entirely and instead leverage imputation to guide feature acqussition (Beebe-Wang et al., 2024). These approaches first identify set of nearest neighbors from the training data based on the currently avail- able features, which are then used to generate an ensemble of imputed instances. For example, AACO (Valancius et al., 2024) evaluates a set of candidate features by computing a loss function (weighted combination of the predictive loss and feature acquisition cost) over the ensemble of imputed instances, using the true labels of the nearest neighbors. Also, the method in Sreebe-Wang et al., 2023) applies a fixed predictor to the ensemble of imputed instances and then uses an explanation technique, i.e., SHAP (Lundberg & Lee, 2017). to acquire the next feature with the highest variance in importance scores for acquisition. Unlike our proposed framework, the approach in (Beebe-Wang et al., 2023) relies on a local explanation technique (e.g.. SHAP) at inference time, which is computationally demanding. In contrast, we used feature rankings derived from explana- tion methods on the training data to train a policy network, thereby enabling efficient and scalable feature acquisition. Conceptually, our approach is closely related to imitation learning approaches (He et al., 201 2012a;b), where the policy network is trained to follow trajectories of another policy, such as greedy one (He et al., 2016a).
al., 2023; Valancius
We employed a decision transformer architecture (Chen et al., 2021) as our policy network and trained it using a two-stage approach. Our method outperformed SOTA techniques, demonstrating that feature impor- tance rankings can be accurately inferred without direct observation.
2. Related work
Generally, the methods the AFA literature have two net- works: a policy network for feature acquisition and a predic- tion network for prediction with available subset of features. The AFA methods mainly differ in training their policy networks, so we only highlight those differences.
The AFA problem can be formulated as a Markov decision process (MDP) (Zubek & Dietterich, 2002; Dulac-Arnole et al., 2011); 1); based on this formulation, there have been
2

---

Active Feature Acquisition Via Explainability-drive Ranking
Masked Input and Target Feature Index Generation
b) Feature Index Prediction and Acquisition
Feature importance makinng order: g- (32, 14, 25,
Initial fenuree index II
Target index
Predicted index
a
a ++
a t+
Target feature index
35
32
X
1
35
32
a ++
a
Predicted feature
index
14
14
11 Nas1
32
32
Causal Transformer (9_)
Causal Transformer (9_))
1x Nea
24
25
p Nygt Xast
14
14
X.,
X
a_
F
X,,
X
a c+]
First stage
Prodictedd fouiure index
feature indix
Ayp Rege Mey ...)
( Anp Sige Kyg Second stage
xx.
a
F-11
X
a
r
X
a
X
[+1
Socond
stage
Fürst stage
<) Inference
ag
az
Causal Transformer (9,))
Causal Transformer (q,))
Causal Transformer (9_)
Nay
Ney
Figure 1. Explainability- driven active feature acquisition framework. a) Our training strategy consists of two stages. This figure shows how the masked inputs and target feature indices are generated during the first and second stages for a given feature importance ranking yp. In the first stage, features are selected based on their ranking order yp In the second stage, features are acquired by the policy network (g=). During the first stage, the next feature in the ranking is the target feature index. However, in the second stage, the target feature the fcature index having the highest ranking order among the ones that are not acquired; because this, the target feature remains the same until it is acquired. b) This part of the figure shows how the policy network qu. based on the decision transformer (Chen et al., 2021). processes the masked inputs during training for both stages. Sequential data with a context length E, set to 2 in this case, is fed into qx- At each time step, (r receives three tokens: the masked input (XM) ) action (az) and reward (r²)) The action token represents the index of the last acquired feature, and the reward is the output of the predictor network. To ensure causality, future tokens are masked while qr predicts the next feature to acquire at any time step. c) This figure illustrates the inference stage for image inputs in the causal transformer model, where predicted fcatures (or patches) are progressively acquired in a series of sequential acquisition steps.
In addition to AFA methods, related approaches from the budget learning literature (Trapeznikov & Saligrama, 2013; Nan & Saligrama, 2017; Ekanayake & Zois, 2024) explore fixed feature acquisition orders, limiting the number of po- tential feature subsets. These methods aim to identify easily classifiable instances, enabling the acquisition of a mini- mal set of features in such cases, thereby reducing overall acquisition costs.
instance-specific feature subsets, aiming either to explain model predictions or to highlight the most informative fea- tures for each input. Additionally, global feature importance methods aim to identify the most relevant static features across a dataset. For example, the Conerete Autoencoder (CAE) (Balin et al., 2019) selects features via an autoen- coder, while SAGE (Covert al., 2020) extends Shapley values to quantify global feature importance. For detailed overview, we refer readers to recent surveys (Samek et al., 2021; Bolón-Canedo et al., 2022; Mesinovic et al., 2023).
With regards to the local explanation literature (Ribeiro et al., 2016; Petsiuk et al., 2018; Jethani et al., 2021; Lundberg & Lee, 2017), various methods focus on quantifying the contribution of individual features to model predictions for each instance. Among these methods, SHAP (Lundberg & Lee, 2017), based on game-theoretic Shapley values (Shap- ley, 1953), is particularly popular. SHAP calculations are computationally intensive, leading to the development of several approximations (Lundberg& & Lee, 2017; Ancona et al., 2019; Jethani et al., 2022; Covert et al., 2023a). Fast- SHAP (Jethani et al., 2022), for instance, provides an effi- cient approximation using deep explainer model. Alterna- tively, INVASE (Yoon et al., 2019) and I.2X (Chen al., 2018) represent approaches that learn to identify relevant
3. Problem description
Let x E Rd represent a d-dimensional input feature vec- tor², and y E (1,2,... C} denote the associated target la- bel, where C' is the number of classes. Additionally, let M S [d] = (1...., d} be the subset of indices indicating the available features, and XM be the masked input vector with these available features. Each j-th feature has an associated cost Cj. and each input x is subject to a budget constraint k.
"Each feature can have different dimension size but ease of ex- position, in here have assumed each feature one dimensional.
3

---

Active Feature Acquisition Via Explainability- Ranking
The objective is to find a predictor Se. parameterized with 0. and a policy network 9r. parameterized with n, such that the following constraint objective is minimized:
aligns with that of (Valancius et al., 2024). The main dis- tinction is that a hard budget constraint is enforced here, whereas the oracle in (Valancius et al., 2024) incorporates feature costs into weighted objective function.
min Exyt Bas-,, (((Jo((M) s.t. Eosk. 0,x
Typically, methods in the literature (Yin et al., 2020; Covert et al., 2023b) assume that features have identical costs and that there is a fixed global budget l for all inputs. Given the available training samples {(x', y (') these methods aim to identify input-specific important features to acquire them sequentially in order of the most informative feature to the least one. To achieve this, they train Or through explo- ration using reinforcement learning (RL) (Yin et al., 2020) or cermmation-thhorette (Covert et al., 2023b) formulations, while simultaneously training the predictor network fo. In this work, we took a different approach by assuming access to feature importance rankings for each training sample. Consequently, instead of treating it as feature exploration problem, we addressed it as feature prediction problem (Figure 1).
JEM
Here, the first expectation is taken over the joint distribu- tion of x, y, and k. Including k in the joint distribution accounts for variability in the budget constraint across dif- ferent instances, reflecting scenarios where the maximum allowable feature acquisition cost is input- input-dependent. The subset M is generated sequentially by the policy network 9.- which determines the next missing feature to acquire, i.e.,
arg max 9- (xx) [d] .M. Subsequently, the predictor fo makes probabilistic predictions for any possible subset M, i.e. fx(xx) E [0, 1|cx1. For brevity, let the output of fo be denoted as y, i.e., y = (x( and the output of qa be denoted as q. i.c., q = (( (XAr)), where a is the softmax function. While softmax is also applied to the output of fe to obtain probabilistic predictions, we omit it from the notation for clarity.
4. Methodology and experimental settings
Oracle policy network. To establish theoretical bench- mark for feature acquisition, we define the oracle policy network q" The oracle policy is assumed to have perfect knowledge of the true importance ranking of features for each instance x. At each acquisition step, this policy se- lects the most important unobserved feature based on this ranking. The oracle policy network can be formally defined as follows: Let M* C [d] be the optimal subset of feature indices selected by the oracle policy, obtained by exhaus- tively evaluating all feasible subsets that satisfy the budget constraint specified in Equation 1. Once M is identified, the oracle constructs an internal ordering
Feature importance ranking. In our method, we assumed access to the feature rankings op for each training sam- ple x', sorted by their importance. While determining the importance of features for each input is challenging, we found that local explanation methods can effectively achieve this goal. We assumed that a model with reason- able task performance would naturally prioritize the most important instance-specifie features, which can be identi- fied using explanation methods. We empirically validated our assumption by demonstrating that when the policy net- work effectively selects features sequentially based on their ranking order during inference for each instance, the pre- dictor achieves superior average performance for given budget of available features, outperforming the current state-of-the-art methods.
M = (i1, ................... SIar-): where M E [d].
where features are sorted greedily in descending order of their importance to the prediction task. The oracle then acquires features sequentially according to this predefined order, selecting one feature at each step until all features in M* are obtained. Notably, the internal ordering does not affect the final outcome, as all features in M are eventually acquired. Hence, this approach differs from a sequential greedy policy, which selects features one at a time based on marginal gains at each step. In this formulation, the oracle policy 9² is not a learnable entity but theoretical construct. assumes perfect knowledge of the true im- portance ranking of features for each instance, which is typically unavailable in real-world scenarios. Essentially, this serves as the ideal baseline and sets an upper bound on the performance achievable by any practical policy network (i.e., 4.). By comparing 4.r against q", we can evaluate the effectiveness of the learned approach in approximating the optimal feature acquisition strategy. This oracle definition
To obtain the ranking order of the features for each training instance, first, we trained classifier using ((x,, , with the standard cross-entropy loss minimization. Then, we run an explanation method (SHAP (Lundberg & Lee, 2017), or LIME (Ribeiro et al., 2016) etc.) to get the feature importance ranking order y", where (1) is the feature index with the highest importance and (d) is the feature index with the least importance for the input x'. So our training set is {xt, v", +
Policy network. - Decision transformer. By approaching the problem as a conditional sequence modeling task. like the "decision transformer" (Chen et al., 2021), we trained 4r. which is a causal transformer model, with the objec- tive of next action/token prediction. We fed dr with se- quential data and sequence length (i.c., context length)
4

---

Active Feature Acquisition Via Explainability-drive Ranking
of l. At each timestep, there are three tokens includ- ing the input, the action and the reward as in Chen et al. (2021). During training, at timestep t, the input is xMa,' which is the i'th sample with many available features and M,- = (1)( spt (D))². Whereas, the action of is the most recently acquired feature index, i.e. of = g and the reward rt is the output of the predictor with the current input, i.c., rj- = y fo (x ). The rewards in RL-based methods (Kachuee et al., 2019; Li & Oliva, 2021) are typically functions of the predictor output; we followed a similar idea, but instead of defining specific function, we directly fed our policy network with the predictor output. So, for given sequence from the timestep t tot+2 + 8--1-11 1, the output of our 4r for the input i is: of = Zz (x),, af r() cct+e-1: 1.re--1-reeee) Fi.+e-1),, where t - indicates all the tokens from the time step to t f We used a mini version of GPT architecture (Radford, 2018) as transformer model. Please refer to the decision transformer paper (Chen et al., 2021) for more details.
arg max e and M, (2), (O). Then, we minimized the same losses in the first stage with the same strategy. In Lg. the index of the feature, which is not acquired yet and having the highest order among the features that are not acquired, is taken as the true label. For example, if yp (1) & ( (1), ('))),)((((()) then 4p (1) is taken as the true label; but if (1) is acquired and (2) is not acquired then (2) is taken as the true label, i.c. get (1) (f* (t)) and (2) & (1), By this second stage, we trained the predictor So to make its prediction with the subset of features M, acquired by OF Also, the policy network qur is trained to predict the feature with the highest ranking order among the features that are not acquired using the input with the imperfect subset of features M, This second stage helps both networks to perform better during inference, where the imperfect subset of features M, can only be used. Note that both the predictor and policy networks are dependent on each other. However, during training, we prevented the gradient flow from one network to another. Therefore, each network has its own independent loss function; because of the dependency, we trained them simultaneously. At = O, there is no feature acquired yet, i.e., Mn 0; so for all i, the outputs of qu are the same at 0. Consequently, at = 0, the same feature must be selected for acquisition across all inputs. In our approach, we initialized each input with the first feature that, on average, holds the highest importance ranking based on the training set. Detailed outlines of both training strategies are provided in the Appendix.
and Li+c-1 = g- (=
Training strategy. To train gr. we minimized the standard cross-entropy loss by considering the index of the next fea- ture that not acquired with the highest importance (i.e., pp'(t (t + 1)) as the true label with the minibatch setting. At each iteration, the loss function is:
No 4
1
Cg =
log(d' ((p))))
(2)
No
Implementation details. During training, we fixed the number of epochs to 200 and 16 for the first and second stage, respectively. We used Adam optimizer (Kingma & Ba, 2014) and a cosine scheduler (Loshchilov & Hutter, 2017). Before starting training, we pre-trained the predic- tor network, as done in (Covert et al., 2023b; Gadgil et al., 2024). We also employed a different augmentation strat- egy proposed in (Hoffer et al., 2020). In addition, as with other methods in the literature (Kachuee al., 2019; Covert et al., 2023b; Gadgil et al., 2024), we shared the backbone between fe and qs. We used this backbone in 4a to get the embedding of the input token. The embedding of action was extracted using a learnable embedding dictionary. For the reward embedding, we applied a linear layer followed by a non-linear activation to transform the output of fe into the embedding dimension. In qr, we set context length l to 4, number of heads and layers 4 and 3. respectively. We would like to clarify that we did not conduct an extensive parameter search in our experiments. The context length parameter was selected based on validation performance on the CIFAR-10 10 dataset, while the remaining parameters were chosen heuristically. These values were then held fixed across all experiments. For l selection, we evaluated the model's performance on the CIFAR-10 validation set,
where No is the batch size, ((+1) is the 't + 1)'th elc- ment of qi. and tù is randomly sampled integer determining the initial time step of sequence fed to the model for the i i'th sample. Simultaneously, we trained the predictor fo also by minimizing the standard cross-entropy loss:
N 1
1
1 No
=
log(Ff.y)
(3)
During the first stage of training, both fo and gx are fed by the input with the features that are acquired based on the ranking order provided by the local explanation method. However during inference, because fr not 100% accurate, the feature subset Mr, generated by yar may not always contain the top t features with the highest ranking order. To train both models to handle this new subset of features not encountered in the first stage, we introduced second stage of training. At the beginning of each iteration of the second stage, we first generated empirrcaal/predited feature acquisition ç order for each x", where (t((++1)
+ 1) =
'Fach sample has its own specific Mg, but we do not specify through superscript i if it is clear from the context. "htps://gituh.ccom com/karpathyliinGP
5

---

Active Feature Acquisition Via Explainability- Ranking
Spambase
CIFAR-10
CIFAR-100
Accuracy %
80
AUROO
Accuracy
10
20
features
of patches ImageNette
of patches
BloodMNIST
Metabric
100
100
100
Accurecy
Aczuany %.
Accuracy
80
80
40
402
20
40
20
patchos
of patches CTGS
of features CKD
80
0.9
0.9
Accurazy
70
AUROC
0.8
AUROO
0.8
0.7
0.7
0.6
0.6
50
0.5
0.5
10
15
20
10
15
20
of foatures
# of features
# of features
Empirical oracle + Our algorithm - ------------ DIME + GDFS - CAE +PL + Center-croppinnn -4 Random selection
Figure Model performance. We compared average classification performance our AFA method with other well-known methods across datasets with varying numbers of features. The results were averaged over three independent runs for the image datasets and nine independent runs for the tabular datasets. These datasets include five tabular datasets, four of which are medical: Spambase, Metabric, CPS, CTGS, and CKD, as well as four image datasets: CIFAR-10, CIFAR-100, BloodMNIST, and ImageNette. For improved readability, plots including standard deviations are provided in the Appendix.
averaging accuracies over the first 20 features. The result- ing mean accuracies were 78. 41%, 78.76%, 79.12%. and 78.12% for e = 1, 2, l 4, and 8, respectively. To maintain a consistent effective batch size at each iter- ation (see Equations 2 and 3), adjusted the batch size No inversely with f. Based on these results, selected f = 4 for all subsequent experiments. In the output of q, we subtracted large constant from the logits of already acquired feature indices before applying the softmax layer to prevent re-acquisition.
Table 1. Summary of datasets used in our experiments. For each dataset, we list the number features (d), classes (C), samples, and image size, along with the utilized patch size when applicable.
Dataset
d
C
Samples
Image size Patch size
Spambase CIFAR-10
57
2
4,601
64
10
60,000
32 4 x 4
CIFAR-100
64
100
60,000
32 x 32
4 x 4
BloodMNIST
196
8
17,092
28 28
Datasets. We conducted experiments on nine datasets (Ta- ble 1): five tabular datasets (Spambase, Metabric, CPS, CTGS, and CKD) and four image datasets (CIFAR-10, C00, BloodMNIST, and ImageNette). For image datasets, we partitioned each image into non-overlapping patches. For detailed descriptions the datasets, we refer readers to the Appendix (A.
ImageNette
196
10
13,395
224 x 224
16 x 16
Metabric CPS
489
6
1,898
8
3
418
CTGS CKD
23
2
2,139
50
2
1,659
Model architectures. To test the robustness of method across different architectures, varied predictor architec- tures. We employed ResNet50 (He et al., 201 6b) for Ima- geNette, ResNet18 (He et al., 201 for the CIFAR-10 and C1FAR-100 datasets, and a custom CNN for the BloodM- NIST dataset. The custom CNN has four convolution layers
with output channels 16, 32, 64, and 64, each followed by ReL.U activation and a max pooling layer. The convolu- tion layers are followed by flattening and linear layers for classification. For the Spambase dataset, we used a multi- layer perception (MLP) consisting of 2 hidden layers with 128 neurons, each followed by ReL.U and dropout layer.
6

---

Active Feature Acquisition Via Explainability-drive Ranking
Table Stage-wise classification performance. The table presents our model's performance after the first and second training stages, averaged over the first 20 features, on the Spambase, CIFAR- 10, CIFAR-I00 100. BloodMNIST, ImageNette, Metabric, CPS, CTGS and CKD datasets. Mean and standard deviation values were calculated across three independent runs for the image datasets and nine independent runs for the tabular datasets. Additionally, we included results from an extended first-stage training (250 epochs) to provide a more rigorous comparison with second-stage performance, illustrating the relative benefits of our two-stage approach versus simply extending training duration. For the Spambase, CTGS and CKD datasets, reported the area under the receiver operating characteristic curve values, while for the remaining datascts, we provided accuracy metrics. For the CPS dataset, the average was calculated over its 8 features, the dataset contains only 8 features.
Spam
CIFARIO
CIFARIOO
BloodMNISTT
ImageNette 10
Metabric
CPS
CTGS
CKD
# of classes:
100
8
First-stage (250) 0.952+.001
75.96 %6 45.91 96%6%6666 79. 79.83 19% 43.25 % 62. 62.52+1227 % 67. 2. 9 0.916+.0002 0.822+.b1
First-stage
0.9511.000 1.000 75.76.0.19 % 46.05.o. 25% 79. 19.25=0.11 9
73.7610.42
% 62. 28.48 % 67.21 10.1%% 0.9161.0000 0.825 000
Second-stage 0.955 15. 46.9920.15% 83.87 % 78.96 % 60.43 % 67.4540.11 1 0.916 0.836
On the medical tabular datasets, we utilized the same MLP architecture with 1024 hidden layer neurons on Metabric, 512 on CKD, 512 on CTGS and 128 on CPS.
Table 3. Model performance using various feature ranking ap- proaches. Comparison of classification performance across five tabular datases using feature rankings derived from INVASE, a learnable instance- feature selection method, and various local explanation methods: TreeSHAP (T-SHAP), LIME, KernelSHAP (K-SHAP), and IME (sampling). The performance metrics are the area under the receiver operating characteristic curve for the hinary- classification datasets and accuracy for the multi-class datasets. Mean and standard deviation values were caleulated across nine independent runs. Results on the Metabric dataset using INVASE omitted due to the high computational cost associated with this ranking method.
Feature importance ranking calculations. We employed various local explanation techniques tailored to the charac- teristics of image and tabular datasets. For image datasets, we employed FastSHAP (Jethani et al., 2022) to generate restance-speciffe feature ranking orders s for each input x², owing to its computational efficiency and ability to han- dle dynamic feature importance changes during training with data augmentations. For tabular datasets, we relied on free-based models, specifically CatBoost (Prokhorenkova et al., 2018), to derive feature ranking orders due to their strong performance on tabular data (Grinsztajn et al., 2022). Instance-specific rankings were computed using TreeSHAP (Lundberg et al., 2020), optimized for tree-based models, via the SHAP package². These rankings were used in our main experiments. Additionally, we explored alternative ranking methods, including INVASE (Yoon et al., 2019), and LIME (Ribeiro et al., 2016), as well two SHAP- based techniques, KernelSHAP (Lundberg & Lee, 2017) and IME (sampling) (Strumbelj & Kononenko, 2010), for tabular datasets.
Spam
Metabric
CPS
CTGS
CKD 2
# of classes:
T-SHAP LIME K-SHAP
0.96+0.001 69. 4..0.11. 91% 07. 07.5+0.13 % 0.92+
0.84+0.07
0.95.2.002
22..
67.120.58 % 0.91 0.82.0.09
0.9620.002
69. 69.610.33
% 67.320-50 58 0.92g0.001
0.83
0.95
69.8g..10%
67.120.61 S. 0.92 0.01 0.83
IN
0.93+0.002
68.4+0.24% R 0.91+0.008
0.83+0.00
employs a concrete selector layer for feature selection, grad- ually discretizing the selection process by lowering the tem- perature parameter. OPI. (Kachuee et al., 2019) is a RI- based method that employs deep Q-learning, using predic- tion uncertainty as the reward signal during training. Due to time constraints, we evaluated OPL. only on the tabu- lar datasets. The baseline methods, center-cropping and random selection, provide simpler comparisons: center- cropping selects patches from the center of the input, while random selection picks patches arbitrarily.
Comparison with the state-of-the-art methods. We evalu- ated our method against several existing approaches includ- ing DIME (discriminative mutual information estimation), GDFS (greedy dynamic feature selection), CAE (concrete autoencoder), OPL (RL-based method) and two baseline methods: center-cropping and random selection. DIME (Gadgil et al., 2024) prioritizes features by estimating their mutual information with the response variable in discrim- inative framework. GDFS (Covert et al., 2023b) uses a greedy strategy to select features based on their conditional mutual information, employing amortized optimization to approximate the greedy policy. CAE (Balın et al., 2019) is an unsupervised, end-to-end differentiable method that
We also evaluated an empirical oracle to approximate the optimal feature acquisition strategy. Given a predictor fe- the ideal feature subset MM, as determined by the oracle policy q", could theoretically be identified through ex- haustive combinatorial search for any given budget k. How- ever, performing such an optimization for every instance across multiple values of & is computationally impractical. Instead, we precomputed the optimal feature acquisition order for each instance using feature importance rankings
7

---

Active Feature Acquisition Via Explainablity-driien Ranking
Table Alignment between model's feature acquisition order and the feature importance rankings. This table presents the percentage overlap between the top T feature indices, ranked by their importance and those acquired by our method for T 10, 15, 20, 25, and on all datasets except CPS, as it only contains 8 features. Mean and standard deviation values were calculated across three independent runs for the image datasets and nine independent runs for the tabular datasets.
Spam 57
CIEAR-10 64
CIFAR-100 64
BloodMNIST 196
ImageNette 196
Metabric 489
CKD 50
CTGS 23
# of features (d):
Top 10 features 77 77.2611.06 % 36.22 %
47.29 12.25 %
40.75 40.75 %
11. 10 % 16.30 %
59.04 % 66.57 % 79.9 4
Top 15 features
82.15 45.83 %
47.94.2.12% 12.10%
47
61.5+1.08
% 69.6 % 91. 0 %
Top 20 features
87.3140.55% 52. 52.43 %
52.59.1.a5% 63.8511.65%
20.7440.07%% %
62.38 % 71.28 % 75.7=n.2
Top 25 features
87.6410.29% 68. 25 %
68.1011.14%
55.60=1.6s%
25.0610.06% %
62.59
% 24.21 %
N/A
Top 30 features
88.1510.17%
2.
10.83
57.82=1.46%
29.07
63.0510.44%
76.81
N/A
deviation bars are included in the Appendix (Figure A1), The reported standard deviations (in Figure A1, Table 2, and Table 3) confirm that our method is robust to both sources of variability. We found that initializing inputs with three features, rather than just one, improves training stability. Based on this, we fixed first three feature acquisition order and obtained the results shown in Figure 2. The second and third features were also selected based on their average importance rankings. Note that fixing the acquisition order for all d features is equivalent to using static global feature selection methods like CAE, which is suboptimal, as our empirical results demonstrate. Therefore, initializing with more than one feature can negatively impact the achiev- able upper bound on performance. However, we found that fixing the acquisition order for a few initial features helps stabilize training. Additionally, since our method relies on the feature ranking order, having a better ranking can lead to improved performance. Our approach can work with any ranking order, including those provided by humans, but we have shown that local model explanation algorithms are effective in providing this order.
Table 5. Nearest neighbor-based feature acquisition baseline. Performance of a baseline inspired by AACO, where features are acquired based on the importance ranking of the nearest training sample. Predictor networks from our method's second stage were reused without retraining.
Spam
Melabric
CPS
CTGS
CKD
classes:
Our method 0.96 001 69.8 41 67.5 13% 0.9220.001 0.84 07
NN
0.95 68.1 76% 67.2 % 0.91 #0.000 0.80
derived from explanation methods. At inference, all fea- tures were acquired in this predetermined order, ensuring that selections followed their instance-speciee importance rankings. Additionally, to standardize initialization, each instance began with the feature that exhibited the highest average ranking across the training set.
5. Results and discussion
Figure 2 demonstrates that our method shows superior, or comparable performance all the datasets. For example, on the ImageNette dataset, with the few number of patches, our method performs well, achieving 63.64% and 74.95% average accuracy with two and five available patches among 196 patches, respectively. Additionally, our model achieved an average AUROC score of 0.8465 on the CKD dataset with 10 features. Furthermore, the superior performance of the empirical oracle across all datasets highlights the rela- tive potential of our approach. To assess the robustness of our method to random weight initialization and the poten- tial variability of local explanation techniques, results were averaged over three independent runs for the image datasets and nine independent runs for the tabular datasets. Specifi- cally, for the tabular datasets, we trained three initial models with different random seeds, producing three distinct fea- ture ranking orders. For each ranking order, our method was trained three times to capture the variability of random network initialization. Due to computational constraints, we only varied the random initialization of the networks on the image datasets. To enhance readability, plots with standard
The average performance after both stages shown for all the datasets in Table 2, highlighting the benefit of the second stage. The second stage provides significant improvement most datasets, except for Spambase, CPS and CTGS that are relatively simpler compared to others, at least in terms of number of classes and features. Specifically, the second stage provides mean classification accuracy increase from 0.94% (on CIFAR-100) to 7.35% (on Metabric). Table 2 also includes results from an extended first-stage training (250 epochs) clarify that the performance gains observed in the second stage are not solely attributable to additional training epochs. The results of the first-stage training with 250 epochs are nearly identical to those with 200 epochs, demonstrating the effectiveness of the second-stage training. However, on BloodMNIST, the additional epochs provide meaningful performance increase. To further assess this, we conducted first-stage training with 300 epochs, yielding results of 79.73. 79.73 19 %, further reinforcing the effectiveness of the second-stage training.
8

---

Active Feature Acquisition Via Explainability-drivee Ranking
We also evaluated the robustness and effectiveness of our method across different feature ranking approaches, includ- ing INVASE (Yoon et al., 2019) and LIME (Ribeiro et al., 2016), as well as two SHAP-basee techniques for estimat- ing feature importance: KernelSHAP (Lundberg & Lee, 2017) and IME (sampling) (Strumbelj & Kononenko, 2010). These results (Table 3) indicate that while our method is robust to different ranking orders, its performance is also dependent on the quality of the ranking order generated by the explainability methods. These results were also aver- aged across nine independent runs, as in Table 2. To further verify the second point and test the dependency of the rank- ing orders' quality on the pre-trained model capacity, we conducted another ablation experiment on the CIFAR-I0 dataset. Specifically, we used ResNet-10, a smaller model compared to ResNet-18, as the pre-trained model for deter- mining the feature ranking order, while retaining ResNet-18 18 as the classification network. We observed that the perfor- mance of our method decreased from 78.44 44. to 78. 78.22% on the test set, and from 79.12% to 78.42% on the valida- tion set. These results confirm that the pre-trained model's capacity impacts the quality of feature ranking order and, consequently, the performance of our method. In addition, we evaluated the effecctiveness of the decision transformer as the policy network by comparing our method's ' performance with different architectures. When the decision transformer was replaced with a ResNet block, the model's 's accuracy de- creased from 78.44 78.44% to 76.83% on the CIFAR-10 dataset and from 46.99% to 46.70% on the CIFAR- dataset. Similarly, substituting the decision transformer with a CNN block reduced the model's accuracy from 83.87% to 78.23% on the BloodMNIST dataset. These results demonstrate the advantage of using a decision transformer as the policy network while highlighting that our method remains effec- tive with alternative architectures. We further evaluated a nearest- neighto (NN) based feature acquisition approach without any policy network, inspired by the AACO method (Valaneius et al., 2024). For each masked test instance, we first identified its nearest neighbor from the training set, then determined the next feature to acquire based on that neigh- bor's feature importance ranking. Specifically, we selected the highedt-ranked feature (according to the neighbor's 'ss rank- ing) that had not yet been acquired for the test instance. The results are presented in Table 5. For this bascline, we used the same predictor networks from our method's sec- ond stage. Additionally, nearest neighbors were identified using raw feature distances as in AACO, which may not be effective for image datasets.
all datasets. This trend indicates that our method's s feature acquisition order increasingly aligns with the feature im- portance rankings more features are considered. While the empirical oracle performances in Figure 2 demonstrate the practical benefits of using explanability- ranking orders in the AFA problem. Table 4 highlights the degree to which our method's acquisition strategy aligns with the fea- ture importance rankings. We also provide example patch acquisition trajectories for both our method and the empir- ical oracle on the ImageNette dataset, illustrated for four classes in the Appendix (Figure A2).
We emphasize the flexibility of our proposed method, which can operate with any given feature ordering. In the absence of definitive ground truth for feature importance rank- ings, we rely on explainability methods to generate these orderings. While such methods are valuable, they may not always yield optimal rankings in all settings (Kumar et al., 2020; Catav et al., 2021). As more accurate explanation techniques emerge, our approach can readily incorporate them to achieve further performance gains. Although our method demonstrates improved accuracy over baseline ap- proaches, it incurs higher training time. However, this can be mitigated through optimizations such as early stopping and mixed-preeee precision (half- prefispre training. In practical domains like medicine, AFA is often performed using pre- trained models customized to specific conditions. In such settings, explainability tools are commonly employed to support interpretability and foster clinical trust, which is an essential criterion for medical AI applications (Hill et al., 2025; Dai et al., 2024; Xue et al., 2024). Our framework is designed to take advantage of these precomputed feature importance rankings, removing the need to recompute ex- planations during training and thereby enabling efficient deployment. Finally, we acknowledge that our method as- sumes uniform feature acquisition costs. While this simpli- fies the framework, it may not reflect real-world conditions. Extending the method to incorporate non-uniform acquisi- tion costs remains important avenue for future research.
6. Conclusion
Our work introduces an active feature acquisition strategy by reframing it as feature prediction task, where the model learns acquire features based on explainability fea- ture importance rankings. Stage-wise results demonstrate that our two-stage training approach improves feature se- lection and classification performance on tabular and image datasets. The findings suggest that our method is robust across various models, datasets, and settings, and that it ex- hibits strong practical applicability in real-world scenarios, including domains such as medicine.
In Table 4, we present the overlap ratios between our method's s acquired feature order and the local explanation bechiiqee--based feature importance rankings across differ- ent datasets. As the number of top features (T) increases from 10 to 30, the percentage overlap generally rises for
9

---

Active Feature Acquisition Via Explainability- Ranking
Impact statement
approach for explaining data. In Meila, M. and Zhang, T. (eds.), Proceedings the 38th International Conference Machine Learning, volume 139 of o Proceedings of Machine Learning Research, pp. 1324-1335. PMLR, 18- 24 Jul 2021. URL https: htpp//oooednngssse ngs mmlr. sress/v133 hwtp:////aaaaaaaaaaaaa..
This paper presents work whose goal is to advance the field of machine learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
Acknowledgments
Chattopadhyay, A., Chan, K. H. R., Haeffele, D., Ge- man, D., and Vidal, R. Variational information pur- suit for interpretable predictions. In The Eleventh In- ternational Conference on Learning Representations, 2023. URL https:/ / http///pppeeeeeeeeeeeeeeeeeeeeep netereeeeeeeeeee/eeeeeee ntw.e/wueete/torr 2a=771sWa-Tm32
This project was supported by grants from the National In- stitute on Aging's Artificial Intelligence and Technology Collaboratories 1P30-AG073105), the American Heart As- sociation (20SFRN35460031) and the National Institutes of Health (R01-HL159620, R01-AG062109, RO1-AG083735, and R01-NS142076).
Chattopadhyay, A., Chan, K. H. R., and Vidal, R. Bootstrap- ping variational information pursuit with large language and vision models for interpretable image classification. In The Tweifth International Conference on Learning Representations, 2024. URL https:: /opereview. net feet/form?200 2- 9bmIbVaA2A.
References
Acevedo, A., Merino, A., Alférez, S., Ángel Molina, Boldú, L., and Rodellar, J. A dataset of microscopic peripheral blood cell images for development of automatic recogni- tion systems. Data in Brief, 30: 105474, 2020.
Chen, J., Song, L., Wainwright, M., and Jordan, M. Learn- ing explain: An cioommation-theorttti perspective on model interpretation. In Dy. J. and Krause, A. (eds.), Proceedings of the 35th International Conference on Ma- chine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 983-892. PMLR, 10-15 Jul 2018.
Ancona, M., Oztireli, C., and Gross, M. Explaining deep neural networks with polynomial time algorithm for shapley value approximation. In Chaudhuri. K. and Salakhutdinov, R. (eds.), Proceedings of the 36th Inter- national Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 272-281.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. Decision transformer: Reinforcement learning via sequence modeling. In Ranzato, M.. Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 15084-15097, Curran Associates, Inc., 2021. URL https: /estps::/prrocs httt:s::/pproceedingg. hewprrooceeeinnn...eii..ee..e.e. cc cc/paper_ htapperefiiieesppppr///////////
PMLR, 09-15 Jun 2019.
Balm, M. F., Abid, A., and Zou, Concrete autoencoders: Differentiable feature selection and reconstruction. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceed- ings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learn- ing Research, pp. 444-453. PMLR, 09-15 Jun 2019. URL https:tp/// //tppo/eeedngs.s../////////s..se mlr. wwllr.presee ptess//97/// balin1 9a. html.
pdf.
Beebe-Wang, N., Qiu, W. and Lee, S.-I. Explanation- guided dynamic feature selection for medical risk pre- diction. In ICMI. 3rd Workshop on Interpretable Machine Learning in Healthcare (IMLH), 2023. URL https: // //openreview wttiiee...eeeeooeee wwt. f i id=1 (eddddiite itffff53v..
Covert, I., Lundberg. S. M., and Lee, S.-I. Under- standing global feature contributions with additive importance measures. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin. H. (eds.), Advances in Neural Information Processing Systems, volume 33, pp. 17212 Curran Associates, Inc., cc httptt//ppcceeed// neurips.. e c75ffb0ccc86d5003322222222222222 pdf.
Bolón-Canedo. V., -etanze, A., Ferän-Ferrández, L.,and and Cancela, B. Feature Selection: From the Past to the Future, pp. 11-34. Springer International Publishing, Cham, 2022. ISBN 978-3-030-93052-33 doi: 10.1007/
2020. URL htwpt:
978-3-030-93052--3222 https:::: /ttpttt::://d...... http://oooooo..../////////////// 1007 00077/97-3--11 3- 030- -00---9052222-2
Covert, I. C., Kim, C., and Lee, S.-I. Learning to estimate shapley values with vision transformers. In The Eleventh International Conference on Learning Representations, 2023a. URL ht https: : /tlLhtttp////ppeeeeieee /pwpnrevieweeeeeeeeoeeeeee net forum?
Catav, A., Fu, B., Zoabi, Y. Meilik, A. L. W., Shomron, N., Ernst, J.. Sankararaman, S., and Biladd-Bacrrchh,, R.
Marginal contribution feature importance an axiomatic
1d=5ktFNZ_ppT..
10

---

Active Feature Acquisition Via Explainability- Ranking
Covert, I. C.,, Qiu, W., Lu, M., Kim, N. Y., White, N. J. and Lee, S.-I. Learning maximize mutual information for dynamic feature selection. In Krause, A., Brunskill, E., Cho, K., Engelhardt, B., Sabato, S., and Scarlett. J. (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 Proceedings of Machine Learning Research, pp. 6424-64477 PMLR, 23- 29 Jul 2023b.
Representations, 2024. URL https: /oppa//////eerieew.....o net/ http:/oormmidddo// peee//forummeidd/jjuuuuuu
Grinsztajn, L., Oyallon, E., and Varoquaux, G. Why do tree-based models still outperform deep learning on typ- ical tabular data? In Thirty-sixth Conference on Neu- ral Information Processing Systems Datasets and Bench- marks Track, 2022. URL https: oppnrevien..e net. forum? www.fforrrugggoo id==p7 _pQ22z.
Curtis, C., Shah, S. P., Chin, S.-F., Turashvili, G., Rueda, O. M., Dunning, M. J., Speed, D., Lynch, A. G., Sama- rajiwa, S., Yuan, Y., Gräf, S., Ha, G., Haffari, G., Bashashati, A., Russell, R., McKinney, S., Langerod, A., Green, A., Provenzano. E., Wishart, G., Pinder, S., Watson, P. Markowetz, F., Murphy, L., Ellis, I., Pu- rushotham, A., Børresen-Dale, A.-L., Brenton, J. D., Tavaré, S., Caldas, C., and Aparicio, S. The genomic and transcriptomic architecture of 2, 000 breast tumours re- veals novel subgroups. Nature, 486(7403):346 352, April URL
Hammer, S. M., Katzenstein, A., Hughes, M. D., Gun- dacker, H., Schooley, R. T., Haubrich, R. H., Henry, W. K., Lederman, M. M., Phair, J. P., Niu, M., Hirsch, M. S., and Merigan, T. C. A trial comparing nucleoside monotherapy with combination therapy in hiv-infected adults with cd4 cell counts from 200 to 500 per cubic millimeter. aids clinical trials group study 175 study team. The New England journal of medicine, 335 15:1081- 90, 1996. URL https: // /api.. semanticscoolar. org / Corpus ID: 40754467.
2012. ISSN 1476-4687, doi: 10. 00.1038/natrre1000 http: / dx. doi. org/ 10 1030010000000000001 /t/////aaueeee..../
He, H., Daumé III, H., and Eisner, J. Cost-sensitive dy- namic feature selection. In International Conference on Machine Learning (ICML) workshop on Inferning: Inter- actions between Inference and Learning, 2012a.
Dai, L., Sheng, B., Chen, T., Wu, Q.. Liu, R., Cai. C., Wu, L., Yang, D., Hamzah, H., Liu, Y., et al. A deep learning system for predicting time to progression of diabetic retinopathy. Nature Medicine, 30(2):584-594, 2024.
He, H., Eisner, J., and Daume, H. Imitation learning by coaching. In Pereira, F., Burges, C., Bottou, L., and Weinberger, K. (eds.), Advances in Neural Information Processing Systems, volume 25, Curran Associates, Inc., neurips..
Deng, J., Dong, W.. Socher, R., Li, L.-J,, Li, K., and Fei-Fei, L. Imagenet: A large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vi- sion and Pattern Recognition, pp. 248-255, 2009. doi: 10.1109/CVP.22009..02 5206648...
2012b. URL cc / 2df pdf.
Dickson, E., Grambsch, P., Fleming, T., Fisher, L.. and Langworthy, A. Cirrhosis Patient Survival Predic- tion. UCI Machine Learning Repository, 1989. DOI: https///ooo.rr/1024444////
He, H., Mineiro, P., and Karampatziakis, N. Active infor- mation acquisition. arXiv, 2016a.
He, K., Zhang, X., Ren, S., and Sun, J. Deep residual learn- ing for image recognition. Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2016b.
Dulac-Arold, G., Denoyer, L., Preux, P., and Gallinari, P. Datum-wise classification: A sequential approach to sparsity. In Gunopulos, D., Hofmann, T., Malerba, D., and Vazirgiannis, M. (eds.), Machine Learning and Knowledge Discovery in Databases, pp. 375-390, Berlin, Heidelberg, 2011. Springer Berlin Heidelberg. ISBN 978-3-642-23788-5..
He, W. and Chen, T. Scalable online disease diagnosis via multi-model--useee actor-critic reinforcement learning. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '22, pp. 4695-4703. Association for Computing Machinery, 2022. ISBN 978 9781450393850, doi: 00.1145/3534678334466777
Ekanayake, S. P. and Zois, D. Sequential acquisition of features and experts for datum-wise classification. In ICASSP 2024 2024 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp- 5225-5229, 2024. doi: 10. 10.119//CASSP44485522222 10447423.
He, W., Mao, X., Ma, C.,, Huang, Y., Hernandez-Lobato, J. M., and Chen, T. Bsoda: A bipartite scalable framework for online disease diagnosis. In Proceed- ings of the ACM Web Conference 2022, WWW '22, pp. 2511-2521, New York, NY, USA, 2022. Associa- tion for Computing Machinery. ISBN 9781450390963 doi: 10.1145/4885447.552222 URL https: : /doi. org/10. 3511221
Gadgil, S., Covert, I. C., and Lee, S.-I. Estimating condi- tional mutual information for dynamic feature selection. In The Twelfth International Conference on Learning

---

Active Feature Acquisition Via Explainabiltyyddrien Ranking
Hill, E. D., Kashyap, P., Raffanello, E., Wang, Y., Moffitt, T. E., Caspi, A., Engelhard, M., and Posner, J. Prediction of mental health risk in adolescents. Nature Medicine, pp. 1-7, 2025.
Krizhevsky, A. Learning multiple layers of fea- tures from tiny images. pp- 32-33, 2009. URL https://wwar cs tow.nto..dduu.. Learning-features--ae 200. pd..
Hoffer, E.,, Ben-Nun, T., Hubara, I., Giladi, N., Hoe- fler, T., and Soudry, D. Augment your batch: Improv- ing generalization through instance repetition. In 2020 IEEE/CVF Conference on Computer Vision and Pat- tem Recognition (CVPR), pp. 8126-8135, 2020. doi: 10.110//CVPR42260.22000..
Kumar, I. E., Venkatasuuraamanian. S., Scheidegger, C., and Friedler, S. Problems with shapley-value saglly-valuu-baseeee explana- tions as feature importance measures. In International conference on machine learning, pp. 5491-5500. PMLR, 2020.
Li, Y. and Oliva, J. Active feature acquisition with genera- tive surrogate models. In Meila, M. and Zhang, T. (eds.), Proceedings of the 38th International Conference on Ma- chine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 6450-6459. PMLR, 18-24 Jul 2021. URL https: httsss///occeeeeeeeeeeeee mlr. ptr.p press/ v139/1i21p html.
Hopkins, M.. Reeber, E., Forman, G., and Suermondt, J. Spambase. UCI Machine Learning Repository, 1999.
DOI: https://bbi.rrg//0.24432//33333
Howard, J. Imagenette: A smaller subset of 10 easily classified classes from imagenet, March 2019. URL https: //tthbb.t com/ wwwtaa/iiaaaaaaaa ittpa///mmmmmeettt...mmmmm
Janisch, J., Pevný, T., and Lisý, V. Classification with costly features using deep reinforcement learning. Proceed- ings of the AAAI Conference on Artificial Intelligence, 33(01))3959-3966 Jul. 2019, doi: 10.1609/aaai.v3331 33013959 URL https: //ojs. aaai. orr./r///////o index. php/ hhtp//AAT//arrticl///iee//
Loshchilov, I. and Hutter, F. SGDR: Stochastic gradient descent with warm restarts. In International Conference on Learning Representations, 2017.
Lundberg, S. M. and Lee, S.-I. A unified approach to inter- preting model predictions. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017.
Jethani, N., Sudarshan, M., Aphinyanaphongs, Y., and Ran- ganath, R. Have we learned explain?: How inter- pretability methods can learn to encode predictions in their interpretations. In Banerjee, A. and Fukumizu, K. (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pp. 1459- 1467. PMLR, 13-15 Apr 2021.
Lundberg, S. M., Erion, G., Chen, H., DeGrave, A., Prutkin, J. M., Nair, B., Katz, R., Himmelfarb, J., Bansal, N., and Lee, S.-I. From local explanations to global under- standing with explainable ai for trees. Nature machine intelligence, 2(1):56-67, 2020.
Jethani, N., Sudarshan, M., Covert, I. Lee, S.-I., and Ran- ganath, R. FastSHAP: Real-time shapley value estimation. In International Conference on Learning Representations, 2022. URL https: /ttp//peeeeiieeeeeeeeeeeeee/eeee n forum? id- (dd2qqG__am 2er22_2GG_VSSSS2 VTV53T.
Mesinovie, M., Watkinson, P., and Zhu, T. Explainable ai for clinical risk prediction: a survey of concepts, methods, and modalities. arXiv, 2023.
Kachuee, M., Goldstein, O., Kärkkäinen, K., and Sar- rafzadeh, M. Opportunistic learning: Budgeted cost-sensitive learning from data streams. In In- ternational Conference on Learning Representations, 2019. URL id=81eOHo
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidje- land, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. Nature, 518(e,408 529-533, 2015.
wwweeeeiiiieeeeeeeeeeeeeeeeeeeee forum?
Nan, F. and Saligrama, V. Adaptive classification for prediction under a budget. Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S.. and Gamett, R. (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL https::/proceedings. new.eeeidddnnnnnnnnneeeeeee cc cc/paper_filedee ////appe__flee//ppppr//21/////11 aBFEOFE00eacdacccc222278982 pdf.
Kharoua, R. E. Chronic kidney disease dataset, 2024. URL https:// /twp///wwww/.////////////// kaggle. www/ dsv 658224.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. arXiv, 2014.
Kingma, D. P. and Welling, M. Auto-encoding varia- tional bayes, 2013. URL https://arxiv..rrr abs/1312. 6114.
12

---

Active Feature Acquisition Via Explainability-driee Ranking
Pereira, B., Chin, S.-F. Rueda, O. M., Vollan, H.-K. M., Provenzano, E., Bardwell, H. A., Pugh, M., Jones, I.,, Russell, R., Sammut, S.-J., Tsui, D. W. Y., Liu, B., Dawson, S.-J., Abraham, J., Northen, H., Peden, J. F., Mukherjee, A., Turashvili, G., Green, A. R., McKinney, S., Oloumi, A., Shah, S., Rosenfeld, N., Murphy, L., Bent- ley, D. R., Ellis, O., Purushotham, A., Pinder, S. E., Børresen-Dafe Dale, A.-L., Earl, H. M., Pharoah, D., Ross, M. T., Aparicio, S., and Caldas, C. Erratum: The somatic mutation profiles of 2,433 breast cancers refine their ge- nomic and transcriptomic landscapes. Nat. Commun., 7 (1):11908, June 2016.
1953. ISBN 9781400881999 9781400881970-018. URL https: ptttss://oo... http://ooooooorr/ 1515/9781400881999 018.
doi: doi:
Shim, H., Hwang, S. J., and Yang. E. Joint active feature acquisition and classification with variable-size set encoding. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Process- ing Systems, volume 31. Curran Associates, Inc., 2018. URL https::/ neuriig e oo//ppeerfiilessppper///////// pdf.
Petsiuk, V., Das, A., and Saenko, K. Rise: Randomized in- put sampling for explanation of black-box models, 2018. URL hwtptttp//// /al:://aaaaa //aaaaiivv..oo abs/ 0742117........
Strumbelj, E. and Kononenko, I. An efficient expla- nation of individual classifications using game the- ory. Journal of Machine Learning Research, 11(1):1- 18, 2010. htthttp: httpt//jmmms httg/.org////ooo httpll/.rrgpapprs////oo papers/vl1/ s strumbel 110a html.
Prokhorenkova, L., Gusev, G., Vorobev, A., Dorogush, A. V., and Gulin, A. Catboost: unbiased boosting with categorical features. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL https://proceeddii neurips. cc cc/paper cc/paper_fileseeeeeeeeeeeeeeseee wttpee///ppee//////////////
Torralba, A., Fergus, R., and Freeman, W. T. 80 million tiny images: A large data set for nonparametric object and scene recognition. IEEE Transactions on Pattern Analysis and Machine Intelligence, 30(11):195-1197. 2008. doi: 10.1 10.1109/TPAM..000..
9-Paper.
pdf.
Trapeznikov, K. and Saligrama, V. Supervised sequen- tial classification under budget constraints. In Car- valho, C. M. and Ravikumar, P. (eds.), Proceedings of the Sixteenth International Conference on Artificial Intelligence and Statisties, volume 31 of Proceedings of Machine Learning Research, pp. 581-589, Scotts- dale, Arizona, USA, 29 Apr-01 May 2013. PMLR. URL WtLLttts://proceeiiiii mlr. press/v33 trapeznikov13a. html.
Radford, A. Improving language understanding by genera- tive pre-training. Preprint, 2018.
Rangrej, S. B. and Clark, J. J. A probabilistic hard attention model for sequentially observed scenes. In 32nd British Machine Vision Confer- ence 2021, BMVC 2021, Online, November 22-25, 2021, pp. 137. BMVA Press, 2021. URL https:
assets/papprs//D22
Valancius, M., Lennon, M., and Oliva, J. Acquisition con- ditioned oracle for nongreedy active feature acquisition. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarlett, J., and Berkenkamp. F. (eds.), Pro- ceedings of the 41st International Conference Ma- chine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 48957-48975. PMLR, 21 Jul 2024.
Ribeiro, M. T., Singh, S., and Guestrin, C. "why should trust you?": Explaining the predictions of any classifier. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Dis- covery and Data Mining, KDD '16, pp. 1135-1144, New York, NY, USA, 2016. Association for Comput- ing Machinery. ISBN 9781450342322.
doi: 10.1 2939672.2939778. URL htttt:::/ddo..orgg/// 1145/2939672.293877
ml wtt.gg.........ss/...... v235 htt///aaaaniussssa..... html.
van Hasselt, H., Guez, A., and Silver, D. Deep re- inforcement learning with double q-learning. Pro- ceedings of the AAAI Conference ligence, 30(1), Mar. 2016. doi: 10.1609/aai..v300 10295. URL bttps:://ojs aaai. org/ index. php/ AAAI httpA///aaaacclleeee/////
Samek, W., Montavon, G., Lapuschkin, S., Anders, C. J., and Müller, K.-R. Explaining deep neural networks and beyond: A review of methods and applications. Proceedings of the IEEE, 109(3):247-278, 2021. doi: 10.1 10.1100//ROC.22013330000
Artificial Intel-
Shapley, L. S. 17. A Value for n-Person Games, pp. 307-318. Princeton University Press, Princeton,
von Kleist, H., Zamanian, A., Shpitser, I., and Ahmidi, N. Evaluation of active feature acquisition methods for static
13

---

Active Feature Acquisition Via Explainabilit--rriee Ranking
feature settings, 2023. URL httpss://axxiiv.orgg abs/2312. 03619. 9.
Xue. C., Kowshik, S. S., Lteif, D., Puducheri, S., Jaso- danand, V. H., Zhou, T., Walia, S., Guney, O. B., Zhang, J. D., Pham, S. T., al. Ai-based differential diag- nosis of dementia etiologies multimodal data. Nature Medicine, 30(10);2977-2988 2024.
Yang. J., Shi, R., and Ni, Medmnist classification de- cathlon: A lightweight automl benchmark for medical image analysis. In IEEE 18th International Symposium Biomedical Imaging (ISBI), pp. 191-195, 2021.
Yang, J., Shi, R., Wei, D., Liu, Z.,, Zhao, L., Ke, B., Pfister, H., and Ni, B. Medmnist v2-a large-scale lightweight benchmark for 2d and 3d biomedical image classification. Scientific Data, 10(1):41, 2023.
Yin, H., Li, Y., Pan, S., Zhang, C., and Tschiatschek, S. Reinforcement learning with efficient active feature ac- quisition. In Learning Meets Combinatorial Algorithms at NeurlPS2020, 2020.
Yoon, J., Jordon, J., and van der Schaar, M. INVASE: instance-wise variable selection using neural networks. In International Conference on Learning Representations, 2019. URL https: http///pppeeeeeeeeeeeeeeeee http:/eentr/roreeeeoe/ao/aa id=BJg_roAck7.
Zubek, B. and Dietterich, T. G. Pruning improves heuris- tic search for cost-sensitive learning. In Proceedings of the Nineteenth International Conference Machine Learning, ICML '02, pp. 19-26, San Francisco, CA, USA. 2002. Morgan Kaufmann Publishers Inc. ISBN 1558608737.
14

---

Active Feature Acquisition Via Explainabiltty-drier Ranking
A. Appendix
A.I. Dataset descriptions
We utilized several datasets in our experiments (Table 1), including ImageNette, CIFAR-10, CIFAR-100, BloodMNIST, and Spambase. ImageNette (Howard, 2019) is a 10-class subset of the ImageNet dataset (Deng et al., 2009). CIFAR-10 and CIFAR-100 (Krizhevsky, 2009) are subsets of the 80 million tiny images dataset (Torralba al., 2008), containing 10 and 100 classes respectively. BloodMNIST (Acevedo al., 2020), derived from the MedMNIST dataset (Yang et al., 2021; 2023), comprises images of individual normal cells collected from individuals without infection, hematologic or oncologic diseases, and free of any pharmacologic treatment at the time of blood collection. The patch sizes are 16 x 16 for ImageNette (makes total of 196 patches, d = 196), 4 x 4 for the CIFAR-10 and CIFAR-100 datasets (d = 64), and 2 x 2 for the BloodMNIST dataset (d = 196). Spambase (Hopkins et al., 1999) is a well-known tabular dataset for classifying spam emails, consisting of 57 features derived from textual data. Additionally, to assess the applicability of our method in real-world scenarios, such as healthcare, we conducted experiments on four medical tabular datasets. As part of the preprocessing. removed ID columns and categorical columns that were not ranking- more than 10% missing values were also excluded, while the remaining missing values were imputed with the mean. In the following and in Table 1, the number of features refers to the count after preprocessing. The Metabric dataset (Curtis et al., 2012; Pereira et al., 2016) contains targeted gene sequencing data from 1,898 breast cancer samples, where we utilized mRNA-level Z-scores, which contains 489 features, to predict the Pam50 gene status that is a multi-class classification task. The cirrhosis patient survival (CPS) dataset (Dickson et al., 1989) includes records from 418 patients, primarily with primary biliary cirrhosis. along with 8 clinical features, with the task of predicting patient survival states categorized as death, censored, or censored due liver transplantation. The AIDS clinical trials group study 175 (CTGS) dataset (Hammer al., 1996) contains 2139 records of patients diagnosed with AIDS, 23 features, with binary classification task to predict whether a patient has died within a specified time period. Lastly, the chronic kidney disease (CKD) dataset (Kharoua, 2024) comprises 1659 patient records with 50 clinical features, and the task is to predict whether patient is diagnosed with chronic kidney disease in a binary classification setting.
binary. Columns with
A.2. Pseudocodes
Below, we provide the pseudocode for our first and second training stages, as well as for the inference stage.
Algorithm 1 Pseudocode for the first-stage training of qa and fa
Require: Training set [(x', y² {') batch size No. context length É, learning rate 7
Pre-train fo on {(x², 8² N using random feature selection
2: Initialize 7s
3: for each epoch do
4: for j = to [N/N,] do
5: Sample minibatch {x', v², N, (recalculate yp for each iteration if random augmentation is applied)
6: Sample random integer t; for each
7: Initialize losses: Lg = 0, L,- =0
8: for 1,- = to do
9:
Define f/ + tx Generate masked input: xiy; where My =1''(1),.... (t)}
10:
11:
Compute 81 r So (x())
12:
Compute q:. o (x M
of baa ) 1e6 x Me), where o/ = (t))
13:
Update Lj + Ls No Σ log(ys )
14:
Update C.y + fug No E log( (. (1(+1))
end for
16: Update parameters 0 0 VaL ग π 7V 7
17:
end for
18:
end for
15

---

Active Feature Acquisition Via Explainablity-ddiien Ranking
Algorithm 2 Pseudocode for the second-stage training of yz and So
Require: Training set {(x', y², s (²)) batch size No. context length E, learning rate 7, So and on from the first stage
1: for each epoch do
2: for j = 1 to [N/No] do
3: Sample minibatch {xx', v², (0"))) bim (recalculate yp for each iteration if random augmentation applied)
4: Generate & for each
5: Sample random integer tú for each i
6: Initialize losses: Lg = 0, Ly 0
7:
for t2 0 to & 1 do
8: Define 14 1 bxr Generate masked imput: xxa
9:
where M; (1"()).....
10:
Compute 814 rig fo (x Si
Compute d (0 (xx
M, auy a
Carty uay) 1e6 x Mg), where az = (1,)
rf
12:
Update Ly +- Ly N. En log( (
13:
Determine the true label for the Or network (denote this true label as V ). The true label is the index of the feature, which is not acquired yet and having the highest SHAP value among the features that are not acquired Lg c Lg No Σ No log( qi
end for
16: Update parameters 0 0 NL Nb
17: end for
18: end for
Algorithm 3 Pseudocode for the inference stage
Require: Input x, context length fe, Qr. and the indices of first three features for the initialization
1: Acquire the first three indices in their order for the first three steps
2: for = 3 to J: do
3: Compute r, f0((
4: Compute q, 0 (gxx (x M m (--------- 1-jr
r mmx((1, -4-++111 1e6 x M,
5: Compute (t() arg max 4r
6: Update the mask M, U (t)
7: end for
8: Predict ýx = fo(x
16

---

Active Feature Acquisition Via Explainablitt-driven Ranking
A.3. Additional results
We present the same performance results shown in Figure now including standard deviations in Figure Al. We also show example patch acquisition trajectories for both method and the empirical oracle in Figure A2. As in Figure A1, method exhibits relatively low standard deviation values, indicating robustness to fluctuations in explanation- based rankings and random weight initialization. Figure A2 further demonstrates that our method tends to acquire patches concentrated near the center of the image, focusing on informative regions aligned with the object's 's shape and structure. In contrast, the empirical oracle, leveraging perfect knowledge the image, selects seattered but informative regions based on object-specific cues, as expected.
Spambase
CIFAR-10
CIFAR-100
80
0.9
Accuracy %
80
Accuracy %
60
AUROC
0.8
40
0.7
20
0.6
40
0
10
20
30
10
20
0
10
20
# of features BloodMNIST
# of patches ImageNette
of patches Metabric
100
100
100
Accuracy %
Accuracy %
80
Accuracy %
80
50
60
60
40
40
20
20
10
20
30
10
15
0
10
20
# of patches CPS
# of patches CTGS
of features CKD
80
1
Accuracy %
70
AUROC
0.8
AUROC
0.8
60
0.6
0.6
50
40
0.4
0.4
0
2
10
20
0
10 of features
20
# of features
of features
- Empirical oracle
Our algorithm
DIME
GDFS
Center-cropping
Random selection
Figure AI. Performance comparison with standard deviations. Mean accuracy (or AUROC) and standard deviation values are shown for method and baselines across nine datasets. Results averaged three independent runs for image datasets and nine independent runs for tabular datasets. Our method demonstrates consistently strong performance with relatively low variance across settings. The empirical oracle, shown in black, represents idealized upper bound. Error bars represent standard deviations.
17

---

Active Feature Acquisition Via Explainabiiit-driven Ranking
patches
patches
patches
patches
patches
Our method
English Springer
Empirical oracle
Our method
Chain Saw
Empirical oracle
Our method
French Horn
Emplrical oracle
Our method
Garbage Truck
Emplrical oracle
Figure A2. Examples of feature/patch acquisition trajectories. Illustrative examples of patch acquisition sequences for four ImageNette classes: English springer, chain saw, frrench horn, and garbage truck. first column shows the original images. while the subsequent columns show the cumulative patches acquired steps 5, 10. 15, 20, and 25, For each example, the top row corresponds to the acquisition trajectory produced by our method, and the bottom row corresponds to the empirical oracle. method tends to acquire patches concentrated in structurally informative regions, whereas the empirical oracle, having full image access, selects scattered but highly discriminative patches.
18


Reviews:
[{"review_id": "TuzL158XqE", "reviewer": "Reviewer_4RZP", "summary": "This paper introduces a novel approach to active feature acquisition (AFA) by reframing the problem as a feature prediction task guided by explainability-driven rankings. Specifically, the authors leverage local explanation methods (e.g., SHAP, LIME, TreeSHAP) to generate instance-wise feature importance rankings. These rankings are used to supervise a decision transformer that learns a policy for acquiring the next most informative feature for each instance.\nEmpirical results across a diverse set of tabular and image datasets (including several medical datasets) show consistent improvements over state-of-the-art AFS methods. Notably, to my knowledge, this is the first known work to directly incorporate local explanation methods as a guiding signal for sequential feature acquisition, offering a more stable and interpretable alternative to traditional reinforcement learning or mutual information-based approaches.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Response to W1 & Q1: We appreciate the reviewer’s comments on the potential variability of post-hoc explanation methods and their impact on the robustness of our learned acquisition policy. To address these concerns, we conducted several robustness evaluations. As shown in Table 3, our method performs consistently across different explanation-based ranking approaches (e.g., SHAP, LIME), demonstrating resilience to variations in feature importance estimations. Additionally, we tested the impact of using a weaker model to compute rankings: when replacing ResNet-18 with ResNet-10 on CIFAR-10, performance declined only marginally (78.61\\% to 78.22\\%), suggesting our method is robust to model capacity differences in the ranking stage. As further suggested, we evaluated the stability of feature importance rankings across random seeds. Specifically, we trained three initial models using different random seeds, resulting in three distinct ranking orders for each explanation method in Table 3. For each ranking order, we trained our policy network three times to account for variability in our model’s initialization. The updated Table 3 now reports the mean and standard deviation across these nine runs, demonstrating stable performance and indicating that our method is robust to fluctuations in explanation-based rankings:\n\n# Updated Table 3\n|  | Spam | Metabric | CPS | CTGS | CKD |\n| --- | --- | --- | --- | --- | --- |\n| T-SHAP | 0.955 ± 0.001 | 69.83 ± 0.41% | 67.45 ± 0.13% | 0.916 ± 0.001 | 0.836 ± 0.07 |\n| LIME | 0.953 ± 0.002 | 69.15 ± 0.18% | 67.06 ± 0.36% | 0.913 ± 0.001 | 0.822 ± 0.09 |\n| K-SHAP | 0.956 ± 0.002 | 69.57 ± 0.33% | 67.32 ± 0.56% | 0.915 ± 0.001 | 0.831 ± 0.005 |\n| IME | 0.954 ± 0.001 | 69.78 ± 0.10% | 67.12 ± 0.61% | 0.916 ± 0.001 | 0.8258 ± 0.1 |\n| INVASE | 0.927 ± 0.002 | - | 68.37 ± 0.23% | 0.912± 0.003 | 0.8305 ± 0.09 |\n| OPL* | 0.889 ± 0.002 | 61.54 ± 0.23% | 63.45 ± 0.80% | 0.864± 0.005 | 0.7003 ± 0.04 |\n\n*We also evaluated an RL-based method (OPL [1]). Since we are unable to update our figures during the rebuttal period, we present this baseline result in table format.\n\nResponse to W2:  Our AFA framework fundamentally differs from end-to-end feature selection methods like INVASE. While INVASE integrates feature selection directly into model training via policy gradients, AFA decouples the feature importance estimation from policy learning. This has two advantages: (a) it enables interpretability aligned with widely used post-hoc explanation tools, and (b) it simplifies training by avoiding the high variance and complexity of reinforcement learning. \n\nResponse to Q2: Thank you also for suggesting the INVASE [2]  method . Our approach is compatible with any ranking order, including those produced by INVASE. Based on your suggestion, we evaluated our method using INVASE-derived rankings and included these results in Table 3 (excluding the Metabric dataset due to computational constraints). Running INVASE on Metabric proved infeasible within the rebuttal window due to its high computational cost.\n\nReferences:\n\n[1] Kachuee, M., Goldstein, O., Kärkkäinen, K., and Sarrafzadeh, M. Opportunistic learning: Budgeted cost-sensitive learning from data streams. In International Conference on Learning Representations, 2019.\n\n[2] Yoon, Jinsung, James Jordon, and Mihaela Van der Schaar. \"INVASE: Instance-wise variable selection using neural networks.\" International conference on learning representations. 2018.", "claims_and_evidence": "Yes. The claims made in the submission are supported by clear and convincing evidence. The experiments are thorough and diverse, covering nine datasets. Performance comparisons with state-of-the-art methods, ablations (e.g., different explanation techniques, architectures), and oracle benchmarks support the central claims.", "methods_and_evaluation_criteria": "Yes. The use of decision transformers, per-instance explainability-based supervision, and standard evaluation on public benchmarks are all appropriate for the AFS setting. No additional explanation necessary here.", "theoretical_claims": "N/A — the paper does not include novel theoretical derivations or formal proofs that require scrutiny.", "experimental_designs_or_analyses": "Yes, the experimental design is sound. The authors use well-established backbones like ResNet variants, a GPT-3 mini version for the transformer, and standard training protocols.", "relation_to_broader_scientific_literature": "The paper builds on two key streams of literature: (1) active feature acquisition, especially via reinforcement learning or mutual information-based methods; and (2) explainability methods like SHAP and LIME.\n\nThis work is distinct in its fusion of these two domains, using post-hoc explanation tools during training to stabilize and guide sequential feature acquisition. I am not aware of any prior attempts that combine local explanation techniques with sequential policy learning in this manner, making the contribution both novel and timely.", "essential_references_not_discussed": "One possibly relevant line of work the authors may consider discussing is INVASE (Yoon et al., ICLR 2018), which also performs instance-wise variable selection using a policy-based network but learns feature importance end-to-end rather than using post-hoc explanations. It would be helpful for the authors to clarify how their use of SHAP/LIME compares in terms of ranking stability and generalization performance.", "other_strengths_and_weaknesses": "Strengths:\n\n 1. Strong combination of explainability and AFA.\n 2. Effective use of decision transformers in a new setting.\n 3. Generalizable two-stage training framework.\n 4. Broad empirical evaluation on both image and medical datasets.\n 5. Well written paper.\n\nWeaknesses\n1. The paper provides a helpful analysis (e.g., Table 4) showing that the learned acquisition policy aligns well with the explainability-based feature rankings, which supports the overall training framework.\nThat said, an important open question is the stability of the explanation methods themselves. Post-hoc techniques like SHAP and LIME can produce variable feature rankings across different runs or model initializations, especially in high-dimensional settings. While the two-stage training strategy may help mitigate this, the paper does not directly evaluate how such variability might affect the robustness of the learned policy.\nIt would strengthen the work to include a ranking stability analysis (e.g., across seeds or models), and assess how sensitive the acquisition policy is to such fluctuations.\n2. Additionally, it would be helpful to contrast this approach with training-time feature selectors like INVASE (Yoon et al., 2018), which do not rely on post-hoc explainability and may produce more stable, model-aligned importance scores. A comparison or discussion here would strengthen the positioning and generalizability of the approach.", "other_comments_or_suggestions": "Line 312: \"due to to\" → Typo, please correct to \"due to\".", "questions_for_authors": "1. How robust is your method to noise or instability in the feature importance rankings generated by SHAP or LIME?\n\n    Would using a weaker model (or different initialization) to compute rankings degrade downstream acquisition policy quality significantly?\n\n2. How does your method compare to training-time feature selection methods like INVASE that don’t rely on post-hoc explainability?\n\n    Could this offer better stability or generalization?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "wP2WEIcoy2", "reviewer": "Reviewer_fG82", "summary": "The authors propose an active feature acquisition (AFA) framework that selects features based on their importance to each individual case. The method leverages local explanation techniques to generate instance-specific feature importance rankings. The authors reframe the AFA problem as a feature prediction task, introducing a policy network grounded in a decision transformer architecture. The authors conducted experiments on multiple datasets and demonstrate that their approach outperforms current state-of-the-art AFA methods in both predictive accuracy and feature acquisition efficiency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Response to W1: We agree that the performance of our AFA framework is influenced by the accuracy of local explanation methods. However, our experiments show that rankings generated by widely used explanation techniques (e.g., SHAP, LIME) consistently match or exceed the performance of rankings derived from existing AFA baselines. This suggests that current explanation methods already provide sufficiently reliable signals for guiding feature acquisition. Moreover, explainability research is advancing rapidly, and we expect that continued improvements in explanation quality will further strengthen the effectiveness and generalizability of our framework.\n\nResponse to W2: We acknowledge the reviewer’s concern regarding the computational cost of explanation methods like SHAP. However, recent advances, such as FastSHAP, have significantly improved efficiency through approximation strategies. Moreover, in practical settings such as medicine, AFA would typically be used with pre-existing, well-trained models tailored to specific conditions. In such scenarios, explanation methods are often already applied to ensure interpretability and clinical trust, which is an important requirement in medical AI [1, 2]. Our framework is designed to take advantage of these precomputed feature importance rankings, avoiding the need to rerun explanation algorithms during training. This makes our approach both practical and computationally efficient in real-world deployments.\n\nResponse to W3: As shown in our ablation experiments, our framework is flexible and can accommodate simpler policy models, albeit with some trade-off in performance. While a decision transformer requires more resources than other architectures that leverage for example, a ResNet block, the decision transformer consistently achieved superior results across tasks, motivating its inclusion as the default choice. To support ease of adoption and reproducibility, we will release all implementation details and code upon acceptance, enabling researchers and practitioners to readily integrate or adapt our method.\n\nResponse to Q1: While RL-based methods are theoretically capable of discovering optimal acquisition policies, in practice they often suffer from instability, high variance, and sensitivity to hyperparameters [3,4,5]. As discussed in [6], greedy-based methods, despite their lack of formal optimality guarantees, frequently outperform RL approaches due to their simplicity and stability during training. Moreover, under certain assumptions about the data distribution, greedy strategies can be provably near-optimal [7], making them a compelling practical alternative. We showed that our framework, which avoids high-variance RL training by using a supervised learning objective with a decision transformer, achieves strong empirical performance while benefiting from interpretability and ease of optimization. Additionally, we evaluated the performance of an RL-based method (OPL [8]) in our experiments (please see Table 3 in our response to Reviewer 1). Our experimental results further suggest that RL-based approaches struggle to match the stability and effectiveness of our method,\n\nReferences: \n\n[1] Hill, E.D., Kashyap, P., Raffanello, E. et al. Prediction of mental health risk in adolescents. Nat Med. 2025 Mar 5. doi: 10.1038/s41591-025-03560-7. \n\n[2] Dai, L., Sheng, B., Chen, T. et al. A deep learning system for predicting time to progression of diabetic retinopathy. Nat Med 30, 584–594 (2024).\n\n[3] Erion, G., Janizek, J. D., Hudelson, C., Utarnachitt, R. B., McCoy, A. M., Sayre, M. R., ... Lee, S. I. (2022). Coai: Cost-aware artificial intelligence for health care. Nature biomedical engineering, 6(12), 1384.\n\n[4] Chattopadhyay, A., Chan, K. H. R., Haeffele, B. D., Geman, D., Vidal, R. (2023). Variational information pursuit for interpretable predictions. arXiv preprint arXiv:2302.02876.\n\n[5] Covert, I. C., Qiu, W., Lu, M., Kim, N. Y., White, N. J., Lee, S. I. (2023, July). Learning to maximize mutual information for dynamic feature selection. In International Conference on Machine Learning (pp. 6424-6447). PMLR.\n\n[6] Gadgil, S., Covert, I., Lee, S. I. (2024). Estimating conditional mutual information for dynamic feature selection. International conference on learning representations..\n\n[7] Chen, Y., Hassani, S. H., Karbasi, A., Krause, A. (2015, June). Sequential information maximization: When is greedy near-optimal?. In Conference on Learning Theory (pp. 338-363). PMLR.\n\n[8] Kachuee, M., Goldstein, O., Kärkkäinen, K., and Sarrafzadeh, M. Opportunistic learning: Budgeted cost-sensitive learning from data streams. In International Conference on Learning Representations, 2019.", "claims_and_evidence": "Looks good to me.", "methods_and_evaluation_criteria": "Looks good to me.", "theoretical_claims": "No theoretical claims.", "experimental_designs_or_analyses": "Looks good to me.", "relation_to_broader_scientific_literature": "N/A", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths:\n1. The paper addresses an important problem in machine learning: efficient feature acquisition in scenarios where data collection is costly or time-consuming.\n2. The approach is intuitive and leverages explainability methods (SHAP, LIME) in a novel way to determine instance-specific feature importance rankings.\n3. The authors reframe the active feature acquisition (AFA) problem as a feature prediction task, which allows them to use a decision transformer architecture for the policy network.\n4. The proposed method outperforms state-of-the-art AFA techniques in both predictive accuracy and feature acquisition efficiency across multiple datasets.\n\nWeaknesses:\n1. The method relies on the accuracy and reliability of local explanation methods. If the explanations are not accurate, the performance of the AFA framework could be affected.\n2. The computational cost of generating feature importance rankings using methods like SHAP can be high, especially for large datasets or complex models.\n3. The complexity of the decision transformer architecture may require significant computational resources and expertise to implement and train.", "other_comments_or_suggestions": "A typo in Ln 291: ImageNette -> ImageNet", "questions_for_authors": "The authors acknowledge that RL-based methods are theoretically capable of finding the optimal policy, but their method outperforms them empirically. It would be good to have a deeper discussion of why this might be the case.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "sE2UJIy9ba", "reviewer": "Reviewer_7dkk", "summary": "This paper tackles the problem of active feature acquisition, a setting where all features might not be available at inference and the model needs to make accurate predictions while minimizing the number of features acquired. The authors propose a framework that dynamically selects instance-specific features based on their importance ranking obtained from local explanation methods. They reformulate the problem as a feature prediction task and introduce a policy network based on the decision-transformer architecture which sequentially predicts the next most-informative feature. This network is trained in a two-stage approach along with a predictor network which makes the predictions given a subset of features that have been acquired. Experiments across both image and tabular modalities demonstrate that this approach outperforms other recent methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Methods and evaluation criteria: The reviewer raises a thoughtful point about the role of the second-stage training. The first stage provides clean supervision through ground-truth feature importance rankings generated by explanation methods. However, during inference, the policy must operate on partially observed inputs and select features sequentially. This setting diverges from the idealized full-information assumption of the first stage. It also introduces two key challenges: (a) the transferability of supervision degrades with incomplete or noisy inputs, and (b) the policy must make decisions under uncertainty introduced by prior imperfect acquisitions. Our second-stage training directly addresses these issues by jointly training the policy and predictor under realistic, policy-driven acquisition settings. To further clarify this benefit, we added a comparison in Table 2 showing that extending first-stage training to 250 epochs yields minimal or no improvement, which underscores the utility of the second stage. Please see the updated Table 2 in our rebuttal to the fourth reviewer. The reviewer also notes that incorporating feature costs could further improve the method. We agree that incorporating feature costs is important for real-world applications, especially in domains like healthcare where test costs vary. While our current method assumes uniform costs, we view cost-aware acquisition as an important direction for future work.\n\nExperimental design or analyses: The reviewer suggests comparing training times due to the high cost of generating explanation-based rankings. Our current implementation trains 2-4x slower than baseline methods. However, we expect significant speed-ups with optimization (e.g., using PyTorch Lightning, early stopping). Additionally, as noted in our response to Reviewer 2, W2, our method can leverage precomputed rankings, which aligns with practical scenarios (e.g., clinical models) where such data already exist. We acknowledge the current limitation and will include a discussion in the final version.\n\nRelation to broader scientific literature: The reviewer would like a deeper discussion of the reliability and computational feasibility of explanation methods. We agree, and have discussed this in response to Reviewer 2 (W1 and W2). In short, while explanation methods have known limitations, our experiments show that the AFA policy is robust to variation across seeds and models. Practical approximations (e.g., FastSHAP) are also available for scaling.\n\nResponse to W1: Dependency on explanation methods may limit scalability due to computational cost. We addressed this in detail in our response to Reviewer 2, W2. To summarize: our approach can reuse precomputed rankings and has room for significant implementation optimizations. This trade-off will be discussed explicitly in the paper.\n\nResponse to W2: Feature costs are not currently incorporated. We acknowledge this limitation and see cost-aware acquisition as a valuable extension. We plan to explore this direction as part of future work.\n\nOther comments: If accepted, we will include visualizations of selected image patches. For tabular data, further work is needed to assess alignment with prior clinical knowledge. We will move the ResNet-10 and decision transformer ablation results to a separate table. Also, we will update the pseudocode and include inference pseudocode.\n\nResponse to Q1: Fixing the initial features stabilizes training by reducing uncertainty in early steps and simplifying the learning of conditional distributions. This is especially helpful when many features are available. Notably, this flexibility is unique to our method compared to other AFA methods.\n\nResponse to Q2: Our framework can be extended to support adaptive stopping criteria (e.g., based on prediction entropy). While not implemented here, we consider this a promising direction for future work.", "claims_and_evidence": "One of the main claims is that local explanation methods can provide a reliable signal for instance-specific feature importance and the experiment results do support that claim well, especially considering the oracle implementation which outperforms the others substantially across most datasets.", "methods_and_evaluation_criteria": "I am a bit skeptical about the two-stage training strategy. Even though the results indicate that the second stage improves performance, I would like more analysis on why the first-stage is not enough for the policy network to learn the top features. \n\nI would have liked to see feature costs getting incorporated into the training as well. Since the authors consider feature costs as a potential constraint, there could be a scenario where a feature is highly informative but has a high cost (in terms or time, money, etc.), so it might be better to acquire a less costly feature instead (e.g. lab test features are costlier than demographic features).", "theoretical_claims": "There are no theoretical claims provided other than maybe the intuition of the oracle policy $q^*$", "experimental_designs_or_analyses": "I think the experimental design makes sense. In terms of the analyses, I would have like a plot of the training times for the different methods since local explanation methods can be expensive to run (see weaknesses).", "relation_to_broader_scientific_literature": "There has been a lot of research work done in the field of feature selection, starting all the way from 1996 [1]. Static feature selection has also been an important subject for a long time [2, 3, 4]. For dynamic feature selection, the literature has focused on formulating it as an RL problem [5, 6] or using mutual information estimations [7, 8]. There are many more works in each of these sub-fields, which indicates that this problem is of great interest to the scientific community. \n \n[1] Geman, Donald, and Bruno Jedynak. \"An active testing model for tracking roads in satellite images.\" IEEE Transactions on Pattern Analysis and Machine Intelligence 18.1 (1996): 1-14\\\n[2] Isabelle Guyon and André Elisseeff. An introduction to variable and feature selection. Journal of Machine Learning Research, 3(Mar):1157–1182, 2003.\\\n[3] Jundong Li, Kewei Cheng, Suhang Wang, Fred Morstatter, Robert P Trevino, Jiliang Tang, and Huan Liu. Feature selection: A data perspective. ACM Computing Surveys (CSUR), 50(6):1–45, 2017.\\\n[4] Balın, Muhammed Fatih, Abubakar Abid, and James Zou. \"Concrete autoencoders: Differentiable feature selection and reconstruction.\" International conference on machine learning. PMLR, 2019.\\\n[5] Jaromír Janisch, Tomáš Pevny, and Viliam Lis y. Classification with costly features as a sequential decision-making problem. Machine Learning, 109:1587–1615, 2020.\\\n[6] Mohammad Kachuee, Orpaz Goldstein, Kimmo Kärkkäinen, Sajad Darabi, and Majid Sarrafzadeh. Opportunistic learning: Budgeted cost-sensitive learning from data streams. In International Conference on Learning Representations, 2018.\\\n[7] Aditya Chattopadhyay, Kwan Ho Ryan Chan, Benjamin D Haeffele, Donald Geman, and René Vidal. Variational information pursuit for interpretable predictions. arXiv preprint arXiv:2302.02876, 2023.\\\n[8] Yang Li and Junier Oliva. Active feature acquisition with generative surrogate models. In International Conference on Machine Learning, pages 6450–6459. PMLR, 2021.", "essential_references_not_discussed": "I think the authors cover most of the important works and references related to the topic of feature selection. I would have liked more discussion on the potential limitations of the local explanation methods (reliability, computational feasibility) , since they form the basis of the proposed framework.", "other_strengths_and_weaknesses": "Strengths:\n1. The paper tackles an important and well-studied task for active feature acquisition. As machine learning is getting used in real-world settings, this will be an important practical consideration to reduce costs and improve efficiency.\n2. The authors integrate different existing works like decision transformers and local explanation methods for this task, which is interesting.\n3. Experiments are performed across a variety of datasets and it looks like they achieve state-of-the-art performance on most of them.\n4. I like the way concept figure 1 is presented as it gives a good high-level idea of how the two-stage approach works.\n\nWeaknesses:\n1. The dependence on local explanation methods also means that the proposed framework also has the same limitations as these methods. Feature attributions methods like SHAP grow exponentially in computation with the number of features. This might reduce scalability of the proposed framework. The authors should show a comparison the total training time with the other methods as well. Since the gains are marginal for some of the datasets, it might be practically more feasible to use another method which is more efficient.\n2. From my understanding, feature costs are not considered during training which is also an essential factor to consider. I would like to see an extension of this method that incorporate feature costs in the training.", "other_comments_or_suggestions": "Suggestions:\n1. Provide visualizations of  the features obtained on the image datasets. Also, for the medical tabular datasets, showing that the features selected correspond with prior knowledge would also be useful. \n2. The results of the ablation experiments on page 8 (using ResNet-10 and replacing the decision transformers) should be shown in a different table\n3. In the training pseudo-code, show how $a^i_{t:t_i'}$ and $r^i_{t:t'_i}$ are computed. Please provide the pseudo-code for the inference stage as well", "questions_for_authors": "1. Initializing the input with three features during inference seems counterintuitive, is there an intuition of why it helps stabilize the training? \n2. Can this method be adapted so that a different number of features are acquired for different instances? One way could be to use a different stopping criteria for each instance (maybe using the prediction entropy).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "OMUq0lBMgZ", "reviewer": "Reviewer_pnLr", "summary": "**Edit:**  Thank you to the authors for their response. Since repeated experiments have been run, an RL baseline, and reproducibility details have been provided, I have raised my score from 2 (weak reject) to 3 (weak accept). There are remaining points which I have detailed in the Rebuttal Comment, if they are addressed adequately I will raise the score to 4 (accept).\n\n**Original Review:**\n\n\nThis paper is about Active Feature Acquisition (AFA) - the test time task of sequentially measuring features on an instance-wise basis to improve predictive performance. This is relevant to all applications where at test-time features are not all jointly available, for example, medicine where for each individual patient a doctor will make different measurements to make a diagnosis. Due to cost or time constraints the aim is to measure as few features as possible.\n\nThe paper proposes a novel method where a policy is trained to predict (for each instance) which feature will be best to acquire next, and it is trained as a classification task, where an oracle gives the best feature and this is used as a label using the cross entropy of the policy's prediction as the loss. The oracle is defined using feature importance ranking methods such as FastSHAP.\n\nThe training is split into two main phases, where the policy uses the oracle features, and the second phase it uses it's own acquisition as inputs (still using the oracle features as labels), so that it is able to train under the conditions which would be experienced during inference.\n\nThe experiments are carried out on multiple real world medical datasets and image datasets. Two SOTA AFA methods are used as baselines as well as additional baselines.\n\nAblations are carried out investigating the importance of the two stage training, as well as sensitivity to the feature importance ranking method used to generate labels.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Theoretical concern and Q5: We acknowledge that our original description of the oracle definition may have been misleading and we clarify the distinction here. Our oracle is not a purely greedy policy. Rather, it first identifies the optimal subset of features that minimizes prediction loss under a given acquisition budget. This is done by exhaustively evaluating all feasible subsets within the budget. Once the optimal subset is identified, features are ranked in a greedy order to define an acquisition trajectory. Importantly, the acquisition proceeds until all features in the subset are obtained. Thus, the internal ordering does not affect the final outcome. This approach is distinct from a sequential greedy policy that chooses features one step at a time based on marginal gains. Our formulation closely aligns with the Acquisition Conditioned Oracle proposed by Valancius et al. (2024), which also aims to account for joint informativeness. The main distinction is that our formulation imposes a hard budget constraint, whereas Valancius et al. incorporate feature costs into a weighted objective. We will reference this work in the paper and discuss the similarities and differences.\n\nResponse to experimental design concerns (Q1–Q4): To improve the robustness of our claims, we conducted additional experiments and updated our results accordingly: \n\n-- (Repeatability and variability): We now report results over three runs (over nine runs on tabular datasets, please see our response to W1 and Q1 of the first reviewer) with different random seeds. Table 2 has been updated to include mean and standard deviation across runs.\n\n-- (First versus second stage comparison): To clarify that performance gains in Table 2 are not due to additional training epochs alone, we extended the first-stage training from 200 to 250. We observed that performance did not significantly improve with longer training, while our second-stage approach yielded clear gains. This confirms that the second stage provides meaningful training benefits beyond simply more epochs. The results of the first-stage training with 250 epochs are nearly identical to those with 200 epochs, demonstrating the effectiveness of the second-stage training. However, on BloodMNIST*, the additional epochs provide a meaningful performance increase. To further assess this, we conducted first-stage training with 300 epochs, yielding results of $79.73 \\pm 0.19$, further reinforcing the effectiveness of the second-stage training. This also highlights the potential for performance improvements through dataset-specific hyperparameter tuning.\n\n-- (Reproducibility details): We used Adam optimizer with a learning rate of 1e-3 for all datasets, except for CIFAR100 and Imagenette, where we used 5e-4 due to larger model sizes. If our paper gets accepted, then we will release the complete code, including all training scripts and configuration details, to ensure reproducibility.\n\n-- (Baseline diversity – Q3): We agree that incorporating more diverse baselines would strengthen our comparisons. Generative AFA methods like ACF (Li and Oliva, 2020) were not included due to their high computational demands and limited scalability to tabular and high-dimensional image datasets. Moreover, several recent works (e.g., Gadgil et al., 2024) show that discriminative methods often outperform generative counterparts in practice. We appreciate the suggestion and plan to include an additional baseline in future work. Meanwhile, we have already evaluated an RL-based method. Please refer to the table in our response to Reviewer 1.\n\n-- (Hyperparameter tuning of baselines): For baselines with published results on datasets like CIFAR10 and Imagenette, we confirmed alignment with reported performances. For datasets without public results, we made reasonable efforts to tune hyperparameters. We will include detailed tuning procedures in the final version. \n\n|  | Spam | CIFAR10 | CIFAR100 | BloodMNIST | ImageNette | Metabric | CPS | CTGS | CKD |\n| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |\n| First-stage (250 epochs) | 0.952 ± 0.001 | 75.96 ± 0.16% | 45.91 ± 0.36% | 79.83 ± 0.19%* | 73.95 ± 0.25% | 62.52 ± 1.27% | 67.23 ± 0.48% | 0.9157 ± 0.0002 | 0.822 ± 0.01 |\n| First-stage | 0.951 ± 0.0002 | 75.76 ± 0.19% | 46.05 ± 0.25% | 79.25 ± 0.15% | 73.76 ± 0.42% | 62.48 ± 1.39% | 67.21 ± 0.15% | 0.9155 ± 0.0004 | 0.824 ± 0.008 |\n| Second-stage | 0.955 ± 0.0001 | 78.44 ± 0.15% | 46.99 ± 0.15% | 83.87 ± 1.05% | 78.96 ± 0.12% | 69.83 ± 0.41% | 67.45 ± 0.13% | 0.9164 ± 0.0001 | 0.836 ± 0.07 |\n\nResponse to missing impact statement - Q6: We have added the suggested statement.\n\nResponse to clarification on center cropping (Comment 2): For the details of center cropping, we kindly refer the reviewer to the papers of DIME and greedy methods.\n\nResponse to comment on Figure 2 (Q7): Please see our response to Q1 of Reviewer 3.\n\nWe also included new references suggested by the reviewer.", "claims_and_evidence": "- The variety of datasets is very good to see.\n- There are no uncertainties provided in the results. Have there not been repeats of the experiments over different seeds? If not these are required to improve the reliability of the results provided.\n- DIME and GDFS are good baselines to include. However, these are greedy methods, showing this method also beats an RL baseline will make the claims stronger. Similarly, beating a generative model for CMI maximization will also make the claim stronger.\n- At the start of Section 5 it is stated that the first three features are fixed to obtain the results in Figure 2. So why do the baselines perform poorly sometimes with 3 or fewer features and not the same as the oracle or the proposed model?", "methods_and_evaluation_criteria": "The method makes sense for the problem at hand. Its very interesting to see a model trained to predict the feature to select as a classification problem. Using the feature importance rankings is also an innovative way to generate the labels.", "theoretical_claims": "There is no theory provided in this paper - which is not a strength or a weakness, just a statement that there are no theoretical claims therefore no errors.\n\nHowever, there is an issue with how the oracle is defined on page 3. Lines 142 - 164 Right. The claim is that at each step the oracle policy selects the most important feature that maximally improves the predictive performance. This is a greedy oracle, and therefore is not guaranteed to be optimal. This will not select features that are jointly informative but individually non-informative (see https://arxiv.org/abs/2302.13960 for example). \n\nFor example we can have features that are only informative if jointly known, and many other very noisy features. The proposed oracle will never select the jointly informative features before the others (unless it already has measured one of them), because a noisy feature will always provide some more **immediate** information, but less information in the long-term. Therefore this greedy oracle is not guaranteed to be optimal, especially with jointly predictive features.\n\nHow does this affect the definition of the oracle? How does it affect how the feature importances are ranked? If a dataset with this property is used, how would the model perform? How does it affect the conclusions drawn from the result in Table 4 if we can't guarantee the oracle is optimal? All greedy methods should \"fail\" on a dataset like this. The empirical results presented are positive, which would imply this property is not present in those datasets. Another experiment or at least a discussion would greatly improve the paper.\n\nAs a starting point:\nStrong performance is also seen in greedy methods like DIME and GDFS in the literature, additionally theory does exist that supports greedy methods being **near-optimal** (https://proceedings.mlr.press/v40/Chen15b.html).", "experimental_designs_or_analyses": "The main issue with this paper is that without theory, the empirical results must be very convincing, and currently they can be improved (the model motivation is clear so this is not a problem).\n\n- Without repeats we can't really know how sensitive this model is to initialization\n- There is no code provided which is not a problem in itself, but the reproducibility details are limited. There are no optimizer parameters for example. \n- There has been limited hyperparameter tuning. This is framed as an advantage - the proposed model performs well with minimal tuning. However, there has also been minimal tuning on the baselines, so we don't know if they have just got poor architecture hyperparameters or optimizer hyperparameters.\n- The result in Table 2 is only meaningful if the first stage of training was continued to acquire its result. As it is currently, there are results from the first stage (200 epochs) then for the second stage (200 epochs of first stage + 16 epochs of second stage). So it is not clear if the result from the second stage is due to training for longer or because it is a useful training change.\n- The dataset selection is good, however the baseline selection is limited. Two AFA methods are used as baselines - GDFS and DIME, and these work in very similar ways, in terms of being greedy methods that are trained by simulating the acquisition process. The results would be improved by including an RL baseline and a Generative Model baseline.", "relation_to_broader_scientific_literature": "The paper is appropriately framed within the AFA literature. Missing references are given next.", "essential_references_not_discussed": "- Partial VAE, Ma et al. 2019 (https://arxiv.org/abs/1809.11142)\n- Dynamic Feature Acquisition with Arbitrary Conditional Flows, Li and Oliva 2020 (https://arxiv.org/abs/2006.07701)\n- Acquisition Conditioned Oracle for Nongreedy Active Feature Acquisition, Valancius et al. 2024 (https://arxiv.org/abs/2302.13960)\n- Sequential Information Maximization: When is Greedy Near-optimal?, Chen et al. 2015 (https://proceedings.mlr.press/v40/Chen15b.html)", "other_strengths_and_weaknesses": "Stengths:\n\n- The idea the paper introduces is a very good one.\n- The paper generally is written nicely.\n- The method is clear, with a good use of diagrams and pseudo-code\n\nWeaknesses:\n\n- There is no impact statement. The paper does not have any ethical issues, so it is not majorly needed. However, it was written as a requirement in the Call for Papers (https://icml.cc/Conferences/2025/CallForPapers): \"Authors are required to include a statement of the potential broader impact of their work, including its ethical aspects and future societal consequences.\"", "other_comments_or_suggestions": "- Line 176 Left: \"these methods aim identifying\" should be \"these methods aim to identify\"\n- How does center cropping keep working after selecting the center? Does it spiral out? It is not clear how it continues to select patches of images? Is it a global fixed ordering as well?", "questions_for_authors": "I like the idea this paper proposes and it's interesting to see it working. I would like to recommend accept, however, currently I do not find the empirical evaluation strong enough for acceptance. I understand I may be wrong, so following the responses and reading other reviewers' thoughts I will re-evaluate. In particular these are the following concerns:\n\n- No repeats of experiments\n- Limited information on reproducing the experiments (no learning rates for example)\n- Lack of RL and generative model baseline\n- Limited hyperparameter tuning of baselines\n- Understanding of the oracle, and how it would work with jointly predictive features. How this might affect feature importances and therefore the training of the proposed model. This can either be with another small synthetic experiment, or with theoretical justification. How does the proposed oracle from this work compare to the oracle described in Acquisition Conditioned Oracle for Nongreedy Active Feature Acquisition, Valancius et al. 2024 (https://arxiv.org/abs/2302.13960), which does consider jointly informative features?\n- Addressing the lack of impact statement, it is enough to use the proposed statement from the call for papers: \"This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a novel approach to active feature acquisition by leveraging explainability methods to guide a decision transformer.
The reviewers agreed on the strong empirical results and the intuitive nature of the method, and its potential to outperform existing techniques.
Decision: Accept (poster)