Learning to Maximize Mutual Information for Dynamic Feature Selection

Let ${\mathbf{x}}$ denote a vector of input features and ${\mathbf{y}}$ a response variable for a supervised learning task. The input consists of $d$ distinct features, or ${\mathbf{x}}=({\mathbf{x}}_{1},\ldots,{\mathbf{x}}_{d})$. We use the notation $s\subseteq[d]\equiv\{1,\ldots,d\}$ to denote a subset of indices and ${\mathbf{x}}_{s}=\{{\mathbf{x}}_{i}:i\in s\}$ a subset of features. Bold symbols ${\mathbf{x}},{\mathbf{y}}$ represent random variables, the symbols $x,y$ are possible values, and $p({\mathbf{x}},{\mathbf{y}})$ denotes the data distribution.

Our goal is to design a policy that controls which features are selected given the currently available information. The selection policy can be viewed as a function $\pi(x_{s})\in[d]$, meaning that it receives a subset of features as its input and outputs the next feature index to query. The policy is accompanied by a predictor $f(x_{s})$ that can make predictions given the set of available features; for example, if ${\mathbf{y}}$ is discrete then predictions lie in the probability simplex, or $f(x_{s})\in\Delta^{K-1}$ for $K$ classes. The notation $f(x_{s}\cup x_{i})$ represents the prediction given the combined features. We initially consider policy and predictor functions that operate on feature subsets, and Section 4 proposes an implementation using a mask variable $m\in[0,1]^{d}$ where the functions operate on $x\odot m$.

The goal of DFS is to select features with minimal budget that achieve maximum predictive accuracy. Having access to more features generally makes prediction easier, so the challenge is selecting a small number of informative features. There are multiple formulations for this problem, including versions with non-uniform feature costs and different budgets for each sample (Kachuee et al., 2018), but we focus on the setting with a fixed budget and uniform costs. Our goal is to handle predictions at inference time by beginning with no features, sequentially selecting features $x_{s}$ such that $|s|=k$ for a fixed budget $k<d$, and finally making accurate predictions for the response variable $y$.

Given a loss function that measures the discrepancy between predictions and labels $\ell(\hat{y},y)$, a natural scoring criterion is the expected loss after selecting $k$ features. The scoring is applied to a policy-predictor pair $(\pi,f)$, and we define the score for a fixed budget $k$ as follows,

where feature indices are chosen sequentially for each $({\mathbf{x}},{\mathbf{y}})$ according to $i_{n}=\pi(\{{\mathbf{x}}_{i_{t}}\}_{t=1}^{n-1})$. The goal is to minimize $v_{k}(\pi,f)$, or equivalently, to maximize our final predictive accuracy.

One approach is to frame this as a Markov decision process (MDP) and solve it using standard RL techniques, so that $\pi$ and $f$ are trained to optimize a reward function based on eq. 1. Several recent works have designed such formulations (Shim et al., 2018; Kachuee et al., 2018; Janisch et al., 2019; Li & Oliva, 2021). However, these approaches are difficult to train effectively, so our work focuses on a greedy approach that is easier to learn and simpler to interpret.

As an idealized approach to DFS, we are interested in the greedy algorithm that selects the most informative feature at each step. This feature can be defined in multiple ways, but we focus on the information-theoretic perspective that the most useful feature has maximum CMI with the response variable (Cover & Thomas, 2012). The CMI, denoted as $I({\mathbf{x}}_{i};{\mathbf{y}}\mid x_{s})$, quantifies how much information an unknown feature ${\mathbf{x}}_{i}$ provides about the response ${\mathbf{y}}$ when accounting for the current features $x_{s}$, and it is defined as the KL divergence between the joint and factorized distributions:

Based on this, we define the greedy CMI policy as $\pi^{*}(x_{s})\equiv\operatorname*{arg\,max}_{i}I({\mathbf{x}}_{i};{\mathbf{y}}\mid x_{s})$, so that features are sequentially selected to maximize our information about the response variable. We can alternatively understand the policy as performing greedy uncertainty minimization, because this is equivalent to minimizing ${\mathbf{y}}$’s conditional entropy at each step, or $\pi^{*}(x_{s})=\operatorname*{arg\,min}_{i}H({\mathbf{y}}\mid{\mathbf{x}}_{i},x_{s})$ (Cover & Thomas, 2012). For a complete characterization of this idealized approach, we also consider that the policy is paired with the Bayes classifier as a predictor, or $f^{*}(x_{s})=p({\mathbf{y}}\mid x_{s})$.

Maximizing the information about ${\mathbf{y}}$ at each step is intuitive and should be effective in many problems. Prior work has considered the same idea, but from two perspectives that differ from ours. First, Chen et al. (2015b) take a theoretical perspective and prove that the greedy algorithm achieves performance within a multiplicative factor of the optimal policy; the proof requires specific distributional assumptions, but we find that the greedy algorithm performs well with many real-world datasets (Section 6). Second, from an implementation perspective, two works aim to provide practical approximations; however, these suffer from several limitations, so our work aims to develop a simple and flexible alternative (Section 4). In these works, Ma et al. (2019) and Chattopadhyay et al. (2022) both require a conditional generative model of the data distribution, which we discuss next.

The greedy policy is trivial to implement if we can directly calculate CMI, but this is rarely the case in practice. Instead, one option is to estimate it. We now describe a procedure to do so iteratively for each feature, assuming for now that we have oracle access to the response distributions $p({\mathbf{y}}\mid{\mathbf{x}}_{s})$ for all $s\subseteq[d]$ and the feature distributions $p({\mathbf{x}}_{i}\mid{\mathbf{x}}_{s})$ for all $s\subseteq[d]$ and $i\in[d]$.

At any point in the selection procedure, given the current features $x_{s}$, we can estimate the CMI for a feature ${\mathbf{x}}_{i}$ where $i\notin s$ as follows. First, we can sample multiple values for ${\mathbf{x}}_{i}$ from its conditional distribution, or $x_{i}^{j}\sim p({\mathbf{x}}_{i}\mid x_{s})$ for $j\in[n]$. Next, we can generate Bayes optimal predictions for each sampled value, or $p({\mathbf{y}}\mid x_{s},x_{i}^{j})$. Finally, we can calculate the mean prediction and the mean KL divergence relative to this prediction, which yields the following CMI estimator:

This score measures the variability among predictions and captures whether different ${\mathbf{x}}_{i}$ values significantly affect ${\mathbf{y}}$’s conditional distribution. The estimator can be used to select features, or we can set $\pi(x_{s})=\operatorname*{arg\,max}_{i}I_{i}^{n}$, due to the following limiting result (see Appendix A):

This procedure thus provides a way to identify the correct greedy selections by estimating the CMI. Prior work has explored similar ideas for scoring features based on sampled predictions (Saar-Tsechansky et al., 2009; Chen et al., 2015a; Early et al., 2016a, b), but the implementation choices in these works prevent them from performing greedy information maximization. In eq. 2, is it important that our estimator uses the Bayes classifier, that we sample features from the conditional distribution $p({\mathbf{x}}_{i}\mid x_{s})$, and that we use the KL divergence as a measure of prediction variability. However, this estimator is impractical because we typically lack access to both $p({\mathbf{y}}\mid{\mathbf{x}}_{s})$ and $p({\mathbf{x}}_{i}\mid{\mathbf{x}}_{s})$.

In practice, we would instead require learned substitutes for each distribution. For example, we can use a a classifier that approximates $f(x_{s})\approx p({\mathbf{y}}\mid x_{s})$ and a generative model that approximates samples from $p({\mathbf{x}}_{i}\mid{\mathbf{x}}_{s})$. Similarly, Ma et al. (2019) propose jointly modeling $({\mathbf{x}},{\mathbf{y}})$ with a conditional generative model, which is implemented via a modified VAE (Kingma et al., 2015). This approach is limited for several reasons, including (i) the difficulty of training an accurate conditional generative model, (ii) the challenge of modeling mixed continuous/categorical features (Ma et al., 2020; Nazabal et al., 2020), and (iii) the slow CMI estimation process. In our approach, which we discuss next, we bypass all three of these challenges by directly predicting the best selection at each step.

