Estimating Conditional Mutual Information for Dynamic Feature Selection
Authors: Soham Gadgil
Ian Connick Covert
Su-In Lee
Venue: ICLR 2024 poster (2024)
Keywords: dynamic feature selection
adaptive
feature selection
mutual information
information theory
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We develop a method for dynamic feature selection by directly predicting the conditional mutual information with the response variable

Abstract: Dynamic feature selection, where we sequentially query features to make accurate predictions with a minimal budget, is a promising paradigm to reduce feature acquisition costs and provide transparency into the prediction process. The problem is challenging, however, as it requires both making predictions with arbitrary feature sets and learning a policy to identify the most valuable selections. Here, we take an information-theoretic perspective and prioritize features based on their mutual information with the response variable. The main challenge is implementing this policy, and we design a new approach that estimates the mutual information in a discriminative rather than a generative fashion. Building on our learning approach, we introduce several further improvements: allowing variable feature budgets across samples, enabling non-uniform costs between features, incorporating prior information, and exploring modern architectures to handle partial input information. We find that our method provides consistent gains over recent state-of-the-art methods across a variety of datasets.

Published as a conference paper at ICLR 2024
ESTIMATING CONDITIONAL MUTUAL INFORMATION FOR DYNAMIC FEATURE SELECTION
Soham Gadgil. Ian Covert', Su-In Lee Paul G. Allen School of Computer Science & Engineering, University of Washington
ABSTRACT
Dynamic feature selection, where we sequentially query features to make accurate predictions with a minimal budget, is a promising paradigm to reduce feature acquisition costs and provide tramsparency into model's predictions. The problem is challenging, however, as it requires both predicting with arbitrary feature sets and learning policy to identify valuable selections. Here, take an information- theoretic perspective and prioritize features based on their mutual information with the response variable. The main challenge implementing this policy, and we design a new approach that estimates the mutual information in discriminative rather than generative fashion. Building on our approach, we then introduce several further improvements: allowing variable feature budgets across samples, enabling non-uniform feature costs, incorporating prior information, and exploring modern architectures to handle partial inputs. Our experiments show that our method provides consistent gains over recent methods across a variety of datasets.
INTRODUCTION
Many machine learning applications rely on high- tigh-ddmensiionl datasets with significant data acqui- sition costs. For example, medical diagnosis can depend on a range of demographic features, lab tests and physical examinations, and each piece of information takes time and money collect (Kachuee et al. 20188] Erion et al. 2022 He et al., (022). To improve interpretability and reduce data acquisition costs, a natural approach
adaptively query features given the current information, so that each prediction relies on only small number of features. This approach is referred to as dynamic feature selection (DFS) and is a promising paradigm considered by several works in recent years (Kachuee et al. 2018 Janishh et al. 2019; Chhttopadhyyy et al. 2022 2023
Covert et al. (023).
Among the existing methods that address this problem, two main approaches have emerged. One idea is to formulate DFS as a Markov decision process (MDP) and use reinforcement learning (RL)
(Dulac-Arnold et al.] [01]] 1 Mnih et al.] 2014 Kachuee et al. 2018 Janisch et al., 2019). This approach theoretically has the capacity to discover the optimal policy, but it faces training difficulties that are common in RL (Henderson et al. 2018). Alternatively, another line of work focuses on
greedy approaches, where Teatures are selected based on their conditional mutual information (CMI) with the response variable (Chen et al. 2015) Ma et al. (201)). While less flexible than RL, the greedy approach is near-optimal under certain assumptions about the data distribution (Chen et al., 2015) and represents a simpler learning problem; as a result, it has been found to perform better than RI in
several recent works (Erion et al. 2022 Chattopadhyay et al. 2023 Covert al. 2023)
Nevertheless, the greedy approach is non-trivial to implement, because calculating the CMI requires detailed knowledge of the data distribution. Many recent works have explored approximating the
CMI using generative models (Ma et al., [2019] Rangrej and Clark 2021 Chattopadhyay et al. 2022 He et al. 2022). but these are difficult to train and lead to a slow CMI estimation process. Instead, two recent works introduced simpler approach. which is to directly estimate the feature index with
maximum CMI (Chattopadhyay et al. 2023) Covert et al. 2023). These methods rely on simpler learning objectives, are faster at inference time, and often provide better predictive accuracy. They can
* Equal contribution. Correspondence to 4sgadgil @en washington.cdi.
Prior works have also referred to the problem as sequential information maxinization Chen al. active variable selection (Ma et al. 2019) and information pursuit (Chattopadhyay al. 2023
2015

---

Published as a conference paper at ICLR 2024
be thought of discriminative alternatives to earlier generative methods, because they directly output the required value rather than modeling the data generation process (Ng and Jordan| 2001). Their main downside is that they bypass estimating the CMI, which can be useful for multiple purposes, such as determining when to stop selecting new features.
Here, our goal is to advance the greedy DFS approach by combining the best aspects of current methods: building on recent work (Chattopadlyay et al. [023] Covert et al.) 2023), we aim to estimate the CMI itself in a discriminative fashion. We aim to do so without requiring additional labels, making strong assumptions about the data distribution, or fitting generative models. We accomplish this by designing suitable learning objective, which we prove recovers the CMI if our model is trained to optimality (Section 4).
Based on our new learning approach, we then explore range of capabilities enabled by accurately estimating the CMI: these include allowing variable per-sample feature budgets, accounting for non-uniform feature costs, and leveraging modern architectures to train models with partial input information. In real-world settings, acquisition costs are not identical for all features, and we can use multiple low-cost features lieu of one high-cost feature to get the same diagnostic performance. Although these capabilities are supported by certain RL methods (Kachuee et al. 2018 Janisch et al.) 2019). they are not supported by prior discriminative approaches that do not estimate the CMI.
We find that our proposal offers a promising alternative to available methods: it enables the capabilities of generative methods while retaining the simplicity of discriminative methods, and shows improved performance in our experiments. The contributions of this work are the following:
1. We develop a learning approach to estimate the CMI in a discriminative fashion. Our method involves training a network to score candidate features based on their predictive utility, and we prove that training with our objective recovers the exact CMI optimality.
2. We generalize our approach to incorporate prior information beyond the main features. Here, we again prove that our procedure recovers modified version of the CMI at optimality.
3. Taking inspiration from adaptive submodular optimization, we show how to adapt our CMI-based approach to scenarios with non-uniform feature costs.
4. We analyze the role of variable feature budgets and how they enable an improved cost-accuracy tradeoff. We show that a single instantiation of our method can be evaluated with multiple stopping criteria, and that a policy with a flexible per-prediction budget performs best.
5. We investigate the role of modern architectures in improving performance in DFS, In particular, we find that for image data, our method benefits from using ViTs rather than standard CNNs.
Our experiments demonstrate the effectiveness of our approach across range of applications, including several tabular and image datasets. We compare our approach to many recent methods, and we find that approach provides consistent gains across all the datasets tested.
PROBLEM FORMULATION
Notation. Let x denote a vector of input features and y response variable for a supervised learning task. The input consists of d separate features x = (x1...., Xa), and we use S C [d] = {I,.... . d} to denote subset of indices and xs = {x; : E S} subset of features. Bold symbols x, y represent random variables, the symbols x. are possible values, and p(x. y) denotes the data distribution.
Our goal is to select features given the currently available information, and do so on a per-instance basis to rapidly arrive at accurate predictions. In doing so, we require predictor f(xs) that makes predictions given any set of available features; for example, if y is discrete then our predictions lie in the simplex, or f(xs) E A K for K classes. We also require selection policy T(xs) e [d], which takes set of features as its input and outputs the next feature index to observe. We next discuss how to design these models, focusing on approach motivated by information theory.
Dynamic feature selection. The goal of DFS is to select features separately for each prediction, and achieve both low acquisition cost and high predictive accuracy. Previous work has explored several approaches to design a selection policy, including training the policy with RL (Kachuee et al. 2018. Janisch al. 2019), imitation learning (He et al. 2012 2016a), and following greedy
2

---

Published as a conference paper at ICLR 2024
2023 We focus here on the latter approach, where the idealized policy selects the feature with maximum CMI at each step, or where we define Intuitively. the CMI represents how much observing x; improves our ability to predict y when a subset of the features æg is already observed, and it is defined as the
policy based CMI (Ma et al. 2019 Covert et al.
x (2xz) arg max; xy; Øg).
following KI. divergence (Cover and Thomas 2012):
xz) xyy:x;
DRt. (p(x:: y (ss) || p(x: Is)p(y [x)).
The idealized selection policy is accompanied by an idealized predictor, which is the Bayes classifier (xs) ply xs) for classification problems, and under certain assumptions
known that this provides performance within a multiplicative factor of the optimal policy (Chen et al. 2015). However, the CMI policy is difficult to implement because it requires oracle access to the data distribution: computing eg. (1) requires both the distributions p(y | xx)) and p(x; (g) for all (S, which presents a challenging modeling problem. For example, some works have approximated X(y;x; æs) using generative models (Ma et al. 2019 Rangrej and Clark 202 He et al. 2022).
while others have directly modeled xt (ær)) (Chattopadhyay et al. 2023: Covert et al. 2023).
When we follow the greedy CMI policy, another question that arises is how many features to select for each prediction. Previous work has focused mainly on the fixed- budget setting, where we stop when | = k for a specified budget & < d (Ma et al. 2019; Rangrej and Clark 2021 Chattopadhyay et al. 2023 Covert et al. 2023). We instead consider variable budgets in this work (Kachuee et al. 2018 Janisch al. 2019). where the goal is achieve high accuracy given low average feature cost. Unlike many works, we also consider distinct, or non-uniform costs for each feature. As we discuss in Section 4, these goals are made easier by estimating the CMI in discriminative fashion.
RELATED WORK
One of the earliest works on DFS is Geman and Jedynak (199)), who used the CMI selection criterion but made simplifying assumptions about the data distribution. Chen et al. (2015) analyzed the greedy CMI approach theoretically and proved conditions where achieves near-optimal performance. More recent works have focused on practical implementations: among them, several use generative models to approximate the CMI (Ma et al. 2019 Rangrej and Clark 2021 Chattopadhyay et al..
2022; He et al. (022). and two works proposed predicting the best feature index in discriminative fashion (Chattopadhyay et al.) 2023 Covert et al. (023). Our work develops a similar discriminative approach, but to estimate the CMI itself rather than the argmax among candidate features.
Apart from these CMI-based methods, other works have addressed DFS as an RL problem (Dulac Armold et al.] 2011 Shim et al.) 2018 Kachuee et al., 2018 Janisch et al. 2019: 2020 Li and Oliva 2021). For example, Janisch et al. (011)) 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. RL theoretically has the capacity to discover better policies than a greedy approach, but it has not been found to perform well in practice, seemingly due to training difficulties that common in RL (Erion et al. 2022] Chattopadhyay et al. 2023]] Covert et al. 2023). Beyond these approaches, other works have instead explored imitation learning by mimicking oracle policy (He et al., 2012)
2016a), minimizing expected misclassification costs at each step (Liyanage et al. 2021b) b). and making selections based on a Bayesian network fit to the joint data distribution (Liyanage et al. 202
Static feature selection has been an important subject in statistics and machine learning for decades; see (Guyon and Elisseefi 2003 Li et al. 2017 Cai al. 2018) for reviews. CMI is also the basis of some static methods, and its estimation has been studied by many works (Fleuret 2004 Peng al.) 20055] Shishkin et al. 2016). Greedy methods perform well under certain assumptions about the data distribution and are popular for simple models (Das and Kempe| 2011 Flenberg et al. 2018)
but feature selection with nonlinear models is more challenging. For neural networks, methods now exist that leverage either group sparse penalties (Feng and Simon, 2017 Tank et al. 2021 | Lembadri
cct al. 202 or differentiable gating mechanisms (Chang et al. 2017 Baln et al. [001 Lindenbaum
al. 202T). Unlike recent DFS methods (Chattopadhyyy et al. 2023 Covert al. 2023 our work
bypasses these techniques by using a simpler regression objective.
Finally, mutual information estimation with deep learning has been an active topic in recent years (Oord et al. 2018; Belghazi et al.] 2018 Poole et al., 2019 Song and Ermon 2019]] (halee et al.) 2022). Unlike prior works, ours focuses the CMI between features and response variable,
3

