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

TL;DR

The paper tackles meta-active feature acquisition (meta-AFA), addressing how to learn acquisition policies that generalize across tasks with retrospective missingness. It introduces Learning-to-Measure (L2M), a two-stage framework that combines sequence-model uncertainty quantification with an uncertainty-guided greedy acquisition policy, trained end-to-end via a differentiable surrogate of conditional mutual information and implemented with a transformer architecture. By formalizing identifiability conditions under missingness mechanisms (e.g., MAR, exclusion restrictions, positivity) and enabling in-context learning, L2M achieves reliable uncertainty estimates and strong acquisition performance on synthetic and real-world tabular benchmarks, particularly when labels are scarce. The work demonstrates scalable, in-context AFA without per-task retraining, with potential impact in healthcare and other domains where feature costs are high and data are incomplete.

Abstract

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. Most prior work addresses acquisition for a single predetermined task, limiting scalability. To address this limitation, we formalize the meta-AFA problem, where the goal is to learn acquisition policies across various tasks. We introduce Learning-to-Measure (L2M), which consists of i) reliable uncertainty quantification over unseen tasks, and ii) an uncertainty-guided greedy feature acquisition agent that maximizes conditional mutual information. We demonstrate a sequence-modeling or autoregressive pre-training approach that underpins reliable uncertainty quantification for tasks with arbitrary missingness. L2M operates directly on datasets with retrospective missingness and performs the meta-AFA 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.

Figures (13)

• Figure 1: Schematic of the meta-AFA problem. The model can acquire lab measurements for multiple tasks, Tasks 1-3. The input is a sequence of past observations together with a query, and the policy $\pi_\theta$ outputs the next greedy acquisition action. The updated query is then used for the following step. After $k$ acquisitions, the predictor $f$ produces the final label prediction.
• Figure 2: (On the left) Pretraining procedure for the predictor $f_\phi$ and policy $\pi_\theta$. The modified transformer encoder model takes as input task-specific context and partially observed states $X_o^{(j)}$ and learns to predict outcomes $Y^{(j)}$ and the optimal greedy action $A^{(j)}$ via a sequence loss. (On the right) The trained transformer model is able to predict labels and optimal actions for new unseen query tasks in-context. The attention mask to enable handling different context lengths used in the Transformer is given in Appendix \ref{['app:architecture']}.
• Figure 3: Left: Improvement in uncertainty estimation relative to baseline task-specific MLP (higher is better). The L2M approach shows progressively increasing gains as more features are acquired in the trajectory. Right: Coverage plots comparing MLP and L2M at various acquisition steps. L2M also shows robust coverage compared to the MLP at various acquisition steps. Error represents standard error across 200 sampled evaluation tasks.To ensure a fair comparison, both methods are evaluated on the same random acquisition trajectories using consistent evaluation tasks and samples.
• Figure 4: Acquisition performance quantified by log loss averaged over tasks derived from various synthetic and real-world datasets. MIMIC-IV demonstrate the ability for a single pretrained L2M model to generalize to diverse tasks with real unseen labels. The acquisition performance also often outperforms task-specific greedy and RL approaches.
• Figure 5: Average improvement in log loss (y-axis) for the GP tasks and MIMIC-LOS for a fixed set of evaluation tasks while varying the number of context samples, and levels of retrospective missingness in each task (x-axis). Task-specific improvements are averaged over 100 (GP) and 50 (MIMIC) tasks. L2M is robust to settings with fewer shots (labeled samples) and with higher rates of feature missingness. L2M also achieves length generalization for the GP task, as it is only trained on sequences of 500 samples.

Theorems & Definitions (14)

• Theorem 3.4
• Definition 4.1: Blocked Policy
• Theorem 4.2: Surrogate optimality for greedy CMI with context
• Remark A.1
• Proposition A.2
• proof
• Lemma A.3: Bayes optimality under cross-entropy
• proof
• Lemma A.4: Risk reduces to (expected) conditional entropy at $\phi^\star$
• proof