Estimating Conditional Mutual Information for Dynamic Feature Selection

We would like to thank the reviewer for closely examining our work and providing their feedback. We have used the space below to respond to the questions raised in the review. Updates to the manuscript are highlighted in red to make them easily identifiable.

“While the incorporation of acquisition costs with CMI estimates is straight-forward, the acquisition costs appear to be an artificial supplementary add-on to the methodology.”

Once we are able to accurately predict the CMI with DIME, extending the method to incorporate non-uniform feature costs may seem like a small add-on, but it makes DIME more useful in real-world settings where features have diverse acquisition costs. Taking the emergency medicine setting as a motivating example, there are many features that can be acquired with significant variation in the temporal and monetary cost. Demographic features like age, sex, and patient history would be cheap and only take only minutes to obtain, whereas lab test features might take days and would be more expensive, and genetic tests could take weeks to conduct. Incorporating non-uniform feature costs allows us to take these differences into consideration, and DIME can use multiple low-cost features in lieu of one high cost feature to get the same diagnostic performance.

“Some related approaches should be mentioned in the related work section and could have been considered as competitors”

Thank you for pointing us to these very interesting works! We carefully read them both, and after studying their proposed methods, we do not think they are closely related to our problem setting. Fumagalli et al. (2023) study global feature importance in the context of streaming data, where new data points come in and the distribution may change over time; in contrast, we focus on feature selection with selections that can differ between data points, and there is no analogous notion of distribution shift. Haug et al. (2020) focus on online feature selection, which is also related to data streams where the distribution can shift over time; our work considers the iterative acquisition of features for each prediction, but this is very different from the notion of incoming data streams with changing distributions. As such, we do not think there is any room to compare DIME with these methods.

“It is not quite clear how prior information is factored into and affects the feature selection policy; this part should be better explained”

Section 4.2 describes how our approach can be modified to include prior information: essentially, $z$ becomes an input to both the predictor and value networks, and they are optimized in the same way as before using the same loss functions. The predictor network is trained with the loss $\ell(f(x_S, z; \theta), y)$ (e.g., cross entropy loss) and the value network is trained with $(v_i(x_S, z; \phi) - \Delta(x_S, x_i, z, y))^2$. Theorem 2 describes how this should theoretically affect the dynamic feature selection procedure.

We demonstrate this approach with the histopathology dataset, using a Canny edge image as prior information. During training, we use separate ViT backbones for the original and the edge images, and we pass these images as inputs to both the value network and the predictor. Specifically, we concatenate the resulting embeddings of the original and edge images from the backbones before either getting the CMI estimates from $v$ or the predictions from $f$, which are implemented as small networks that take concatenated embeddings as their input. We can update this section in the manuscript to better explain how the prior information is passed to the networks.

“An exact algorithmic representation should be used to provide clarity on the interplay between the networks, the CMI estimation, the use of prior information, and how the features are finally selected in various scenarios.”

For more information about how the networks are trained, please see Algorithm 1 in Appendix C, which shows how the value network and predictor network are jointly optimized. We are happy to specify how the algorithm incorporates prior information, as well as add another one describing how features are selected at inference time. We will also add a reference to these algorithms in the main text, which we forgot to include in the original submission.

Thank you for taking the time to read our response and adjust your score, we appreciate it! We will consider the differences between our method and data streams.

We would like to thank the reviewer for closely examining our work and providing their feedback. We have used the space below to respond to the questions raised in the review. Updates to the manuscript are highlighted in red to make them easily identifiable.

“As far as I understand, the authors propose a method to learn CMI improvement. It would be nice to compare this method to the ones, where CMI is approximated in the direct statistical way. For instance, there is a lot of approaches like MIM, MIFS, mRMR, JMI, CMICOT, etc. See for a brief overview of them in [Shishkin, A., et al. “Efficient high-order interaction-aware feature selection based on conditional mutual information”, NeurIPS 2016]. Could you please compare your approach with such statistical methods? Including, in experiments. Is it true that the gain we obtain from learning CMI improvement over statistical methods is cost effective? (learning consumes more computation => on the same level of computation costs, we can get more from statistical methods, right?)”

The reviewer is correct that DIME estimates each feature’s CMI with the response variable given a set of observed features. However, our problem formulation is different from all the methods that the reviewer points to: these methods perform static feature selection, which means that the features are selected just once and used for all predictions, whereas DIME performs dynamic feature selection, which means that we select features separately for each prediction. To be precise, these static methods are concerned with the quantity $I(\mathbf{y}; \mathbf{x}_i \mid \mathbf{x}_S)$, while DIME aims to estimate $I(\mathbf{y}; \mathbf{x}_i \mid x_S)$ with $x_S$ being a set of observed features for the current prediction. Estimating these two quantities requires very different methodologies, and it is not clear if the procedures used by the suggested methods are applicable here or if they are computationally feasible to be run repeatedly for each selection performed for each prediction.