---

Published as a conference paper at ICLR 2024
method can condition arbitrary feature sets, and estimate many CMI terms via a single predictive model rather than training separate networks to estimate each mutual information term.
PROPOSED METHOD
In this section, we introduce our method to dynamically select features by estimating the CMI in a discriminative fashion. We then discuss how to incorporate prior information into the selection process, handle non-uniform feature costs, and enable variable budgets across predictions.
4.1 ESTIMATING THE CONDITIONAL MUTUAL INFORMATION
We parameterize two networks to implement our selection policy. First, we have a predictor network [xxs; 0), e.g., classifier with predictions in AK Next, we have value network uxx; 6) Rd designed to estimate the CMI for each feature, or u; ( () R (yy; X; xts). These are implemented with zero-masking to represent missing features, and we can also pass the mask as a binary indicator vector. Once they are trained, we make selections according to arg max, U/ (x5; Φ), and we can make predictions at any time using fxx; 0). The question is how to train the networks, given that prior works required generative models to estimate the CMI (Ma et al. 2019 Chattopadhyay et al. (022).
Our main insight train the models jointly but with their own objectives, and design an objective for the value network that recovers the CMI at optimality. Specifically, we formulate regression problem whose target the incremental improvement in the loss when incorporating a single new feature. We train the predictor to make accurate predictions given any feature set, or
min Exy Eng [e (f(xsi 0), y)},
0
and we simultaneously train the value network with the following regression objective,
min Eg (c: o) XXx.:i.))))
where we define the loss improvement as A(rg, v) = ((f(xs; 0). gz v). The regression objective in inLeq. (3) is motivated by the following property, which shows that if assume an accurate predictor f(x;: 8) (i.e., the Bayes classifier), the value network's labels are unbiased estimates of the CMI (proofs in[Appendix A).
Lemma 1. When we use the Bayes classifier ply x5) as a predictor and & is cross entropy loss, the incremental loss improvement is an unbiased estimator of the CMI for each (I5, x;) pair:
Ey [A[rg y)] = I (y:x; rs).
(4)
Based on this result and the fact that the optimal predictor does not depend on the selection policy (Covert et al., 2023), we can make the following claim about jointly training the two models. We assume that both models are infinitely expressive (e.g., very wide networks) that they can achieve their respective optimizers.
Theorem When E cross entropy loss, the objectives eq. (2) and eg. (3) are jointly optimized by predictor f(xs; 0*) =pyy| xs) and value network where Us (Its; 0*) (y; xt)) for E [d].
This allows us to train the models in an end-to-end fashion using stochastic gradient descent, as depicted in Figure I] In Appendix A we prove a similar result for regression problems: that the policy estimates the reduction in conditional variance associated with each candidate feature. Additional analysis in Appendiix B shows how suboptimality in the classifier can affect the learned CMI estimates; however, even if the learned estimates 1a (xs; Φ) are imperfect in practice, we expect good performance because the policy replicates selections that improve the loss during training.
Several other steps are important during training, and these are detailed in Appendix C aalng with pseudo-code for the training algorithm and the inference procedure. First, like several prior methods, we pre-train the predictor with random feature sets before beginning joint training (Rangrej and Clark 2021 Chattopadhyay et al. 2023 Covert et al. 2023), Next, we generate training samples (Zgs Te, y) by executing the current policy with a random exploration probability e E [0. 1]. which can be decayed throughout training. Finally, sometimes share parameters between the models, particularly when they are large networks; this helps in our experiments with image data, which use
either CNNs or ViTs (Dosovitskiy al. 2020).
4

---

