Towards Robust Active Feature Acquisition

Typical discriminative models predict a target $y$ using all $d$-dimensional features $x\in\mathbb{R}^{d}$. AFA, instead, actively acquires feature values to improve the prediction. It typically starts from an empty set of features and sequentially acquires more features $x_{i}$ until the prediction is sufficiently accurate or it exceeds the given acquisition budget. The goal of an AFA model is to minimize the following objective

where $\mathcal{L}(\hat{y}(x_{o}),y)$ is the prediction error between the groundtruth target $y$ and the prediction $\hat{y}$ using the acquired features $x_{o}$, $\mathcal{C}(o)$ measures the total cost of acquiring a subset of features $o\subseteq\{1,\ldots,d\}$, and $\alpha$ balances these two terms.

However, directly optimizing (1) is not trivial, since it involves optimizing over a combinatorial number of possible subsets. Many heuristic approaches have been developed to approximately solve this problem. For example, in [4], the authors propose to take into account the cost of features when selecting an attribute for building a decision tree so that the final tree will have a minimum total cost. [5] utilizes a naive Bayes classifier to handle the partially observed features, where the unobserved features are simply ignored in the likelihood objective. They then assess the utility of each unobserved feature by their expected reduction of the misclassification cost. At each acquisition step, the feature with highest utility is acquired. Authors in [6] instead leverage a margin-based classifier. An instance is classified by retrieving the nearest neighbors from training set using the partially observed features, and the utility of each unobserved feature is calculated by the one-step ahead classification accuracy. Following the same greedy solution, EDDI [7] utilizes the modern generative models to handle partially observed instances. Specifically, they propose a partial VAE to model the arbitrary marginal likelihoods $p(x_{o})$ (target variable $y$ is concatenated into $x$ and modeled together). Inspired by the experimental design approaches [8], they assess the utility of each unobserved feature with their expected information gain to the target variable $y$, i.e.,

The feature with highest utility is acquired at each step. Similar to EDDI, Icebreaker [9] proposes to use a Bayesian Deep Latent Gaussian model to capture the uncertainty of unobserved features and to assess their utilities. They further extend the problem to actively acquire additional information during training.

Greedy approaches are easy to understand, but they are also inherently flawed, since they are myopic and unaware of the long-term goal of obtaining multiple features that are jointly informative. Instead of acquiring features greedily, the AFA problem has been formulated as a Markov Decision Process (MDP) [10, 11]. Therefore, reinforcement learning based approaches can be utilized, where a long-term discounted reward is optimized. In the MDP formulation of AFA problem, the state is the current observed features, the action is the next feature to acquire, and the reward contains the final prediction reward and the cost of each acquired feature. In [1] and [2], a special action indicating the termination of the acquisition process is also introduced. The agent will stop acquiring more features when it selects the termination action. Specifically, we have

where the state, $s$, consists of the current acquired feature subset, $o\subseteq\{1,\ldots,d\}$, and their values, $x_{o}$. The action, $a$, is either one of the remaining unobserved features, $u=\{1,\ldots,d\}\setminus o$, or the termination action, $\phi$. When a new feature, $i$, is acquired, the current state transits to a new state following $o\xrightarrow{i}o\cup i$, $x_{o}\xrightarrow{i}x_{o}\cup x_{i}$, and the agent receives the negative acquisition cost of this feature as a reward. If the termination action is selected (i.e., $a=\phi$), the agent makes a prediction based on all acquired features, $x_{o}$, and receives a final reward as $-\mathcal{L}(\hat{y}(x_{o}),y)$.

Given the above MDP formulation, several RL approaches have been explored. [10] fits a transition model using complete data, and then uses the $\text{AO}^{*}$ heuristic search algorithm to find an optimal policy. [11] utilizes Fitted Q-Iteration to optimize the MDP. [12] and [13] instead employ a imitation learning approach coached by a greedy reference policy. JAFA [2] jointly learns an RL agent and a classifier, where the classifier is deemed as the environment to calculate the reward.

