Active Feature Acquisition with Generative Surrogate Models

A typical machine learning paradigm for discriminative tasks is to learn the distribution of an output, $y$ given a complete set of features, $x\in\mathbb{R}^{d}$: $p(y\mid x)$. Although this paradigm is successful in a multitude of domains, it is incongruous with the expectations of many real-world intelligent systems in two key ways: first, it assumes that a complete set of features has been observed; second, as a consequence, it also assumes that no additional information (features) of an instance may be obtained at evaluation time. These assumptions often do not hold; human agents routinely reason over instances with incomplete data and decide when and what additional information to obtain. For example, consider a doctor diagnosing a patient. The doctor usually has not observed all possible measurements (such as blood samples, x-rays, etc.) for the patient. He/she is not forced to make a diagnosis based on the observed measurements; instead, he/she may dynamically decide to take more measurements to help determine the diagnosis. Of course, the next measurement to make (feature to observe), if any, will depend on the values of the already observed features; thus, the doctor may determine a different set of features to observe from patient to patient (instance to instance) depending on the values of the features that were observed. Hence, each patient will not have the same subset of features selected (as would be the case with typical feature selection). Furthermore, acquiring features typically involves some cost (in time, money and risk), and intelligent systems are expected to automatically balance the cost and improvement on performance. In order to more closely match the needs of many real-world applications, we propose an active feature acquisition (AFA) model that not only makes predictions with incomplete/missing features, but also determines the next feature that would be the most valuable to obtain for a particular instance.

As noted in (Shim et al., 2018), the active feature acquisition problem may be formulated as a Markov decision process (MDP), where the state is the set of currently observed features and the action is the next feature to acquire. A special action indicates whether to stop the acquisition process and make a final prediction. After acquiring its value and paying the acquisition cost, the newly acquired feature is added to the observed subset and the agent proceeds to the next acquisition step. Once the agent decides to terminate the acquisition, it makes a final prediction based on the features acquired thus far. For example, in an image classification task (Fig. 2), the agent would dynamically acquire pixels until it is certain of the image class. The goal of the agent is to maximize the prediction performance while minimizing the acquisition cost.

The key insight of this work is that the dynamics model for the AFA MDP is based on the conditionals of the features: $p(x_{j}\mid x_{o})$, where $x_{j}$ is an unobserved feature selected for acquisition and $x_{o}$ are the previously acquired features. Thus, we develop a model-based approach through generative modeling of all conditional dependencies. Equipped with the surrogate model, our method, Generative Surrogate Models for RL (GSMRL), essentially combines model-free and model-based RL into a holistic framework.

GSMRL rectifies several short-comings of a model-free scheme such as JAFA (Shim et al., 2018). 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. To reduce the sparsity of the rewards and simplify the credit assignment problem for potentially long episodes (Minsky, 1961; Sutton, 1988), we leverage a surrogate model to provide intermediate rewards by assessing the information gain of a newly acquired feature, which quantifies how much our confidence about the prediction improves by acquiring this feature. In addition to sparse rewards, an AFA agent must also navigate a complicated high-dimensional action space (Dulac-Arnold et al., 2015), and must manage multiple roles as it has to: implicitly model dependencies, perform a cost/benefit analysis, and act as a classifier. To lessen the burden, we also propose using the surrogate model to provide side information that assists the agent. The side information shall explicitly inform the agent of: 1) uncertainty and imputations for unobserved features; 2) an estimate of the expected information gain of future acquisitions; 3) uncertainty of the target output. This allows the agent to easily assess its current uncertainty and helps the agent ‘look ahead’ to expected outcomes from future acquisitions.

In this work, we also propose the first (to the best of our knowledge) unsupervised AFA task, which we coin active instance recognition (AIR). Here we consider the case where there is not a single target variable, but instead the target of interest may be the remaining unobserved features themselves. That is, rather than reducing the uncertainty with respect to some desired output response (that cannot be directly queried and must be predicted), the task is to query as few features as possible that allows the agent to correctly uncover the remaining unobserved features. For example, in image data AIR, an agent queries new pixels until it can reliably uncover the remaining pixels (see Fig. 2). AIR is especially relevant in surveying tasks, which are broadly applicable across various domains and applications. Most surveys aim to discover a broad set of underlying characteristics of instances (e.g., citizens in a census) using a limited number of queries (questions in the census form), which is at the core of AIR. Policies for AIR would build a personalized subset of survey questions (for individual instances) that quickly uncovered the likely answers to all remaining questions.

