Non-Myopic Active Feature Acquisition via Pathwise Policy Gradients

Abstract: Active feature acquisition (AFA) considers prediction problems in which features are costly to obtain and the learner adaptively decides which feature values to acquire for each instance and when to stop and predict. AFA can be formulated as a partially observable Markov decision process (POMDP), which naturally admits a sequential decision-making perspective. In this paper, we present non-myopic pathwise policy gradients (NM-PPG), a new AFA method built around this formulation. We introduce a continuous relaxation of the acquisition process that enables pathwise gradients through the full acquisition trajectory, avoiding the high variance of standard score-function policy gradients while allowing end-to-end optimization of a non-myopic acquisition policy. To better align training with deployment, we further develop a straight-through rollout scheme that follows hard feature acquisitions in the forward pass while backpropagating through the corresponding soft relaxation in the backward pass. We stabilize optimization with entropy regularization and staged temperature sharpening. Experiments on both synthetic and real-world datasets demonstrate that NM-PPG yields superior performance relative to state-of-the-art AFA baselines.

• The paper introduces NM-PPG, a novel approach leveraging continuous relaxations and Gumbel-Softmax for adaptive non-myopic feature acquisition.
• It formulates active feature acquisition as a POMDP to balance prediction loss with acquisition costs, enabling instance-wise, context-aware feature selection.
• Empirical results on synthetic and real-world datasets demonstrate NM-PPG’s robust accuracy-cost trade-offs outperforming traditional myopic and RL-based methods.

Non-Myopic Active Feature Acquisition via Pathwise Policy Gradients

Problem Formulation and Motivation

The paper "Non-Myopic Active Feature Acquisition via Pathwise Policy Gradients" (2605.05511) develops a principled approach to Active Feature Acquisition (AFA), focusing on settings where features are costly and acquisition strategies must be adaptive and instance-wise. Traditional static feature selection is insufficient because it cannot exploit instance-specific information about which features are informative, especially when informative features vary or are context-dependent. AFA is naturally framed as a Partially Observable Markov Decision Process (POMDP), in which acquisition decisions are sequential and a learner must determine both which feature to acquire next and when to stop acquisition and predict.

The objective is to minimize the expected sum of prediction loss and acquisition cost:

$\min_{\theta} \mathbb{E}_{\mathbf{x},\mathbf{y} \mathbb{E}_{\pi_{\theta} [ \ell (f_{\phi}(\mathbf{x}(m_{t_\theta(\mathbf{x})})),\mathbf{y}) + \alpha c(m_{t_\theta(\mathbf{x})})$

where $f_\phi$ is a predictor on masked inputs, $\pi_\theta$ is the acquisition policy, $m_t$ is the current observation mask, and $\alpha$ controls the trade-off. This formulation is fully general to arbitrary cost structures and permits adaptive, non-myopic sequential acquisition.

Pathwise Policy Gradient Approach

AFA-POMDPs are computationally challenging; prior approaches include model-based, model-free (mainly RL-based), and hybrid techniques. RL-based non-myopic methods optimize long-horizon cost-minimization objectives but suffer from high variance due to score-function methods (e.g., policy gradient theorem), while non-RL approaches use heuristics or jointly informative feature groups and are biased relative to true adaptive cost minimization.

The core contribution is NM-PPG (Non-Myopic Pathwise Policy Gradients), which uses a continuous relaxation of the discrete acquisition process, implemented via Gumbel-Softmax reparameterization. This allows pathwise gradients to propagate through the entire acquisition trajectory rather than a single step. The relaxed acquisition mask ($\tilde m_t$) and conditional feature distribution ($\tilde r_t$) evolve through deterministic dynamics, supporting backpropagation and end-to-end optimization over the policy network.

Hard feature acquisitions are enforced in the forward pass (straight-through rollout), aligning training with test-time deployment, while gradients propagate via corresponding soft masks in the backward pass. Entropy regularization and staged temperature sharpening stabilize optimization.

Algorithmic Framework

NM-PPG is trained on fully observed datasets, split into train/validation/test splits. Each batch proceeds through a staged soft-temperature schedule, wherein the Gumbel-Softmax relaxation sharpens toward the discrete acquisition process as training progresses. The predictor is first pre-trained on random masks, followed by fine-tuning on masks produced by the learned policy rollouts. Early stopping uses validation performance. At inference, the trained policy is evaluated greedily with blocked logits for acquired features.

Empirical Analysis and Numerical Results

NM-PPG is benchmarked against a comprehensive suite of AFA baselines (myopic: DiFA, GDFS, DIME; heuristic non-myopic: AACO, SEFA; RL-based non-myopic: GSMRL, OL) across 12 datasets spanning synthetic, tabular, medical, and image domains, with various cost structures and imbalanced scenarios.

Results demonstrate:

• On synthetic datasets (Cube-NM, Syn1, Syn3) explicitly constructed for non-myopic acquisition, NM-PPG achieves superior accuracy-cost trade-offs and recovers context-first acquisition paths required for optimality.
• On high-dimensional real-world tasks (e.g., MNIST), NM-PPG remains stable and performant while RL-based baselines suffer from instability.
• On many real-world tabular and medical datasets (e.g., NHANES Mortality, Diabetes), NM-PPG either matches or improves upon myopic baselines, particularly when non-myopic reasoning yields clear gains.

  Figure 1: Results for all methods across synthetic (top row) and real-world (2 bottom rows) datasets; performance by accuracy or F1-score as appropriate, over acquisition cost.

NM-PPG is explicitly consistent with myopic baselines when the dataset structure does not reward non-myopic planning, but outperforms them in presence of joint feature informativeness or contextually expensive features. The approach robustly addresses the relaxation gap that hinders prior continuous relaxations.

Ablation, Acquisition Behavior, and Training Dynamics

Ablation studies confirm the necessity of the straight-through rollout mechanism and entropy regularization; without them, optimization loses stability and alignment with deployment behavior. Acquisition trajectory visualizations (e.g., Cube-NM, Syn1, Syn3) show NM-PPG replicating the optimal context-first policy and adaptive branching based on observed context, in contrast to myopic methods which prefer locally informative but globally sub-optimal features.

Figure 2: Ablation study comparing NM-PPG with variants lacking straight-through rollouts or entropy regularization; both components are critical for robust performance.

Figure 3: Acquisition trajectories on Cube-NM with $n_c=5$ and $\sigma=0.1$, showing NM-PPG consistently acquires context features early and adapts acquisition sequence accordingly.

Figure 4: Training curves for NM-PPG, demonstrating stable loss, acquisition cost, and overall objective convergence across staged soft-temperature annealing.

Practical and Theoretical Implications

NM-PPG bridges the gap between principled sequential decision-making and practical, stable optimization in AFA. By enabling pathwise gradients over non-myopic acquisition trajectories and aligning training with deployment, it overcomes critical shortcomings of both RL-based and heuristic non-myopic strategies. Empirical results highlight that model-free, end-to-end non-myopic policy optimization is computationally tractable and practically useful for high-dimensional and heterogeneous cost domains.

Theoretically, NM-PPG matches the optimality conditions of solving the truncated AFA-POMDP, with rigorous guarantees regarding the continuous relaxation. The method naturally generalizes to variable cost, mixed data modalities, and arbitrary loss functions. The approach provides a template for policy-gradient optimization in other sequential decision settings with discrete and costly actions.

Future Directions

Advances may focus on integrating NM-PPG with more flexible, uncertainty-aware predictors for partial observations (e.g., conditional generative models), deeper theoretical analysis of relaxation bias under more challenging regimes, and deployment of NM-PPG in domains with long temporal horizons (e.g., longitudinal health records, active vision). Extensions to structured feature selection, adaptive expert selection, or hierarchical acquisition can further improve performance in real-world settings.

Conclusion

NM-PPG provides a rigorous framework for non-myopic AFA, yielding robust and adaptive acquisition policies that achieve near-optimal accuracy-cost trade-offs across diverse datasets. Its pathwise gradient architecture and straight-through rollout mechanism address both the variance problem of RL-based methods and the bias of heuristic non-myopic approaches. The method demonstrates that non-myopic decision-making can be practically deployed and reliably optimized for active information acquisition in settings with complex cost structures and adaptive feature dependencies.