Although MDPs are broad enough to encapsulate the active acquisition of features, there are several challenges that limit the success of a naive reinforcement learning approach. In the aforementioned MDP, the agent pays the acquisition cost at each acquisition step but only receives a reward about the prediction after completing the acquisition process. This results in sparse rewards leading to credit assignment problems for potentially long episodes [14, 15], which may make training difficult. In addition, an agent that is making feature acquisitions must also navigate a complicated high-dimensional action space, as the action space scales with the number of features, making for a challenging RL problem [3]. Finally, the agent must manage multiple roles as it has to: implicitly model dependencies (to avoid selecting redundant features); perform a cost/benefit analysis (to judge what unobserved informative feature is worth the cost); and act as a classifier (to make a final prediction). Attaining these roles individually is a challenge, doing so jointly and without gradients on the reward (as with the MDP) is an even more daunting problem. To lessen the burden on the agent, and assuage the other challenges, GSMRL [16] proposes a model-based alternative. The key observation of GSMRL is that the dynamics of the above MDP can be modeled by the conditional dependencies among features. That is, the state transitions are based on the conditionals: $p(x_{j}\mid x_{o})$, where $x_{o}$ is current observed features and $x_{j}$ is an unobserved feature. Base on this observation, a surrogate model that captures the arbitrary conditionals, $p(x_{u},y\mid x_{o})$, is employed to assist the agent. Specifically, GSMRL defines an intermediate reward using the information gain of the acquired feature, $x_{i}$, at the current acquisition step, i.e., $r_{m}(s,i)=H(y\mid x_{o})-H(y\mid x_{o},x_{i}).$ The entropy terms are estimated using the learned surrogate model. Furthermore, the expected information gain for each candidate acquisition is provided to the agent as side information, i.e.,

In addition to $\mathcal{U}_{j}$, the surrogate model can also provide the current prediction $\hat{y}$, the prediction probability, $p(y\mid x_{o})$, the imputed values of unobserved features and their uncertainties, $p(x_{u}\mid x_{o})$, as auxiliary information. Armed with the auxiliary information and the intermediate reward, GSMRL alleviates the challenge of a model-free approach and obtains state-of-the-art performance for several AFA problems. Given their established excellence, we use GSMRL as the base model for our robust AFA framework.

AFA acquires features actively to improve the prediction of a target variable, while some application do not have an explicit target; instead, the features are acquired to improve our understanding to the instance. In GSMRL [1], the authors propose a task named AIR, where an agent acquires features actively to reconstruct the unobserved features. A similar model-based RL approach is used for AIR, where a dynamics model $p(x_{u}\mid x_{o})$ captures the state transition. The intermediate reward and auxiliary information can be similarly derived by replacing $y$ with $x_{u\setminus i}$ (the unobserved features excluding the current candidate $i$). A special case of AIR is to acquire features in k-space for accelerated MRI reconstruction. [17] and [18] have explored this application using deep Q-learning and policy gradient respectively. Several non-RL approaches [19, 20] have also been proposed.

ML models are typically trained with a specific data distribution, however, when deployed, the model may encounter data that is outside the training distribution. For those out-of-distribution inputs, the prediction could be arbitrarily bad. Therefore, detecting OOD inputs has been an active research direction. One approach relies on the uncertainty of prediction. The prediction for OOD inputs are expected to have higher uncertainty. However, deep classifiers tend to be overly confident about their prediction, thus a postprocessing step to calibrate the uncertainty is needed [21, 22]. Bayesian neural networks (BNNs), where weights are sampled from their posterior distributions, have also been used to quantify the uncertainty [23]. However, the full posterior is usually intractable. To avoid the complex posterior inference of BNNs, deep ensemble [24] and MC dropout [25] are proposed as two approximations, where multiple independent forward passes are conducted to obtain a set of predictions. To avoid the expensive multiple forward passes, DUQ [26] propose to detect OOD inputs with a single deterministic model in a single pass. It learns a set of centroid vectors corresponding to the different classes and measures uncertainty according to the distance calculated by an RBF network [27] between the model output and its closest centroid. Recently, Gaussian processes have been combined with deep neural networks to quantify the prediction uncertainty, such as SNGP [28] and DUE [29].