Published as a conference paper at ICLR 2024
Step
Step k
25
us
U
[ly: 2 jass
9
1(9.0 v)
Arrsus 25
1(ū. v)
Diagram of training approach. At each selection step ni the value network (ats; () predicts the CMI for all features, and a single feature v; is chosen for the next prediction [ssuu 0). The prediction loss is used update the predictor (see cq. (2)). and the loss improvement is used to update the value network (see eq. (3) The networks are trained jointly with SGD.
Figure 1:
4.2 INCORPORATING PRIOR INFORMATION
further direction is utilizing prior information obtained before beginning the selection process. We view such prior information as separate random variable z, and can be either an exogenous input, a subset of features that are available with no associated cost, or even noisy or low-resolution version of the main input x (Ranzato) 2014 Ba et al.) 2014 2015). Situations of this form arise in multiple applications, c.g., a patient's demographic features in a medical diagnosis setting.
Given such prior information, selections should be based on [(y; X; &gg 2). which captures how informative x; is given knowledge of both X, and z. Next, the predictions made any time given by p(y æ2s. because z provides information that can improve our predictions. As for our proposed CMI estimation procedure from Section 4.1 is straightforward to modify. The two models must take the prior information additional input, and we can train them with modified versions of eqqqs (2) and
idealized DFS policy must be modified follows. First, the
min ExyzlEs [e (f(xs, z; 0),
min ExyzLE,
(vi Φ) A(xs, x,, y,))
(5)
where A(xs, I4, 2, v) l(f(((( z, (). - (ff(us z; 0), y) is the incremental loss improvement. We then have the following result for jointly training the two models.
Theorem When & cross entropy loss, the objectives in ey. (5) are jointly optimized by a predictor f 2; ply æs.2) 2) and value network where (z
-(y;x, 2),2) 2 for all [d].
The same implementation details discussed in Section 4.1 apply here, and Section 5.2 discusses how the prior information z is incorporated into the value and predictor networks practice.
4.3 ALLOWING A VARIABLE FEATURE BUDGET
Given our approach for estimating each feature's CMI, natural question is how trade off informa- tion with feature acquisition costs. We now consider two challenges related to features costs: (1) how handle non-uniform costs between features, and (2) when to stop collecting new features.
Non-uniform costs. For the first challenge, consider medical diagnosis motivating example. Diagnoses can be informed by heterogeneous data, including demographic variables, questionnaires, physical examinations, and lab tests; each measurement requires different amount of time or money (Kachuee et al.) 2018 Erion et al. 2022), and feature costs must be balanced with the information they provide. We consider here that each feature has cost C 0 and that costs are additive.
There are multiple ways to trade off cost with information, but we take inspiration from adaptive sub- modular optimization, where costs are accounted for via the ratio between the expected improvement and cost. Here, this suggests that our selections should be arg max, I(y:x; rs)/c;. For adaptive submodular objectives, this criterion guarantees netr-optimal performance (Golovin and Krause 201T]; the DES problem is known to not be adaptive submodular (Chen et al. 2015), which means that cannot offer performance guarantees, but find that this approach works well in practice.

---

Published as a conference paper at ICLR 2024
Variable budgets. Next, we consider when to stop acquiring new features. Many previous works focused on the
setting, where we adopt a budget for all predictions (Chen
al.
2015 Ma et al. 2019 Rangrej and Clark 2021: Covert et al. 2023). This can be viewed as
stopping criterion, and it generalizes to non-unifory costs: we can keep collecting new information long Dies G I k. Alternativlyy, itt al. 2023). where selection terminates once the predictions have low uncertainty. For classthication problems, a natural approach to stop collecting features when 11 (y (x)) S m.
adopt a seriddence-oostrainnee setup (Chattopadhyay
In general, it is unclear whether we should follow a budget- or confidence-constraince approach, or whether another option offers a better cost-accuracy tradeoff. We resolve this by considering the optimal performance achievable by non non- greedy policies, and we present the following insight: that policies with per-prediction constraints are Pareto-dominated by those that satisfy their constraints on average. We state this claim here in simplified form, and we defer the formal version to Appendix A
Proposition 1. (Informal) For any feature budget k, the best policy achieve this budget on average achieves lower loss than the best policy with a per-prediction budget constraint. Similarly, for any confidence level m, the best policy to achieve this confidence on average achieves lower cost than the best policy with a per-prediction confidence constraint.
Intuitively, when designing a policy to achieve given average cost or confidence level, it should help to let the policy violate that level for certain predictions. For example, for patient whose medical condition is inherently uncertain and will not be resolved by any number of tests, it is preferable from cost-accuracy perspective to stop early rather than run many expensive tests. [Proposition I suggests that we should avoid adopting budget or confidence constraints and instead seek the optimal unconstrained policy, but because we assume that we only have CMI estimates (Section 4. 4.1,, we opt for a simple alternative: we adopt penalty parameter À >> O, we make selections at each step according to I (y; x; xs)/c;, and we terminate the algorithm when max; / (y; X, xs)/c; < A.
Following this approach, we see that single instantiation of our model can be run with three different stopping criteria: we can use the budget k, the confidence m, or the penalty A. contrast, prior methods that penalize feature costs required training separately for each penalty strength (Janisch) et al. 2019). Although our penalized criterion is not the optimal one alluded to in inPrrpooitiinn I| we find empirically that it consistently improves performance relative fixed budget.
5 EXPERIMENTS
We refer to our approach as DIME (discriminative mutual information estimation) and explore two data modalities evaluate its performance. Our tabular datasets include two medical diagnosis tasks, which represent natural and valuable use cases for DFS; we also use MNIST, which was considered in prior works (Ma et al. 2019; Chattopadhyay et al. 2022; 2023). As for our image datasets, we include these because they are studied in several earlier works and represent challenging DFS problems (Karayev et al., 2012, Mnih and Gregor 2014 Janisch et al. 2019 Rangrej and Clark) 2021 We also design an experiment to test the CMI estimation accuracy in Appendix G
In terms of baselines, we compare DIME to both static and dynamic feature selection methods. As a strong static baseline, we compare to supervised Concrete Autoencoder (CAE) (Balin et al. (2019). which outperformed several dynamic methods in recent work (Covert et al., 2002)). ). We also compare to two more static baselines that statistically estimate the CMI feature selection, mRMR (Peng et al. 2005) and CMICOT (Shishkin et al., 2016). For dynamic baselines, we consider multiple approaches: first, we compare to the recent discriminative methods that directly predict the CMI's argmax (Chattopadhyay et al. 2023 Covert et al. 2023), which we refer to as Argmax Direct. Next, consider EDDI (Ma et al. 2019). generative approach that uses partial variational autoencoder (PVAE) to sample unknown features. Because EDDI does not scale as well beyond tabular datasets, we compare probabilistic hard attention (Hard Attention) (Rangrej and Clark 2021). a method that adapts EDDI to work with image data. Finally, we also compare the tabular datasets to two RL approaches: classification with costly features (CwCF) (Janisch et al., (020) and opportunistic learning (OL) (Kachuee et al. 2018), which both design rewards based on features' predictive utility. Appendix F provides more information about the baseline methods.
"Code available at http://gtttttutec /ttppp///giitthbbbcccmmm

---

Published as a conference paper at ICLR 2024
MNIST
Intubation
ROSMAP
0.85
AAc
0.6
0.80
AUROC
0.80
Top
0.75
0.2
0.70
0.65
15
10
15 Fealures EDDI
25
Avg
DIME (Penalized)
-.... CAE
Argmax max Direct
-
OL
CNCF
....
MRMA
CMICOT
Evaluation with tabular datasets for varying feature acquisition budgets. Results averaged across trials, and shaded regions indicate the standard error for each method.
Figure 2:
5.1 TABULAR DATASETS
Our first medical diagnosis task involves predicting whether a patient requires endotracheal intubation for respiratory support in emergency medicine setting (d 112) (Covert et al. 2023). The second uses cognitive, demographic, and medical history data from two longitudinal aging cohort studies (ROSMAP, Bennett 2012a b) to predict imminent dementia onset (d 46). In both scenarios, it is difficult to acquire all features due to time constraints, making DFS a promising paradigm. The third is the standard MNIST dataset (LeCun al. 1998), which we formulate as a tabular problem by treating each pixel as feature (d 784). Across all methods, we use fully connected networks with dropout to reduce overfitting (Srivastava et al. 2014), and the classification performance measured using AUROC for the medical tasks and top-1 accuracy for MNIST. Appendix D provides more details about the datasets, and Appendix E provides more information about the models.
Uniform feature costs. We first consider the scenario with equal costs for all features, and Figure 2 shows results with a range feature budgets. DIME with the penalized stopping criterion achieves the best results among all methods for both medical diagnosis tasks. It also performs the best on MNIST, where we achieve above 90% accuracy with an average of ~ 10/784 features (1.27%). Among the baselines, Argmax Direct is the strongest dynamic method, CAE is competitive static method that outperforms both CMICOT and mRMR, and EDDI usually does not perform well. The competitive, expected from results in prior work (Rangrej and Clark 2021] Chattopadhyay et al. 2023 Covert et al. 2023). DIME generally shows the greatest advantage for moderate numbers of features, and the gap reduces as the performance saturates. The variability between trials for the basclines is typically low, as shown by our uncertainty estimates; for DIME, we show results from five independent trials in Appenddix G because it difficult to ensure identical budgets with separate training runs.
RL approaches (CwCF, OL)
Intubation
ROSMAP
0.85
0.85
AUROC
AUROC
0.75
0.75
0.70
0.60
0
100
150
200
200
400
600
800
1000
Cost (Relative, 1-10)
Cost (Seconds)
DIME
(Inference Costs)
DIME (Uniform Costs)
EDDi
CWCF
Figure 3: Evaluation with non-uniform feature costs for medical diagnosis tasks. Costs are relative for Intubation and expressed in seconds for ROSMAP. The results show the classification performance for varying levels of average feature acquisition cost.
Non-uniform feature costs. Our CMI estimation approach lets us incorporate non-uniform feature costs into DIME, and we demonstrate this with the Intubation and ROSMAP datasets. For ROSMAP, we use costs expressed as the time required to acquire each feature, and for Intubation we use relative costs estimated by board-certified physician (Appendix D). For comparisons,
EDDI and
7

---

Published as a conference paper at ICLR 2024
Imagenette
ImageNet-100
1.0
0.8
0.8
Acc
0.7
0.5
0.6
0.5
0.4
0.4
0.3
0.3
12
14
10
15
20
25
30
Avg Patches
Avg Patchees Budget (Resmet)
(VIT)
Budget (MT)
Penalized (Resnet)
Figure 4: Evaluation of DIME on image datasets with different vision architectures.
CwCF to accommodate non-uniform feature costs. We also compare to two ablations our approach: (1) using uniform costs during training but the true costs during inference (Inference Costs), which tests DIME's robustness to changing feature costs after training; and (2) using uniform costs during both training and inference (Uniform Costs), which demonstrates the importance of using correct costs. All methods are compared here with the budget-constrained stopping criterion.
Figure 3 shows the results with non-uniform feature costs. DIME outperforms EDDI and CwCF by substantial margin, echoing the earlier results, and reflecting the improved CMI estimation with our discriminative approach. Comparing to the variations of DIME, using the true non-uniform costs during both training and inference outperforms both variations, showing that considering costs when making selections is important. The version that uses costs only during inference slightly outperforms ignoring costs on ROSMAP, indicating a degree of robustness to changing feature costs between training and inference, but both versions give similar results on Intubation.
Additional results. Due to space constraints, we defer several additional results to Appendix G
5.2
IMAGE DATASETS
Next, we applied our method to three image classification datasets. The first two are subsets of ImageNet (Deng et al., 2009), one with 10 classes (Imagenette. Howard) and the other with 100 classes (ImageNet- 100, Ambityga). The third is histopathology dataset (MHIST, comprising hematoxylin and eosin -H&a)e-ttanned images obtained by extracting diagnostically- relevant tiles from whole-slide images (WSIs). The task is predict the histological pattern as either a benign or precancerous lesion. WSIs have extremely high resolution and are infeasible for direct use in any classification task, making them a potential use case for DIME to identify important patches. The images in all three datasets are 224 x 224, and we view them as d = 196 patches of size 16 x 16. We explore different architectures for the value and predictor networks, namely ResNets (He et al. 2016b) and Vision Transformers (ViTs) (Dosovitskiy et al. 2020). We use a shared backbone in both cases, with each component having its own output head. Classification performance is measured using top-1 accuracy for both ImageNet subsets, and with AUROC for the histopathology task.
al. (02I),
Imagenette
Image Image Net-
Histopathology
0.8
0.9
0.8
0.80
0.6
0.7
AUROC
0.75
0.4
0.65
0.2
0.60
0.3
2.5
5.0
7.5
10.0
12.5
0
10
20
20
40
60
Aug # Patches
Avg Patches
Avg * Patches
(Penalized)
Argmax- Direct
Figure 5: Evaluation with image datasets for varying numbers average patches selected. Results averaged across 5 trials, and shaded regions indicate the standard error for each method.
8

---

Published as a conference paper at ICLR 2024
Original
Edge Image
Histopathology
0.80
AUROC
0.70
DINE (Budgee)
(Ponolized)
DIME predictor sketch)
0.60
20
40
60
Avg Patches
Figure 6: Evaluation of DIME with prior information for histopathology classification. Left: Original MHIST image. Center: Canny edge image. Right: Results for varying number of average patches.
Exploring modern architectures. As an initial experiment, Figure compares the ResNet and ViT architectures for the Imagenette and ImageNet- 100 datasets. We use DIME to conduct this analysis because its discriminative approach lets us seamlessly plug in any network architecture. Across both datasets, DIME's penalized version outperforms the budget-constraine ViTs significantly outperform ResNets. This can be attributed to the ViT's self-attention mechanism: this architecture is better suited to handle information spread across an image, a property that has made ViTs useful in other applications with partial inputs (Naseer et al. 2021 Jain et al. 2021 Salman et al. 2022). Given their superior performance, we use VITs as backbones in subsequent experiments where possible. This includes DIME, CAE and Argmax Direct, but Hard Attention lacks this flexibility because it uses a recurrent module to process subsets of image regions.
version, and we find that
Next, Figure 5) compares DIME to the baselines across multiple feature budgets. DIME with the penalized stopping criteria outperforms the baselines for all feature budgets, with the largest gains observed for Imagenette. Notably, we achieve nearly 97% accuracy on Imagenette with only ~15 1996 patches (7.7%). The Argmax Direct baseline is competitive with DIME, but the Hard Attention baseline shows a larger drop in performance.
Incorporating prior information. Next, we explored the possibility of incorporating prior infor- mation into DIME's selection process for the histopathology dataset. To simulate informing our selections with a less exact but easily acquirable version of the tissue, we use the Canny edge image as a sketch (Canny) 1986), which can help generate more valuable selections than a blank image. We use separate VIT backbones for the original and edge images, and concatenate the resulting embeddings before estimating the CMI or making class predictions (see /Appendix E for details). Figure 6 shows example images, along with the results obtained with DIME for various feature budgets. The results show that the prior information is incorporated successfully, leading to improved performance for both the penalized and budget-constrained versions. To verify that the improvement is not solely due to the predictive signal provided by the edge image, we conduct an ablation where the sketch is integrated into the predictor only for a pre-trained and otherwise frozen version of DIME. This middle ground improves upon no prior information, but it generally performs well below the version that uses prior information both when making selections and predictions.
CONCLUSION
This work presents DIME, new DFS approach enabled by estimating the CMI in discriminative fashion. Our approach involves learning value and predictor networks, trained in an end-to-end fashion with straightforward regression objective. From theoretical perspective, we prove that our training approach recovers the exact CMI at optimality. Empirically, DIME accurately estimates the CMI and enables an improved cost-accuracy tradeoff, exceeding both the generative and discriminative methods it builds upon (Chattopadhyay et al. range of tabular and image datasets and demonstrates the potential to implement several additions to the classic greedy CMI selection policy: these include allowing non-uniform features costs, variable budgets, and incorporating prior information. The results also show that DIME is robust to higher image resolutions, scales more classes, and benefits from modern architectures. Future work may focus on promising applications like MRIs and region-of of-inttrestt selection within WSIs, using DIME to initialize RL methods, and otherwise accelerating improving DIME'S training.
Covert et al.) 2023). Our evaluation considers a

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENT
We thank the members of the Lee Lab for helpful discussions. This work was supported by the National Science Foundation (CAREER DBI-1552309 and 683 DBI-1759487) and the National Institutes of Health (R35 GM 128638 and ROI (G061132).
REPRODUCIBILITY STATEMENT
The code to train DIME has been submitted as part of the supplementary material. It includes a README that describes the steps to be taken for training and evaluating models with the dif- ferent datasets. Pseudocode for DIME's s training algorithm along with details about our training implementation are provided Appendix C Proofs for Lemma Theorem 1 are provided in Appendix A.T. The proof for [heoremm iss provided in Appendix A.2] The adaptation of DIME for regression problems, along with the corresponding theorem and proof, are provided in (pppeddix A.3] Descriptions of all datasets used in experiments, along with ways access them, are provided in Appendix D All the datasets are publicly available except for the intubation dataset, which is private due to patient privacy concerns.
REFERENCES
David A Bennett, Julie A Schneider, Zoe Arvanitakis, and Robert S Wilson. Overview and findings from the religious orders study. Current Alzheimer Research, (6):628-645, 2012a.
David A Bennett, Julie A Schneider, Aron S Buchman, Lisa L Barnes, Patricia A Boyle, and Robert Wilson. Overview and findings from the rush memory and aging project. Current Alzheimer Research, 9(6):6466633 2012b.
Khaled Abdeljawad, Krishna C Vemulapalli, Charles J Kahi, Oscar W Cummings, Dale C Snover, and Douglas K Rex. Sessile serrated polyp prevalence determined by a colonoscopist with a high lesion detection rate and an experienced pathologist. Gastrointestinal Endoscopy, 8133:5117-524 2015.
Ambityga.
Imagenet100.
https: /ww. httpss/ww.kaggge.coo/aaaaaea http://aalee..ooattttttt////aaaa
imagenet100
Jimmy Ba, Volodymyr Mnih, and Koray Kavukcuoglu. Multiple object recognition with visual
attention. arXiv preprint arXiv: 1412.7755, 2014.
Jimmy Ba. Russ R Salakhutdinov, Roger B Grosse, and Brendan J Frey. Learning wake-sleep recurrent attention models. Advances in Neural Information Processing Systems, 28, 2015.
Muhammed Fatih Balm, Abubakar Abid, and James Zou. Concrete autoencoders: Differentiable feature selection and reconstruction. In International Conference on Machine Learning, pages 444-453, PMLR, 2019.
Nicasia Beebe-Want Wang, Alex Okeson, Tim Althoff, and Su-In Lee. Efficient and explainable risk assessments for imminent dementia an aging cohort study. IEEE Journal of Biomedicul and Health Informatics, 25(7):2409-2420) 2021.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference on Machine Learning, pages 531-540. PMLR, 2018
Wieland Brendel and Matthias Bethge. Approximating CNNs with -ecal--e models works surprisingly well on ImageNet. arXiv preprint arXiv: 104.00000 2019.
Jie Cai, Jiawei Luo, Shulin Wang, and Sheng Yang. Feature selection in machine learning: A new perspective. Neurocomputing, 300:70-79, 2018.
John Canny. A computational upproach to edge detection. JEEE Transactions on Pattern Analysis and Machine Intelligence, (6):679-698) 1986.
10

---

Published as a conference paper at ICLR 2024
Chun-Hao Chang, Ladislav Rampasek, and Anna Goldenberg. Dropout feature ranking for deep learning models. arXiv preprint arXiv:1712.08645,, 2017.
Aditya Chattopadhyay, Stewart Slocum, Benjamin Haeffele, Rene Vidal, and Donald Geman. In- terpretable by design: Learning predictors by composing interpretable queries. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022.
Aditya Chattopadhyay, Kwan Ho Ryan Chan, Benjamin D Haeffele, Donald Geman, and René Vidal. Variational information pursuit for interpretable predictions. arXiv preprint arXiv:2302.02876
2023.
Yuxin Chen, Hamed Hassani, Amin Karbasi, and Andreas Krause. Sequential information max- imization: When is greedy near-optimal? Conference on Learning Theory, pages 338-363. PMLR, 2015.
T.M. Cover and J.A. Thomas. Elements of Information Theory. Wiley, 2012. ISBN 97811 185857777777.
Ian Covert, Wei Qiu, Mingyu Lu, Nayoon Kim, Nathan White, and Su-In Lee. Learning to maximize mutual information for dynamic feature selection. arXiv preprint arXiv:2331..055 2023.
Abhimanyu Das and David Kempe. Submodular meets spectral: Greedy algorithms for subset selection, sparse approximation and dictionary selection. arXiv preprint arXiv: 102.3995. 2011.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: large-scale hier- archical image database. In 2009 IEEE Conference Computer Vision and Pattern Recognition, pages 148-255 Ieee, 2009.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11922 2020.
Gabriel Dulac-Aroold, Ludovic Denoyer, Philippe Preux, and Patrick Gallinari. Datum-wise classifi- cation: a sequential approach to sparsity. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 375-390, Springer, 2011.
Ethan R Elenberg, Rajiv Khanna, Alexandros G Dimakis, and Sahand Negahban. Restricted strong convexity implies weak submodularity. The Annals of Statistics, 36(6B):353993366 2018.
Gabriel Erion, Joseph Janizek, Carly Hudelson, Richard B Utarnachitt, Andrew M McCoy, Michael Sayre, Nathan J White, and Su-In Lee. CoAl: Cost-aware artificial intelligence for health care. Nature Biomedical Engineering, 6(12):1384, 2022.
Alton B Farris, Joseph Misdraji, Amitabh Srivastava, Alona Muzikansky, Vikram Deshpande, Gre- Sessile serrated adenoma: challenging discrimination
gory Y Lauwers, and Mari from other serrated colonic polyps. The American Journal of Surgical Puthology, 32(1):30-35, 2008.
Jean Feng and Noah Simon. Sparse-input neural networks for high-dimensional nonparametric regression and classification. arXiv preprint 1rX15:1711.0552,, 2017.
François Fleuret. Fast binary feature selection with conditional mutual information. Journal of Machine Learning Research, 5(9), 2004.
Donald Geman and Bruno Jedynak. An active testing model for tracking roads in satellite images. IEEE Transactions on Pattern Analysis and Machine Intelligence, 18(1):1-14, 1996.
Katharina Glatz, Bobbi Pritt, Dieter Glatz, Arddt Hartmann, Michael J O' OrBei, and Hagen Blaszyk. multinational, internet-based assessment of observer variability the diagnosis of serrated colorectal polyps. American Journal of Clinical Pathology. 127(6))938- 979-9459 2007.
Daniel Golovin and Andreas Krause. Adaptive submodularity: Theory and applications in active learning and stochastic optimization. Journal of Artificial Intelligence Research, 42:427-488 2011.
11

---

Published as a conference paper at ICLR 2024
Isabelle Guyon and André Elisseeff. An introduction to variable and feature selection. Journal of Machine Learning Research, 31Mar):1157-1182, 2003.
He He, Hal Daumé III, and Jason Eisner. Cost-sensitive dynamic feature selection. In ICML Inferning Workshop, 2012.
He He, Paul Mineiro, and Nikos Karampatziakis. Active information acquisition. arXiv preprint arXiv:1602.02181, 2016a.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE Conference Computer Vision and Pattern Recognition, pages 770-778, 2016b.
Weijie He, Xiaohao Mao, Chao Ma, Yu Huang, José Miguel Hemandezz-Lbatto and Ting Chen. BSODA: bipartite scalable framework for online disease diagnosis. In Proceedings of the ACM Web Conference 2022, pages 2511-2521,
Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In Proceedings of the AAAI Conference Artificial Intelligence, volume 32, 2018.
FastAl Jeremy Howard. imagenette
The Imagenette dataset.
bttps://githu..o fastai/
Saachi Jain, Hadi Salman, Eric Wong, Pengchuan Zhang, Vibhav Vineet, Sai Vemprala, and Alek- sander Madry. Missingness bias in model debugging. In International Conference on Learning Representations, 2021.
Jaromír Janisch, Tomáš Pevný, and Viliam Lisy. Classification with costly features using deep rein- forcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3959-3966, 2019.
Jaromír Janisch, Tomáš Pevný, and Viliam Lisy. Classification with costly features sequential decision-making problem. Machine Learning, 109:1587-1615, 2020.
Mohammad Kachuce, Orpaz Goldstein, Kimmo Kärkkäinen, Sajad Darabi, and Majid Sarrafzadeh. Opportunistic learning: Budgeted cost-sensitive learning from data streams. In International Conference on Learning Representations, 2018
Sergey Karayev, Tobias Baumgartner, Mario Fritz, and Trevor Darrell. Timely object recognition. Advances in Neural Information Processing Systems, 25, 2012.
Omer Khalid, Sofyan Radaideh, Oscar W Cummings, Michael O'Brien, John R Goldblum, and Douglas K Rex. Reinterpretation of histology of proximal colon polyps called hyperplastic in 2001. World Journal of Gasiroenterology. 15(30):3767, 2009.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to document recognition. Proceedings of the IEFE, 86(11):2288-2222 1998.
Ismael Lemhadri, Feng Ruan, and Rob Tibshirani. Lassonet: Neural networks with feature sparsity. In International Conference on Artificial Intelligence and Statistics, pages 10-18. PMLR, 2021.
Jundong Li, Kewei Cheng, Suhang Wang, Fred Morstatter, Robert P Trevino. Jiliang Tang, and Huan Liu. Feature selection: A data perspective. ACM Computing Surveys (CSUR), 50(6):1-55, 2017.
Yang Li and Junier Oliva. Active feature acquisition with generative surrogate models. In International Conference Machine Learning, pages 1450-6459. PMLR, 2021.
Ofir Lindenbaum, Uni Shaham, Erez Peterfreund, Jonathan Svirsky, Nicolas Casey, and Yuval Kluger. Differentiable unsupervised feature selection based on a gated Laplacian. Advances in Neural Information Processing Systems, 34: 14:1550-1549 2021.
Yasitha Warahena Liyanage. Haphney-Stavronn Zois, and Charalampos Chelmis. Dynamic instance- wise classification in correlated feature spaces. IEEE Transactions on Artificial Intelligence, 2(6):
537-548, 2021a.
12

