Active Feature Acquisition with Generative Surrogate Models

1 INTRODUCTION . A typical machine learning paradigm for discriminative tasks is to learn the distribution of an output , y given a complete set of features , x ∈ Rd : p ( y | 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 , not each patient will 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 the return on improvement of the task 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 what next feature would be the most valuable to obtain for a particular instance . In this work , we formulate the active feature acquisition problem 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 . We also introduce a special action to indicate whether to stop the acquisition process and make a final prediction . Reinforcement learning is then utilized to optimize the MDP , and the agent learns a policy for selecting which next feature to acquire based on the current state . 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 . 1 ) , 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 . 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 . The surrogate model captures the arbitrary conditional distribution p ( y , xu | xo ) , where y is the target variable and u , o ⊆ { 1 , . . . , d } are arbitrary subsets of all d-dimensional features . Note that the surrogate model must be able to capture arbitrary conditionals ( for subsets u , o ) since the acquired features will vary from instance to instance . We propose using the surrogate model to calculate intermediate rewards by assessing the information gain of the newly acquired feature , which quantifies how much our confidence about the prediction improves by acquiring this feature . In addition to producing intermediate rewards , we also propose using the surrogate model to provide side information that assists the agent . First , in order to inform the agent of the current information held in observed features , we pass uncertainty on the target through p ( y | xo ) . Second , to inform the agent about potential values for unobserved features , we pass imputed values by sampling x̂u ∼ p ( xu | xo ) . Lastly , to inform the agent about the expected utility of acquisitions , we pass an estimate of the expected information gain of acquisitions i for the target variable , i.e. , H ( y | xo ) − Ep ( xi|xo ) H ( y | xi , xo ) . We note that the expected information gain can be used to directly build a greedy policy , where the next feature to acquire is the one maximizes the expected information gain ( Ma et al. , 2018 ; Gong et al. , 2019 ) . In contrast , our agent learns a non-greedy policy to maximize the long-term returns and use the greedy approach as a ‘ prior ’ policy to guide our agent . In summary , our agent actively acquires new feature and pays the acquisition cost until it decides to terminate the acquisition process and make a final prediction . Meanwhile , the surrogate model calculates the information gain of the acquired feature as an intermediate reward and provides side information to assist the agent in assessing its current uncertainty and help it ‘ look ahead ’ to expected outcomes from future acquisitions . When the acquisition process is completed , the environment provides a final reward based on the agent ’ s prediction . Note that the environment does have access to the ground-truth target y to evaluate the reward , but can not reveal it to the agent . Equipped with the surrogate model , our method , denoted as GSMRL , essentially combines model-free and modelbased RL into a holistic framework . Above we discussed AFA for supervised tasks , where the goal is to acquire new features to predict a target variable y . In some cases , however , there may not be a single target variable , but instead the target of interest may be the remaining unobserved features themselves . That is , rather than reduce the uncertainty with respect to some desired output response ( that can not be directly queried and must be predicted ) , we now propose active instance recognition ( AIR ) , where 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 survey 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 Algorithm 1 : Active Feature Acquisition with GSMRL 1. load pretrained surrogate model M , agent agent and prediction model fθ ( · ) ; 2. instantiate an environment with data D and surrogate model M : env = Environment ( D , M ) ; xo , o , done , reward = env.reset ( ) ; // o = ∅ , done=False , reward=0 while not done do aux = M .query ( xo , o ) ; // query M for auxiliary information // aux contains the prediction ŷ ∼ p ( y | xo ) and output likelihoods , // the imputed values x̂u ∼ p ( xu | xo ) and their uncertainties , // and estimated utilities Ui for each i ∈ u ( equation 4 ) . action = agent.act ( xo , o , aux ) ; // act based on the state and auxiliary info xo , o , done , r = env.step ( action ) ; // take a step based on the action // if action indicates termination : done=True , r=-L ( ŷ ( xo ) , y ) // else : done=False , r=rm ( s , action ) − αC ( action ) , o = o ∪ action reward = reward + r ; // accumulate rewards end prediction = M .predict ( xo , o , aux ) or fθ ( xo , o , aux ) ; // make a final prediction of survey questions ( for individual instances ) that quickly uncovered the likely answers to all remaining questions . To adapt our GSMRL framework to AIR , we set the target variable y equal to x and modify the surrogate model accordingly . Our contributions are as follows : 1 ) We propose a way of building surrogate models for AFA problem that captures the state transitions with arbitrary conditional distributions . 2 ) We 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 extend the active feature acquisition problem to an unsupervised case where the target variables are the unobserved features themselves . Our RL agent can be adapted to this problem with simple modifications . 4 ) We achieve state-of-the-art performance on both supervised and unsupervised tasks . 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 . 2 METHODS . 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 . 2.1 AFA AND AIR WITH GSMRL . Consider a discriminative task with features x ∈ Rd and target y . Instead of predicting the target by first collecting all the features , we perform a sequential feature acquisition process in which we start from an empty set of features and actively acquire more features . There is typically a cost associated with features and the goal is to maximize the task performance while minimizing the acquisition cost , i.e. , minimize L ( ŷ ( xo ) , y ) + αC ( o ) , ( 1 ) where L ( ŷ ( xo ) , y ) represents the loss function between the prediction ŷ ( xo ) and the target y . Note that the prediction is made with the acquired feature subset xo , o ⊆ { 1 , . . . , d } . Therefore the agent should be able to predict with arbitrary subset as inputs . C ( o ) represents the acquisition cost of the acquired features o . The hyperparameter α controls the trade-off between prediction loss and acquisition cost . For unsupervised tasks , the target variable y is equal to x ; that is , we acquire features actively to represent the instance with a selected subset . In order to solve the optimization problem in equation 1 , we formulate it as a Markov decision process as done in ( Shim et al. , 2018 ) : s = [ o , xo ] , a ∈ u ∪ φ , r ( s , a ) = −L ( ŷ , y ) I ( a = φ ) − αC ( a ) I ( a �= φ ) . ( 2 ) The state s is the current acquired feature subset o ⊆ { 1 , . . . , d } and their values xo . The action space contains the remaining candidate features u = { 1 , . . . , d } \ o and a special action φ that in- dicates the termination of the acquisition process . To optimize the MDP , a reinforcement learning agent acts based on the observed state and receives rewards from the environment . When the agent acquires a new feature i , the current state transits to a new state following o i−→ o ∪ i , xo i−→ xo ∪ xi , and the reward is the negative acquisition cost of this feature . Note xi is obtained from the environment ( i.e . we observe the true ith feature value for the instance ) . When the agent terminates the acquisition and makes a prediction , the reward equals to the negative prediction loss using current acquired features . Since the prediction is made at the end of the acquisition , the reward of the prediction is received only when the agent decide to terminate the acquisition process . This is a typical temporal credit assignment problem for RL algorithms , which could affect the learning of the agent ( Minsky , 1961 ; Sutton , 1988 ) . In order to remedy this issue , we propose to leverage a generative surrogate model to provide intermediate rewards for each acquisition . The surrogate model estimates the state transitions with arbitrary conditional distributions p ( y , xu | xo ) for arbitrary subsets u and o . We propose using the surrogate model to assess the intermediate reward rm for a newly acquired feature i . The intermediate rewards are inspired by the information gain to the target variable rm ( s , i ) = H ( y | xo ) − γH ( y | xo , xi ) , ( 3 ) where γ is a discount factor for the MDP . In appendix A , we show that our intermediate rewards will not change the optimal policy . In addition to intermediate rewards , we propose using the surrogate model to also provide side information to assist the agent , which includes the current prediction and output likelihood , the possible values and corresponding uncertainties of the unobserved features , and the estimated utilities of the candidate acquisitions . The current prediction ŷ and likelihood p ( y | xo ) inform the agent about its confidence , which can help the agent determine whether to stop the acquisition . The imputed values and uncertainties of the unobserved features give the agent the ability to look ahead into and future and guide its exploration . For example , if the surrogate model is very confident about the value of a currently unobserved feature , then acquiring it would be redundant . The utility of a feature i is estimated by its expected information gain to the target variable : Ui = H ( y | xo ) − Ep ( xi|xo ) H ( y | xi , xo ) = H ( xi | xo ) − Ep ( y|xo ) H ( xi | y , xo ) , ( 4 ) where the surrogate model is used to estimate the entropies . The utility essentially quantifies the conditional mutual information I ( xi ; y | xo ) between each candidate feature and the target variable . A greedy policy can be easily built based on the utilities where the next feature to acquire is the one with maximum utility ( Ma et al. , 2018 ; Gong et al. , 2019 ) . Here , our agent takes the utilities as side information to help balance exploration and exploitation , and eventually learns a non-greedy policy . When the agent deems that acquisition is complete , it makes a final prediction based on the acquired feaures thus far . The final prediction may be made using the surrogate model , i.e. , p ( y | xo ) , but it might be beneficial to train predictions specifically based on the agent ’ s own distribution of acquired features o , since the surrogate model is agnostic to the feature acquisition policy of the agent . Therefore , we optionally build a prediction model fθ ( · ) that takes both the current state xo and the side information as inputs ( i.e . the same inputs as the policy ) . The prediction model can be trained simultaneously with the policy as an auxiliary task ; weight sharing between the policy and prediction function helps facilitate the learning of more meaningful representations . Now we have two predictions , from the surrogate model and the prediction model respectively . The final reward −L ( ŷ , y ) during training is the maximum one using either predictions . During test time , we choose one prediction based on validation performance . An illustration of our framework is presented in Fig . 3 . Please refer to Algorithm 1 for the pseudo-code of the acquisition process with our GSMRL framework . We will expound on the surrogate models for different settings below .

Summary: In this work, a reinforcement learning (RL) approach is proposed to solve the active feature acquisition (AFA) problem (as well as the active instance recognition problem). Comparing to existing RL approaches for AFA, the main difference of the proposed approach is to introduce a generative model (utilizing the existing ACFlow model) to learn the transition function, in order to provide additional rewards and auxiliary information. The proposed approach is evaluated on MNIST and UCI datasets, which can outperform two existing baselines.