Distribution Guided Active Feature Acquisition

TL;DR

This work introduces Distribution Guided Active Feature Acquisition (DistAFA), a framework enabling prediction-time interaction with the environment to acquire informative features under cost constraints. It combines a single, flexible ACFlow model to capture arbitrary conditional distributions $p(x_u\mid x_o)$ with greedy and RL-based acquisition strategies, notably the Generative Surrogate Models for RL (GSMRL) that provide intermediate rewards and side information to the agent. It further enhances interpretability through goal-based explainable AFA and improves robustness via a partially observed OOD detector integrated into acquisition. The approach demonstrates state-of-the-art performance across diverse tasks, offers scalable hierarchical action spaces, and provides principled guarantees on reward shaping within its RL setup. This has practical implications for real-world decision systems that must actively query features while balancing cost, accuracy, interpretability, and reliability under distribution shift.

Abstract

Human agents routinely reason on instances with incomplete and muddied data (and weigh the cost of obtaining further features). In contrast, much of ML is devoted to the unrealistic, sterile environment where all features are observed and further information on an instance is obviated. Here we extend past static ML and develop an active feature acquisition (AFA) framework that interacts with the environment to obtain new information on-the-fly and can: 1) make inferences on an instance in the face of incomplete features, 2) determine a plan for feature acquisitions to obtain additional information on the instance at hand. We build our AFA framework on a backbone of understanding the information and conditional dependencies that are present in the data. First, we show how to build generative models that can capture dependencies over arbitrary subsets of features and employ these models for acquisitions in a greedy scheme. After, we show that it is possible to guide the training of RL agents for AFA via side-information and auxiliary rewards stemming from our generative models. We also examine two important factors for deploying AFA models in real-world scenarios, namely interpretability and robustness. Extensive experiments demonstrate the state-of-the-art performance of our AFA framework.

Figures (35)

• Figure 1: Acquisition process for supervised and unsupervised AFA tasks. Top: acquisition process where one pixel value is acquired at each step and the green masks indicate the unobserved features. Bottom: the prediction probabilities and averaged inpaintings for supervised and unsupervised tasks respectively.
• Figure 2: An illustrative example of the grouped action space, where 6 features are grouped into 3 clusters. The grayed circles represent the current observed features (or fully observed groups) and are not considered as candidates anymore. The dashed line shows one acquisition at the current step, which acquires the feature $g_2^{(1)}$. The corresponding circles will be grayed after this acquisition step.
• Figure 3: Estimating the clustering posterior for a partially observed instance. We first impute the unobserved features based on the surrogate model, and then get the cluster assignment probabilities for each imputed sample based on a fully observed clustering model. The estimated posterior is the average of those cluster assignment probabilities.
• Figure 4: Schematic illustration of our robust AFA framework.
• Figure 5: Example of acquired features and prediction probabilities. The green masks indicate the unobserved features.