---

Published as a conference paper at ICLR 2024
Yasitha Warahena Liyanage. Huphhee--Starooute Zois, and Charalampos Chelmis. Dynamic instance- wise joint feature selection and classification. IEEE Transactions on Artificial Intelligence, 2(2): 169-184, 2021b.
Chao Ma, Sebastian Tschiatschek, Konstantina Palla, Jose Miguel Hernandez-Lobato. Sebastian Nowozin, and Cheng Zhang. EDDI: Efficient dynamic discovery of high-value information with partial VAE. In International Conference Machine Learning. pages 4234 4243. PMLR, 2019.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning, pages 1791-1799. PMLR, 2014.
Volodymyr Mnih, Nicolas Heess, Alex Graves, et al. Recurrent models of visual attention. Advances in Neural Information Processing Systems, 27, 2014.
Muhammad Muzammal Naseer, Kanchana Ranasinghe, Salman H Khan, Munawar Hayat, Fahad Shahbaz Khan, and Ming- Ming-Hsuan Yang. Intriguing properties of vision transformers. Advances in Neural Information Processing Systems, 34, 2021.
Andrew Ng and Michael Jordan. On discriminative generative classifiers: A comparison of logistic regression and naive Bayes. Advances in Neural Information Processing Systems, 14, 2001.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint 1 2018.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary Vito,, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in PyTorch. 2017.
Hanchuan Peng, Fuhui Long, and Chris Ding. Feature selection based on mutual information criteria of max-dependency, max-relevance, and min-redundancy. IEEE Transactions on patlern analysis and machine intelligence, 27(8): 1221-123222 2005.
Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker. On variational bounds of mutual information. In International Conference on Machine Learning, pages 5171- 5180. PMLR, 2019.
Samrudhdhi B Rangrej and James J Clark. A probabilistic hard attention model for sequentially observed scenes. arXiv preprint arXiv:2111.07534 2021.
Aurelio Ranzato. On learning where look. arXiv preprint
2014.
Hadi Salman, Saachi Jain, Eric Wong, and Aleksander Madry. Certified patch robustness via smoothed vision transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 15137-15147, 2022.
Yuval Shalev, Amichai Painsky, and Irad Ben-Gal. Neural joint entropy estimation. IEEE Transactions on Neural Networks and Learning Systems, 2022.
Hajin Shim, Sung Ju Hwang, and Eunho Yang. Joint active feature acquisition and classification with variable-size set encoding. Advances in Neural Information Processing Systems, 31, 2018.
Alexander Shishkin, Anastasia Bezzubtseva, Alexey Drutsa, Ilia Shishkov. Ekaterina Gladkikh, Gleb Gusev, and Pavel Serdyukov. Efficient high-order Interaction-awar feature selection based on conditional mutual information. Advances in neural information processing systems, 29, 2016.
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estimators. arXiv preprint arXiv:1910.06222, 2019.
Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The Journal of Machine Learning Research,
2014.
Alex Tank, Ian Covert, Nicholas Foti, Ali Shojaie, and Emily Fox. Neural Granger causality. IEEE Transactions Pattern Analysis and Machine Intelligence, 44(8):426744279, 2021.
13

---