Our contributions are as follows: 1) We reformulate the AFA problem as a generative modeling task and build surrogate models that capture the state transitions with arbitrary conditional distributions. 2) We develop methodology to leverage the surrogate model to provide intermediate rewards as training signals and to provide auxiliary information that assists the agent. Our framework represents a novel combination of model-free and model-based RL. 3) We propose the first unsupervised active feature acquisition task where the target variables are the unobserved features themselves. 4) We achieve state-of-the-art performance on both supervised and unsupervised tasks in the largest scale AFA study to date. 5) We open-source a standardized environment inheriting the OpenAI gym interfaces (Brockman et al., 2016) to assist future research on active feature acquisition. Code will be released upon publication.

In this section, we first describe our GSMRL framework for both active feature acquisition (AFA) and active instance recognition (AIR) problems. We then develop our RL algorithm and the corresponding surrogate models for different settings. We also introduce a special application that acquires features for time series data.

Active Learning Active learning (Fu et al., 2013; Konyushkova et al., 2017; Yoo & Kweon, 2019) is a related approach in ML to gather more information when a learner can query an oracle for the true label, $y$, of a complete feature vector $x\in\mathbb{R}^{d}$ to build a better estimator. However, our methods consider queries to the environment for the feature value corresponding to an unobserved feature dimension, $i$, in order to provide a better prediction on the current instance. Thus, while the active learning paradigm queries an oracle during training to build a classifier with complete features, our paradigm queries the environment at evaluation to obtain missing features of a current instance to help its current assessment.

Feature Selection Feature selection (Miao & Niu, 2016; Li et al., 2017; Cai et al., 2018), ascertains a static subset of important features to eliminate redundancies, which can help reduce computation and improve generalization. Feature selection methods choose a fixed subset of features $s\subseteq\{1,\dots,d\}$, and always predict $y$ using this same subset of feature values, $x_{s}$. In contrast, our model considers a dynamic subset of features that is sequentially chosen and personalized on an instance-by-instance basis to increase useful information.

Active Feature Acquisition Instead of predicting the target passively using collected features, previous works have explored actively acquiring features in the cost-sensitive setting. Active perception is a relevant sub-field where a robot with a mounted camera is planning by selecting the best next view (Bajcsy, 1988; Aloimonos et al., 1988; Cheng et al., 2018; Jayaraman & Grauman, 2018). In this work we consider general features, and take images as one of many data sources. For general data, Ling et al. (2004), Chai et al. (2004) and Nan et al. (2014) propose decision tree, naive Bayes and maximum margin based classifiers, respectively, to jointly minimize the misclassification cost and feature acquisition cost. Ma et al. (2018) and Gong et al. (2019) acquire features greedily using mutual information as the estimated utility. Zubek et al. (2004) formulate the AFA problem as a MDP and fit a transition model using complete data, then they use the AO* heuristic search algorithm to find an optimal policy. Rückstieß et al. (2011) formulate the problem as a partially observable MDP and solve it using Fitted Q-Iteration. He et al. (2012) and He et al. (2016) instead employ a imitation learning approach guided by a greedy reference policy. Shim et al. (2018) utilize Deep Q-Learning and jointly learn a policy and a classifier. The classifier is treated as an environment that calculates the classification loss as the reward. ODIN (Zannone et al., 2019) presents an approach to learn a policy and a prediction model using augmented data with a Partial VAE (Ma et al., 2018). In contrast, GSMRL uses a surrogate model, which estimates both the state transitions and the prediction in a unified model, to directly provide intermediate rewards and auxiliary information to an agent.

