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.

Response to experimental design concerns (Q1–Q4): To improve the robustness of our claims, we conducted additional experiments and updated our results accordingly:

-- (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.

-- (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.

-- (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.

-- (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.

-- (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.

Response to missing impact statement - Q6: We have added the suggested statement.

Response 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.

Response to comment on Figure 2 (Q7): Please see our response to Q1 of Reviewer 3.

We also included new references suggested by the reviewer.

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.

Experimental 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.

Relation 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.

Response 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.

Response 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.

Other 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.

Response 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.

Response 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.

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.

Response 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.

Response 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.

Response 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,

References:

[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.

[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).

[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.

[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.

[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.

[6] Gadgil, S., Covert, I., Lee, S. I. (2024). Estimating conditional mutual information for dynamic feature selection. International conference on learning representations..

[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.

[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.

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:

Updated Table 3

*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.

Response 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.

Response 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.

References:

[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.

[2] Yoon, Jinsung, James Jordon, and Mihaela Van der Schaar. "INVASE: Instance-wise variable selection using neural networks." International conference on learning representations. 2018.