Published as a conference paper at ICLR 2024
Jerry Wei, Arief Suriawinata, Bing Ren, Xiaoying Liu, Mikhail Lisovsky, Louis Vaickus, Charles Brown, Michael Baker, Naofumi Tomita, Lorenzo Torresani, et al. A petri dish for histopathology image analysis. In Artificial Intelligence in Medicine, pages 11-24. Springer, 2021.
Newton ACS Wong, Linda P Hunt, Marco R Novelli, Neil A Shepherd, and Bryan Warren. Observer agreement in the diagnosis of serrated polyps of the large bowel. Histopathology. 55(1):63-66, 2009.
14

---

Published as a conference paper at ICLR 2024
A PROOFS
In this section, we re-state and prove our results from the main text.
A.1 ESTIMATING CONDITIONAL MUTUAL INFORMATION
Lemma 1. When we use the Bayes classifier ply xs) as a predictor and cross entropy loss, the incremental loss improvement is unbiased estimator of the CMI for each (xs, xi) pair:
By,x,/= [Ax Ily: rs).
Proof. Consider the expected eross entropy loss given the prediction ply as):
Eylng p(ppl [sss
Σ
ply y xs) logp(y = y as)
y=1
Hy xs).
Next, consider the loss given the prediction p(y I5, taken in expectation across x; and y's conditional distribution ply, X; ws)::
Ey x/,x [[(p(y] xs)x = Exilas Halleryyyylee 1y [pp(y &5, = 2)). =Ex, (H T3s X; = 2;)]
Hy Iy X;).
We therefore have the following expectation for the incremental loss improvement:
Ey.x/es ( [A(s, y)] Ey.x.lxs ((p(y xs)s E(p(y xs), Hy xs)) Hy (g,
T(y; X; xs))
Thus, the loss improvement A y) an unbiased estimator of the CMI [(;; (s))
Theorem When E cross entropy loss, the objectives eq. (2) eg. (3) are jointly optimized by
predictor f(zs; 0*) = p(y æg) and value network where (;(4s; o xy; for E [d].
2023 our proof considers both models' optimal predictions for Proof. Similar to (Covert et al.,
each input. Beginning with the predictor, consider the output given the input æs. The selections were made given only æs, observing this input conveys no information about y the remaining features s. The expected loss is therefore
Eylas ((f(s))) 0).)]
Assuming that is cross entropy loss, we can decompose the expected loss as follows:
By jmg (1(f(x;;))))))))
Σ ply xs) log (xss:
Σ ply = xs) log ply = [sss) fy(as;;) 0)
ply y xs))
y=1
H x3) + DKL (ply 1xs) |l f();(0)
15

---

Published as a conference paper at ICLR 2024
Due to the non-negative KL divergence term, make this argument for any imput ø,, so we say that the optimal predictor is fas; (") = p(y ws)
that the optimal prediction is p(y zz). We
for all zz. Notably, this argument does not depend the selection policy: it only requires that the policy has no additional information about the response variable or unobserved features.
Next, consider the value network while assuming that the optimal predictor f(x;; ((x);0)).
Given an input 2g, we once again have no further information about or x (s))ss so the expected loss taken across the distribution ply, x; mg) as follows:
x jay [(u(zs:)) o) - Aats,x,,))
The expected loss can then be decomposed.
By xilzs (o(Ds;)) - (rss, y),
yx xi [(v(ars;)) Byy las ( Xi, y)
+ Var (A(Es, x;) xs),
which reveals that the optimal prediction is Ey.x/ (zs, x,, y)]. Following Lemma we know that this equal xy;; Iz). And because we can make this argument for any øs, we conclude that the optimal value network is given by v(rs: x*) = T(y; x; xs).
A.2 PRIOR INFORMATION
Before proving Theorem 2]
we first present preliminary result analogous to Lemma
Lemma 2. When the Bayes classifier ply xx, predictor and E is entropy loss, the incremental loss improvement is an unbiased estimator of the CMI for each (Is, 2,x, tuple:
Byy xis [A Xi, v)] (y; X; z).
Proof. The proof follows the same logic as Lemma and after incorporating the additional feature x;. The only difference is that each expectation must also condition on z so the terms to analyze are
where we consider the expected loss before
By [ xs. y)] las. Is, x,),
We now prove the main result for incorporating prior information.
Theorem 2. When & is cross entropy loss, the objectives in inleg. (5) are jointly optimized by a predictor
f( 2; 0*) p(y æs,2) 2) and value network where vg 2; 0²) = T(y:x; xs, 2) for all i [d].
Proof. The proof follows the same logic as Theorem decompose the expected loss as follows:
For the predictor with input ss, we can
Eylng.:(((f(Is: 2: x; H(y xs. 2) + DKI 25.2) || f(zs.
This shows that the optimal predictor is [( 0*) ply =), Next, assuming the optimal predictor, the value network's expected loss can be decomposed as follows:
Ey.x,lss. (o(as, 2;) A(xs, x,,,)
Eyxijzs. [(0(1) 2; Ey Ey.x/zs. ( [A(ss, Xi, x,,
+ Var (A(zs, y) 2g,2).
Based on this, Lemma limplies that the optimal value network is u< 2;0²) = I(y;x; 12 2).
16

---

Published as a conference paper at ICLR 2024
A.3 REGRESSION VERSION
Before proving our main result for regression models, we first present a preliminary result analogous Lemma
Lemma 3. When we use the conditional expectation Ely xs] as predictor and & is mean squared error, the incremental loss improvement is an unbiased estimator of the expected reduction in conditional variance for each (xg, xi) pair:
Ey as (A(Ts, y), = Var(y x.) - Exiles [Var(y [xs; x:)]
= Varr(Ely x.) rs).
Proof. Consider the expected loss given the prediction Ely æs]:
E Eylss
æg) - yy) = Var(y xg). (Bly
Next. consider the loss given the prediction Ely xss taken in expectation across x, and y's conditional distribution p(y, x; xs):
Ey,x, lus
(Ely
Is. xx}
y)
Ex las bas (Ely xss xx,xxxx]] ->
Ex lag [Var(y) x;)]
We therefore have the following expectation for the incremental loss improvement:
By las (xss x,, y)] Var(y (ss) Exatas [Var(y as, x;)]. Using the law of total variance, simplify this difference follows:
Ey,x, lay (Ars x,, y)] = Var (Eyy
x,] xs)
This provides measure similar to CMI: it quantifies to what extent different plausible values of affect our best estimate for the response variable.
We now present main result for regression models.
Theorem When e mean squared error, the objectives eg. (2) and eg. (3) predictor fxrs; 0*) = Ely æss] and value network where for i [d] we have
(3) are jointly optimized by
v(æs: x²) = Var(E x5 xs).
Proof. We follow the same proof technique in Theorem The expected loss for the predictor with input øg can be decomposed as follows,
Eyles ((s((xs))) Eylzs ( y)²]
Eylzs (f(xs;)) - Ely =ss))]] + Var(y xs),
which shows that the optimal predictor network is f(zg; = Ely xs]- Assuming we use the optimal predictor, the expected loss for the value network then be decomposed as
17

---

