AFABench: Active Feature Acquisition Benchmark

• AFABench is a framework for benchmarking active feature acquisition, where models select a subset of features dynamically at test time under a hard budget.
• It features a modular architecture with standardized data loaders, policy interfaces, and evaluation protocols supporting static, greedy, RL-based, and oracle methods.
• Empirical findings reveal trade-offs among predictive accuracy, computational cost, and lookahead efficacy across synthetic and real-world datasets.

AFABench is a generic framework for benchmarking Active Feature Acquisition (AFA), introduced to enable fair, reproducible, and extensible comparison of methods that dynamically select which features to acquire at test time under a hard budget and uniform per-feature cost (Schütz et al., 20 Aug 2025). In AFA, the learner trades predictive performance against acquisition cost by querying only a subset of informative features for each instance, a setting motivated by monetary cost, latency, and privacy constraints. AFABench is described as the first benchmark framework for AFA; it combines a modular software architecture, synthetic and real-world datasets, a standardized evaluation protocol, and implementations spanning static, greedy, reinforcement learning-based, and oracle-based acquisition policies.

1. Problem formulation

Active Feature Acquisition studies the setting in which a model cannot assume full observation of every test instance. The formal setup begins from a fully observed training set

$$\mathcal{D}=\{(x_i,y_i)\}_{i=1}^N,\quad x_i\in\mathbb{R}^d,\; y_i\in\mathcal{Y},$$

and seeks both a predictor $f:\mathbb{R}^b\to\mathcal{Y}$, able to handle arbitrary subsets of $b$ features, and an acquisition policy $\pi$ that, given previously observed feature indices $S\subset [d]$, selects the next feature $i\notin S$ (Schütz et al., 20 Aug 2025). With a hard budget $b$ and uniform per-feature cost, the joint objective is to minimize expected prediction loss after sequential acquisition: $$\min_{f,\pi}\;\mathbb{E}_{(x,y)\sim p}\bigl[\ell\bigl(f(x_{a_{1:b}}),y\bigr)\bigr],$$ where $a_n=\pi(x_{a_{1:n-1}})$ and $|a_{1:b}|=b$.

This formulation places AFABench at the intersection of cost-sensitive learning, sequential decision-making, and partial-information prediction. The benchmark is motivated by the claim that, despite a growing literature on greedy information-theoretic methods, reinforcement-learning approaches, and oracle approximations, there had been no unified suite for comparing these methods on identical data, costs, and budgets. AFABench is therefore designed not primarily as a new acquisition algorithm, but as an infrastructure for standardizing empirical comparison.

2. Benchmark architecture and execution model

AFABench is organized into three core modules: a dataset module, a policy interface, and an evaluation engine (Schütz et al., 20 Aug 2025). The dataset module provides loaders, applies uniform preprocessing such as train/validation/test splits and masking protocols during predictor pre-training, and exposes a standard interface .load() -> (train, val, test) where each split is represented as $f:\mathbb{R}^b\to\mathcal{Y}$0 with $f:\mathbb{R}^b\to\mathcal{Y}$1. The policy interface is built around an abstract base class AFAPolicy with two methods: .train(self, train_data, val_data) and .select(self, observed_x, mask). Built-in registries support extension through calls such as register_policy("MyMethod", MyMethodClass).

The evaluation engine standardizes the benchmarking loop. Given a dataset, a set of policies, and budgets $f:\mathbb{R}^b\to\mathcal{Y}$2, it first pre-trains a shared predictor $f:\mathbb{R}^b\to\mathcal{Y}$3 on fully masked data; then, for each policy and each test instance, it simulates sequential acquisitions until budget $f:\mathbb{R}^b\to\mathcal{Y}$4 is reached; it records the selected features $f:\mathbb{R}^b\to\mathcal{Y}$5, the prediction $f:\mathbb{R}^b\to\mathcal{Y}$6, and the cumulative reward; and finally aggregates metrics into standard accuracy-versus-budget curves with error bars. This shared execution path is intended to isolate differences in acquisition logic from differences in implementation.