Model-based and Model-free RL Reinforcement learning can be roughly grouped into model-based methods and model-free methods depending on whether they use a transition model (Li, 2017). Model-based methods are more data efficient but could suffer from significant bias if the dynamics are misspecified. On the contrary, model-free methods can handle arbitrary dynamic system but typically requires substantially more data samples. There have been works that combine model-free and model-based methods to compensate with each other. The usage of the model includes generating synthetic samples to learn a policy (Gu et al., 2016), back-propagating the reward to the policy along a trajectory (Heess et al., 2015), and planning (Chebotar et al., 2017; Pong et al., 2018). In this work, we rely on the model to provide intermediate rewards and side information. We compare this strategy to other model-based approaches in section 5.

In this section, we evaluate our method on several benchmark environments built upon the UCI repository (Dua & Graff, 2017) and MNIST dataset (LeCun, 1998). We compare our method to another RL based approach, JAFA (Shim et al., 2018), which jointly trains an agent and a classifier. We also compare to a greedy policy EDDI (Ma et al., 2018) that estimates the utility for each candidate feature using a VAE based model and selects one feature with the highest utility at each acquisition step. As a baseline, we also acquire features greedily using our surrogate model that estimates the utility following (6), (8) and (9). We use a fixed cost for each feature and report multiple results with different $\alpha$ in (1) to control the trade-off between task performance and acquisition cost. We cross validate the best architecture and hyperparameters for baselines. Architectural details, hyperparameters and sensitivity analysis are provided in the Appendix. In this work, we conduct the largest scale AFA study to date. Previous works have typically considered smaller datasets (both in terms of the number of features and the number of instances). We instead consider a broad range of datasets with more instances and higher dimensionality. In terms of comparisons, previous works often compare to naively simple baselines, such as a random acquisition order. In this work, we compare our GSMRL to the state-of-the-art models with both greedy policy and non-greedy RL policy.

Classification We first perform classification on the MNIST dataset. We downsample the original images to $16\times 16$ to reduce the action space to accommodate baselines such as EDDI that have trouble scaling (see Sec. D in the appendix for details on full MNIST).

Fig. 5 illustrates several examples of the acquired features and their prediction probability for different images. We can see that our model acquires a different subset of features for different images. Notice the checkerboard patterns of the acquired features, which indicates our model is able to exploit the spatial correlation of the data. Fig. 2 shows the acquisition process and the prediction probability along the acquisition. We can see the prediction become certain after acquiring only a small subset of features. The test accuracy in Fig. 10 demonstrates the superiority of our method over other baselines. It typically achieves higher accuracy with a lower acquisition cost. It is worth noting that our surrogate model with a greedy acquisition policy outperforms EDDI. We believe the improvement is due to the better distribution modeling ability of our surrogate model so that the utility and the prediction are more accurately estimated. We also perform classification using several UCI datasets. The test accuracy is presented in Fig. 7. Again, our method outperforms baselines under the same acquisition budget.

Regression We also conduct experiments for regression tasks using several UCI datasets. We report the root mean squared error (RMSE) of the target variable in Fig. 7. Similar to the classification task, our model outperforms baselines with a lower acquisition cost.

Time Series To evaluate the performance with constraints in action space, we classify over time series data where the acquired features must follow chronological ordering. The datasets are from the UEA & UCR time series classification repository (Bagnall et al., 2017). For GSMRL and JAFA, we clip the probability of invalid actions to zero; for the greedy method, we use a prior to bias the selection towards earlier time points. Please refer to Appendix B.4 for details. Fig. 9 shows the accuracy with different numbers of acquired features. Our method achieves high accuracy by collecting a small subset of the features.

Medical Diagnosis We evaluate the AFA performance for medical diagnosis. We use the Physionet challenge

2012 dataset (Goldberger et al., 2000) to predict the in-hospital mortality. Since the classes are heavily imbalanced, we use weighted cross entropy as training loss and the final rewards. For evaluation, we report the F1 scores in Fig. 11. Compared to baselines, our model achieves higher F1 with lower acquisition cost.

Unsupervised Next, we evaluate our method on unsupervised tasks where features are actively acquired to impute the unobserved features.

