Learning-To-Measure: In-Context Active Feature Acquisition

Active feature acquisition (AFA) is a sequential decision-making problem where the goal is to improve model performance for test instances by adaptively selecting which features to acquire. In practice, AFA methods often learn from retrospective data with systematic missingness in the features and limited task-specific labels. To address this limitation, we introduce Learning-to-Measure (L2M), which consists of i) reliable uncertainty quantification over unseen tasks, and ii) an uncertainty-guided feature acquisition agent that maximizes conditional mutual information. We demonstrate an autoregressive pre-training approach that underpins reliable uncertainty quantification and acquisition for tasks with arbitrary missingness. L2M operates directly on datasets with retrospective missingness and performs the task in-context, eliminating per-task retraining. Across synthetic and real-world tabular benchmarks, L2M matches or surpasses task-specific baselines, particularly under scarce labels and high missingness.

Active Feature Acquisition (AFA) is a process where AI models sequentially gather missing information to improve predictions, much like a doctor ordering tests one by one. The goal is to figure out the most efficient order to collect these missing pieces to make an accurate diagnosis or prediction. Usually, AI learns how to do this by studying historical data. However, real-world records are often messy and incomplete due to how they were originally collected. To overcome this, we introduce a novel framework that learns directly from this fragmented history. Using the messy historical data as guidance, the model calculates which specific piece of missing information will clear up the most confusion for new instances, and chooses to gather that data next. We introduce a training procedure that learns over many such examples, letting our agent adapt on the fly to brand-new problems without needing to be retrained. When tested across various settings, our approach consistently matched or outperformed traditional models, proving especially powerful in challenging scenarios where data is highly incomplete and historical data is limited.