Under the hood, all reinforcement learning-based policies use a common MDP environment via TorchRL. The state is $f:\mathbb{R}^b\to\mathcal{Y}$7, the action is $f:\mathbb{R}^b\to\mathcal{Y}$8, and the reward $f:\mathbb{R}^b\to\mathcal{Y}$9 is method-specific. A plausible implication is that this common environment makes AFABench particularly suitable for controlled ablation across RL reward designs and policy parameterizations.

3. Datasets and the AFAContext stress test

AFABench includes synthetic and real-world tasks and emphasizes diversity of domain and feature dimensionality (Schütz et al., 20 Aug 2025). The overview states that the dataset module provides loaders for six real-world and two synthetic tasks, while the listed benchmark summary names the following datasets: CUBE, AFAContext, MNIST, FashionMNIST, Diabetes, PhysioNet, and MiniBooNE. The listed dataset characteristics are as follows: CUBE has 700 / 150 / 150 train/validation/test instances, 20 features, and 8 classes; AFAContext has 700 / 150 / 150, 30 features, and 8 classes; MNIST and FashionMNIST each have 42,000 / 9,000 / 9,000 with 784 features and 10 classes; Diabetes has 64,443 / 13,810 / 13,810 with 45 features and 3 classes; PhysioNet has 8,400 / 1,800 / 1,800 with 41 features and 2 classes; and MiniBooNE has 91,044 / 19,509 / 19,511 with 50 features and 2 classes. For PhysioNet, the reported performance metric is F1-score. All real-world tasks use uniform-cost per feature and a hard budget. In the image datasets, each pixel is treated as an independent feature.

Preprocessing is standardized. Real-world datasets undergo normalization, and shared predictor pretraining uses random feature masking drawn from $b$0 for MNIST-type datasets and $b$1 elsewhere. This uniform preprocessing is central to the benchmark’s claim of fair comparison.

AFAContext is the benchmark’s purpose-built synthetic task for testing lookahead. It is explicitly designed to expose the limitations of greedy selection. Features 1 and 2 form a one-hot context indicator $b$2, with exactly one entry equal to 1. Features $b$3 and $b$4 define two disjoint groups of ten “CUBE-style” features. If $b$5, block A contains the three class-informative features drawn from

$b$6

while block B is noise drawn from $b$7; if $b$8, the roles of blocks A and B are reversed. Features 23–30 are independent noise with $b$9. All features cost 1 unit. The intended consequence is that an optimal policy must first query $\pi$0, even though those features provide zero immediate predictive gain, in order to discover which block contains signal. Greedy methods that optimize short-term improvement are expected to avoid the context bits and waste budget on uninformative features.

4. Acquisition policy families

AFABench implements representatives from static, greedy, non-greedy, and oracle categories (Schütz et al., 20 Aug 2025). This breadth is essential to its role as a benchmark, because AFA methods differ not only in predictive performance but in computational profile, lookahead depth, and assumptions about feature–label structure.

The static family includes PT-S and CAE-S. PT-S ranks features by permutation importance: for each feature $\pi$1, it measures $\pi$2, the drop in validation accuracy when values of $\pi$3 are permuted, and then selects the top-$\pi$4 features by $\pi$5. CAE-S uses a Concrete Autoencoder to learn a differentiable mask $\pi$6 for reconstructing $\pi$7, adapted here to the supervised loss $\pi$8. These methods do not adapt feature choices instance by instance.

The greedy family is unified by a conditional mutual information criterion: $\pi$9 EDDI-GG trains a partial VAE to approximate $S\subset [d]$0 and $S\subset [d]$1, estimating conditional mutual information analytically in the Gaussian case or by Monte Carlo. GDFS-DG estimates expected one-step lookahead loss reduction,

$S\subset [d]$2

which matches the conditional mutual information optimum when $S\subset [d]$3 is Bayes. DIME-DG uses two networks $S\subset [d]$4, with $S\subset [d]$5, and a joint objective intended to drive $S\subset [d]$6 toward the true conditional mutual information.

The reinforcement learning family treats acquisition as an MDP with state $S\subset [d]$7, action $S\subset [d]$8, transition $S\subset [d]$9, and budget constraint $i\notin S$0. JAFA-MFRL uses a sparse terminal reward,

$i\notin S$1

whereas ODIN-MBRL uses a dense per-step reward,

$i\notin S$2