As for the static feature selection methods mentioned by the reviewer, we did not include these initially because we instead use a more recent state-of-the-art static baseline: the Concrete Autoencoder (Balin et al., 2019). For completeness, we tried comparing DIME with mRMR and CMICOT, and the results are now shown in Figure 2. We can confirm that these methods both underperform DIME. CMICOT does better than mRMR, but both methods underperform the Concrete Autoencoder, our original static baseline. Note that we chose mRMR and CMICOT among the suggested methods because CMICOT has been shown to outperform the others and both of them had public implementations that were easy to incorporate into our problem.

“Returning back to [Shishkin, A., et al. “Efficient high-order interaction-aware feature selection based on conditional mutual information”, NeurIPS 2016], we can see that greedy approach of selecting features one by one is not the best way. The ideal approach is to select the best subset of features (which is computationally and dimensionally infeasible). But we at least can select features not one-by-one, but taking into account interactions between selecting features (like in CMICOT). Is there any way to incorporate that approach into the work of the authors?”

As described in the previous response, a key difference between CMICOT and DIME is whether features are selected just once (static) or separately for each prediction (dynamic). In our case, DIME can only observe features once they are acquired, so features must be selected one at a time. However, in a sense DIME’s training procedure does consider feature interactions, because the CMI conditions on the subset of features that are already observed. For example, if two features are perfectly correlated, after observing one of them, the other’s CMI would drop to zero because acquiring it will provide no additional information about the response variable.

We would like to thank the reviewer for closely examining our work and providing their feedback. We have used the space below to respond to the questions raised in the review. Updates to the manuscript are highlighted in red to make them easily identifiable.

“I believe that some relevant work on dynamic feature selection that is neither RL-based nor greedy-based is missing”

We agree that formulating the problem as a POMDP and using a Bayesian network to capture feature dependencies are valid approaches, and we are happy to include other references in the related work. The suggested works are now included in the updated manuscript, with descriptions based on our best understanding of those works.

“I also believe that the algorithms proposed in the previous references as well as (Janisch et al; 2019) are designed to handle non-uniform costs and variable feature budgets. I think this need to be clearly stated in the paper, since at the moment, it is suggested that only DIME can accommodate such settings”

We don’t wish to suggest that DIME is the only method that allows non-uniform costs and variable feature budgets. In the original submission, we cite Janisch et al. (2019) and Kachuee et al. (2018) as examples of methods that consider these, see for example the following quote: “For example, Janisch et al. (2019) formulate DFS as a MDP where the reward is the 0-1 loss minus the feature cost, while considering both variable budgets and non-uniform feature costs.” Notice that we also compare DIME to EDDI and CwCF in Figure 3 where we consider non-uniform costs.

However, in the context of greedy approaches for dynamic feature selection, it’s worth noting that most previous works focus on the fixed-budget setting with uniform costs. DIME is therefore unique in enabling these additions to the greedy approach, and this is possible only because it more accurately estimates the CMI in a discriminative fashion. If it is helpful, we can update the manuscript to make this point clearer.

“The performance of DIME is very close to Argmax Direct. What is the benefit of using DIME instead of Argmax Direct in that sense?”

The performance of Argmax Direct approaches that of DIME for two tabular datasets (Intubation and ROSMAP), but for the other 4/6 datasets we tested, DIME substantially outperforms Argmax Direct for almost all feature budgets. Furthermore, DIME addresses real-world challenges that Argmax-Direct cannot address: first, feature acquisition costs are not always identical for all features (e.g., a genetic test takes longer than obtaining demographic features), and second, allowing variable per-sample feature budgets enables quicker diagnosis for certain samples (e.g., it may become immediately clear that a patient with a neck injury needs intubation). Argmax Direct does not estimate the CMI, which means that it cannot be adjusted when using non-uniform feature costs, and the CMI values cannot be incorporated into a stopping criterion (Section 4.3). By estimating the CMI with DIME, we are able to do both of these; the former is crucial to support for some applications, and the latter leads to clear performance improvements under a given average budget. To help readers appreciate these points, we will emphasize them in our revised conclusion.

Thank you for your response and for engaging in the discussion period! We appreciate your comments and have made a couple additional changes to the introduction.

1. We now specify that greedy methods have been found to outperform RL methods in several specific works, rather than making a more general claim. We also mention what we believe is the main reason RL methods currently don't work as well: policy optimization is a much more difficult learning problem. In the DFS context, RL incurs sparse rewards and must explore an exponential state space, so even with an appropriate reward like classification accuracy it's understandably difficult to train. With DIME, learning the greedy CMI policy involves fitting one classifier model (the predictor network) and one regression model (the value network), and training these is nearly as simple as standard supervised learning.

2. We added further clarification that some existing RL methods consider non-uniform costs and variable feature budgets (Janisch et al, Kachuee et al). We added this in the discussion of our contributions, and we clarified that these capabilities are new for greedy CMI-based methods.

Thanks again for engaging in the discussion, and we appreciate you adjusting your score!