For our purpose, it is helpful to recognize that the greedy policy can be viewed as the solution to an optimization problem. Section 3 provides a conventional definition of CMI as a KL divergence, but this is difficult to integrate into an end-to-end learning approach. Instead, we now consider the one-step-ahead prediction achieved by a policy $\pi$ and predictor $f$, and we determine the behavior that minimizes their loss. Given the current features $x_{s}$ and a selection $i=\pi(x_{s})$, the expected one-step-ahead loss is:

The variational perspective we develop here consists of two main results regarding this expected loss. The first result relates to the predictor, and we show that the loss-minimizing predictor can be defined independently of the policy $\pi$. We formalize this in the following proposition for classification tasks, and our results can also be generalized to regression tasks (see proofs in Appendix A).

This property requires that features are selected without knowledge of the remaining features or the response variable, which is a valid assumption for DFS, but not in scenarios where selections are based on the full feature set (Chen et al., 2018; Yoon et al., 2018; Jethani et al., 2021). Now, assuming that we use the Bayes classifier $f^{*}$ as a predictor, our second result concerns the selection policy. As we show next, the loss-minimizing policy is equivalent to making selections based on CMI.

With this, we can see that the greedy CMI policy defined in Section 3 is equivalent to minimizing the one-step-ahead prediction loss. Next, we exploit this variational perspective to develop a joint learning procedure for a policy and predictor network.

Instead of estimating each feature’s CMI to identify the next selection, we now develop an approach that directly predicts the best selection at each step. The greedy policy implicitly requires solving an optimization problem for each selection, or $\operatorname*{arg\,max}_{i}I({\mathbf{y}},{\mathbf{x}}_{i};x_{s})$, but since we lack access to this objective, we now formulate an approach that directly predicts the solution. Following a technique known as amortized optimization (Amos, 2022), we do so by casting our variational perspective on CMI from Section 4.1 as an objective function to be optimized by a learnable network.

First, because it facilitates gradient-based optimization, we now consider that the policy outputs a distribution over feature indices. With slight abuse of notation, this section lets the policy be a function $\pi(x_{s})\in\Delta^{d-1}$, which generalizes the previous definition $\pi(x_{s})\in[d]$. Using this stochastic version of the policy, we can now formulate our objective function as follows.

Let the selection policy be parameterized by a neural network $\pi({\mathbf{x}}_{s};\phi)$ and the predictor by a neural network $f({\mathbf{x}}_{s};\theta)$. Let $p({\mathbf{s}})$ represent a distribution with support over all subsets, or $p(s)>0$ for all $|s|<d$. Then, our objective function $\mathcal{L}(\theta,\phi)$ is defined as

Intuitively, eq. 5 represents generating a random feature set ${\mathbf{x}}_{s}$, sampling a feature index according to $i\sim\pi({\mathbf{x}}_{s};\phi)$, and then measuring the loss of the prediction $f({\mathbf{x}}_{s}\cup{\mathbf{x}}_{i};\theta)$. Our objective thus optimizes for individual selections and predictions rather than the entire trajectory, which lets us build on Proposition 1-2. We describe this as an implementation of the greedy approach because it recovers the greedy CMI selections when it is trained to optimality. In the classification case, we show the following result under a mild assumption that there is a unique optimal selection.

If we relax the assumption of a unique optimal selection, the optimal policy $\pi({\mathbf{x}}_{s};\phi^{*})$ simply splits probability mass among the best indices. A similar result holds in the regression case, where we can interpret the greedy policy as performing conditional variance minimization.

Proofs for these results are in Appendix A. We note that the function class for each model must be expressive enough to contain their respective optima, and that the result holds for any $p({\mathbf{s}})$ with support over all subsets.