with a pretrained PVAE for simulated rollouts. ODIN-MFRL disables model-based rollouts and trains via PPO. OL-MFRL uses an uncertainty-based reward,

$i\notin S$3

where $i\notin S$4 is the class-probability vector under dropout. All RL agents share a set-invariant encoder for $i\notin S$5, use DQN or PPO via TorchRL, set $i\notin S$6, and operate with fixed episode length $i\notin S$7.

The oracle-based method AACO approximates the globally optimal subset by solving

$i\notin S$8

using $i\notin S$9-nearest neighbors density estimates and random subset sampling. In practice, the method uses a heuristic that picks features one at a time to meet the hard budget.

5. Evaluation protocol and reported metrics

For each policy $b$0 and budget $b$1, AFABench reports accuracy, or F1 for PhysioNet, as

$b$2

where $b$3 is the set selected by $b$4 (Schütz et al., 20 Aug 2025). Because acquisition cost is uniform, cost is identified directly with the budget $b$5. For reinforcement learning methods, the benchmark also reports a reward curve,

$b$6

The principal visualization is the trade-off curve of performance versus budget, namely accuracy versus $b$7, and F1 versus $b$8 for PhysioNet.

Experimental protocol is standardized across methods. All experiments average over 5 random splits and 3 policy-retraining seeds, with error bars given as $b$9 standard deviation. The benchmark therefore evaluates not only central tendency but also sensitivity to data partitioning and training stochasticity. This emphasis on repeated runs is particularly relevant for RL-based methods, whose empirical variance is often substantial.

The benchmark’s evaluation engine also records the selected subset $$\min_{f,\pi}\;\mathbb{E}_{(x,y)\sim p}\bigl[\ell\bigl(f(x_{a_{1:b}}),y\bigr)\bigr],$$0, the resulting prediction $$\min_{f,\pi}\;\mathbb{E}_{(x,y)\sim p}\bigl[\ell\bigl(f(x_{a_{1:b}}),y\bigr)\bigr],$$1, and cumulative reward during simulation. This design makes AFABench useful for both aggregate benchmarking and post hoc analysis of acquisition behavior, such as whether a method systematically prioritizes context features, high-variance features, or globally important features.

6. Empirical findings, runtime trade-offs, and practical use

The reported results distinguish sharply between settings that reward myopic acquisition and settings that require lookahead (Schütz et al., 20 Aug 2025). On the synthetic CUBE task, discriminative greedy methods—specifically GDFS-DG and DIME-DG—dominate, while PT-S is described as surprisingly competitive and RL and oracle methods bring little gain. On AFAContext, by contrast, greedy methods collapse because they ignore the context bits; RL-based methods, especially ODIN-MFRL and OL-MFRL, excel by learning lookahead, while AACO fails because of its one-at-a-time heuristic.

Across the real-world datasets MNIST, FashionMNIST, Diabetes, PhysioNet, and MiniBooNE, GDFS-DG and DIME-DG consistently rank at or near the top. EDDI-GG trails, with the explanation attributed to PVAE limitations. PT-S remains strong on tabular health and particle-identification tasks, which are characterized in the overview as having low inter-instance variability. RL methods generally underperform on real data, and the stated reasons are that real datasets seldom require deep lookahead and RL training is more fragile and expensive.

Runtime comparisons reinforce these differences. RL training times are reported as tens of thousands of seconds, with evaluation remaining cheap at less than 100 seconds. Greedy methods require 1,000–4,000 seconds for training and about 20 seconds for evaluation. AACO is reported at approximately 10 seconds of training time but about 15,700 seconds for evaluation. Static methods train in less than 500 seconds and evaluate in about 20 seconds. The benchmark therefore makes visible a three-way trade-off among predictive performance, lookahead capability, and compute.

AFABench is distributed as open-source code at https://github.com/Linusaronsson/AFA-Benchmark, and the overview includes a Python quick-start built around AFABenchmark, register_policy, and load_dataset. The framework is presented as modular and intended for easy integration of new methods and tasks. A plausible implication is that AFABench serves not only as a scoreboard for existing AFA algorithms but also as an experimental substrate for studying when dynamic acquisition is genuinely necessary, when static selection suffices, and when non-myopic decision-making justifies its computational overhead.