Active feature acquisition via explainability-driven ranking

Introduction

In many real-world scenarios—most notably in medicine—acquiring a complete set of data for a machine learning model is expensive, time-consuming, or even risky. For instance, a doctor wouldn't order every possible lab test for a patient if a few key metrics could provide a confident diagnosis.

Traditional "static" feature selection picks a fixed subset of features for every instance, which is suboptimal because the most important feature for one patient might be irrelevant for another. Active Feature Acquisition (AFA) aims to solve this by choosing features dynamically, one by one. However, existing AFA methods often rely on Reinforcement Learning (RL), which can be notoriously difficult to train, or complex information-theoretic calculations that struggle to scale.

This paper introduces a novel AFA framework that reframes the problem from a complex RL "exploration" task into a straightforward "feature prediction" task. By using local explainability methods (like SHAP) to rank feature importance for individual instances, the authors train a Decision Transformer to sequentially acquire the most informative features, significantly outperforming current state-of-the-art methods.

Methodology

The authors propose a two-stage framework that treats feature acquisition as a sequential prediction problem:

1. Generating the "Teacher" Rankings: First, a standard classifier is trained on the full dataset. Local explanation techniques (like SHAP or LIME) are then used to generate an importance ranking of features for every single instance in the training set. This creates a "ground truth" sequence of which features should have been acquired first for each specific case.
2. The Policy Network (Decision Transformer): Instead of using RL to "discover" which features are good through trial and error, the framework employs a Decision Transformer architecture. This network is trained to predict the next best feature to acquire based on the current "state" (the set of already acquired features and their values).
3. Two-Stage Training:
   • Stage 1: The model learns to mimic the ideal rankings provided by the explainer.
4. Stage 2: The model is fine-tuned to handle its own potential mistakes. Because the policy won't be 100% accurate during inference, Stage 2 trains the predictor and policy networks to remain robust even when the acquired feature subset is slightly "off-order" compared to the ideal ranking.

Expand

Source, p. 3

Key Insights

• Explainability as a Supervisor: The core breakthrough is realizing that local explainers—typically used for post-hoc analysis—can serve as an excellent "oracle" for training a policy. This bypasses the need for the "exploration" phase in RL, which is often a major source of training instability.
• Reframing AFA as Sequence Modeling: By using a Decision Transformer, the authors treat feature acquisition like a language modeling task. The "context" is the set of features already known, and the "next token" is the next feature to acquire.
• The "Oracle" Gap: The paper demonstrates that an "ideal" policy (one that knows the SHAP values perfectly) significantly outperforms all existing state-of-the-art methods. The learned policy in this framework gets remarkably close to this theoretical upper bound.

Experiments and Results

The researchers evaluated their approach on a diverse mix of nine datasets, spanning both tabular data (Spambase, Metabric, CPS, CTGS, CKD) and image data (CIFAR-10, CIFAR-100, BloodMNIST, ImageNette). For images, "features" were defined as non-overlapping patches of the image.

The proposed method consistently outperformed state-of-the-art baselines like DIME, GDFS, and OPL across several dimensions:

• Accuracy vs. Budget: Across all datasets, the model achieved higher accuracy (or AUROC) with fewer features compared to baselines. For example, on ImageNette, the model reached high accuracy with just 5-10 patches out of 196.
• Efficiency Gains: On the Metabric breast cancer dataset, the second-stage training improved classification accuracy by 7.35% over the first stage alone.
• Robustness: The method proved robust across different architectures (ResNet18, ResNet50, MLPs) and different underlying explainers (SHAP, LIME).
• Structural Focus in Images: Qualitative results showed that the model naturally learns to acquire patches concentrated on the object's core structure (e.g., the head of a dog) rather than random background pixels.

Expand

Source, p. 6

Limitations and Future Work

While highly effective, the framework has a few notable limitations:

• Uniform Costs: Currently, the model assumes every feature costs the same to acquire. In reality, an MRI costs much more than a routine blood pressure reading. Incorporating non-uniform acquisition costs is a primary avenue for future research.
• Training Overhead: While inference is fast, the training process requires pre-calculating explanations for the entire training set, which can be computationally expensive for very large datasets.
• Explainer Dependency: The policy's performance is ultimately capped by the quality of the local explainer used to generate the initial rankings. If the explainer is biased or inaccurate, the policy will inherit those flaws.

@inproceedings{guney2025active,
  author = {Osman Berke Guney and Ketan Suhaas Saichandran and Karim Elzokm and Ziming Zhang and Vijaya B Kolachalama},
  title = {{Active feature acquisition via explainability-driven ranking}},
  booktitle = {International Conference on Machine Learning},
  year = {2025},
  url = {https://openreview.net/forum?id=J8YRdm39jn},
}