Published as a conference paper at ICLR 2024
xx (o(Is; Φ) ((Es,
x))
xx [(v(zs; Φ) xx (1 Xi, Xi, x
Var (A(Ts; Xx, y)] as))
which shows that the optimal value network prediction is Ey xy,x//zs [ Igs X x)) conclude that the optimal value network is therefore u(rs; x²) = Var [8 Ess Xi] z5).
Lemma lets us
A.4 ALLOWING A VARIABLE FEATURE BUDGET
We first re-state our informal claim, and then introduce notation required to show formal version.
Proposition 1. (Informal) For any feature budget k, the best policy to achieve this budget on average achieves lower loss than the best policy with a per-prediction budget constraint. Similarly, for any confidence level m, the best policy to achieve this confidence on average achieves lower cost than the best policy with per-prediction confidence constraint.
In order to account for policy's s stopping criterion, we generalize our earlier notation so that policies are functions the form v (2ts) E {0) U [d], and we say policy terminates (or stops selecting new features) when ı (ars) O. Given an input x, we let S'T,, x)) S [d] denote the set of indices selected upon termination. The cost for this prediction is c(TT, x) LiEs(arz)) C₁ and there is also notion of
expected loss {(T, x) that we define as follows:
[TT x) = By (ns ( (f(((()))))
(6)
For example, if is cross entropy loss and {(H, x) Hy Is): due to this interpretation of the expected loss, we refer to constraints on {(TT x) confidence constraints. For example. Chattopadhyay et al. (2023) suggests selecting features until H(y xg) < m for a confidence level m. In comparing policies, we must consider the tradeoff between accuracy and feature cost, and we have two competing objectives the average loss and the average cost:
the Bayes classifier f(xg p(y xs), we have
(TT) = E[((TTx))
<(πT =E[c(T, x)].
Now, there are three types of policies we wish to compare: (1) those that adopt budget constraint for each prediction, (2) those that adopt confidence constraint for each prediction, and (3) those with no constraints. These classes of selection policies are defined as follows:
1. (Budget-constrained) These policies adopt a budget k that must be respected for each input x. That is, have c( < k for all This can be ensured by terminating the policy when the budget is exactly satisfied (Ma et al. 2019 Rangrej and Clark 2021 Covert et al. (2023) or when there are no more candidates that will exceed the budget. Policies of this form are said to belong to the set I[g-
2. (Confidence constrained) These policies adopt a minimum confidence m that must be respected for each input m. That is, we must have ((z, S m for all ær. Technically, we may not be able to guarantee this for all predictions due to inherent uncertainty, so we can instead keep making predictions as long as the expected loss exceeds m (Chattopadhyay al. 2023). Policies of this form are said to belong to the set IIm-
3. (Unconstrained) These policies have no per- per-prediction constraints on the feature cost or expected loss. These policies are said to belong to the II, where have II, II and II,,, C II.
With these definitions in place, we now present more formal version of our claim.
Proposition 2. (Formal) For any average feature cost k. the best unconstrained policy achieves lower expected loss than the best budget-constrant policy:
min Sk
{(x)) S min f(TT).
(8)
sell:
18

---

Published as a conference paper at ICLR 2024
Similarly for any average confidence level the best unconstrained policy achieves lower expected cost than the best confidence-construn policy:
min rell:(()sm
<(T) S min c(x).
aellm
In other words, for any desired average feature confidence level, it cannot help adopt that level as per-prediction constraint. The best policy to achieve these levels on average can violate the constraint for some predictions, and result provide either lower average or expected loss.
Proof. The proof of this claim relies on the fact that IT. S II and IT S IT. It is easy to see that 11. S { E :I: <(π)= S k}. This implies the inequality ineqq because the right-hand side takes the minimum over a smaller set of policies. Similarly, it is easy to see that 1m S { 11 ((x) m). which implies the inequality injeq. (9)
B PREDICTOR SUBOPTIMALITY
Consider a feature subset Vg where the ideal prediction from the Bayes classifier is p(y [g). but the learned classifier instead outputs q æs). The incorrect prediction result in skewed loss, which then provides incorrect labels to the value network v(rs; Φ). Specifically, the expected loss assuming many data points yx such that xs xs becomes
Eyles ((all xs). (y as))
DKL (ply (g) 1 [x)))
(10)
The loss is therefore higher average than it should be given the Bayes classifier, with the extra loss being equal to the KL divergence between the ideal and actual predictions. However, this does not imply that v(rs;)) Φ) systematically overestimates the CMI, because its labels are based on the expected loss reduction.
Consider that the above situation with incorrect predictions occurs not only for æs, but also for all values of x; that is, the classifier outputs q(y Es, xz) rather than p(y Igs xz) for each value T, Now, the expected loss reduction is the following:
Ey.x, (Ars, X., x). I(y;x; (s)) + DKL. (Kl ass) |l gl (ss))
- Ex,1 [DKL(P(y xs,x,) g(y x:)]]
(11)
This implies that given infinite data and a value network that perfectly optimizes its objective, the learned CMI estimates are biased according to difference in KL divergence terms. Notably, the difference can be either positive or negative, And intuitively, the bias shrinks to zero as the classifier approaches p(y .s)) for all predictions.
the CMI estimates can be incorrect in either direction.
19

---

Published as a conference paper at ICLR 2024
C TRAINING ALGORITHM
Algorithm I summarizes our learning approach, where we jointly train the predictor and value networks according the objectives in eq. (2) and (3) We implemented in PyTorch (Paszke et al., 2017) using PyTorch Lightning Note that the algorithm requires a dataset of fully observed x samples with corresponding labels y.
Algorithm 1: Training algorithm
Input: Data distribution p(x, y). learning rate y, budget ks exploration E (0, 1), costs E R Output: Predictor f(xs; f), value network v(xs; Ø)
// Prepare models initialize u(xs; Φ). pre-train f(xs; 0) with random masks
while converged do
// Initialize aiee variables
initialize = () Lg 0, f. =0
sample z, ~ p(x,y)
// Initial prediction calculate Uprev = f(x
update La - L. + (Uprev y)
while Eies Q: S do
// Determine next selection
calculate / v(ts: Φ)
set = arg msx T//e, with probability e, else sample from [d] S
// Update predictor loss update S + SUj calculate f(zg;
update Ca + La + (()
// Update value loss calculate A [(pper 8) f(j. v) update Cg + Ls (1, A)
set Sprev
end
/ Gradient step
update +- 0 -
end
httpps//www.prrrrrhhhgghttttt in..iiggttnnniiii
20

---

Published as a conference paper at ICLR 2024
Next. Algorithm shows how features selected inference time. For simplicity. show only the penalized stopping criterion, which requires a penalty term > 0. To implement either the budget confidence constraints described in Section 4.3 we only need to change the stopping criterion.
Algorithm 2: Inference algorithm
Input:
Test instance (z, y). predictor f(xs; 0), value network u(xs; Ø), penalty parameter >
Output: Prediction
// Initialize feature set initialize S = {}
for i = 1= do
// Est mmtte current CMT
calculate / = o(
// Check stopping criterion
max T < X then
I break end
// Determine next selection
set arg max li/a update S SUj
end
// Make prediction calculate = (z);0)
To incorporate prior information into Algorithm T Jand Algorithm 2 (denoted by z), simply update / = u(rs 2; Φ) and ý fIrs, 2; 0) during both training and inference.
Algorithm is simplified to omit several details that we implement practice and these details discussed below.
Masked pre-training. When pre-training the predictor s(xs; 0), sample feature subsets as follows: we first sample a cardinality {O..... d} uniformly at random, and we then sample the members of the subset random. This distribution ensures even coverage of different subset sizes [S), whereas treating each feature's 's membership as an independent Bernoulli variable biases the subsets towards a specific size.
Minibatching, As is conventional in deep learning, we calculate gradients in parallel for multiple inputs. Allorithm I] this means that we take gradient steps calculated over multiple data samples y) and multiple feature budgets.
Learning rate schedule, Rather than train with a fixed learning rate ? > we reduce its value over the course of training. To avoid setting a precise number of epochs for each dataset, we decay the learning rate when the loss reaches a plateau, and we perform early stopping when the learning rate is sufficiently low. The initial learning rate depends on the architecture, but use values similar to those used for conventional training (e.g., VITs require a lower learning rate than CNNs or MLPs).
Annealing exploration probability. Setting large value for helps encourage exploration, but inference time we set ε = 0. To avoid the mismatch between these settings, we anneal e towards zero over the course of training. Specifically, we train the model with sequence of values, warm-starting each model with the output from the previous value.
Parameter sharing. As mentioned in Section 4 sometimes share parameters between the predictor and value network. We implement this via a shared backbone, e.g., sequence of self- attention layers in VIT (Dosovitskiy al.) 2020). The backbone initialized via the predictor pre-training with random masks, and it is then used for both / (x 0) and o(x;; Φ) with separate output heads for each one.
21

---

Published as a conference paper at ICLR 2024
Scaling value network outputs. To learn the optimal value network outputs, it is technically sufficient to let the network make unconstrained, real-valued predictions. However, given that the non-negative, or xy; xg) 2 0 for all (z,, it sensible to constrain the predictions: for example, we can apply a softplus output activation. Similarly, we know that the true CMI values are upper bounded by the current prediction entropy Hyy xs) (Cover and Thomas 2012). These simultancous bounds can be summarized as follows:
true CMI values
0 S xy; (..))
Hy xs)).
To enforce both inequalities, we apply a sigmoid operation the unconstrained value network prediction u(zs; Φ), and we multiply this by the empirical prediction entropy from fæzs; 8). An ablation showing the effect of this approach in Figure 9]
Prior information. We found that issue with using prior information (as discussed in Section 4.2 overfitting z. This perhaps unsurprising, particularly when is high-dimensional. because the same input is used repeatedly with different feature subsets Xş and the same label To mitigate this, we applied the following simple fix: for the separate network that processes the prior variable z, we detached gradients when using the learned representation to make classifier predictions, so that gradients are propagated only for the value network's 's CMI predictions. An ablation demonstrating this approach is in Figure 15
Inference time. At inference time, we follow a similar procedure as in Algorithm Tbbt with = so that always make the most valuable selection. Algorithm demonsttaates the pseudo-code for inference. In terms of stopping criteria for making a prediction, we explore multiple approaches, discussed in (Section 4.3 (I) a budget-constrained approach with parameter k, (2) confidence constrained approach with parameter m, and (3) penalized approach with parameter A. Our results generated by evaluating single learned policy with several values for each of these parameters. The range of reasonable values for the confidence parameter m and penalty parameter À depend on the dataset, so these are tuned by hand.
Feature grouping. Several of our datasets involve grouped features: for example, we group pixels in the image datasets into patches, and our medical diagnosis datasets have grouped one-hot indicators for categorical variables (Intubation, ROSMAP). To implement this grouping structure in our method, we simply predict the CMI for each group, and then calculate the value network's 's objective based on the loss improvement after revealing the group's 's values.
D DATASETS
This section provides details about the datasets used in our experiments. The size of each dataset is summarized in Table
Table Summary of datasets used in our experiments.
Dataset
# Features
# Feature Groups
# Classes
# Samples
MNIST
784
10
60,000
Intubation
112
35
65,515
ROSMAP
46
43
13,438 2
ImageNette
50,176
196
10
13,395
ImageNet-100 50,176
196
100
135,000
Histopathology 50,176
196
3152
MNIST. This is the standard digit classification dataset (LeCun et al.) 1998). We downloaded it with PyTorch and used the standard train and test splits, with 10,000 training samples held out as a validation set.
Intubation.
This is a privately curated dataset from a university medical center, gathered over 13-year period (2007-2020). Our goal is to predict which patients require respiratory support upon arrival in the emergency department. We selected 112 pre-hospital clinical features including dispatch information (injury date, time, cause, and location), demographic information (age, sex),
22

---

Published as a conference paper at ICLR 2024
and pre-hospital vital signs (blood pressure, heart rate, respiratory rate). The outcome is defined based on whether a patient received respiratory support, including both invasive (intubation) and non-invasive (BiPap) approaches. We excluded patients under the age of 18, and because many features represent one-hot encodings for categorical variables, we grouped them into 35 feature groups. Feature acquisition costs were obtained by having a board-certified emergency physician estimate the relative cost of obtaining each feature. The dataset is not publicly available due to patient privacy concerns.
ROSMAP. The Religious Order Study (ROS) and Memory Aging Project (MAP) (A Bennett et al. 2012a b) are complementary epidemiological studies that enroll participants to study dementia. ROS is a logitudinal study that enrolls clergy without known dementia from across the United States, including Catholic nuns, priests, and brothers aged 65 years and older. Participants agree to annual medical and psychological evaluation and pledge their brain for donation. MAP is longitudinal study that enrolls participants encompassing a wider community from 40 continuous care retirement facilities around the Chicago metropolitan area. Participants are without known dementia and agree to annual clinical evaluation and donation of brain, spinal cord and muscle after death. While entering the study, participants share demographic information (e.g. age, sex) and also provide their blood samples for genotyping. At each annual visit, their medical information is updated and they take series of cognitive tests, which generate multiple measurements over time. This results in 46 different variables, grouped into 43 feature groups to account for one-hot encodings. The task is to predict dementia onset within the next three years given the current medical information and no prior history of dementia. In total, the data contains 3,194 individuals with between 1 and 23 annual visits. Following the preprocessing steps used in (Beebe-Wang et al.) 202I), we applied a four-year sliding window over each sample, thereby generating multiple samples per participant. Each sample split into an input window consisting of the current year visit t and a prediction window of the next three years (l + 1, + 2, + 3). To avoid overlap between the training, validation, or testing sets, we ensured that all samples from a single individual fell into only one of the data splits. Feature acquisition costs expressed in terms of time taken were borrowed from (Beebe- Wang et al.) (002)) for the cognitive tests and rough estimates were assigned to the remaining features using prior knowledge. We discarded the genotypic feature (APOE e4 allele) from the feature set since it is highly predictive of dementia and is difficult to assign an appropriate cost. The dataset can be accessed at httpt: //de:: niagads org/ cwwprts.h//////o//o/oooooo religious Peligious-orders- senders--tudy memory- and aging project--eectttcoc rosmap/|
Imagenette and ImageNet-100. These are both subsets of the standard ImageNet dataset (Deng et al. 2009). Imagenette contains 10 classes and was downloaded using the Fast.ai deep learning library (Howard), ImageNet- 100 contains 100 classes and was downloaded from Kaggle (Ambityga). and in both cases we split the images to obtain train, validation and test splits. Images were resized to 224 x 224 resolution for both architectures we explored, ResNets (He et al., 2016b) and ViTs (Dosovitskiy et al. 2020).
MHIST. The MHIST (minimalist histopathology) (Wei et al. 2021) dataset comprises 3.152 hema- toxylin and eosin (H&E)-stained Formalin Fixed Paraffin-Embeddee (FFPE) fixed-size images of colorectal polyps from patients at Bartmouth-Hitcheere Medical Center (DHMC). The task is to perform binary classification between hyperplastic polyps (HPs) and sessile serrated adenomas (SSAs), which is a challenging prediction task with significant variation in inter- anter-patholgiss. agree- ment (Abdeljawad et al. 2015 Farris et al. 2008 Glatz al. 2007 Khalid et al., 2009) WWng et al. 2009). HPs are typically benign, while SSAS are precancerous lesions that can turn into cancer if not treated promptly. The fixed-size images were obtained by scanning 328 whole-slide images and then extracting regions of size 224 x 224 representing biagnostically-reeeve
regions of interest for HPs or SSAs. For the ground truth, each image was assigned a gold-standard label determined by the majority vote of seven board-certified gastrointestinal pathologists at the DHMC. The dataset can be accessed by filling out the form at hhtps:://// hipp://bb.. //mmi//ii rds. wea.s.gittan github........ 10/MHIST/
23

---

