🤖 machine learning

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

This paper introduces Learning-to-Measure (L2M), a meta-learning framework that enables in-context active feature acquisition across diverse tasks by leveraging autoregressive pre-training for uncertainty quantification and a greedy agent maximizing conditional mutual information, thereby eliminating the need for per-task retraining while outperforming baselines under scarce labels and high missingness.

Original authors: Yuta Kobayashi, Zilin Jing, Jiayu Yao, Hongseok Namkoong, Shalmali Joshi

Published 2026-06-02

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yuta Kobayashi, Zilin Jing, Jiayu Yao, Hongseok Namkoong, Shalmali Joshi

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a doctor trying to diagnose a patient, but you don't have their full medical history. You know their age and gender (the "baseline" features), but you don't know their blood pressure, cholesterol, or if they have a specific allergy. In the real world, getting these extra details costs money, takes time, or might even be risky (like a painful biopsy).

The Problem: The "Guessing Game" of Data
Traditional machine learning models usually assume they get all the data at once. But in reality, data is often messy. Some records are missing because the doctor didn't order a test, or the patient couldn't afford it. This is called "retrospective missingness."

Old methods for handling this are like a student who tries to memorize the answer key for one specific exam. If they face a slightly different exam, they fail. They also often try to "fill in the blanks" of missing data with guesses (imputation), which can lead to bad decisions because the guess might be wrong.

The Solution: Learning-To-Measure (L2M)
The paper introduces a new framework called Learning-To-Measure (L2M). Think of L2M not as a student memorizing one exam, but as a super-intelligent detective who has read thousands of different mystery novels.

Here is how it works, using simple analogies:

1. The "Universal Detective" (Meta-Learning)

Instead of training a separate detective for every single case (which is slow and requires lots of data), L2M trains one single model on a massive, diverse library of "practice cases."

The Analogy: Imagine a detective who has solved 10,000 different types of mysteries (some with missing fingerprints, some with missing alibis, some with only one witness). Because they've seen so many variations, they don't need to re-learn how to solve a new case from scratch. They just look at the new clues and say, "Ah, this looks like Case #4,502. I know exactly which clue to ask for next."
The Paper's Claim: This allows the model to work on new, unseen tasks without needing to be retrained for each specific job.

2. The "Uncertainty Compass" (Reliable Uncertainty)

When a detective is unsure, they know why they are unsure. L2M is special because it doesn't just guess; it has a built-in "Uncertainty Compass."

The Analogy: If you ask a normal model, "Is this a cat or a dog?" it might say "Dog" with 99% confidence, even if it's missing the picture of the animal's face. L2M says, "I think it's a dog, but I'm only 60% sure because I can't see the ears."
The Paper's Claim: The model uses a technique called sequence modeling (like how a language model predicts the next word in a sentence) to predict the outcome. This gives it a very reliable sense of how "confused" it is.

3. The "Strategic Clue Hunter" (Active Feature Acquisition)

This is the core magic. The model doesn't just guess; it actively decides which piece of information to buy next to reduce its confusion the most.

The Analogy: Imagine you are playing a game of "20 Questions." A bad player asks random questions. A smart player asks, "Is it an animal?" because that cuts the possibilities in half immediately.
How L2M does it: It looks at its "Uncertainty Compass." If it's confused about whether a patient has a heart condition, it checks: "If I get the blood test, will I be less confused? If I get the X-ray, will I be less confused?" It picks the one that gives the biggest "bang for the buck" (the most information for the cost).
The Paper's Claim: It uses a mathematical trick to make this decision process smooth and fast, allowing it to learn the best strategy for asking questions across all its practice cases.

4. Why It's Better Than the Old Way

The paper tested L2M against older methods on real-world data (like hospital records) and made-up data.

The Result: When data is scarce (few labeled examples) or when data is very messy (lots of missing pieces), L2M wins.
The Analogy: If you have a very small puzzle with many missing pieces, an old method might try to force a piece into the wrong spot. L2M looks at the whole picture, realizes it's missing a specific corner piece, and asks for that specific piece first, solving the puzzle faster and more accurately.

Summary

Learning-To-Measure is a system that learns from a huge variety of past experiences to become an expert at deciding what information to collect next.

It doesn't need to be retrained for every new job.
It knows when it is confused and uses that feeling to guide its next move.
It works best when data is missing or expensive to get, making it a "smart shopper" for information.

The authors show that this approach is more robust and accurate than previous methods, especially in tricky situations where data is incomplete.

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

1. Problem Definition

The paper addresses Active Feature Acquisition (AFA), a sequential decision-making problem where an agent adaptively selects which features to acquire for a specific test instance to improve predictive performance while minimizing acquisition costs. This is particularly relevant in domains like medical diagnostics, where acquiring certain features (e.g., imaging, biopsies) incurs significant financial costs or risks.

A critical limitation of existing AFA methods is their handling of retrospective missingness. Historical training data often contains systematic missing features due to prior data collection policies (e.g., clinical protocols that prioritize lab tests before imaging). Standard methods often fail to account for this missingness structure, leading to biased acquisition policies and poorly calibrated uncertainty estimates. Furthermore, most existing approaches are designed for single, predetermined tasks, requiring per-task retraining and relying on complex latent-variable models that struggle with scalability and reliability in high-dimensional settings.

The authors formalize the Meta-AFA problem: training a single agent across a distribution of tasks with diverse data distributions and missingness patterns, enabling it to generalize to unseen tasks at inference time without task-specific retraining.