This approach has two key advantages over the CMI estimation procedure from Section 3.2. First, we avoid modeling the feature conditional distributions $p({\mathbf{x}}_{i}\mid{\mathbf{x}}_{s})$ for all $(s,i)$. Modeling these distributions is a difficult intermediate step, and our approach instead aims to directly output the optimal index. Second, our approach is faster because each selection is made in a single forward pass: selecting $k$ features using the procedure from Ma et al. (2019) requires $\mathcal{O}(dk)$ scoring steps, but our approach requires only $k$ forward passes through the policy network $\pi({\mathbf{x}}_{s};\phi)$.

Furthermore, compared to a policy trained by RL, the greedy approach is easier to learn. Our training procedure can be viewed as a form of reward shaping (Sutton et al., 1998; Randløv & Alstrøm, 1998), where the reward accounts for the loss after each step and provides a strong signal about whether each selection is helpful. In comparison, observing the reward only after selecting $k$ features provides a comparably weak signal to the policy network (see eq. 1). RL methods generally face a challenging exploration-exploitation trade-off, but learning the greedy policy is simpler because it only requires finding the locally optimal choice at each step.

Our objective in eq. 5 yields the correct greedy policy when it is perfectly optimized, but $\mathcal{L}(\theta,\phi)$ is difficult to optimize by gradient descent. In particular, gradients are difficult to propagate through the policy network given a sampled index $i\sim\pi({\mathbf{x}}_{s};\phi)$. The Reinforce trick (Williams, 1992) is one way to get stochastic gradients, but high gradient variance can make it ineffective in many problems. There is a robust literature on reducing gradient variance in this setting (Tucker et al., 2017; Grathwohl et al., 2018), but we propose using a simple alternative: the Concrete distribution (Maddison et al., 2016).

An index sampled according to $i\sim\pi(x_{s};\phi)$ can be represented by a one-hot vector $m\in\{0,1\}^{d}$ indicating the chosen index, and with the Concrete distribution we instead sample an approximately one-hot vector in the probability simplex, or $m\in\Delta^{d-1}$. This continuous relaxation lets us calculate gradients using the reparameterization trick (Maddison et al., 2016; Jang et al., 2016). Relaxing the subset $s\subseteq[d]$ to a continuous vector also requires relaxing the policy and predictor functions, so we let these operate on a masked input $x$, or the element-wise product $x\odot m$. To avoid ambiguity about whether features are zero or masked, we can also pass the mask as a model input.

Training with the Concrete distribution requires specifying a temperature parameter $\tau>0$ to control how discrete the samples are. Previous works have typically trained with a fixed temperature or annealed it over a pre-determined number of epochs (Chang et al., 2017; Chen et al., 2018; Balın et al., 2019), but we instead train with a sequence of $\tau$ values and perform early stopping at each step. This removes the temperature and number of epochs as important hyperparameters to tune. Our training procedure is summarized in Figure 1, and in more detail by Algorithm 1.

There are also several optional steps that we found can improve optimization:

• •

  Parameters can be shared between the predictor and policy networks $f({\mathbf{x}};\theta),\pi({\mathbf{x}},\phi)$. This does not complicate their joint optimization, and learning a shared representation in the early layers can in some cases help the networks optimize faster (e.g., for image data).

• •

  Rather than training with a random subset distribution $p({\mathbf{s}})$, we generate subsets using features selected by the current policy $\pi({\mathbf{x}};\phi)$. This allows the models to focus on subsets likely to be encountered at inference time, and it does not affect the globally optimal policy/predictor: gradients are not propagated between selections, so both eq. 5 and this sampling approach treat each feature set as an independent optimization problem, only with different weights (see Appendix D).

• •

  We pre-train the predictor $f({\mathbf{x}};\theta)$ using random subsets before jointly training the policy-predictor pair. This works better than optimizing $\mathcal{L}(\theta,\phi)$ from a random initialization, because a random predictor $f({\mathbf{x}};\theta)$ provides no signal to $\pi({\mathbf{x}};\phi)$ about which features are useful.

We first applied our method to three medical diagnosis tasks derived from an emergency medicine setting. The tasks involve predicting a patient’s bleeding risk via a low fibrinogen concentration (bleeding), whether the patient requires endotracheal intubation for respiratory support (respiratory), and whether the patient will be responsive to fluid resuscitation (fluid). See Appendix B for more details about the datasets. In each scenario, gathering all possible inputs at inference time is challenging due to time and resource constraints, thus making DFS a natural solution.