Intuitively, generative models are expected to be able to detect OOD inputs using the likelihood [30], that is, OOD inputs should have lower likelihood than the in-distribution ones. However, recent works [31, 32] show that it is not the case for high dimensional distributions. This oddity has been verified for many modern deep generative models, such as PixelCNN, VAE and normalizing flows. To rectify this pathology, [33] proposes to train multiple generative models in parallel and utilize the Watanabe-Akaike Information Criterion to identify OOD data. [34] proposes to use likelihood ratio of an autoregressive model to remove the contributions from background pixels. [35] propose a simple typicality test for a batch of data. Furthermore, [36] propose to utilize an estimator for the density of states on several summary statistics of the in-distribution data. It then evaluates the density of states estimator (DoSE) on testing data and marks those with low support as OOD. MSMA [37] empirically observes that the norm of scores conditioned on different noise levels serves as very effective summary statistics in the DoSE framework, and thus utilizes the multi-scale score matching network [38] to calculate those statistics.

As described in Sec. 2.1, the AFA problem can be interpreted as a MDP, where the action space at each acquisition step contains the current unobserved features. For certain problems, the action space could be enormous. For example, in the aforementioned health care example, the action space could contain an exhaustive list of possible inspections a hospital can offer. Dealing with large action space for RL is generally challenging, since the agent may not be able to effectively explore the entire action space. Several approaches have been proposed to train RL agent with a large discrete action space. For instance, [3] proposes a Wolpertinger policy that maps a state representation to a continuous action representation. The action representation is then used to look up $k$-nearest valid actions as candidates. Finally, the action with the highest Q value is selected and executed in the environment. [39] proposes a sequentialization scheme, where the action space is transformed into a sequence of $\mathcal{B}$-ary decision code words. A pair of bijective encoder-decoder is defined to perform this transformation. Running the agent will produce a sequence of decisions, which are subsequently decoded as a valid action that can be executed in the environment.

Similar to [39], we also formulate the action space as a sequence of decisions. Here, we propose to utilize the inherited clustering properties of the candidate features. Given a set of features, $\{x_{1},\ldots,x_{d}\}$, we assume features can be clustered based on their informativeness to the target variable $y$. That is, there might be a subset of features that are decisive about $y$ and another subset of features that are not relevant to $y$. This assumption holds true for many real-life tasks. For example, a music recommender system might use a questionnaire to collect information from a user. Some questions about musicians or song genres are closely related to the target, while some questions about addresses might not be relevant at all. Based on this intuition, we propose to assess the informativeness of the candidate features using their mutual information to the target variable, $y$, i.e., $I(x_{i};y)$, where $i\in\{1,\ldots,d\}$. The mutual information can be estimated using the learned arbitrary conditionals of the surrogate model

where the expectation is estimated using a held-out validation set. Given the estimated mutual information, we can simply sort and divide the candidate features into different groups. For the sake of implementation simplicity, we use clusters with the same number of features. We can further group features inside each cluster into smaller clusters and develop a tree-structured action space as in [39], which we leave for future works. Note that the clustering is not performed actively for each instance; instead, we cluster once for each dataset and keep the cluster structure fixed throughout the acquisition process.

It is worth noting that the mutual information $I(x_{i};y)$ is not the only choice for clustering features. Alternative quantities, such as the mutual information $I(x_{i};x_{j})$ or a metric $d(x_{i},x_{j})=H(x_{i},x_{j})-I(x_{i};x_{j})$, can be used together with a hierarchical clustering procedure to group candidate features. However, these alternatives need to be estimated for each pair of candidate features, which incurs a $O(d^{2})$ computational complexity, while the mutual information, $I(x_{i};y)$, only has $O(d)$ complexity.