Published as a conference paper at ICLR 2024
E
MODELS
Here, we briefly describe the types of models used for each dataset. The exploration probability for all models is set to 0.05 at the start with an annealing rate of 0.2.
Tabular datasets. For all the tabular datasets, we multilayer perceptrons (MIPs) with two hidden layers and ReLU non-linearity. We use 128 neurons in the hidden layers for the ROSMAP and Intubation datasets, and 512 neurons for MNIST. The initial learning rate is set to 10-3 at the start and we also use dropout with probability 0.3 in all layers to reduce overfitting (Srivastava et al. 2014). The value and predictor networks use separate but identical network architectures. The networks are trained on a NVIDIA RTX 2080 Ti GPU with 12GB of memory.
Image datasets: ResNet. We use shared ResNet-50 backbone (He et al. 2016b) for the predictor and value networks. The final representation from the backbone has shape 7 x 7, and the output heads for each network are specified as follows. The predictor head contains a Conv - Batch Norm -> ReLU sequence followed by global average pooling and a fully connected layer. The value network head consists of an upsampling block with a transposed convolutional layer, followed by a 1 x convolution and sigmoid to scale the predictions (see Appendix C). The learning rate starts at 10 and the networks are trained on a NVIDIA RTX 2080 Ti with 12GB of memory.
Image datasets: VIT. We use shared ViT backbone (vit_sma 1l 11_11pppaae patchatchl 6 6_224) (Doooviskkiy ct al. 2020) for the predictor and value networks. We use ViT and ResNet backbones having a similar number of parameters for fair comparison: ResNet-50 has approximately 23M parameters, and ViT-Small has 22M parameters. The predictor head contains a linear layer applied to the class token, and the value network head contains a linear layer applied to all tokens except for the class token, followed by sigmoid function. When incorporating prior information, a separate ViT backbone is used for both the predictor and value networks to generate an embedding, which is then concatenated with the masked image embedding to get either the predicted CMIs or the class prediction. The learning rate starts at 10-5, and the networks are trained on NVIDIA Quadro RTX 6000 GPU with 24GB of memory.
F BASELINES
Here, we provide more details here on our baseline methods.
Concrete autoencoder. This is a static feature selection method that optimizes a differentiable selection module within neural network (Balin et al.) (019).. The layer can be added at the input of any architecture, so we use this method for both tabular and image datasets. The original work suggested training with an exponentially decayed temperature and hand-tuned number of epochs, but we use a different approach to minimize the tuning required for each dataset: we train with a sequence of temperature values, and we perform early stopping for each one based the validation loss. We return the features that are selected after training with the lowest temperature, and we evaluate them by training model from scratch with only those features provided.
This is static feature selection method that scores features based on their mutual information with the response variable. To identify joint interactions between several features, the authors build a scoring function using a min-max optimization objective (Shishkin et al., 2016). To make the method practically feasible, binary representations of features are used. To adapt this method to our setting, for each feature budget, we use CMICOT to select the best subset and then fit a classifier on the selected features to get the final performance. We ensure that the classifier used has the same architecture as the one used for DIME.
CMICOT.
mRMR. This is static feature selection method that identifies a subset of features from a larger set that maximizes the relevance to the target variable while minimizing redundancy among selected features (Peng et al. 2005). In other words, it aims to find a set of features that are individually informative for predicting the target variable and, at the same time, not highly correlated with each other. Similar to CMICOT, we fit a separate classifier on the subset identified for each feature budget.
EDDI. This is DFS method that uses generative model to sample the unobserved features (Mal ct al. 2019). We implement a PVAE to sample the unknown features, and these samples are used to estimate the CMI for candidate features each selection step. We separately implement classifier
24

---

Published as a conference paper at ICLR 2024
that makes predictions with arbitrary feature sets, similar the one obtained after masked pre-training in Algorithm I] We use this method only for our tabular datasets, as the PVAE is not expected to work well for images, and the computational cost at inference time is relatively high due to its iteration across candidate features.
Probabilistic hard attention. This method extends EDDI to work for images by imputing unobserved features within a diwddeensiooal, learned feature space (Rangrej and Clark (021]. To ensure that the method operates on the same image regions DIME, we implemented feature extractor that separately computes embeddings for non-overlapping 14 x 14 patches, similar to a ViT (Dosovitskiy al. 2020) or bag-of-feature model (Brendel and Bethge| 2019 Specifically, our extractor consists of 16 x 16 convolutional layer, followed by a series of I x 1 convolutions. The features from each patch are aggregated by recurrent module, and we retain the same structure used in the original implementation.
Argmax Direct. This is DFS method that directly estimates the feature index with maximum CMI. It is based on two concurrent works whose main difference is how gradients are calculated for the selector network (Chattopadhyay et al., 2023 Covert et al. 2023 for simplicity, we only use the technique based on the Concrete distribution from Covert et al. (202)). As a discriminative method, this baseline allows us to use arbitrary architectures and is straightforward to apply with either tabular or image datasets.
Classification with costly features (CwCF). This is an RL-based approach that converts the DFS problem into MDP, considering the expected utility of acquiring feature given its cost and impact on the classification accuracy (Janisch et al.) 2019). A variant of deep Q-learning is implemented as the RL solver, and the method is adapted work with sparse training datasets with missing features. We used the budget-constrained implementation, we added support for feature grouping, and we use the same architectures as DIME for the Q-network. We train and evaluate separate models for each target budget. We limit comparison to tabular datasets because the method does not scale well and high performance is not expected for images. CwCF initially suffered in our evaluation with the medical diagnosis datasets because we use AUROC to measure performance: CwCF produces hard classifications rather than predicted probabilities, so it can suffer on a ranking- based metrics like AUROC that are more meaningful for clinical prediction tasks. To account for this discrepancy, we followed an approach from Erion et al. (2022) and used the Q-values for classes as proxies for pseudo predicted probabilities, since the Q values for "class" action can be interpreted as a score of how confident the model is in predicting that class. This resulted in improved performance for the ROSMAP and the Intubation datasets. However, still observe that CwCF underperforms compared to the other methods.
Opportunistic learning (OL). This is another RL-based approach to solve DFS (Kachuee et al. 2018). The model consists of two networks: a Q-network that estimates the value associated with each action, where actions correspond features, and P-network responsible for making predictions. These are similar to the value and predictors network used in DIME, so we use the same architectures as our approach, and OL shares parameters between the P- and Q-networks. We use the implementation from (Covert et al., (023)) which modifies the method by preventing the prediction action until the pre-specified budget is met, and supports pre-defined feature groups. We limit comparison to tabular datasets since the method does not scale well and high performance is not expected for images.
G ADDITIONAL RESULTS
Here, we present the results of several additional experiments and ablations on the different datasets.
Figure 7| shows the prediction calibration of the predictor network by plotting DIME's performance at different levels of confidence for a specific budget of k 15, along with the density of the samples at those contidence levels across multiple datasets. This shows that the predictor network is well calibrated and does not systematically overestimate underestimate its predicted probabilities. Proper calibration is important to achieve accurate loss values, because these are then used to train the value network (see eq. (3)p.
Figure 8) shows the calibration of the predicted CMIs by the value network for both a tabular and an image dataset by plotting the difference in entropy or losses against the predicted CMI. A linear trend
25

---