2. Methodology: Learning-To-Measure (L2M)

The authors propose Learning-To-Measure (L2M), a meta-learning framework that performs AFA in-context using a pre-trained sequence model (Transformer). The framework consists of two main stages:

2.1. Theoretical Foundation

The objective of AFA is to maximize the Conditional Mutual Information (CMI) between the target variable $Y$ and a candidate feature $X_j$ given the currently observed features $X_t$ .
$I_\tau(Y; X_j | X_t = x_t) \triangleq \mathbb{E}_{X_j|x_t} [D_{KL}(P_\tau(Y | X_j \cup x_t) \parallel P_\tau(Y | x_t))]$

The paper establishes conditions under which CMI is identifiable from retrospectively missing data using a causal framework (Potential Outcomes). Under assumptions of Missing at Random (MAR), No Direct Effect of missingness on the target, and Positivity (sufficient coverage of acquisition actions), the CMI can be estimated using only the observed data where the feature was actually measured.

2.2. Architecture and Training

L2M replaces explicit generative modeling with autoregressive sequence modeling.

Model Architecture: A Transformer-based sequence model with two heads:
1. Target Head ( $f_\phi$ ): Predicts the posterior predictive distribution (PPD) of the target $Y$ given partially observed inputs and a context of historical samples.
2. Policy Head ( $\pi_\theta$ ): Selects the next feature to acquire.
In-Context Learning: The model takes a sequence of context samples (historical data with known labels and missingness patterns) and query samples (partially observed inputs) as input. It learns to infer the task-specific data-generating mechanism and missingness structure directly from the context.
Surrogate Optimization: Since exact CMI maximization is intractable, L2M optimizes a surrogate objective: the expected one-step-ahead predictive loss (e.g., negative log-likelihood) of the target head after acquiring a candidate feature.
- The policy is trained to minimize this loss, which theoretically corresponds to maximizing CMI under the Bayes-optimal predictor.
Differentiable Policy Learning: To enable gradient-based optimization for the discrete action of feature selection, the authors employ a straight-through Gumbel-Softmax relaxation. This allows the policy to sample actions during the forward pass while backpropagating gradients through a continuous relaxation.
Blocked Policies: During training, the policy is constrained ("blocked") to only select features that are available in the retrospective training data ( $R_j=1$ ), ensuring the model learns valid acquisition strategies given the missingness constraints.

2.3. Training Procedure

The pretraining involves sampling diverse tasks from a prior (synthetic or semi-synthetic) with varying feature distributions, label mechanisms, and missingness patterns. The model is trained end-to-end on sequences of context and query samples, optimizing the joint loss of the predictor and the policy.

3. Key Contributions

Meta-Learning Formulation: The paper formalizes the Meta-AFA problem, enabling a single agent to learn acquisition policies that generalize across time-invariant tasks with diverse data distributions and retrospective missingness patterns, eliminating the need for per-task retraining.
Sequence Modeling for Uncertainty and Decision-Making: L2M combines uncertainty estimation and decision-making via a scalable Transformer architecture. It provides reliable uncertainty estimates for partially observed inputs and leverages these estimates to acquire features that minimize expected loss. The authors demonstrate that this formulation enables the meta-learning of both greedy and planning-based (lookahead) AFA policies.
Robustness to Data Scarcity and Missingness: Empirical results show that L2M outperforms task-specific baselines, particularly in regimes with scarce labeled data and high rates of feature missingness. The approach avoids the pitfalls of deterministic imputation by explicitly propagating uncertainty through the sequence model.

4. Experimental Results

The authors evaluate L2M on synthetic and real-world tabular benchmarks (Metabric, MiniBooNE, MIMIC-IV, MNIST).

Uncertainty Quantification: L2M demonstrates superior uncertainty calibration compared to task-specific MLP baselines. Coverage plots show L2M remains closer to the ideal diagonal line across acquisition trajectories, indicating more reliable predictive distributions.
Acquisition Performance: L2M achieves lower log loss (better predictive performance) than random acquisition and consistently outperforms task-specific greedy (GDFS, DIME) and reinforcement learning (DQN) baselines across various datasets.
Generalization: A single L2M model pretrained on synthetic tasks generalizes to unseen real-world tasks (e.g., MIMIC-IV) and out-of-distribution kernel settings (Matern vs. RBF), maintaining performance without fine-tuning.
Lookahead Capabilities: The framework supports lookahead policies (k-step planning). Experiments show that lookahead variants can outperform greedy strategies, particularly in tasks with feature synergies where submodularity assumptions fail.
Efficiency: The pathwise gradient estimator used for policy learning is shown to be more sample-efficient and stable than zeroth-order methods like REINFORCE.

5. Significance and Claims

The paper claims that L2M represents a significant step toward foundation models for sequential decision-making in data-scarce and high-missingness environments. By leveraging in-context learning, L2M eliminates the need for complex latent-variable approximations and per-task retraining, offering a principled, scalable, and calibrated approach to active feature acquisition.

The authors position this work as a proof of concept. They explicitly state that while the method shows promise in reducing unnecessary measurements (e.g., in clinical triage), it is not yet ready for deployment. Real-world application would require prospective validation, careful specification of costs and utilities, and human oversight. The work opens avenues for extending meta-learning to broader data modalities and richer task distributions, while acknowledging current limitations regarding time-varying dynamics and the reliance on the MAR assumption.