Given the grouped action space, $\mathcal{A}=\{g_{k}\}_{k=1}^{K}$, with $K$ distinct clusters, we develop a hierarchical policy to select one candidate feature at each acquisition step. $g_{k}=\{g_{k}^{(1)},\ldots,g_{k}^{(N)}\}\subseteq\{1,\ldots,d\}$ represents the $k_{th}$ group of features of size $N$, where $\forall k\neq k^{\prime},\ g_{k}\cap g_{k^{\prime}}=\emptyset$ and $\cup_{k=1}^{K}g_{k}=\{1,\ldots,d\}$. The policy factorizes autoregressively by first selecting the group index, $k$, and then selecting the feature index, $n$, inside the selected group, i.e.,

The actual feature index being acquired is then decoded as $g_{k}^{(n)}$. As the agent acquires features, the already acquired features are removed from the candidate set. We simply set the probabilities of those features to zeros and renormalize the distribution. Similarly, if all features of a group have been acquired, the probability of this group is set to zero. With the proposed action space grouping, the original $d$-dimensional action space is reduced to $K+N$ decisions. Please refer to Fig. 1 for an illustration.

In Sec. 2.3, we introduce several advanced techniques to detect out-of-distribution inputs. However, those approaches require fully observed data. In an AFA framework, data are partially observed at any acquisition step, which renders those approaches inappropriate. In this section, we develop a novel OOD detection algorithm specifically tailored for partially observed data. Inspired by MSMA [37], we propose to use the norm of scores from an arbitrary marginal distribution $p(x_{o})$ as summary statistics and further detect partially observed OOD inputs with a DoSE [36] approach. MSMA for fully observed data is built by the following steps:

1. (i)

   Train a noise conditioned score matching network $s_{\theta}$ [38] with $L$ noise levels by optimizing

   $$\frac{1}{L}\sum_{i=1}^{L}\frac{\sigma_{i}^{2}}{2}\mathbb{E}_{p_{data}(x)}\mathbb{E}_{\tilde{x}\sim\mathcal{N}(x,\sigma_{i}^{2}I)}\left[\left\|s_{\theta}(\tilde{x},\sigma_{i})+\frac{\tilde{x}-x}{\sigma_{i}^{2}}\right\|_{2}^{2}\right].$$

   (7)

   The score network essentially approximates the score of a series of smoothed data distributions $\nabla_{\tilde{x}}\log q_{\sigma_{i}}(\tilde{x})$, where $q_{\sigma_{i}}(\tilde{x})=\int p_{data}(x)q_{\sigma_{i}}(\tilde{x}\mid x)dx$, and $q_{\sigma_{i}}(\tilde{x}\mid x)$ transforms $x$ by adding some Gaussian noise form $\mathcal{N}(0,\sigma_{i}^{2}I)$.

2. (ii)

   For a given input $x$, compute the L2 norm of scores at each noise level, i.e., $s_{i}=\|s_{\theta}(x,\sigma_{i})\|$.

3. (iii)

   Fit a low dimensional likelihood model for the norm of scores using in-distribution data, i.e., $p(s_{1},\ldots,s_{L})$, which is called density of states in [36] following the concept in statistical mechanics.

4. (iv)

   Threshold the likelihood to determine whether the input $x$ is OOD or not.

In order to deal with partially observed data, we modify the score network to output scores of arbitrary marginal distributions, i.e., $\nabla_{\tilde{x}_{m}}\log q_{\sigma_{i}}(\tilde{x}_{m})$, where $m\subseteq\{1,\ldots,d\}$ represents an arbitrary subset of features. The training objective (7) is modified accordingly to