We use fully connected networks for all methods, and we use dropout to reduce overfitting (Srivastava et al., 2014). Figure 2 (top) shows the results of applying each method with various feature budgets. The classification accuracy is measured via AUROC, and the greedy method achieves the best results for nearly all feature budgets on all three tasks. Among the baselines, several static methods are sometimes close, but the CMI estimation method is rarely competitive (Ma et al., 2019). Additionally, OL provides unstable and weak results. The greedy method’s advantage is often largest when selecting a small number of features, and it usually becomes narrower once the accuracy saturates.

Next, we conducted experiments using three publicly available tabular datasets: spam classification (Dua & Graff, 2017), particle identification (MiniBooNE) (Roe et al., 2005) and diabetes diagnosis (Miller, 1973). The diabetes task is a natural application for DFS and was used in prior work (Kachuee et al., 2018). We again tested various numbers of features, and Figure 2 (bottom) shows plots of the AUROC for each feature budget. On these tasks, the greedy method is once again most accurate for nearly all numbers of features. Table 1 summarizes the results via the mean AUROC across $k=1,\ldots,10$ features, further emphasizing the benefits of the greedy method across all six datasets. Appendix E shows larger versions of the AUROC curves (Figure 4 and Figure 5), as well as plots demonstrating the variability of selections within each dataset.

The results with these datasets reveal that, perhaps surprisingly, dynamic methods can be outperformed by static methods. Interestingly, this point was not highlighted in prior works where strong static baselines were not tested (Kachuee et al., 2018; Janisch et al., 2019). For example, OL is not competitive on these datasets, and the two versions of the CMI estimation approach are not consistently among the top baselines. Dynamic methods are in principle capable of performing better, so the sub-par results from these methods underscore the difficulty of learning both a selection policy and a prediction function that works for multiple feature sets. In these experiments, our approach is the only dynamic method to do both successfully.

Next, we considered two standard image classification datasets: MNIST (LeCun et al., 1998) and CIFAR-10 (Krizhevsky et al., 2009). Our goal is to begin with a blank image, sequentially reveal multiple pixels or patches, and ultimately make a classification using a small portion of the image. Although this is not an obvious use case for DFS, it represents a challenging problem for our method, and similar tasks were considered in several earlier works (Karayev et al., 2012; Mnih et al., 2014; Early et al., 2016a; Janisch et al., 2019).

For MNIST, we use fully connected architectures for both the policy and predictor, and we treat pixels as individual features; we therefore have $d=784$. For CIFAR-10, we use a shared ResNet backbone (He et al., 2016b) for the policy and predictor networks, and each network uses its own output head. The $32\times 32$ images are coarsened into $d=64$ patches of size $4\times 4$, so the selector head generates logits corresponding to each patch, and the predictor head generates probabilities for each class.

Figure 3 shows our method’s accuracy for different feature budgets. For MNIST, we use the previous baselines but exclude the CMI estimation method due to its computational cost: it becomes slow when evaluating many candidate features. We observe a large benefit for our method, particularly when making a small number of selections. Our greedy method reaches nearly 90% accuracy with just 10 pixels, which is roughly 10% higher than the best baseline and considerably higher than prior work (Balın et al., 2019; Yamada et al., 2020; Covert et al., 2020). OL yields the worst results, and it also trains slowly due to the large number of states.

For CIFAR-10, we omit several baseline comparisons due to their computational cost. We use the CAE, which is our most competitive static baseline, as well as two simple baselines: center crops and random masks of various sizes. For each method, we plot the mean and 95% confidence intervals determined from five trials. Our greedy approach is slightly less accurate with a very small number of patches, but it reaches significantly higher accuracy when using 6-20 patches. Finally, Figure 3 (bottom) also shows qualitative examples of our method’s predictions after selecting 10 out of 64 patches, and Appendix E shows similar plots with different numbers of patches.