We use negative MSE as the reward for GSMRL and JAFA. The greedy policy calculates the utility following (9). For low dimensional UCI datasets, our method is comparable to baselines as shown in Fig. 9; but for the high dimensional case, as shown in Fig. 12, our method is doing better. Note JAFA is worse than the greedy policy for MNIST. We found it hard to train the policy and the reconstruction model jointly without the help of the surrogate model in this case. See Fig. 2 for an example of the acquisition process.

We now conduct a series of ablation studies to explore the capabilities of our GSMRL model.

Model-based Alternatives Our GSMRL model combines model-based and model-free approach into a holistic framework by providing the agent with auxiliary information and intermediate rewards. Here, we study different ways of utilizing the dynamics model. As in ODIN (Zannone et al., 2019), we utilize class conditioned generative models to generate synthetic trajectories.

The agent is then trained with both real and synthetic data (PPO+Syn). Another way of using the model is to extract a semantic embedding from the observations (Kumar et al., 2018). We use a pretrained EDDI to embed the current observed features into a 100-dimensional feature vector. An agent then takes the embedding as input and predicts the next acquisition (PPO+Embed). Figure 13 compares our method with these alternatives. We also present the results from a model-free approach as a baseline. We see our GSMRL outperforms other model-based approaches by a large margin.

Surrogate Models Our method relies on the surrogate model to provide intermediate rewards and auxiliary information. To better understand the contributions each component does to the overall framework, we conduct ablation studies using the MNIST dataset. We gradually drop one component from the full model and report the results in Fig. 15. The ‘Full Model’ uses both intermediate rewards and auxiliary information.

We then drop the intermediate rewards and denote it as ‘w/o rm’. The model without auxiliary information is denoted as ‘w/o aux’. We further drop both components and denote it as ‘w/o rm & aux’. From Fig. 15, we see these two components contribute significantly to the final results. We also compare models with and without the surrogate model. For models without a surrogate model, we train a classifier jointly with the agent as in JAFA. We plot the smoothed rewards using moving window average during training in Fig. 15. The agent with a surrogate model not only produces higher and smoother rewards but also converges faster.

Dynamic vs. Static Acquisition Our GSMRL acquires features following a dynamic order where it eventually acquires different features for different instances. A dynamic acquisition policy should perform better than a static one (i.e., the same set of features are acquired for each instance), since the dynamic policy allows the acquisition to be specifically adapted to the corresponding instance. To verify this is actually the case, we compare the dynamic and static acquisition under a greedy policy for MNIST classification.

Similar to the dynamic greedy policy, the static acquisition policy acquires the feature with maximum utility at each step, but the utility is averaged over the whole testing set, therefore the same acquisition order is adopted for the whole testing set. Figure 16 shows the classification accuracy for both EDDI and GSM under a greedy acquisition policy. We can see the dynamic policy is always better than the corresponding static one. Furthermore, our GSM with a static acquisition can already outperform dynamic EDDI.

Greedy vs. Non-greedy Acquisition Our GSMRL will terminate the acquisition process if the agent deems the current acquisition achieves the optimal trade-off between the prediction performance and the acquisition cost.

To evaluate how much the termination action affects the performance and to directly compare with the greedy policies under the same acquisition budget, we conduct an ablation study that removes the termination action and gives the agent a hard acquisition budget (i.e., forcing the agent to predict after some number of acquisitions). We can see (Fig. 17) GSMRL outperforms the greedy policy under all budgets. Moreover, we see that the agent is able to correctly assess whether or not more acquisitions are useful, since it obtains better performance when it dictates when to predict with the termination action.

In this work, we reformulate the dynamics of the AFA MDP as a generative modeling task among features. We leverage a generative surrogate model to capture the state transitions across arbitrary feature subsets. The surrogate model also provides auxiliary information and intermediate rewards to assist the agent. Our GSMRL model essentially combines model-based and model-free approaches. We conduct a large scale study to evaluate our model on both supervised and unsupervised AFA problems. Our model achieves state-of-the-art performance on both problems. In future work, we will explore AFA in spatial-temporal setting with continuously indexed features.