Published as a conference paper at ICLR 2024
showcases that the CMIs predicted by the value network align well with the difference in either the entropy or the loss. Since we do not have ground truth CMI values to evaluate the accuracy of our value network, this serves viable alternative for real- world datasets, and we can verify that the CMI predictions correctly represent the expected reduction in either loss or prediction entropy.
Figure 9 shows the effect of constraining the predicted CMIs using the current prediction entropy. as described in [Apendixx CJ Without the constraint, there are some samples with unrealistically high CMI estimates that are greater than the prediction entropy. After applying the sigmoid activation on the value network predictions, this issue is corrected.
Figure shows multiple trials for the penalized policy on the tabular datasets. This provides simple way to represent variability between trials when we cannot precisely control the budget between separately trained policies. Similarly, FFiuuee 12 and Figure 13 show multiple trials while considering non-uniform feature costs, and Figure shows multiple trials for the image datasets. The relative results stay the same across all trials, as the performance variability between trials is generally small.
Figure 14 shows the confidence distribution of full input predictions in the tabular datasets. Across all datasets, we observe that the model has high confidence in many of the samples, but there are some that remain uncertain even after observing all the features. This provides motivation for using the penalized approach, because confidence-constraingg the entire feature acquisition budget only to remain at high uncertainty.
approach could suffer here by expending
Figure 13 shows the effectiveness of detaching gradients for the predictor network of the sketch z in the histopathology dataset, as described in Appendix C The penalized policy performs significantly better when we propagate gradients only for CMI predictions, which we attribute to reduced overfitting to the prior information.
Figure 16 compares the penalized stopping criterion with the confidence- and budget-constraiin versions for the image datasets. Similarly. Figures 10 11 include the budget-constrained approach in comparisons with the baselines. The penalized stopping criterion introduced in Secciin 4.3 consistently achieves the best results.
Table 2 compares the training times (in hours) of DIME with the baselines. It is difficult to compare exact training times across methods because each has hyperparameters that be tuned to make it converge faster, but we compared them under the hyperpurameters used to generate our results. We observe that DIME trains slightly faster than Argmax Direct for both the tabular and image datasets. Compared to EDDI and Hard Attention, we see that these generative methods are actually somewhat faster to train; however, because they must generate a large number of samples at inference time and iterate over candidate features, their evaluation time is far slower. The total training and inference time for EDDI is 48.26 hours with MNIST, and for Hard Attn is 25.45 hours on Imagenette, whereas the evaluation time for DIME is negligible compared to its training. As expected, the RL method CwCF is very slow to run on MNIST, taking almost 29 hours to train models for all target budgets.
Finally, Figure 17 shows example feature selections for the MNIST dataset, and Figure 18 shows feature selection frequencies for ROSMAP, to verify the distinct selections made across predictions. This capability differs from static selection methods like the CAE (Balin al. 2019). which select the same features for all predictions.
Table 2: Training times for each method (in hours).
Method
MNIST
Imagenette
DIME
4.88
27.14
Argmax Direct
6.02
43.67
EDDI
0.39
Hard Atm
18.25
CwCF
28.91
26

---

Published as a conference paper at ICLR 2024
7000
6000
0.8
a Sammpley
15000
0.7
4000
6.66
3000
0.5
2000
0.4
1000
Confidence
Contisencc
ROSMAP
ROSMAP
1750
0.9
0.6
TofSenpler
1250
1.000
0.6
500
0.5
0.5
0.7
0.5
0.9
0.5
0.5
0.7
0.3
0.9
1.0
Intuba tion
Intubation
3000
0.9
2500
0.8
2000
1500
0.6
1000
0.4
0.5
0.7
0.5
0.5
Confirenc
Confidence
Imagenette
Imagenette
1.0
2500
0.8
Samples
0.6
0.2
Confiiincc
Figure 7: Evaluation of prediction calibration for a fixed budget of A: = 15. The left column shows the prediction calibration of the predictor network by plotting the accuracy for different confidence levels. The right column shows the distribution of confidences across all samples.
27

---

Published as a conference paper at ICLR 2024
MNIST
Imagenette
0:30
Entropy Loss
Entropy
0.30
0.25
Metric Difference
Metric Difference
0.20
0.05
0.00
0.00
0:15 Predicted
Figure 8: Evaluation of CMI calibration. The x-axis shows different values for the predicted CMI throughout the selection process, and the y-axis shows the reduction either loss or entropy after the corresponding feature is selected.
Imagenette w/ Entropy Trick
Imagenette w/o Entropy Trick
2.00
2.00
1.75
1.50
1.50
Predicted CHI
Predicted CMI
1.00
1.00
0.75
0.50
0.50
0.25
0.00
0.00
Prediction Entropy
Prediction Entropy
Figure 9: Predicted CMIs with and without the entropy trick to scale value network outputs.
Intubation
ROSMAP
0.90
0.92
0.9
0.90
0.8
0.85
0.7
0.86
0.80
0.82
0.4
0.80
0.3
0.70
10
15
20
25
5
10
15
20
25
2
10
Avg Featuress
Avg # Features
Avg # Features
(Penalized)
(Budget)
Argmax
Figure 10: Multiple trials using the penalized policy and the budget constraint for tabular datasets. DIME with penalized policy remains the best method across five independent trials.
28

---

Published as a conference paper at ICLR 2024
Imagenette
Net-100
Histopathology
1.0
0.9
0.8
0.85
0.6
AUROC
0.75
0.4
0.65
0.4
0.3
15
30
20
60
Aug Patches
Avg Patches
Avg * Patches
DIME (Penalized)
DIME (Budget)
CAE
Argmax Direct
Hard Attn
Figure 11: Multiple trials using DIME'S penalized policy and the budget constraint for image datasets. DIME with penalized policy remains the best method across five independent trials.
ROSMAP Trial
ROSMAP Trial 2
ROSMAP Trial 3
0.90
0.90
0.85
0.85
AUROC
AUROC
0.80
0.80
0.75
0.75
0
250
500
750
1000
500
1000
1500
500
1000
1500
Cost (Seconds)
Cost (Seconds)
Cost (Seconds)
ROSMAP Trial
ROSMAP Trial 5
0.90
AUROC
AUROC
0.80
0.70
500
1000
1500
2000
500
1000
1500
Cost, (Seconds)
Cost (Seconds)
DIME (Budget)
- DIME (Budget, Inference Costs)
DIME (Budget, No Costs)
EDDI
Figure 12: Multiple trials when using non-uniform feature costs for the ROSMAP dataset.
29

---

Published as a conference paper at ICLR 2024
Intubation
0.90
0.90
0.90
AUROC
0.80
0.80
0.80
0.75
0.75
0.70
0.70
0.70
50
100
150
200
50
100
150
200
50
100
150
200
Cost (Relative. 1-10)
Cost (Relative. 1-10)
Cost (Relative. 1-10)
Intubation Trial 4
Intubation Trial 5
0.90
0.85
0.85
AUROC
AUROC
0.80
0.80
0.75
0.75
0.70
Cost (Relative, 1-10)
Cost (Relative, 1-10)
(Budget)
DIME (Budget, Inference Costs)
DIME [Budget, No Costs)
Figure 13: Multiple trials when using non-uniform feature costs for the intubation dataset.
Intubation
ROSMAP
10000
3000
8000
Semeles
1250
6000
Same
2000
4000
1000
2000
250
0.4
0.6
0.8
0.5
0.6
0.7
1.0
0.5
0.6
0.7
0.8
0.9
1.0
Confidence
Confidence
Confidence
Figure 14: Confidence distribution on full-input predictions across the tabular datasets.
Histopathology
0.85
0.80
AUROC
0.75
DIME (Penal
0.70
DIME (Budg et)
DIME (Penalized, detach)
0.65
DIME (Budget, No detach)
20
40
60
Avg # Patches
Ablation of stop-gradients trick when using prior information for the histopathology dataset (ppendix C).
Figure 15:
30

---

Published as a conference paper at ICLR 2024
Imagenette
ImageNet-100
Histopathology
0.8
0.85
0.9
0.7
Top-1. Arc
0.6
0.75
0.7
0.70
0.4
0.6
0.3
0.60
2.5
7.5
10.0
20
Avg # Patches
Avg Patches
Aug # Patches
(Penalized)
(Budget)
(Confidence)
Figure 16: Comparison of the budget-constrained..ooffe com. and penalized stopping criteria.
7
2
7
8
6
3
7
9
4
Figure 17: Examples of MNIST feature selections across several samples, with budget
ROSMAP
1.0
2
0.8
# Features
0.6
9
0.4
10
12
0.2
14
0.0
12-
15-
18
21
24-
27
30
33
36
39
42
Feature
x
Figure 18: Feature selection frequency for ROSMAP. Each entry shows the fraction of samples in
which a feature is selected when we use the budget-constrained stopping criterion for the specified number of features.
31

---

Published as a conference paper at ICLR 2024
Intubation
ROSMAP
0.90
0.90
0.85
0.85
AUROC
0.80
AUROC
0.80
0.75
0.75
0.70
0.70
0.65
5
15
20
25
2
8
Avg # Features
Avg Features
DIME (Penalized)
Argmax Direct
OL
mRMR
CMICOT
FIRES
CAE
EDDI
CwCF
Evaluation with tabular datasets for varying feature acquisition budgets. Results averaged across 5 trials, and shaded regions indicate the standard error for each method.
Figure 19:
32


Reviews:
[{"review_id": "iPrbTdvOLV", "reviewer": "Reviewer_iU6i", "summary": "This paper studies the problem of dynamic feature selection (i.e., adaptively acquire features based on previously acquired ones), where the ultimate goal is to classify each data instance using a small subset of the available features. A greedy approach, termed DIME, is proposed, where feature selection is based on the conditional mutual information (CMI) with the response variable. The paper focuses on estimating CMI in a discriminative fashion, and then trains two networks to implement the feature selection policy. Furthermore, the paper discusses how to handle non-uniform costs across features, prior information and variable feature budgets. A number of experiments on tabular and image datasets are provided that showcase the performance of DIME in contrast to prior work on the area of dynamic feature selection.", "questions": "(1) I believe that the following set of references also study dynamic feature selection, but from a different perspective than the references presented in the paper. Specifically, they formulate the problem as POMDP and study the theoretical properties of the optimum solution, which enables them to design fast algorithms for dynamic feature selection. They also give bounds on the expected number of features needed to achieve a specific accuracy level. Finally, the latter paper uses mutual information to also drive dynamic feature selection. I think that after reviewing these papers, the authors would agree with me that this is relevant work and they should consider citing them for completeness.\n \n- Liyanage, Y.W., Zois, D.S. and Chelmis, C., 2021. Dynamic instance-wise joint feature selection and classification. IEEE Transactions on Artificial Intelligence, 2(2), pp.169-184. \n\n- Liyanage, Y.W., Zois, D.S. and Chelmis, C., 2021. Dynamic Instance-Wise Classification in Correlated Feature Spaces. IEEE Transactions on Artificial Intelligence, 2(6), pp.537-548. \n\n(2) 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.\n\n(3) The performance of DIME is very close to Argmax Direct. What is the benefit of using DIME instead of Argmax Direct in that sense?\n\n(4) My major concern is the really surprisingly bad performance of RL-based methods compared to the greedy methods, including DIME. The reason is that I see the RL-based methods as a way of implementing dynamic programming, which is the optimal thing to do, and should be better than greedy methods. Of course, RL has its own issues, but I am still wondering if the performance improvement observed in the experiments is due to using more powerful neural networks to approximate the selection process and not the greedy nature of the method. In addition, it may be because the RL-based methods that DIME is compared with use different criteria than mutual information (excluding the missing references above). Finally, I believe that the RL-based methods consider the cost of features during their decision making process. Is it possible that the results presented in the paper have misspecified this parameter? May be an ablation study  would help me better justify the above surprising behavior of RL-based methods compared to DIME\n\n(5) I noticed in the Appendix that features are in groups. Does DIME use the group structure somehow?\n\n(6) It is unclear to me from the description in the paper if DIME uses a hard feature budget constraint for each data instance or an average constraint or something else. May be you can clarify this?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "+ The paper studies a very important and timely problem, i.e., dynamic feature selection with a minimal budget.\n+ A new greedy algorithm is proposed that estimates CMI in a discriminative fashion.\n+ The paper discusses how to handle common issues happening in practice, e.g., non-uniform costs, variable feature budgets, etc.\n+ Theoretical results are provided that lead to an unbiased estimator of CMI.\n+ The paper is well-written and effort has been put to clarify its contributions.\n+ The performance of the proposed algorithm is illustrated on a number of datasets and compared with existing methods.", "weaknesses": "- I believe that some relevant wont on dynamic feature selection that is neither RL-based nor greedy-based is missing (see references below)\n- The performance improvement of DIME wrt to one of the baselines (i.e., Argmax Direct) is small (especially for the tabular datasets). \n- The paper claims that DIME only accounts for non-uniform costs and variable feature budgets, but I am pretty sure that some of the RL-based methods handle such settings too.\n- The really bad performance of the RL-based methods compared to the greedy methods, including DIME, is quite surprising to me and I am wondering if this is some kind of typo.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HWk317YILR", "reviewer": "Reviewer_j5dM", "summary": "The work is devoted to the problem of dynamic feature selection. The authors propose to estimate – learn – Conditional Mutual Information (CMI) of the next feature given the past selected ones. The learning is based on the incremental loss improvements of a given predictor (simultaneously learned with the CMI). It is proven that, for cross entropy loss, the optimal learning point will be CMI. The paper also contains experiments on several datasets that demonstrate better performance of the proposed method in comparison to several baselines.", "questions": "1.\tAs 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?)\n\n\n2.\tReturning 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?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "-\tClear problem setup\n-\tLarge volume of experiments", "weaknesses": "-\tComparison with baselines (see questions)\n-\tIt looks like incorporating prior information and variable feature selection constitute 2/3 of the overall work contribution, but they seem more-or-less incremental things (i.e., given such a task a practitioner would easily reproduce Sections 4.2 and 4.3 having Section 4.1). More non-trivial insights would be nice to see in the main contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zMCg8Yk6OH", "reviewer": "Reviewer_s6Fj", "summary": "The authors propose DIME a dynamic feature selection (DFS) method that is based on two neural networks, aiming to maximize prediction accuracy while balancing the cost of acquiring features. Specifically, a value network is designed to estimate the conditional mutual information (CMI), $I(y;x_i∣x_S)$, and a prediction network is used to make predictions based on the currently selected features. Both networks are jointly trained. If the acquisition costs for features are known, the objective aims to select the feature that maximizes $I(y;x_i∣x_S)/c_i$, where $c_i$ is the acquisition cost of $i$'th feature. In the experiments, DIME is compared against both static and dynamic feature selection methods on various (tabular and image) datasets. The results show that DIME outperforms the competitors in terms of predictive performance and effectiveness (i.e., number of selected features).", "questions": "See weaknesses 3.) and 4.) from above.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1.) The research investigates various policies for feature selection, such as RL, imitation learning, and a greedy policy based on CMI.\n\n2.) The introduction of the Predictor Network and Value Network, which jointly work to estimate Conditional Mutual Information (CMI) is interesting and represents a novel approach.\n\n3.) The theoretical parts of the paper provide a solid foundation for the approach.\n\n4.) The experimental evaluation is quite comprehensive.", "weaknesses": "1.) 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.\n\n2.) Some related approaches should be mentioned in the related work section and could have been considered as competitors\n\nIncremental permutation feature importance (iPFI): towards online explanations on data streams (Machine Learning 2023) \nF Fumagalli, M Muschalik, E Hüllermeier, B Hammer\n\nLeveraging model inherent variable importance for stable online feature selection (KDD 2020)\nJ Haug, M Pawelczyk, K Broelemann, G Kasneci \n\n3.) It is not quite clear how prior information is factored into and affects the feature selection policy; this part should be better explained.\n\n4.) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)