AFABench

A benchmark for Active Feature Acquisition (AFA) methods

Compare state-of-the-art algorithms for sequential feature selection in scenarios where acquiring features is costly. Includes implementations of multiple AFA methods, standardized datasets, and automated evaluation pipelines.

Installation

uv is the only external dependency.

# Clone repository
git clone https://github.com/Linusaronsson/AFA-Benchmark.git
cd AFA-Benchmark

# Install dependencies with uv
uv sync

Quickstart

Local execution is not recommended for reproducing the full benchmark because the pipeline generates a large number of jobs. If you still want to run it locally with 8 CPU cores, execute this command from the repo root:

uv run snakemake \
    -s extra/workflow/snakefiles/orchestration/pipeline.smk \
    all \
    --configfile \
      extra/workflow/conf/eval_hard_budgets/all.yaml \
      extra/workflow/conf/methods/all.yaml \
      extra/workflow/conf/method_sets/all.yaml \
      extra/workflow/conf/method_options/all.yaml \
      extra/workflow/conf/pretrain_mappings/all.yaml \
      extra/workflow/conf/soft_budget_params/all.yaml \
      extra/workflow/conf/unmaskers/all.yaml \
      extra/workflow/conf/classifier_names/all.yaml \
      extra/workflow/conf/datasets/all.yaml \
    --config \
      device=cpu \
    --jobs 8

To reproduce the full benchmark results, use SLURM instead. See the reproducing full results tutorial for the exact commands.

Features

• Accessible configuration using hydra
• Reproducible pipeline using snakemake.
• Modular design: rerun specific parts of the pipeline as needed.
• Extensible framework: add custom datasets and AFA methods.

Limitations

• Supports only classification tasks; regression tasks are not yet implemented.

What is Active Feature Acquisition?

Active Feature Acquisition (AFA) addresses scenarios where,

• Features are expensive to obtain (medical tests, surveys, sensors),
• Real-time decisions must be made with partial information,
• Budget constraints limit which features you can acquire.

Example: Medical diagnosis where each test costs money and time. AFA methods intelligently decide which tests to order next based on previous results, aiming for accurate diagnosis with minimal cost. See the following survey for details: AFA Survey.

Implemented Methods

Method	Paper	Strategy	Greedy?
EDDI	link	Generative estimation of CMI	Greedy
GDFS	link	Discriminative estimation of CMI	Greedy
DIME	link	Discriminative estimation of CMI	Greedy
JAFA	link	Model-free RL	Non-greedy
OL	link	Model-free RL	Non-greedy
ODIN	link	Model-free RL	Non-greedy
ODIN	link	Model-based RL	Non-greedy
AACO	link	Oracle-based	Non-greedy
PT	link	Global feature importance	N/A
CAE	link	Global feature importance	N/A

Datasets

Dataset	Type	Modality	Train Size	Val Size	Test Size	# Features	# Groups	# Classes
CUBE	Synthetic	Tabular	600	200	200	20		20
CUBE-NM	Synthetic	Tabular	600	200	200	55	51	8
MNIST	Real World	Image (tabularized)	36,000	12,000	12,000	784	784	10
FashionMNIST	Real World	Image (tabularized)	36,000	12,000	12,000	784	784	10
Diabetes	Real World	Tabular	55,237	18,412	18,413	45	45	3
PhysioNet	Real World	Tabular	7,200	2,400	2,400	41	41	2
MiniBooNE	Real World	Tabular	78,038	26,012	26,014	50	50	2
ACTG175	Real World	Tabular	1,283	427	429	23	23
CKD	Real World	Tabular	240	80	80	24	24	2
BankMarketing	Real World	Tabular	27,126	9,042	9,043	16	16	2
Imagenette	Real World	Image	5,681	3,788	3,925	150,528	196	10

Project structure

• afabench: Main package.
• docs: Documentation.
• extra: Saved methods, data, logs and so on, non-source code files.
  • conf: This is where all the script configuration files are. Each configuration file corresponds to a class in config_classes.py.
  • data: Where miscellaneous files for datasets (e.g., CSVs, custom generated costs, etc.) are stored.
  • workflow: Snakemake workflows for running the full pipeline.
  • output: Folder where outputs from the snakemake pipeline are stored.
• scripts/: Folder of scripts, many of which are called from the snakemake pipeline.
• test: Tests.
  • src: Tests related to library code in afabench.
  • scripts: Tests related to specific scripts in scripts.

Tutorials

Learn more in our tutorials:

• Reproducing full results
• Pipeline explanation
• Adding a new dataset
• Adding a new method

Development

We encourage researchers to fork this repository and implement their own methods. Take a look at the tutorials to get started.

To follow repo conventions, run

uv run just qa

which will

• format code with ruff format
• do linting and formatting with ruff check --fix
• type checking with basedpyright --warnings
• run tests with pytest

Citation

If you use this benchmark in your research, please cite,

@inproceedings{Schütz2026,
    author = {Sch{\"u}tz, Valter and Wu, Han and Rezvan, Reza and Aronsson,
              Linus and Haghir Chehreghani, Morteza},
    title = {AFABench: A Generic Framework for Benchmarking Active Feature
             Acquisition},
    year = {2026},
    booktitle = {Proceedings of the 32nd ACM SIGKDD Conference on Knowledge
                 Discovery and Data Mining},
    url = {https://doi.org/10.1145/3770855.3817493},
}