where $\mathbb{I}_{m}$ represents a $d$-dimensional binary mask indicating the partially observed features, $\odot$ represents the element-wise product operation, and $p(m)$ is the distribution for generating observed dimensions. Similar to the fully observed case, we compute the L2 norm of scores at each noise level, i.e., $s_{i}=\|s_{\theta}(x\odot\mathbb{I}_{m},\mathbb{I}_{m},\sigma_{i})\odot\mathbb{I}_{m}\|$, and fit a likelihood model in this transformed low-dimensional space. The likelihood model is also conditioned on the binary mask $\mathbb{I}_{m}$ to indicate the observed dimensions, i.e., $p(s_{1},\ldots,s_{L}\mid\mathbb{I}_{m})$. Given an input $x$ with observed dimensions $m$, we threshold the likelihood $p(s_{i},\ldots,s_{L}\mid\mathbb{I}_{m})$ to determine whether the partially observed data $x_{m}$ is OOD or not. To train the partially observed MSMA (PO-MSMA), we generate a mask for each input data $x$ at random. The conditional likelihood over norm of scores is estimated by a conditional autoregressive model, for which we utilize the efficient masked autoregressive implementation [40].

One benefit of our proposed PO-MSMA approach is that a single model can be used to detect OOD inputs with arbitrary observed features, which is convenient for detecting OOD inputs along the acquisition trajectories. Furthermore, sharing weights across different tasks (i.e., different marginal distributions) could act as a regularization (as discussed in [16]), thus the unified score matching network can potentially perform better than separately trained ones for each different conditional, which we will investigate in future works.

Above, we introduce our proposed action space grouping technique and a partially observed OOD detection algorithm. Combining those components, we can now actively acquire features for problem with a large action space and simultaneously detect OOD inputs using the acquired subset of features. However, the OOD detection performance might be suboptimal, since the agent is not informed of the detection goal and merely focuses on predicting the target variable $y$. The features that are informative for predicting the target might not be informative for distinguishing OOD inputs.

In order to guide the agent to acquire features that are informative for OOD detection, we propose to utilize the likelihood, $p(s_{1},\ldots,s_{L}\mid\mathbb{I}_{m})$, as an auxiliary reward, which encourages the agent to acquire features more closely resemble the in-distribution ones, and thus reduces the false positive detection. Alternatively, if reducing the false negative is desired, we could use the negative likelihood $-p(s_{1},\ldots,s_{L}\mid\mathbb{I}_{m})$ as the auxiliary reward.

In summary, our robust AFA framework contains a dynamics model, an OOD detector and an RL agent. The dynamics model captures the arbitrary conditionals, $p(x_{u},y\mid x_{o})$, and is utilized to provide auxiliary information and intermediate rewards. It also enables a simple and efficient action space grouping technique and thus scales AFA up to applications with large action spaces. The partially observed OOD detector is used to distinguish OOD inputs alongside the acquisition procedure and also used to provide an auxiliary reward so that the agent is encouraged to acquire informative features for OOD detection. The RL agent takes in the current acquired features and auxiliary information from the dynamics model and predicts what next feature to acquire. When the feature is actually acquired, the agent pays the acquisition cost of the feature and receives an intermediate reward from the dynamics model. When the acquisition process is terminated, the agent makes a final prediction about the target, $y$, using all its acquired features and receives an reward about its prediction. It also receives an reward from the OOD detector about the likelihood of the acquired feature subset in the transformed space (i.e., the norm of the scores). Please refer to Algorithm 1 for additional details and to Fig. 2 for an illustration.

In GSMRL [1], the acquisition procedure is terminated when the agent selects a special termination action, which means each instance could have different number of features acquired. Although intriguing for practical use, it introduces additional complexity to assess OOD detection performance, since we need to separate two possible causes of a detection failure, i.e., not sufficient acquisitions and not effective acquisitions. To simplify the evaluation, we instead specify a fixed acquisition budget (i.e., the number of acquired features). The agent will terminate the acquisition process when it exceeds the specified acquisition budget. However, it is possible to incorporate a termination action into our framework.