Stochastic Encodings for Active Feature Acquisition
Authors: Alexander Luke Ian Norcliffe
Changhee Lee
Fergus Imrie
Mihaela van der Schaar
Pietro Lio
Venue: ICML 2025 poster (2025)
Keywords: Active Feature Acquisition
Dynamic Feature Selection
Primary area: deep_learning->everything_else
TL;DR: We present a new method using stochastic encoders for Active Feature Acquisition, the test time task of iteratively observing features to improve the current prediction.

Abstract: Active Feature Acquisition is an instance-wise, sequential decision making problem. The aim is to dynamically select which feature to measure based on current observations, independently for each test instance. Common approaches either use Reinforcement Learning, which experiences training difficulties, or greedily maximize the conditional mutual information of the label and unobserved features, which makes myopic acquisitions. To address these shortcomings, we introduce a latent variable model, trained in a supervised manner. Acquisitions are made by reasoning about the features across many possible unobserved realizations in a stochastic latent space. Extensive evaluation on a large range of synthetic and real datasets demonstrates that our approach reliably outperforms a diverse set of baselines.

Stochastic Encodings for Active Feature Acquisition
Alexander Norcliffe'
Changhee Lee? Fergus Imrie
Mihaela van der Schaar
Pietro Liò
Abstract
ray might be prioritized. Active Feature Acquisition (AFA I is an inference time task, where the features are not assumed to be all available at once. Instead, on an instance-wise basis, model sequentially acquires features based on the existing observations to best aid long-term prediction. A common approach is to use Reinforcement Learning (RL) (RückstieB] et al. 2013 Shim et al., 2018), since this is a natural so- lution sequential decision making problem. However, RL suffers from training difficulties such as sparse reward, exploration exploitation, and the deadly-triad (Henderson) ct al. 2018: Erion et al., 2022] Van Hasselt et al. 2018). An alternative approach is to select features that greedily maximize the conditional mutual information (CMI) (Chen et al. 2015a b). This has a significant drawback: CMI does not capture the effects of unobserved features that can be acquired at later stage. This results in myopic decision making that favors immediate predictive power over sets of jointly informative features, or features that are highly informative of which feature to acquire next. Additionally, we argue that CMI is not even guaranteed to be the best short-term objective from the perspective of the 0-1 Loss (making single decisive prediction). Since maximizing CMI is equivalent to minimizing entropy, and this can be achieved by making unlikely classes even more unlikely, rather than distinguishing between more probable outcomes. We further explore the drawbacks of CMI in Section 4
Active Feature Acquisition is an instance-wise, sequential decision making problem. The aim is to dynamically select which feature to measure based on current observations, independently for each test instance. Common approaches either use Reinforcement Learning. which experiences training difficulties, or greedily maximize the con- ditional mutual information of the label and unob- served features, which makes myopic acquisitions. To address these shortcomings, we introduce la- tent variable model, trained in supervised man- ner. Acquisitions are made by reasoning about the features across many possible unobserved re- alizations a stochastic latent space. Extensive evaluation on a large range of synthetic and real datasets demonstrates that our approach reliably outperforms a diverse set of baselines.
1. Introduction
The standard supervised learning paradigm is to learn a pre- dictive model using a training dataset of features and labels, such that the model can make accurate predictions on unseen test inputs. A fundamental assumption is that, test time, all features are jointly available. However, this assumption does not always hold. Consider the example of doctor di- agnosing patient (Kachuee et al. 2019b a). Initially, there is little to no information available and, while there are many tests that could be conducted, the doctor will choose which ones to carry out based on their current understanding of the specific patient's condition. For instance, if a patient has pain in their leg, and the doctor suspects fracture, a leg X-
Motivated by the shortcomings of RL and CMI maximiza- tion, we introduce a novel AFA approach, which we call Stochastic Encodings for Feature Acquisition (SEFA), that departs from existing methods in several key ways. First, we shift the acquisition problem from reasoning in com- plex feature space to a latent space. We use an information bottleneck-inspired term (Tishby et al. 1999) to regularize the latent space such that decisions are made using label- relevant information only and not noise associated with the features. Second, we use stochastic encoders, allowing us to acquire features by considering their effect across di- verse range of possible latent realizations. This allows us to capture the effects of features that are yet to be observed, re- sulting in acquisitions that are non-greedy by design. Third, our acquisition objective places more focus on labels with higher predicted likelihood, leading to acquisitions that help
Department of Computer Science and Technology, University of Cambridge, Cambridge, United Kingdom 'Department of Artifi cial Intelligence, Korea University, Seoul, Korea 'Department of Statistics, University of Oxford, Oxford, United Kingdom "Department of Applied Mathematics and Theoretical Physics, University of Cambridge, Cambridge. United Kingdom. Corre- spondence to: Alexander Norcliffe <alin2 felin2@cam... cam.ac.uks
Proceedings of the 4nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 20255 by the author(s).
'This problem has also been referred to as Dynamic Feature Selection (DFS).

---

Stochastic Encodings for Active Feature Acquisition
to disambiguate between the most likely classes. Finally, to avoid the difficulties posed by RL., we do not train model to make acquisitions directly. Instead, we train with pre- dictive loss and make acquisitions by maximizing custom objective in a suitably regularized latent space. Our contribu- tions are as follows: (1) We re-examine the CMI acquisition objective and provide theoretical reasoning and concrete examples of its sub-optimality, (2) We introduce SEFA, our novel AFA approach motivated by the limitations of RL (3) We evaluate SEFA on multiple synthetic and real-world datasets, including cancer classi- fication tasks. Comparing against various state-of- AFA baselines, we see that SEFA consistently outperforms these methods. Extensive ablations further demonstrate each novel design choice is required for the best performance.
modeling (Peis al. 2022 He et al. 2022 Li al. 2020
Li & Oliva 2020).
A Alternative Solutions. Sensitivity- solutions make selections based on how sensitive the label is to a given fea- ture (Kachuee et al. 2017 2018). However, since missing values are filled with zero and measuring a feature is dis- continuous, the gradient does not reliably represent the true sensitivity. Imitation learning has been applied (Valancius et al. 2024 He et al.) 2016), however, this requires access to an oracle or to construct one. Prior to deep learning, deci- sion trees were used, with features acquired at each branch of a tree if the feature is unobserved (Xu et al. [2012] 2013 Kusner et al., 2014 Trapeznikov & Saligramaa 2013 Xu et al. 2014). This has also been generalized to ensembles of decision trees (Nan et al. 2015 2016).
and CMI maximization.
2. Related Work
3. Active Feature Acquisition
Reinforcement Learning. The most common AFA ap- proach is to frame the problem as a Markov Decision Pro- cess and train a policy network with RL to decide which
Problem Setup. In standard C-way classification, we have a dd-dimensional feature vector given by the random variable X E X with realization x = (I1, T2,.... . (2))..,;)). and a label given by Y E [C] with realization y. Ordinarily, we assume all features are observed; however, more generally, we wish to allow arbitrary feature subsets as valid inputs. Therefore, let represent a missing feature value and X' = -I (X; (+)). We denote input with feature subset S C [d]. as Xs, where Z5 = 2; if i E S, and ISS = if i S. Given a training set, Draan = {(xs, y)m Dy the AFA task is to train a model that takes a test instance with arbitrary observations xo, and iteratively acquires new features. The model's long-term acquisition goal is to acquire a sequence of features S* to maximize the confidence in its prediction whilst minimizing the number of acquired features:
feature to acquire next Dbulac-Amold et al. 2011 RückstieB
al.
2013
Shim et al. 2018 Janisch al.) 2019; Mnih
let al.
2014 Kachuee al. 2019a). The RI approach read-
ily extends to temporal setting where features and labels can change over time (Kossen et al.) 2023] Yin et al.) 2020). Whilst a natural solution to AFA, RL suffers from training difficultics, and various advances in the RL field have been applied to account for this. For example, using generative models to augment datasets (Zannone et al.) 2019). provid- ing mutual information as additional input to the policy (Li & Oliva 202I), using gradient information in the training process Ghosh & Lan 2023), and reward shaping (Peng
ct al. 2018).
Conditional Mutual Information Maximization. Condi- tional mutual information tells us how much we can learn about one variable by measuring a second, whilst already knowing third. Greedy CMI maximization is a common AFA approach, due to its grounding in information the- ory. However, as we demonstrate in Section 4 it inherently makes short-term acquisitions and is prone to making acqui- sitions that do not distinguish between likely labels. Among existing approaches, networks can be trained to directly predict CMI (Gadgil al. 2024). or policy networks can be specially trained to maximize CMI without ever calcu- lating it (Chattopadhyay et al.) 2023) Covert et al. 2023) Generative models are a second way to estimate CMI by taking Monte Carlo estimates over conditional distributions, (Ma et al. 2019: Chattopadhyay et al.) 2022, Rangrej & Clark| 202 Eaaly et al. (2016). This approach suffers from
S
arg max S
max PModa clxous) 2/51| cE|C]
subject to S.
Here, X balances how much we optimize for confident prediction compared to limiting the number of acquisitions, and B is a given feature budget. These parameters are highly domain-dependent. For example, in medicine, where the stakes are high, we have large B and low X. There is a high tolerance for acquiring features if we make confident predictions. In this paper, we assume we are able to acquire all features if necessary and, unless otherwise stated, mod- els are evaluated by tracking their predictive performance during entire acquisition starting with features.
A Aquusition in Practice. The standard approach to AFA is to construct an acquisition objective function R: X'' [d] + Rgo- that scores each feature conditioned on existing ob- servations, and to acquire the feature that maximizes this: i" arg (4 R(xo, i). The objective defined
associated generative modeling challenges, producing poor estimates of CMI, thus adding the limitations. Improved performance can be achieved with advances in generative
2

---

Stochastic Encodings for Active Feature Acquisition
by the method. For instance, CMI methods use the CMI: (x,ı(xo,1) i) T(;;; Yjxo), telling us how much mea- suring X, will reduce the entropy of Y conditioned on Xo. RL methods use the output of policy or Q network, trained directly on the sequential feature acquisition prob- lem: (xo.Xx. i) = Qx(xo); Following an acquisition, we update the observed feature set to be O U and repeat the acquisition process.
takeaway from these two propositions is that considering possible values of other unobserved features is necessary for optimality and, if the objective is chosen well, sufficient. Whilst we used the indicator example to show this, more generally it applies to settings with features that are jointly informative but individually uninformative CMI does not acquire jointly informative features because marginalizes out their interdependencies with the label.
The second drawback of CMI is that, even short-term objective, it is not guaranteed to be the best objective for identifying the single most likely class. CMI maximiza- tion is equivalent to minimizing entropy, and to show why this is not guaranteed to be optimal, consider two distri- butions over three classes: H[[0.5, 0.5, 0.0]) = 0.693 and H ([0.7, 0.15, 0.15)) 0.819. The first distribution has lower entropy, but the second identifies the most likely class. The insight is that it is possible to maximize CMI by re- ducing some class probabilities to be as low as possible. However, we often desire to identify the most likely class, rather than ruling out options that may have already had a low probability. Therefore, an effective acquisition objec- tive will place focus on distinguishing between the most likely labels. We provide another example in Appendix G
4. Limitations of CMI Maximization
Here we more closely examine the shortcomings of greedy CMI maximization for AFA, to gain an understanding of why CMI maximization can be sub-optimal and how this can be addressed. Whilst grounded in theory and extensively applied, suffers from two drawbacks.
First, greedy CMI maximization makes myopic acquisitions, which in some scenarios is guaranteed to be sub-optimal. We prove this with an example. Consider feature vector with d ++ 1 features, the first d of which are binary, and the last taking an integer value from 1 to d: X E {o, 1]dd x [d]. The final feature acts as an indicator, informing us which of the other d features gives the label, y Fazi The optimal strategy is to first choose the indicator then its designated feature, requiring two acquisitions. However, the expected number of acquisitions to arrive at the same prediction by (proof in Appendix F. The theoretical insight into why CMI fails is because possible future observations are not considered in the present decision since they are marginalized out, p( y x p(z,,,;;, #4, y (1xx)ddz, Each acquisition is made if there are no subsequent acquisitions and therefore the indicator is not chosen first. This is not specific to CMI, but any scoring that marginalizes out unobserved features.
As an aside, from the perspective of Rényi entropies (Rényj 1961). CMI maximization minimizes the Shannon entropy H1 (p) = = -, ps log(pt). Whereas if the aim is to identify the single most likely class, minimizing the min-entropy Ho) = - log max; Pi, might be more appropriate. How- ever, this would still suffer from making myopic acqui- sitions. More importantly, the theoretical foundation no longer holds: minimizing Shannon entropy is equivalent to maximizing the KL divergence; however, minimizing Ho is not equivalent to maximizing the Dao divergence.
greedily maximizing CMI is 3
Proposition 4.1. Any acquisition objective that uses the marginal p(æ,,y) y) to score feature will not select the indi- cator first.
5. Method
To address the limitations of RL and CMI for AFA, we propose a novel method, called Stochastic Encodings for Feature Acquisition (SEFA). We provide a block diagram of our method in Figure I showing both how the model makes predictions and calculates the acquisition objective. We describe the architecture and training in Section 5.1 and the acquisition objective in Section 5.2
The proof is straightforward: With no other features, the indicator and label are independent, p(za+1, y) = p(a ()y(). It therefore impossible to measure its effect on the label without considering possible values of other features, regardless of how the effect is measured. RL meth- ods do not suffer from this, since during training different scenarios are seen and the effects distilled into the parame- ters. Building on this, an adjusted objective that uses CMI but also considers possible unobserved feature values can solve the indicator problem under greedy maximization. Proposition 4.2. Greedy maximization of Epixe|xo) xx) I(X;; Y|xo, xu) is an optimal strategy for the indicator problem, where U
5.1. Architecture and Training
SEFA uses an encoder-predictor architecture with interme- diate stochastic latent variable Z E 2. Predictions are given by p9.0 (9.o (xp())))) Po(y)z).
Encoder. In order to use neural networks to implement
Throughout the paper, unless specified, entropy refers Shan- non entropy.
We prove this in Appendix F Note we do not use this as our acquisition objective, since this is intractable. The key
3

---

Stochastic Encodings for Active Feature Acquisition
Prediction
Stochastic Encoder 1
14g
Po(Y 1
Stochastic Encoder 2
Sample
Po 2
Predictor
Pa 3/2)
Pe (z(xs)
Stochastic Encoder 3
P₂ C(2)
A cquisition
-166
-1.6
r(c,z,1) = 0.49
3.11
-0.9
+0.9
-1.0
2.3
2.3
-14
2(, 0.30
1.89
1.1
Predictor
Po c/a)
0.4
0.4
40.5
-0.5
-0.8
-0.8
r(c,2,3) = 0.21
1.33
-044
-0.4
Backpropagation
-0.4
-0.9
-0.9
Normalize
Euclidean Norms
VzP clz)
Reshape
Figure 1. Block diagram of SEFA. Illustrated using three features and four latent components per feature. The presence absence of feature value is indicated with binary mask vector m. Prediction and acquisition scoring with one latent sample is given together with example numerical values for acquisition.
SEFA, the input vector must not change size. Therefore, to account for missing values, we impute missing feature values and use a binary mask, M {0, 13d, as additional input to the encoder to signal if value is real or imputed. SSS F4, mg,
avoiding potential issues associated with training using RL. We also impose an information bottleneck regularization term on the latent space [(Xs; (Tishby et 1999). minimizing how much information about the features is encoded in the latent variable. The regularization term is implemented using a variational upper bound (Alemi et al. 2017), as the expected KL divergence between pa (z X and a prior p(z). We use N 1) as the prior so that the KL divergence has a closed-form expression. This multiplied by 3 to balance predictive performance with regularization. Additionally, train SEFA make predictions arbitrary feature subsets, at each training step we first uniformly sample removal probability, and then each feature has that probability of being removed. This gives the loss
if i E S and sgs if i S. Where * is for continuous features and a new
mg, = 0
category for categorical features.
For reasons that become clear when we describe the acqui- sition objective (see Section 5.2). we factorize the latent space, so that each feature is separately encoded to latent components (see Figure [
(z(2) = II Po. (zg, ISS,s m..)).
Here G; indexes the latent components that feature i is responsible for encoding. For example, if = 3 then
L
E Pauu 1) PSatbanck
E
log
(1)
Gi = {1, 2, 3), G2 = (4,5, 5,6,6) 6) etc. Each encoder an MLP. fū, X, (0, 13 R B that outputs mean and diagonal standard deviation of normal distribution.
(Z + (p x )lm(2)
5.2. Acquisition
We make predictions on individual latent sam- ples with predictor network given by an MIP, l IRddd Ac. that predicts a probability distribution over C' classes.
Predictor.
Once the SEFA architecture is trained, we can make acqui- sitions using our novel acquisition objective
Training. The SEFA architecture is trained a supervised manner, minimizing the predictive negative log- likelihood,
R(xo, Σ po c(xo) EE|C]
E pu(s(xo)
r(c, 2, (2)

---

Stochastic Encodings for Active Feature Acquisition
where r: [C] 2 [d]
- R2o is a scoring function
5.3. Understanding the Latent Space
r(c, Σj |gg 1|2
Benefits. The remaining novelty of SEFA is the use of the latent space to calculate the acquisition objective. However, it is possible to calculate the objective entirely in the feature space. Here we would train predictive model P5 ()())))) and a generative model pø(x(xs), and the objectíve is
g = VY = clz).
(3)
18,,
Recall Gi indexes the components of the latent gradient that feature encodes. Below we go through this objective, understanding the purpose of each technical component:
R(xo, i) = Σ P
E F(c, x. pa(x(xo)
Scoring Function as an Importance Measure. For a given latent representation z, and given class c, we can use the gradient VuP& c'z), to determine how important each component of is for making that prediction. Gra- dients are an established feature attribution technique a local explainer (Baehrens et al.) 2010 (Simonyan et al.] 2014 Ribeiro et al. 2016),, our novelty is to use them in the latent space, not the feature space.
cE|C]
In this case, each component of the gradient maps directly to one feature, so the scoring function simplifies to
19c1 F(c, x) = Σ3 10yl
g VxP3 c|x).
We use the latent space instead of the feature space because it allows us to calculate the acquisition objective using rep- resentations of the features, rather than the features them- selves. The benefits are fole:fold: (1) Gradients in the latent space are more meaningful and comparable, since all latent components are continuous, and at a similar scale (if there is appropriate regularization). Whereas features can be categorical or at different inherent scales. (2) The information bottleneck regularization removes feature-level noise. Therefore, in the latent space we calculate the ac- quisition objective using only label-relevant information, instead of taking gradients with noisy feature values. (3) We do not need to train a generative model, and therefore avoid associated complexities such as continuous and categorical variables, multi-modal densities, and highly correlated fea- tures. Our ablations in Section and Appendix C verify the use of the latent space over the feature space for caleulating the acquisition objective.
In order to convert the gradients of z to feature scores, we calculate the length of the gradient in feature's associated latent components, |g 22- Note that this is only possible because we used feature- encoders rather than one fully connected encoder. Finally, we normalize scores to treat each latent sample equally, removing the effect of the total gradient length. See Figure T for an example calculation.
Stochastic Encodings. r(c, z. i) only tells us how impor- tant feature i's latent components are for predicting class for single z. However, demonstrated in Section 4 considering possible values of other unobserved features is necessary for optimality, and lacking in CMI. By using stochastic encoders and taking an expectation of r(c, z, i) over pa xo), multiple possible latent realizations (includ- ing those associated with different unobserved feature val- ues) are taken into account in the current decision. We can analogously think of this inner loop as conducting Monte Carlo tree search (Coulom) (006) through possible future observations and measuring which feature on average is most important for predicting the likelihood of class c. To fully sample the possible realizations, multiple latent sam- ples are taken; we empirically verify this and the use of stochastic encoders in Section and Appendix C
Independent Latent Components. One of the model- ing choices to make the acquisition objective tractable is to encode each feature separately. However, this means mea- suring one feature does not affect the latent distribution of another feature, even if they are highly correlated. Naively, this decoupling should harm the performance. However, the complex interdependencies can be accounted for by the predictor network Pz(y)z). For example, even though measuring feature I does not affect feature 2's latent distri- bution, the predictor network can learn to treat feature 2's latent samples differently based on the value of feature I's latent samples (whose distribution does change after mea- surement). Therefore, the gradients change, and measuring feature 1 does affect the score for feature 2 (and all other features), despite the latent distribution not changing.
Probability Weighting. The inner loop of the objec- tive scores feature i, based on its expected importance for predicting the probability of class c, over many possible unobserved latent realizations. To aggregate across classes, rather than treating each class equally, we take a weighted sum of the scores, using the current predicted probabilities (Y = clxo). This places more focus on the more likely classes, overcoming the issue with CMI maximization being possible by making low probabilities lower (see Section 4 Empirically, we verify this in ablations in Appendix C
The predictive part of the loss in |I is given by the negative log-likelihood log (xp()) Po(y/z). It is standard to take one latent sample to calculate this. However, we take multiple samples for two reasons. (1) The predictor network must see enough
Revisiting the Loss.
Supervised Training. Since our acquisition objective is hand-crafted, we can train with a supervised loss, avoiding training difficulties associated with RI.
5

---

Stochastic Encodings for Active Feature Acquisition
latent samples during training learn to adapt based on different samples during acquisition. (2) If only one sample is taken for given xs, the predictive loss encourages each sample to output the same prediction. Either the encoder po(2)xs) reduces the latent standard deviation, or the pre- dictor ( ()) predicts the same label distribution for the different latent samples. Either way, there a lack of diver- sity in the latent space predictions. If multiple samples are taken, this allows individual samples' predictions to vary more. This will not improve the predictions, but improves the acquisition since more diverse set of possible realiza- tions are used in the inner loop of the objective. We verify this empirically in Section and Appendix
acquire 211 first and then acquire the remaining relevant features. Table I shows how many features each model ac- quires until all features relevant to particular instance are selected (including T11). SEFA achieves this in the fewest acquisitions and is close to optimal on all three datasets. Estimating CMI using generative models (ACFlow, EDDI, and VAE) performs worse than the fixed ordering, show- ing that inaccurate estimation of CMI worsens the issues already associated with its greedy maximization. EDDI, in particular, nearly performs as poorly as random selections and consistently performs poorly across all experiments, since it is only trained to indirectly predict from X5 and thus produces inaccurate estimates of CMI and p(y|xs).
We investigate which features are acquired by the best four models for Syn 3 (Figure 2 SEFA consistently chooses 211 first and then continues to make optimal acquisitions, almost achieving the optimal performance of 5 (Table I In contrast, DIME acquires I7 first, since this has the highest mutual information initially, despite not being the best for long-term acquisitions. Therefore, when ø11 < 0, DIME does not start acquiring features 3-6 until acquisition 3. GDFS performs similarly, since it is also trained to max- imize CMI. Opportunistic RL tends make noisy acqui- sitions, seen by the red trajectories, demonstrating how it suffers from training difficulties. See Appendix [A] for equivalent diagrams and analysis for Syn 1 and Syn 2.
6. Experiments
Here we evaluate SEFA against various deep AFA baselines. We consider a range of synthetic, tabular, image, and medi- cal datasets. For reproducibility, we provide full experimen- tal details in Appendix K including hyperparameter choices and training procedures, and full dataset details in Appendix H The code for our method and experiments is available at https: //tp://gtttuuu.u com/ a now.ooffff/eerr
We consider six different state-of offte-of-thheart the-art base- lines: Opportunistic Learning (Kachuee et al. (019a) and GSMRL &L Oliva 2021) as RL baselines, GDFS (Covert 2024) as greedy CMI maximization methods, and ACFlow (Li & Oliva 2020) and EDDI (Ma et al.] 2019) as generative models for CMI max- imization. We also use three vanilla baselines: random ordering of features with an MLP predictor, an MLP with fixed global ordering of features, and a VAE (Kingma Diederik P and Welling, Max) 2014), which has a sepa- rate predictive and generative model to estimate the CMI. Further details about the baselines are given in Appendix]] !!
Baselines.
et al.
2023 and DIME (Gadgil et al.
Ablations. To provide further insight into why SEFA per- forms well, we conduct ablations on the synthetic datasets (Table D.. We investigate the impact of each of our design choices: removing the latent space regularization (3 = 0); using one latent sample during acquisition; using one la- tent sample during training; using a deterministic encoder (and therefore no expectation over latent samples); calcu- lating the acquisition objective in the feature space rather than the latent space and not normalizing the acquisi- tion scores. Note removing probability weighting has no impact binary classification; we prove this and ablate this component in Appendix C Removing any of the novel components impacts SEFA's performance. Calculating the acquisition objective in latent space is the most important de- sign choice, followed by using multiple acquisition samples and a stochastic encoder. To better understand the perfor- mance differences, in Appendix B we examine acquisition heat maps and carry out sensitivity analyses on 3, number of acquisition samples, number of train samples, and number of latent components per feature.
6.1. Synthetic Datasets
We begin by constructing three synthetic classification tasks (denoted Syn 1-3) based on the synthetic experiments used by Yoon et al. (211)), where we know the optimal instance- wise feature ordering. These are binary classification tasks with eleven features sampled from standard normal. Three logits are calculated from the first ten features, defined as:
6
f 02,
l2
1 4.2,
8- 10 sin
+ [z8| + xg te - 224
6.2. Datasets with Unknown Feature Orderings
The binary label is sampled with r(Y 1) = 11/(1= + e Syn I uses & if x11 < 0 and la otherwise. Syn 2 uses & if if < 0 and & otherwise. Syn 3 uses &g if 1 < 0 and lg otherwise. In all cases 211 determines which features are important to the prediction, so the optimal strategy is to
Here, consider multiple synthetic and real-world datasets where the correct feature ordering is not known a priori.
"This was achieved by using VAE for px(x(x)) and MLP for p3 p)(()ss).
6

---

Stochastic Encodings for Active Feature Acquisition
Table Number of aequisitions acquire the correct features the synthetic datasets, the lower the better. We provide the and one standard emror
Model
Syn
Syn 2
Syn
ACFlow
7.730 0.139
7.527 = 0.254
9.194 + 0.278
DIME
4.079 0.064
4.581
0.217
5.667 + 0.038
EDDI
9.197 0.203
9.214
9.794 0.441.0.188 0.186
Fixed MLP
6.009 0.000
5.996 0.000
7.999
GDFS
4.568 +: 0.219
1.114.
5.587 0.201
GSMRL
5.570 0.127
6.227
8.199 0.066
Opportunistic RL
4.201 + 0.041
1.846
5.850 - 0.072
Random
9.484 0.006
9.499 0.005
9.987 + 0.008
VAR
6.589
0.088
6.667
+ 0.140
7.888
+ 0.065
SEPA (ours)
4.017 0.003
4.099
0.009
0884
0.026
Ablation
Syn
Syn 2
Syn
4.520 0.082
4.578 +: 0.109
5.716 + 0.104
Acquisition Sample
4.683 0.030
4.862 0.035
5.700
0.027
1 Train Sample Deterministic Encoder
4.421 + 0.174
4.713 0.160
5.188 + 0.107
4.593 0.236
4.773 0.228
5.744
Feature Space Calculation
5.111 0.070
5.461 0.128
5.977 + 0.052
WO Normalization
4.036 0.009
4.104 1: 0.008
5.101 1 0.015
SEFA (full)
4.017 + 0.003
4.099 + 0.009
5.084 + 0.026
SEFA (ours) X1j 0.0
Opportunistic RI X11 0.0
DIME X11 <.0.0
GDFS X11 5.00
11
11
10-
10-
10-
10
-0.8
Acquisition Proportions
Feature
-0.6
5
2
SEFA X11 0.0
Opportunistic X11
DIME X11
GDFS
1.0
10
10
0.2
2
Acquisition
Figure 2. Acquisition heat maps and trajectories for Syn 3. Individual trajectories are plotted in red, with the acquisition proportions each step as a heat map. Green boxes show optimal strategy, while the vertical black line denotes the minimum number features required (5).
To evaluate, we start with zero features and calculate the evaluation metric at every step during acquisition. We use
active methods outperform the fixed ordering, except EDDI which suffers from the lack of an inbuilt predictive model.
AUROC for binary classification tasks and accuracy for multi-class. We report the average metric during acquisition in Table 2 and plot the curves for SEFA, DIME, GDFS, Opportunistic RL, and the fixed MLP ordering in Figure which tend be the best performing models.
Real Tabular. Next, we consider three real tabular datasets. Bank Marketing (Moro et al.) 2014). California Housing (Pace & Barry 1997) and MiniBooNE (Roe et al., 2005 Roe 2010). The Bank Marketing dataset is a binary classifi product
cation task, predicting if eustomer subscribes based on marketing data. California Housing consists of features about houses in California districts and the label is the median house price. We converted this into four-war classification by bucketing the labels into four equally sized bins. The MiniBooNE dataset particle physics binary classification task to distinguish between electron- neutrinos and muon muun-neutrinos. In all cases, SEFA has both the highest average evaluation metric and maintains the best evalua-
Cube. We start with the Cube Synthetic Dataset (RückstieB et al. 2013 Shim et al. 2018 Zannone al. 2019). The task is right-way classification with twenty features. The feature vector is normally distributed around the corners of a cube, with the cube occupying three different dimensions for each class. Irrelevant features are normally distributed around the center. SEFA has the highest average accuracy, and consistently maintains the highest acquisition curve. All

---

Stochastic Encodings for Active Feature Acquisition
Table 2. Average evaluation metries during acquisition. Higher values are better; report the mean and standard error.
Model
Cube
Bank Marketing
California Housing
MiniBooNE
ACFlow
0.899 0.001
0.823 + 0.013
0.558 1 0.010
0.881 0.018
DIME
0.901 0..01
0.907 1 0.002
0.661 1 0.002
0.951 0.001
EDDI
0.764 0.005
0.705 = 0.010
0.412 + 0.013
0,842 + 0.009
Fixed MLP
0.883 0.001
0.909 0.001
0.658 + 0.002
0.954 0.000
GDFS
0.900 + 0.000
0.907 + 0.001
0.653 I 0.002
0.949 + 0.000
GSMRL
0.891 0.001
0,879 = 0.006
0.638 + 0.003
0,946 0.001
Opportunistic RI.
0.901 0.000
0.910 0.000
0.657 + 0.001
0.953 0.000
Random
0.699 0.001
0.816 1: 0.003
0.569 1 0.003
0.912 0.001
VAE
0.901 0.001
0.878 1: 0.002
0.633 1 0.005
0.925 0.002
SEFA (ours)
0.904
0.001
0.919 + 0.001
0.676 + 0.005
0.957
0.000
Model
MNIST
Fashion MNIST
METABRIC
TCGA
ACFlow
0.667 = 0.003
0.652 + 0.002
0.542 + 0.006
0.711 + 0.009
DIME
0.731 0.002
0.703 + 0.002
0.670 0.007
0.805 0.002
EDDI
0.572 1. 0.003
0.604 1 0.001
0.563 1 0.011
0.634 1 0.005
Fixed MLP
0.708 0.001
0.690 = 0.001
0.685 + 0.003
0.799 - 0.004
GDFS
0.732 0.001
0.692 = 0.002
0.671 - 0.005
0.797 0.002
GSMRI.
0.701 0.002
0.683 + 0.001
0.665 +: 0.002
0.781 + 0.003
Opportunistic RL
0.740 0.000
0.708 0.000
0.706 + 0.004
0.838 0.002
Random
0.661 0.001
0.648 + 0.001
0.647 + 0.005
0.753 + 0.003
VAE
0.716 0.001
0.685 1
0.690 0.004
0.800 0.003
SEFA (ours)
0.761 0.001
0.721 + 0.000
0.709 + 0.003
0.843 + 0.002
Cube
Bank Marketing
California Housing
MiniBooNE
0.75
0,523
0.70
0..5
Accuroe
AUROC
0.903
0.50
AUROC
0.873
0.50
0.50
3.5
0.603
0.55
0.85
0,02
0.30
0.46
Acquisttion
Arquisitoo
Acrguisttion
Aoguisttion
MNIST
Fashion MNIST
METABRIC
TCGA
0.5
4.7
0.0
Accuracy
Accuracy
Accurary
6.6
1.3
0.2
8.40
4.2
12
15 20
10
14
24
18
14
15
Acquisition No..
Acquisiios No.
Acquiiioo
No.
Acquistion. No.
SEPA (muns)
DIME
Fixwd MLP
GDPS
Opportianattcc
Figure 3. Evaluation metries plots, starting from the first to the final acquisition across all datasets. Zoomed-in curves shown in the bottom right corner of each plot.
tion metric through the acquisition curve, in particular Bank Marketing and California Housing. Opportunistic RL, DIME, and GDFS perform approximately as well as each other across the real data. Interestingly, on MiniBooNE, the fixed ordering is the second best, despite other methods actively acquiring features. Again, the generative models underperform due inaccurate CMI estimation.
all methods by significant margin, both in terms of aver- age acquisition performance and the acquisition curve being consistently the highest throughout the acquisition.
6.3. Cancer Classification
Finally. look at SEFA in the context of medicine. We consider two cancer classification tasks. The first is METABRIC (Curtis et al. 2012 Pereira et al., 2016). where the task is to predict the PAM50 status of breast cancer subjects from gene expression data. The six classes are Lu- minal A, Luminal B, HER2 Enriched, Basal Like, Claudin Low, and Normal Like. The second dataset uses The Cancer Genome Atlas (TCGA) (Weinstein et al. 2013). The goal is to predict the location of a tumor based on DNA methylation data. The average accuracies are given in Table 2 and the
Image Classification. Next we consider MNIST (LeCun
al. 1998) and Fashion MNIST (Xiao et al. 2017), and acquire up to twenty pixels. Here, the fixed ordering is inadequate, and the active methods perform better. Oppor- tunistic RL outperforms DIME and GDFS, demonstrating RL still an effective method for AFA despite training difficulties, whereas the problems associated with CMI max- imization are more fundamental. Again, SEFA outperforms
8

---

Stochastic Encodings for Active Feature Acquisition
Breast
Lung
Prostate
Liver
2NP750
2NF7700
2NF7000
1.0
Acquisition Proportions
TMEM
TMEMI
TMEMESAA
TMEM
STSGALI
STSCALL
STSCALI
SISGALI
SERPINII
SEKPINSI
SERPINSI
SERPIN SI
0.8
POUSES
POURER
POUSEN
POURES
EXPEALL
EYBLALL
EXPLALI
LYCLALI
LTBR
LTBR
LTBR
LTER
-0.6
LOCaRSSR2
LOCHESE2
LOC
KLE'S
KLPS
KLES
KLTS
HOXA
HOXA
HOXAS
HOXAS
0.4
GSDMC
GSOME
GSDMC
CSDMC
CREAZ
CRUAZ
CRUAZ
CRLAZ
EPRA
GPAM
GPAN
GPEI
0.2
DNASEILS
DNASELLS
DINGSELL
DINSSEILS
DEPS
DEPG
DEPU
DEPG
C7eet1
Cheert
C7ee11
Cleet1
2
3
5
6
2
5
Acquistions
Figure 4. TCGA acquisition heat maps and trajectories for four tumor locations. We show the first six acquisitions. Notable acquisitions are highlighted with green boxes and discussed in Section 6.3 Despite the associated noise of medical dataset, the trajectories and heat maps show instance-wiss wise acquisitions.
acquisition curves in Figure 3 On METABRIC, SEFA and Opportunistic RL perform similarly, outperforming all other baselines. On TCGA, SEFA significantly outperforms all baselines, with Opportunistic RL strong second, signifi- cantly outperforming DIME and GDFS, further demonstrat- ing CMI is a flawed AFA objective.
Note this not an exhaustive list, and often it is helpful to acquire a feature that may not be associated with one type of cancer, if it is with another, since that feature helps to differentiate between them. Therefore, some acquisitions with a high proportion may not have an associated citation.
To further validate the acquisitions of SEFA, we visualize the trajectories and heat maps for four cancer types in Figure and provide scientific literature supporting the acquisi- tions made. Firstly, we can see that the acquisitions made are instance-win since the underlying acquisition heat maps change between classes. For example, HOXA9 (feature 9) is selected at acquisition 4, 5, and 6 for lung and prostate cancer, but not breast or liver. The first feature selected is almost always ST6GAL1 (feature 18), which is known to be upregulated in a number of cancers including breast, prostate, pancreatic, and ovarian (Garnham al. 2019). For breast cancer, DEF6 (feature 2) is often acquired second, which has been identified to be correlated with metastatic behavior of breast cancer (Zhang et al.) 2020). For lung and liver cancers, DNASEIL3 (feature 3) is often acquired sec- ond; this gene has been identified as a potential biomarker in liver and lung cancer (as well as breast, kidney, and stom- ach) (Deng al.] 2021). For prostate cancer, the second feature acquired tends be SERPINB1 (feature 17), which is linked to prostate cancer (Lerman et al. 2019). For breast cancer, the third acquisition tends to be C7orf51 1 (feature 1), which is altered in triple negative invasive breast cancer (Brown 2016). For lung and liver cancers, SEFA typically acquires PON3 (feature 15) as the third acquisition. It has been shown that PON3 is largely restricted to solid tumors such those in liver, lung, and colon cancer (Schweikert et al. 2012). For prostate cancer, the most common third feature acquired is GRIA2 (feature 7), which has been found to correlate with the recurrence and prognosis of prostate cancer (Alwadi et al. 2022).
7. Conclusion
This paper considered Active Feature Acquisition, the test time task of actively choosing which features to observe to improve prediction. We introduced SEFA, novel approach for AFA by calculating the acquisition objective using samples from a suitably regularized stochastic latent space, moving away from previous solutions based on RL and CMI maximization. SEFA regularly outperformed pre- vious methods across a range of tasks, and we validated acquired features in the scientific literature. Our ablation study demonstrated that each component of SEFA leads to performance gains.
Limitations. Currently SEFA applies to classification tasks but not to regression tasks. This is because SEFA requires separation of class probabilities during acquisition and this notion is not well defined for continuous labels. We view this as an interesting avenue for future work. One possible solution is to have two prediction heads, one that predicts the continuous label as regression task, and one that predicts discretized version of the label a classification task. The classification head can be used for acquisition, and the regression head for prediction.
Additionally, due requiring multiple latent samples, SEFA has larger memory requirements at inference time than RL baselines, depending on how many samples used. How- ever, CMI maximization methods with generative models also require multiple samples at inference time, so this is not new limitation for AFA models.

---

Stochastic Encodings for Active Feature Acquisition
Acknowledgements
Baehrens, D., Schroeter, T., Harmeling, S., Kawanabe, M., Hansen, K., and Müller, K.-R. How Explain Individ- ual Classification Decisions. The Journal of Machine Learning Research, 11:1803-831, 2010.
The results in the TCGA experiment are based upon data generated by the TCGA Research Network: https:// www. www.cancergoooo gov/tcga We thank the ICML review- for their comments and suggestions. We also thank Lucie Charlotte Magister and lulia Duta for providing feed- back. Alexander Norcliffe is supported by a GSK plc grant. Changhee Lee is supported by the National Research Foun- dation of Korea (NRF) grant funded by the Korea govern- ment (MSIT) (No. R0-2002-003386002 and the Institute of Information & Communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT), Artificial Intelligence Graduate School Program (No. RSS-2019-1119077 Korea University) and the Artifi- cial Intelligence Star Fellowship Support Program to nurture the best talents (No. 2025-02304.28..
Brown, J. P. Inmunohistochemiica and Genomic Analysis of Ductal Carcinoma in Situ of the Human Breast. Phl) thesis, King's College London, 2016.
Chattopadhyay, A., Slocum, S., Haeffele, B. D., Vidal, R., and Geman, D. Interpretable by Design: Learning Pre- dictors by Composing Interpretable Queries. IEEE Trans- actions Pattern Analysis and Machine Intelligence, 45 (6):7430-7443, 2022.
Chattopadhyay, A., Chan, K. H. R., Haeffele, B. D., Geman, D., and Vidal, R. Variational Information Pursuit for Interpretable Predictions. In International Conference on Learning Representations, 2023.
Impact Statement
Chen, Y., Hassani, S. H., Karbasi, A., and Krause, A. Se- quential Information Maximization: When is Greedy Near-optimal In Conference Learning Theory, pp. 338-363. PMLR, 2015a.
Our paper is concerned with Active Feature Acquisition, the test time task of acquiring features to iteratively improve model's 's predictions on given data instance. Applications range from medical diagnosis to polling a population. We believe, on the whole, these applications pose a positive benefit to society. Naturally, since this is general task (any task where features are not all available immediately), malicious applications do exist. For example, iteratively harvesting personal data to send targeted misinformation. However, this work does not focus on those applications, and since this area of machine learning research is still in its relative infancy, we do not envisage this occurring for the foreseeable future. An important consideration is the possibility of this work being used in a high-stakes setting or scenario with ethical considerations, but producing incorrect or suboptimal results. In the medical scenario, doctor might miss an important test to diagnose patient, or conduct a painful/dangerous but unnecessary test. Currently, this work is not in a position to be deployed, so this is not a concern yet. However, if it were to be deployed, this problem can be mitigated by being used as a tool by domain experts to aid them in their decision making, and not replacing them.
Chen, Y., Javdani, S., Karbasi, A., Bagnell, J., Srinivasa, S., and Krause, A. Submodular Surrogates for Value of Information. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 29, 2015b.
Coulom, R. Efficient Selectivity and Backup Operators in Monte-Carlo Tree Search. In International conference on computers and games, pp. 72 83.. Springer, 2006.
Covert, I. C.C., Qiu, W., Lu, M., Kim, N. Y., White, J., and Lee, S.-I. Learning to Maximize Mutual Information for Dynamic Feature Selection. In International Conference on Machine Learning, pp. 1424-6447. PMLR, 2023.
Cui, W., Bai, X., Bai, Z.,, Chen, F,, Xu, J., Bai, W., and Xi, Y. Exploring the expression and clinical significance of the mir-140-3p 140-3p-hoxa9 axis in colorectal cancer. Journal of Cancer Research and Clinical Oncology, 150(2):47, 2024.
Curtis, C., Shah, S. P., Chin, S.-F., Turashvili, G., Rueda, O. M., Dunning, M. J., Speed, D., Lynch, A. G., Samara- jiwa, S., Yuan, Y., et al. The genomic and transcriptomic architecture of 2,000 breast tumours reveals novel sub- groups. Nature, 486(7403):346-352) 2012.
References
Alemi, A. A., Fischer, I., Dillon, J. V., and Murphy, K. Deep Variational Information Bottleneck. In International Con- ference on Learning Representations, 2017. 2017.
Deng, Z.,, Xiao, M., Du, D., Luo, N., Liu, D.,, Liu, T., Lian, D., and Peng. DNASEILL as Prognostic Biomarker Associated with Immune Cell Infiltration in Cancer. On- coTargets and Therapy, pp. 201,2017,2011
Alwadi, D.,, Deoraj, A., Felty, Q., and Roy, D. Discovery of recurrence and prognosis associated genes in prostate can- cer. Cancer Research, hupplement::1442--1 2022.
G., Denoyer, L., Preux, P., and Gallinari, P. Ratum-Wiss Classification: A Sequential Approach to
10

---

Stochastic Encodings for Active Feature Acquisition
Sparsity. In Machine Learning and Knowledge Discovery in Datahases: European Conference, ECML. PKDD 2011, Athens, Greece, September 5-9, 2011. Proceedings, Part I 111. pp. 375-390. Springer. 2011.
Kachuee, M., Darabi, S., Moatamed, B., and Sarrafzadeh, M. Dynamic Feature Acquisition using Denoising Au- toencoders. IEEE Transactions Neural Networks and Learning Systems, 30(8):2252-222) 2262, 2018.
Early, K., Fienberg, S. E., and Mankoff, J. Test time fea- ture ordering with FOCUS: interactive predictions with minimal user burden. International Joint Conference on Pervasive and Ubiquitous Computing, 2016.
Kachuee, M., Goldstein, O., Kärkkäinen, K., and Sar- rafzadehm, M. Opportunistic Learning: Budgeted Cost- Sensitive Learning from Data Streams. In International Conference on Learning Representations, 2019a.
Kachuee, M., Karkkainen, K., Goldstein, O., Zamanzadeh, D., and Sarrafzadeh, M. Cost-Sensitive Diagnosis and Learning Leveraging Public Health Data. arXiv preprint arXiv: arXiv:/902.00 /902.00022100 102, 2019b.
Erion, G., Janizek, J. D., Hudelson, C.. Utarnachitt, R. B., McCoy, A. M., Sayre, M. R., White, N. J., and Lee, S.-I. CoAI: Cost-Aware Artificial Intelligence for Health Care.
Nature Biomedical Engineering, 6:1384, 2022.
Kingma. D. P. and Ba, J. Adam: A Method for Stochastic Optimization. In International Conference on Learning Representations, 2015.
Gadgil, S., Covert, C., and Lee, Estimating Conditional Mutual Information for Dynamic Feature Selection. In International Conference on Learning Representations, 2024.
Kingma, Diederik P and Welling, Max. Auto-Encoding Vari- ational Bayes. In International Conference Learning Representations, 2014.
Garnham, R., Scott, E., Livermore, K. E., and Munkley, J. ST6GAL1: A key player in cancer. Oncology Letters, 18 (2):983-989, 2019.
Kossen, J., Cangea, C., Vértes, E.,, Jaegle, A., Patraucean, V., Ktena, I., Tomasev, N., and Belgrave, D. Active Ac- quisition for Multimodal Temporal Data: A Challenging Decision-Making Task. Transactions on Machine Learn- ing Research, 2023. ISSN 2835-8856.
Ghosh, A. and Lan, A. DiFA: Differentiable Feature Ac- quisition. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 7705-7713, 2023.
He, H., Mineiro, P., and Karampatziakis, N. Active Infor- mation Acquisition. arXiv preprint arXiv:1602.02181, 2016.
Kusner, M., Chen, W., Zhou, Q., Xu, Z. E., Weinberger, K., and Chen, Y. Feature-Cost Sensitive Learning with Submodular Trees of Classifiers. In Proceedings of the AAAI Conference Artificial Intellignce, volume 28, 2014.
He, W., Mao, X., Ma, C.,, Huang, Y., Hernandez-Lobato, M., and Chen, T. BSODA: a Bipartite Scalable Frame- work for Online Disease Diagnosis. In Proceedings of the ACM Web Conference 2022, pp. 2511-2521, 2022.
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient- based learning applied to document recognition. Proceed- ings of the IEEE, 86(1 1):2272222224 1998.
Henderson, P., Islam, R., Bachman, P., Pineau, J., Precup, D., and Meger, D. Deep Reinforcement Learning that Matters. In Proceedings of the AAAI conference on artifi- cial intelligence, volume 32, 2018
Lerman, I., Ma, X., Seger, C., Maolake, A., Garcia- Hernandez, M. d. 1. L., Rangel-Moreno, J., Ackerman, J., Nastiuk, K. L., Susiarjo. M., and Hammes, S. R. Epige- netic Suppression of SERPINB1 Promotes Inflammation- Mediated Prostate Cancer Progression. Molecular Cancer Research, 17(4):84--85) 2019.
loffe, S. and Szegedy, C. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In International Conference on Machine Learning, pp. 448-456. pmir, 2015.
Li, W. Nakano, H., Fan, W. Li, Y., Sil, P., Nakano, K., Zhao, F., Karmaus, P. W., Grimm, S. A., Shi, M., et al. DNASEIL3 enhances antitumor immunity and sup- presses tumor progression in colon cancer. JCI Insight, 8 (17), 2023.
Janisch, J., Pevný, T., and Lisy, V. Classification with Costly Features using Deep Reinforcement Learning. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 33, pp. 3959-39666 2019.
Li, Y. and Oliva, J. Active Feature Acquisition with Genera- tive Surrogate Models. In International Conference on Machine Learning, pp. 6450-6459. PMLR, 2021.
Kachuee, M., Hosseini, A., Moatamed, B., Darabi, S., and Sarrafzadch, M. Context-aware feature query to improve the prediction performance. In 2017 IEEE Global Confer- ence on Signal and Information Processing (GlobalSIP), pp. 838 842. IEFE, 2017.
Li, Y. and Oliva, J. B. Dynamic Feature Acquisition with Arbitrary Conditional Flows. arXiv preprint MXXv:2006.077000 2020.

---

Stochastic Encodings for Active Feature Acquisition
Li, Y., Akbar, S., and Oliva, J. ACFlow: Flow Models for Arbitrary Conditional Likelihoods. In International Con- ference Machine Learning, pp. 5831-5844 PMLR, 2020.
Pereira, B., Chin, S.-F., Rueda, O. M., Vollan, H.-K. M., Provenzano, E., Bardwell, H. A., Pugh, M., Jones, L.,, Russell, R., Sammut, S.-J., et al. The somatic mutation profiles of 2,433 breast cancers refine their genomic and transcriptomic landscapes. Nature communications, 7(1): 1-16, 2016.
Ma, C., Tschiatschek, S., Palla, K., Hernández-Lobato, M., Nowozin, S., and Zhang, C. EDDI: Efficient Dy- namic Discovery of High- Value Information with Partial VAE. International Conference on Machine Learning, 2019.
Rangrej. S. B. and Clark, J. A Probabilistic Hard Attention Model for Sequentially Observed Scenes. arXiv preprint arXiv:2111.07544 2021.
Mnih, V., Heess, N., Graves, A., et al. Recurrent Models of Visual Attention. Advances in Neural Information Processing Systems, 27,2011 2014.
Rényi, A. On Measures of Entropy and Information. In Proceedings of the fourth Berkeley symposium mathe- matical statistics and probability, volume I: contributions to the theory of statisties, volume 4, pp. 547-562. Univer- sity of California Press, 1961.
Moro, S., Cortez, P., and Rita, A data-driven approach to predict the success of bank telemarketing. Decision Support Systems, 62:22-33 31, 2014.
Ribeiro, M. T., Singh, S., and Guestrin, C. "Why Should I Trust You?" Explaining the Predictions of any Classifier. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, pp. 1135-1144) 2016.
Nan, F., Wang, J., and Saligrama, V. Feature-Budgeted Random Forest. In International Conference on Machine Learning, pp. 1983-1991. PMLR, 2015.
Nan, F.. Wang. J., and Saligrama, V. Pruning Random Forests for Prediction on Budget. Advances in Neural Information Processing Systemx, 29, 2016.
Roe, B. MiniBooNE particle identification. Machine Learning Repository, 2010.
UCI
DOE:
Pace, R. K. and Barry, R. Sparse Spatial Autoregressions. Statistics and Probability Letters, 1997.
Roe, B. P., Yang, H.-.,, Zhu, J., Liu, Y., Stancu, I,, and McGregor, G. Boosted decision trees alternative to artificial neural networks for particle identification. Nuclear Instruments and Methods in Physics Research Section A: Accelerators. Spectrometers, Detectors and Asxociated Equipment, 543( 2005.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang, E.,
DeVito, Z., Lin, Z., Desmaison, A., Antiga, L., and Lerer, A. Automatic differentiation in PyTorch. NeurlPS Work- shop on Automatic Differentiation, 2017.
Pedregosa, F., Varoquaux, G., Gramfort, A., Michel, V., Thirion, B., Grisel, O., Blondel, M., Prettenhofer, P., Weiss, R., Dubourg, V., Vanderplas. J., Passos, A., Cour- napeau, D., Brucher, M., Perrot, M., and Duchesnay, E. Scikit-learn: Machine Learning in Python. Journal of Machine Learning Research, 12:2825-283 2011.
RückstieB, T., Osendorfer, C., and van der Smagt, P. Mini- mizing Data Consumption with Sequential Online Feature Selection. International Journal of Machine Learning and Cybernetics, 4:235-243. 243,-224,,255 2013.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal Policy Optimization Algorithms. arXiv preprint arXiv:1707.06347 2017.
Peis, I., Ma, C., and Hernández-LLobato, M. Missing Data Imputation and Acquisition with Deep Hierarchical Mod- els and Hamiltonian Monte Carlo. Advances in Neural Information Processing Systems, 35:35839-35851, 2022.
Schweikert, E., Devarajan, A., Witte, I., Wilgenbus, P., Amort, J., Förstermann, U., Shabazian, A., Grijalva, V., Shih, D., lated in cancer tissues and protects against mitochondrial superoxide-mediater cell death. Cell Death & Differenti- ation, 19(9): 1549-150, 2012.
Peng. Y.-.., Tang, K.-F., Lin, H.-T., and Chang, E. REFUEL: Exploring Sparse Features in Deep Reinforcement Learn- ing for Fast Disease Diagnosis. In Advances in Neural Information Processing Systems, volume 31, 2018.
R., et al. PON3 is upregu-
Penna-Martinez, M., Epp, F., Kahles, H., Ramos-Lopez, E.,, Hinsch, N., Hansmann, M.-L., Selkinski, I., Grünwald, F,, Holzer, K., Bechstein, W. O., et al. FOXEI Association with Differentiated Thyroid Cancer and its Progression. Thyroid, 24(5):845-83) 2014.
Shenoy, U. S., Basavarajappa, D. S., Kabekkodu, S. P., and Radhakrishnan, R. Pan-cancer exploration of oncogenic and clinical impacts revealed that HOXA9 a diagnostic indicator of tumorigenesis. Clinical and Experimental Medicine, 24(1):1344 2024.
12

---

Stochastic Encodings for Active Feature Acquisition
Shim, H., Hwang, S. J., and Yang, E. Joint Active Fea- ture Acquisition and Classification with Variable-Size Set Encoding. Advances in Neural Information Processing Systems, 31, 2018.
Yin, H., Li, Y., Pan, S. J., Zhang, C., and Tschiatschek, S. Reinforcement Learning with Efficient Active Feature Acquisition. arXiv preprint 2rXiv:2011.00822 2020.
Yoon, J., Jordon, J., and van der Schaar, M. INVASE: Instance-wise Variable Selection using Neural Networks. In International Conference Learning Representations, 2019.
Simonyan, K., Vedaldi, A., and Zisserman, A. Deep In- side Convolutional Networks: Visualising Image Clas- sification Models and Saliency Maps. In International Conference on Learning Representations, 2014.
Zannone, S., Hemández-LLobato, M., Zhang, C., and Palla, K. ODIN: Optimal Discovery High- Value INformation Using Model-based Deep Reinforcement Learning, In ICML Real-world Sequential Decision Making Workshop, 2019.
Tishby, N., Pereira, C., and Bialek, W. The information bottleneck method, Allerton Conference on Communi- cation, Control and Computing, volume 37, pp. 368-377, 1999.
Trapeznikov, K. and Saligrama, V. Supervised Sequential Classification Under Budget Constraints. In Artificial Intelligence and Statistics, pp. 581-589 PMIR, 2013.
Zhang, L.. Wang. C., and Ma, M. Lnerna pou3f3 promotes cancer cell proliferation, migration, and invasion in renal cell carcinoma by downregulating Inerna gas5. Kidney and Blood Pressure Research, 46(5):613-61) 2021.
Valancius, M., Lennon, M., and Oliva, J. Acquisition Condi- tioned Oracle for Nongreedy Active Feature Acquisition. In International Conference on Machine Learning, 2024.
Zhang, Q., Zhao, G.-S., Cao, Y., Tang, X.-F., Tan, Q.-L.,
Lin, L., and Guo, Q.-N. Increased DEF6 Expression Correlated with Metastasis and Poor Prognosis in Hu- man Osteosarcoma. Oncology Letters, 1092)1629-1160, 2020.
Van Hasselt, H., Doron, Y., Strub, F., Hessel, M., Sonnerat, N., and Modayil, J. Deep Reinforcement Learning and the Deadly Triad. arXiv preprint arXiv: /8/2.02648, 2018.
Weinstein, N., Collisson, E. A., Mills, G. B., Shaw, K. R., Ozenberger, B. A., Ellrott, K., Shmulevich, I., Sander, C., and Stuart, M. The Cancer Genome Adla Pan-Cancer Analysis Project. Nature genetics, 111101;111--1120, 2013.
Wieczorek, A., Wieser, M., Murezzan, D., and Roth, V. Learning Sparse Latent Representations with the Deep Copula Information Bottleneck. International Conference Learning Representations, 2018.
Xiao, H., Rasul, K., and Vollgraf, R. Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learn- ing Algorithms. arXiv preprint arXiv: 077407747, 2017.
Xu, Z., Weinberger, K., and Chapelle, O. The Greedy Miser: Learning under Test-time Budgets. International Conference on Machine Learning, 2012.
Xu. Z., Kusner, M., Weinberger, K., and Chen, M. Cost- Sensitive Tree of Classifiers. In International Conference Machine Learning, pp. 133-141, PMLR, 2013.
Xu, Z., Kusner, M. J., Weinberger, K. Q., Chen, M., and Chapelle, O. Classifier Cascades and Trees for Mini- mizing Feature Evaluation Cost. Journal of Machine Learning Research, 15(1):2113-2144 2014.
Yamada, Y., Lindenbaum, O., Negahban, S., and Kluger, Y. Feature Selection using Stochastic Gates. In International Conference on Machine Learning, 2020.
13

---

Stochastic Encodings for Active Feature Acquisition
A. Additional Synthetic Heat Maps and Trajectories
To complement the synthetic experiments presented Section 6 we provide the heat maps and trajectories for Syn in Figure 5 and Syn in Figure In agreement with Table SEFA clearly performs best on both Syn and Syn 2. both cases, 211 acquired first, informing the model where it needs look next. All features are acquired by the theoretical minimum with the exception of minority of trajectories. Opportunistie RL and DIME have small but noticeable portion of sub-optimal trajectories Syn I when proportion of required feature acquisitions are made after the theoretical minimum of since initially and ø5 are selected. Additionally, GDFS regularly selects I11 late into the acquisition process. On Syn 2, the three baselines do not place all attention on ø11 initially. In fact, Opportunistic RL and GDFS mostly acquire 27 first since provides the best immediate predictive signal. When #11 2 O, the baselines tend to acquire relevant features in the theoretical minimum, albeit in sub-optimal orders. However, we see that when æ11 < O, this is not the case with many required acquisitions being made after the minimum of 3, since 27 has been selected first.
O. GDFS performs particularly poorly on Syn 1, when æ11 < 0 a high
SEFA (ours) X11 < 0.0
Opportunistic RL X11 0.0
DIME X11 0.0
CDFS Xx1 0.0
1.0
10
10
10
9-
9-
9
B-
0.4 Acquisition Proportions
O.0
2-
Feature
0.6
SEFA (ours) X11 2 0.0
Opportunistic RL X11 2 0.0
DIME X11 2 0.0
GDFS X11 0.0
9-
Acquisition
Figure Acquisition heat maps trajectories on Syn Individual trajectories are plotted in red, with the acquisition proportions each step a heat map behind. We use green boxes to highlight the optimal strategy and a vertical black line to show the minimum number features required or 5).
SEFA (ours) 0.0
Oppartunistic RI. X11 0.0.
DIME X11 < 0.0
GDFS 0.0.
11
11
1.0
V
-.88
Acquisition Proportions
Feature
0.6
SEFA (ours) X11
0.0
Opportunistic RL 0.0
DIME 0.0
GDFS
0.2
2
3
5
3
4
6
2
2
5
Acquisition
Figure 6. Acquisition heat maps and trajectories on Syn 2. Individual trajectories are plotted in red, with the acquisition proportions at each step heat map behind. We green boxes highlight the optimal strategy and vertical black line show the minimum number of features required (3 or 5).
14

---

Stochastic Encodings for Active Feature Acquisition
B. Synthetic Ablations and Sensitivity Analysis
Heat maps and Trajectories. We supplement the synthetic ablations Table by studying the acquisition heat maps and trajectories when we: calculate the acquisition objective in the feature space; use a deterministic encoder and use one acquisition sample of the latent space. We plot these for Syn 1-3 in Figures 7 removing each of our proposed components degrades acquisition performance, confirming Table [ All three reduced versions of SEFA, in all cases, select relevant features after the theoretical minimum since they regularly select 4:1 late into the acquisition. Using the feature space almost never selects first. Acquiring with one latent sample leads to trajectories that approximately sample uniformly among all features relevant to a given synthetic task. Confirming that we need to take many acquisition samples to see feature's 's effect on a diverse range of possible latent realizations.
g and All three figures show that
SEFA (full) 0.0
Peature Space = 0.0
Deterministic Encoder 0.0
Anq Sample <0.0
11
11
1.0
7
6-
-0.0
Acquisition Proportions
Feature
-06
3
5
2
SEFA (full) XI1
Feature Space
Deterministic Encoder X11
Ang Sample X
1.1
M
11
10-
10
10
9-
6-
0.2
5-
31
5
3
Acquisition
Figure 7. Acquisition heat maps and trajectories on Syn 1 ablations. Individual trajectories are plotted in red, with the acquisition proportions each step as heat map behind. We use green boxes highlight the optimal strategy and vertical black line show the minimum number of features required (3 or 5).
SEFA (full) 0.0
Feature Space 0.0
Deterministic Encoder 0.0
Ang Sample X11 <0.00
11
11
11
-0.0
Acquisition Proportions
Feature
0.6
3
5
2
4
SEFA (full) X11
Feature Space X11
Deterministic Encoder X11
Acg Sample X11
11
10
10
10.
10
0.2
2
3
5
2
Acquisition
Figure 8. Acquisition heat maps and trajectories on Syn 2 ablations. Individual trajectories are plotted in red, with the acquisition proportions each step as heat map behind. use green boxes highlight the optimal strategy and vertical black line show the minimum number of features required (3 or 5).
15

---

Stochastic Encodings for Active Feature Acquisition
SEFA (full) 0.0
Feature Space <..0
Deterministic Encoder
0,00
Anq Sample N11 <0.00
11
1.0
0.0
Acquisition Proportions
Feature
-0.6
5
SEFA (full) 0.0
Feature Space 2.0
Deterministic Encoder 0.0
Ang Sample =0.0
10-
0.2
2
5
3
2
2
5
Acquisition
Figure 9. Acquisition heat maps and trajectories on Syn 3 ablations. Individual trajectories are plotted in red, with the acquisition proportions each step heat map behind. We green boxes highlight the optimal strategy and vertical black line show the minimum number of features required (5).
Sensitivity Analysis of B. To further explore the importance of a well-regularized latent space, conduct sensitivity analysis the hyperparameter 3, keeping all other hyperparameters the same. Higher 3 leads to the encoders removing more information about the features. We plot the number of acquisitions required to select all relevant features on the synthetic datasets in Figure IO. For all datasets, as expected, if 3 is too high, the latent space is too heavily regularized. There enough label information the latent space, so decisions made there lead sub-optimal acquisitions. Equally, by not regularizing the latent space enough, there is nothing explicitly enforcing the latent space to remove irrelevant information about the features, also leading sub-optimal acquisitions.
Synthetic
Synthetic 2
Synthetic 3
4.6
5.4
Required No. Acquisitinns
4.5
Required No. Acquistiios
5.2
Required No. Aequisitions
6.50
6.25
5.0
4.4
6.00
4.0
4.3
4.6
4.2
5.50
4.4
4.1
5.25
4.0
4.0
5.00
10-4
10-3
10-1
10-3
10-2
10-1
10a
10
10-6
10-4
10-4
10-4
10-2
10-1
10a
10
10-6
10-4
10-4
10-3
10-2 10-1 100
B
B
Figure 10. The number of acquisitions to select the relevant features for different values on the synthetic tasks. The x-axis logarithmic and includes
Sensitivity Analysis of Number of Acquisition Samples. To further investigate the importance of using multiple acquisition samples to sample the full latent diversity, we run a sensitivity analysis on the synthetic tasks. We plot the number of acquisitions required to select relevant features Figure II As expected, if not enough samples are used, the number acquisitions required is larger. We use 200 acquisition samples in our experiments, which is low enough for fast acquisition and high enough that performance has plateaued.
16

---

Stochastic Encodings for Active Feature Acquisition
Synthetic
Synthetic 2
Synthetic 3
4.9
Required No. Acquisitions
Roquired No. Acguisitions
4.8
Required No. Aoguisitions
5.6
6.5
4.6
4.5
5.4
4.3
4.4
5.3
4.2
133
4.1
4.2
4.0
10a
No. Acquisition Samples
No. Acquisition Samples
No. Acquisition Samples
Figure 11. The number of acquisitions to select the relevant features for different numbers of acquisition samples on the synthetic tasks. The x-axis logarithmic.
Sensitivity Analysis of Number of Train Samples. To further investigate the importance of using multiple training samples to encourage diverse latent space, we run a sensitivity analysis on the synthetic tasks. We plot the number of acquisitions required to select all relevant features Figure [2 For Syn and Syn 2, we see that performance tends to improve with the number of samples as expected. For Syn 3, we see the best performance is achieved with 100 samples, which is the number used in experiments.
Synthetic
Synthetic 2
Synthetic 3
5.6
Required No. Acquisitions
Required No. Acquisitions
4.
Required No. Acquisition
5.5
4.6
5.2
42
10a
10a
10a
10a
10P
10a
10a
No. Training Samples
No. Training Samples
No. Training Samples
Figure 12. The number of acquisitions to select the relevant features for different numbers of training samples on the synthetic tasks. The X-axis is logarithmic.
Sensitivity Analysis of Number of Latent Components. To investigate the sensitivity of SEFA's 's performance to the number of latent components per feature, we run a sensitivity analysis on the synthetic tasks. We plot the number of acquisitions required to select all relevant features in Figure 13] We see the performance remains relatively constant, showing SEFA is fairly robust to this hyperparameter. We found between 5-10 is good value such that there enough capacity for rich representations of each feature, but not enough overfit.
17

---

Stochastic Encodings for Active Feature Acquisition
Synthetic
Synthetic 2
Synthetic 3
4.30
4.30
5.8
Required No. Acquistions
4.25
Required Acquisttions
4.25
Required No. Acquistiion
5.7
4.20
0.20
6.6
4.10
4.10
5.4
4.05
4.05
5.3
4.00
0.00
8.2
3.95
3.95
5.1
3.90
3.90
5.0
No. Latent Components
No. Latent Components
No. Latent Components
Figure 13. The number of acquisitions to select the relevant features for different numbers of latent components per feature on the synthetic tasks.
C. Real Data Ablations
To further demonstrate that each novel model component leads to performance gains, we also carry out ablations on the real datasets. Additionally, here we investigate the final novelty introduced, probability weighting, where weight the scores during acquisition by the predicted probabilities po c(xo)- We investigate the of this technique by removing the weight and taking mean, treating each class equally. This was not meaningful the synthetic ablations because this does not affect binary classification tasks. To see this, recall how features are scored in (2):
R(xo.
Σ PY c(xo)
E po(x)xo)
r(c, z,
cE|C]
Writing this in the binary case (class labels are either or 1) gives
R(xo,1) i) po 0(xo)
E
20 PY Y 1(x0)
E po(x(xo)
r(1, z, i).
pa p=(x/xx)
(Y = 1|z) = 1- (Y 0jz), VePa(Y = 1/z) = VVPo = 0jz). therefore r(0, z, i) = r(1, z, i), since the gradients point in opposite directions, and taking Fuclidean norms and normalizing is agnostic to the negative sign. Therefore
R(xo, po. = =[xx)) pa(2/xo)
r(0,z, (Y 1[xo)
r(0, z, i),
px(x
R(xo, (po,o(Y 0|xo) + pe = =1xo))
E
r(0, 2, i),
pe(2|xo)
R(xo.i)
E r(0,z,i)
E p=(2(x))
r(1, z,
p(z(xo)
The weighting is removed in the binary case, proving probability weighting only affects the multi-class setting. We the ablations the real datasets well Cube, where the precise ordering is not known priori. We provide average evaluation metrics during acquisition in Table 3 Calculating the acquisition objective in feature space consistently leads to the worst result (except for Cube and METABRIC where it is second worst), demonstrating that calculating the objective in latent space is a crucial component to the performance of SEFA, even if we give up the explicit ability to model conditional distributions between
features. Following this, using deterministic encoder or one acquisition sample the next worst result, showing the second novelty of using stochastic encoders to consider a diverse set of possible unobserved latent realizations is key factor contributing the performance of SEFA. Removing probability weighting, as expected, has no effect on binary classification tasks (Bank Marketing and MiniBooNE), it does however have significant effect on MNIST, Fashion MNIST, METABRIC and TCGA where there are multiple classes. The performance is also marginally worse without probability weighting on California Housing, and does not change on Cube. In two cases, removing regularization improves performance, but in six out of eight cases and for the synthetic datasets, non-zero 3 improves the acquisition performance of SEFA. Using one train sample consistently leads to slightly worse performance, matching the results on the synthetic datasets. Removing the normalization in the calculation has minimal effect: improves performance on Cube,
18

---

Stochastic Encodings for Active Feature Acquisition
MNIST, and TCGA (the change is marginal in these cases) and worsens performance Fashion MNIST, showing the exact form for is not important in the acquisition objective as using stochastic latent encodings probability weighting.
An additional ablation that does not apply to the synthetic datasets is the use of copula transform (WO Copula). It transforms the continuous features using the empirical cumulative distribution function, rather than raw feature values. It does not affect the synthetic datasets because those features are sampled from a standard normal. Note the transform is specific to SEFA due to its information bottleneck regularization (see Appendix [.2p. Similar to removing normalization, on some datasets the performance improves slightly by removing the Copula (Cube, MNIST, METABRIC), and for others the performance worse (Bank Marketing, MiniBooNE, Fashion MNIST). The difference is marginal compared using the feature space, deterministic encoders, or removing probability weighting, the main novelties of our work.
We plot acquisition curves for key subset of ablations in Figure 14 We do not show the feature space calculation trajectory since it is consistently the worst (see Table 3 Deterministic encoders (green) regularly perform noticeably worse than using stochastic encoders. Removing probability weighting (orange) also leads to significantly worse performance on four of the six multi-class datasets. Using multiple latent acquisition and training samples also required for the best performance.
Table Average acquisition metries on ablations. We give the mean and standard error. Ablations that outperform SEFA are underlined.
Ablation
Cube
Bank Marketing
Califomia Housing
MiniBooNF
8
0.906 1 0.000
0.917
0.001
0.680 1 0.005
0.957 1 0.000
Acquisition Sumple
0.889
0.001
0.912 0.001
0.667 1: 0.004
0.952 + 0.000
Train Sample
0.901 = 0.002
0.912 = 0.003
0.676 + 0.006
0.953 + 0.000
Deterministic Encoder
0.900
0.000
0.893 0.006
0.667 1 0.005
0.944 1 0.002
Feature Space Calculation
0.897 0.002
0.880 = 0.003
0.600 + 0.004
0.914 = 0.004
WO Copula
0.905 = 0.000
0.917 0.002
0.675 + 0.002
0.952 + 0.001
WO Normalization
0.905 0.000
0.919 + 0.001
0.676 + 0.005
0.957 0.000
WO Prob Weighting
0.904 = 0.001
0.919 - 0.001
0.674 + 0.005
0.937 + 0.000
SEFA (full)
0.904 0.001
0.919 0.001
0.676 + 0.005
0.937 + 0.000
Ablation
MNIST
Fashion MNIST
METABRIC
TCGA
0.759 0.000
0.719 0.000
0.706 0.003
0.842 0.002
Acquisition Sumple
0.728 = 0.001
0.699 1 0.000
0.694 + 0.003
0.827 + 0.002
Train Sample
0.741 0.002
0.707 = 0.001
0.705 + 0.003
0.834 + 0.002
Deterministic Encoder
0.741 0.002
0.704 0.001
0.702 + 0.004
0.838 + 0.002
Feature Space Calculation
0.697 1 0.001
0.634 1: 0.001
0.695 1 0.001
0.795 1 0.001
WO Copula
0.762 0.001
0.717 0.001
0.710 0.002
0.843 0.002
WO Normalization
0.762 1 0.00T
0.718 0.001
0.709 1 0.002
0.844 1 0.002
WO Prob Weighting
0.751 = 0.00T
0.694 0.001
0.698 + 0.003
0.832 - 0.001
SEFA (full)
0.761 + 0.001
0.721 + 0.000
0.709 + 0.003
0.843 + 0.002
Cube
Bank Marketing
California Housing
MiniBooNE
0.7
0.32
Ancuracy
ATTROC
Aucurssv
A3300
0.82
0.92
1.0
0.32
5:50
0.5
6:70
0.32
0.75
0.00
0.30
Acquistiion No.
Acquisition
Acquisition No.
Acquiiition
MNIST
Fashion MNIST
METABRIC
TCGA
0.2
Accurti
Accurccy
Accuracy
Accumccy
6.33
0.3
3.50
0.2
22
25
20
10
15
3.3
10. 13
24
10
10
8
20
12
24
16
10
20
10
12
0
10
13
14
15
10
20
Acquistion No.
Acquisition No.
Acquisition No.
Acquisition No.
SERA (tull)
WD Prob Weightting
Deterninistic Eucodor
1 Acquisttio Sample
Train Sample
Figure 14, Evaluation metries starting from the first to the final acquisition for the ablations. To distinguish curves, provide zoomed-in in versions of the plots in the bottom right corner of each one.
19

---

Stochastic Encodings for Active Feature Acquisition
D. Additional Experiment: Robustness to Noise
To investigate how robust SEFA is to noisy features, we repeat the Syn 1-3 experiments; however, after calculating the label, we add normally distributed noise to the features, with increasing standard deviations. The results are given in Table All models get worse with more noise. SEFA remains the best model at all levels of noise. This is likely due to the latent space regularization removing noise between the feature space and the latent space.
Table 4. Number of acquisitions to acquire the correct features under increasing levels of noise, the lower the better. We provide the mean one standard error.
Syn
σ a= 0.0
a 0.1
0.2
= 0.4
ACFlow
7.730 0.139
7.754 = 0.199
7.755 + 0.275
8.255 = 0.131
DIME
4.079 +: 0.064
4.688 + 0.211
4.703 + 0.276
5.368 - 0.267
EDDI
9.197 0.203
8.988
9.216 0.136
9.382 0.245
Fixed MI.P
6.009 + 0.000
7.321 + 0.378
6.009 = 0.000
GDFS
4.568 0.219
4.566 0.1887
4.543
5.583 0.283
GSMRL
5.570 1: 0.127
1.
5.778 0.120
6.858 = 0.269
Opportunistic RL
4.201 - 0.041
4.347 0.091
4.720 0.142
5.488 0.084
Random
9.484 0.006
9.495 0.010
9.495
0.010
9.495 0.010
VAE
6.589
= 0.088
6.866 = 0.041
7.005
+ 0.086
7.095 0.078
SETA (ours)
4.017 1 0.003
4.100 1 0.004
4.207
0.008
4.406
0.004
Syn
0.0
0.1
0.2
0.4
ACFlow
7.527 + 0.254
7.366 + 0.325
7.927 0.363
7,974 0.072
DIME
4.581 1.81+000 0.217
4.830 = 0.037
5.093 + 0.218
5.015 0.124
EDDI
.214+:0.41
9.550 0.342
9.079
9.571 0.225
Fixed MLP
5.996 0.000
0.373
6.095 0.
6.494 0.315
GDFS
4.484 0.159
4.981 0.211
5.655 1 0.386
6.804 0.384
GSMRL
6.227 0.185
0.146
6.525 0.158
7.286 0.088
Opportunistic RL
4.846 1 0.021
5.042 1 0.007
5.108 1 0.064
0.064
5.246 1 0.041
Random
9.499 0.005
9.504 0.007
9.504 0.007
9.50M 0.007
VAE
6.667 = 0.140
6.874 - 0.161
7.057 + 0.134
7.204 0.094
SEFA (ours)
4.099 + 0.009
4.247 1 0.042
4.388
0.045
4.778
0.170
Syn
0.0
0.1
0.2
0.4
ACFlow
9.194 0,278
9.105 0.218
9.184
0.248
9.284
0,231
DIME
5.667 = 0,038
5.771 = 0.222
6.180 + 0.161
6.213 0.180
EDDI
9.794 = 0.186
9.969 = 0.156
9.906
0.212
9.710
Fixed MI.P
7.999 = 0.000
7.999 + 0.000
7.999 0.000
19999=0.00
GDFS
5.587 = 0.201
6,839 = 0.338
7.179 0.134
7,814
GSMRL
8.199 0.067
8.326 0.233
8.347
8.627
Opportunistic RL
5.850 1 0.072
5.859
1 0.063
6.135 1 0.114
0.114
6.672 1 0.105
Random
9.987 + 0.008
9.991
1 0.006
9.991 0.006
9.991 0.006
VAE
7,888
= 0.065
7.968
= 0.068
0.
8.469 0.064
SEFA (ours)
5.084
0.026
5.448 0.042
5.689 0.061
6.207 0.058
E. Additional TCGA Trajectories
To further augment the TCGA analysis in Section 6 we provide the heat maps and trajectories across all tumor locations in Figure [5] We see the selections are instance-wise orderings since different classes have different heat maps and trajectories (as well as different trajectories emerging for the same tumor location). Due to the nature of the task and data, there is still associated noise. Further to the justification in the main paper, we see that in many cases DNASEIL (feature 3) is selected after ST6GAL1 (feature 18). This is because it has been linked to: bladder cancer, breast cancer, gastric carcinoma, liver cancer, lung adenocarcinoma, lung squamous cell carcinoma, ovarian cancer, cervical squamous cell carcinoma, head-neck squamous cell carcinoma, pancreatic adenocarcinoma, and kidney renal clear cell carcinoma (Deng et al. 2021). Additionally, it has been linked to colon cancer progression (Li et al., 2023), and was found to be downregulated in prostate adenocarcinoma and uterine corpus endometrial carcinoma (Degg et al. 202I). DNASEIL3 is regularly selected for all of these tumor locations, even first occasionally, since it is strong predictor its own. We likely see the selection appearing for brain and bone marrow acquisitions with supporting literature are: kidney cancer acquiring POU3F3 (feature 16) (Zhang et al., 2021), KAAGI (feature 10) (the name is Kidney Associated Antigen 1), and HOXA9 (feature 9) (Shenoy et al. 2024); 4); colon cancer acquiring HOXA9 CCi et al. 2024) and thyroid cancer acquiring FOXE1 (feature 5) this is not an exhaustive list.
way rule out these other likely locations after selecting ST6GALI. Additional notable
et al. 2014). Note
20

---

Stochastic Encodings for Active Feature Acquisition
Breast
Lung
Kidney
Brain
Ovary
2N8792
ZNF750
ZAET50 ZNFIRI
231750
2NE181
2NFFSR
ZNFORD
THEMINEA
TNEMI
TOEM
STEGAL:
SIECALL
1.0
POUNT
PORERR
POLDIO
POWIE
POUSTO
PONT
FONS
N
PONO
PONT
PONO
LXPLALI
LYPEAL1
LSPLALL
LYPLAL
ESPLALL
LECK
LIRK
LIBK
LOC283312
KLL'S
KLPPS
KEFS
KLPS
KANG:
KAACI
KAGGE
KAAGE
RAAGE
HOCKS
KOLAS
HONAY
POMAS
BOXAR
GPAE2
GURE1
GPARL
CPRIL
GPREL
POXEI
FOXE
FOXEI
POOR
POITL
EFS
RFS
DNASELLS
DNASELLS
DNASEILD
DNASELLS
Clabb
ChexE1
Clarbt
ChedS
Chuebt
0.8
Endometrium
Head and Neck
Central Nervous System
Thyroid
Prostate
ZNFF63
2NFIS3
ZNFIRO
ZITEGR
IMENIOSA
"NEESE OUA
THEMIPSA
TNEMLE2A
STEGAL:
STRGAL
STSCALL
STOGALL
SISCALL
SERFINDI
SURPIND:
SERFINII
PONS
LYPLALE
LYPLALI
LYPLALL
PXPLAL
ESPLALL
LIBR
LISR
LESR
LIRK
LOCHETES
LOCAES
LOCIBI2EG
LOC283002
LOCHESORG
KAAG1
KANGI
KANGI
HOXAS
KOLAS
HOXAR
KORAS
HOXAS
OSONE
CSENI
GEDMI
CSOMC
GSODIC
0.6
Acquisition Proportions
POXEI
FOXE
POXEI
POXTT
DORE
EPS
EPS
EPS
EPS
BNASE:L3
DNASEILS
DNASELLL
DNASEILL
DNASELLS
Features
DEPS
Colon
Stomach
Bladder
Liver
2NF746
ZNPERS
STEGAGD
STARAAL
STEGALL
STIGIAL
SEGGALI
SURRINIS
STRRPINII
SE SERFINDI
SURPINDE
SHRPPINT
POUSES
POUSES
POUSPS
POURPS
POUSPS
PONS
0.4
LIBH
LIBR
LIDH
CTDR
LOCHID192
3PC 392
LOCIESTRO
LOC2SS302
LOCHESIRG
KLPS
KEPS
KEPS
KLLS
KILSS
OSONI
CSENC
CEOMC
CEOMC
GEDMA
GRIAZ
GRIAS
GRIAS
GRIAD
GREAS
GPISS
GPRRT
OPRI
POSE:
DNASBIL3
DNASH1L3 DEPG
DNASESLI
DNASEILL
DIAAEIL
DAFE
DETG
DOSS
DETS
Cheet::
CTeefSl
Bone Marrow
Pancress
ZNF720 ZNP3RS
2NP183
THEMIOGA
"MMESI
DOUSP3
DOUSES
PON3
FONS
LOC393382
KLRS
KEPS
KANG:
RAAGI
HOSAD
MOLAD
GRISZ
GRIAS
OPSE:
CORS1-
POREI
POXES
EPS
Chalb.
ChecG
Acquistions
Figure Acquisition trajectories TCGA across all classes. Individual trajectories are given red, with the heat map of acquisition proportions each step shown behind. Notable acquisitions with high proportions are highlighted with green boxes.
21

---

Stochastic Encodings for Active Feature Acquisition
F. Indicator Example
Here we elaborate on our indicator example, a simple case where CMI fails. First, we demonstrate that CMI fails and then we show that, by considering possible unobserved feature values in the calculation, we can recover the optimal policy.
Recall the example, have features X {0. 1)dh x [d], i.e the first d dimensions are binary and the final feature is an indicator. The label is given by the value of the feature indexed by the indicator Z In the absence of any of the first d features, the indicator and label are independent p(y, Id+1) ply)p(za+1). Substituting this into the definition of mutual information gives
X8+11 = Dxi. (Y,, X4+1 IllP(YlPP))))) ( 1)) = DRI.(p( (P))))(()((xxs
= 0.
Now consider the mutual information for the other features. Due to the symmetry of the problem. the mutual information for one of these features the same for all others. The mutual information can be more usefully written as
T(Y; x.) = H()) - (( (ra)p((;;)d
The entropy of the label is log 2 since there is equal chance of being 0 or Again, using the symmetry of the system, the entropy of Y if X, is the same as if X, 1, only calculate for one case. When X, = 0, the probability of Y = 0 is x 1 + - 1. since in 1/ cases it takes the exact value of X, based on the value of the indicator, and in ليك cases Y given by different unknown feature value. This gives p(Y 0|X; 0)
d+1 The expression for binary entropy, 2d -plo((p) (1 p) log(1 - p) is maximized by p = 0.5, giving log 2. Since p(Y = 0,X,, 0) > 0.5. the entropy is lower than log 2 in this case. Exploiting the symmetry of the system we conclude that f H(Y a )p(,;)ddr < log 2, and therefore I(); > 0.
Therefore, the indicator is never chosen first, which is a sub-optimal strategy. It can be shown, but is not necessary, that the indicator will be chosen second. A sketch of the reasoning is that now that the value of one feature is known, the indicator and the label are correlated. Therefore, there is non-zero CMI, which turns out to be larger than for the other features, and once the indicator is chosen, the correct feature is the only feature afterward with non-zero CMI. So this strategy will acquire the correct features in 3 selections 1=ه of the time (random feature, indicator, correct feature) and in 2 selections Y of the time (correct feature, indicator). Thus, the expected number of acquisitions for this strategy is
1
d
2 + =
d
So as d gets large, the expected number of required acquisitions approaches 3.
Now consider the adjusted solution of using an information- objective that considers the values of other features. Recall Proposition 4.2. we propose Epixe (x)) [X,;; Y ||xo, (x) recovers the optimal strategy, where Xy is the vector of all other unobserved features. We prove that this will lead to an optimal strategy below.
Initially there are no features, so the acquisition objective is f [X M|xx (xx))((x))d )dxg. Writing this in terms of entropies gives
/ (H(Y - / | (Y1)), (1, xx))(x)))))))) )p(x;) Jdz, p(xu )dxy.
The entropy when all features are known is zero, so for any i the objective simplifies to
f HY (((())))((x))))
If we consider the first d features, we can again apply symmetry to calculate this quantity for one feature feature 1 and apply to all of them. In d-1 cases the entropy is zero, since we will have all of the information required. However if d Id+1 1, then H HY xxy) = log 2, since we don't know the value of feature and therefore Y has equal likelihood of being 0 or 1, this happens in 1 cases this quantity is log2 for the first d features.
For the indicator, p(Y 0fx(r) is the proportion of the first d features that are 0 for given sample xy. All features are independent with probability 0.5 of being 0, so the expression is given by binomial distribution with d trials
d
h - (1 log
h )
log (1 5)
k
2d
d
O
22

---

Stochastic Encodings for Active Feature Acquisition
It is not immediately clear that this is larger than the quantity - for the other features. The first thing we can do is calculate this quantity when d = 3, which gives 0.477, and this is larger than log2 0.231. And the next thing is to notice that this quantity is increasing with d, since as d gets larger there will be more probability mass at k: = 1 As d -> oo the binomial distribution becomes Gaussian with mean and variance 4. so z will approximately be distributed normally with mean 2 and standard deviation 2/2 Therefore this quantity asymptotes towards log 2.
3
Therefore for d 2 3, this objective will choose the indicator first, and not the other features (for 2 all features are scored the same, and for d = 1 the indicator is not the optimal choice). After choosing the indicator, the second selection is trivial. The relevant feature has non-zero CMI, all other features are independent of the label conditioned on the indicator so they have zero CMI. Therefore the correct feature is chosen. This strategy's 's expected number of acquisitions is 2, which is less than 3 - 3.
This example illustrates that by considering the possible realizations in the calculation, and not marginalizing them out, we can make long-term acquisitions. Note we do not use this specific quantity in our paper since involves an additional expectation over unobserved values as well as the expectation inside the CMI. which intractable. This does not even account for the difficulty in estimating the conditional distributions in feature space.
G. Entropy Example
In Section we claimed that CMI maximization can lead to acquisitions that focus on making low probabilities lower, rather than distinguishing between possible answers. Here we provide a concrete example of this occurring.
Consider a feature vector with two features, X E (1,2,3)2, and a label with three possible classes, Y E (1, To generate the label from the features, we start with vector with three zeros which value in v to reduce by 8. And the value of 22 tells us which value in to increase by 6.5. For example if x 2]. then v = 8.0, 6.5, 0.0). p(y]x x) is then given by the softmax of v. With these rules, we can calculate all possible feature values and probabilities in Table
[0.0, 0.0, 0.0). The value of x; tells us
Table 5. All feature values and corresponding y probabilities given by the problem in Appendix C
21
3t2
1/2
>(Y
1/x)
p(Y 2/x)
r(Y 3/x)
1
-1.5
0.0
0.0
0.10037
0,44982
0,44982
1
-8.0
6.5
0.0
0.00000
0.99850
0.00150
1
3
-8.0
0.0
6.5
0.00000
0,00150
0,99850
2
6.5
8.0
0.0
0.99850
0,00000
0,00150
2
0.0
-1.5
0.0
0,44982
0.10037
0,44982
2
3
0.0
-8.0
6.5
0.00150
0.00000
0.99850
3
6.5
0.0
-8.0
0.99850
0.00150
0.00000
3
0.0
6.5
8.0
0.00150
0. 0.99850
0,00000
3
0.0
0.0
1..5
0.444982
0.44982
0.10037
By taking the average, or by seeing that all components of v are treated equally across all feature values, the marginal of the label p(y) = [0.33, 0.33, 0.33), and the entropy is (Y) = log 3. Based on this, we can then calculate the marginal distributions and the entropies of the marginals given in Table
Table 6. Marginal distributions and entropies calculated using Table 5
r(Y = 1 xx)
r(Y 2/x)
p(Y = 3)x)
[(Y
H(Y) - H(
Missing Missing Missing
0.03346
0.48327
0.48327
0.81652
0.28210
0.48327
0.03346
0.48327
0.81652
0.28210
0.48327
0.48327
0.03346
0.81652
0.28210
Missing
0.69912
0.15044
0.15044
0.82016
0.27845
Missing
2
0.15044
0.69912
0.15044
0.82016
0.27845
Missing
3
0.15044
0.15044
0.69912
0.82016
0.27845
Feature is consistently more useful for identifying the most likely class, since it can always identify a class with likelihood 0.699, and feature is better at identifying which class is least likely. However, feature 1 consistently has a lower entropy, so it would be selected before feature 2 by CMI.
23

---

Stochastic Encodings for Active Feature Acquisition
H. Dataset Details
Here we provide all the details about each dataset, including sizes, number of features, and how to access the real datasets.
Synthetic. The synthetic experiments are based on (Yoon et al. 2019), where we know the features that are predictive. and we know that there is heterogeneous order. The datasets are binary datasets where the feature vector has eleven independent features drawn from standard normal. There are three possible logits:
6
i = -r1
=
1.2x 4.2,
Eg = 10 (i.(0)) + |ay| + 2y + - #10 -2.4
1m3
Then, for a given logit value, the label is sampled from a Bernoulli distribution with probability p(Y = 1) = 1/(1111111111 + =e)) We construct three datasets:
Synthetic I: If æ11 < 0 we use &,, otherwise l
Synthetic 2: æ11 < 0 we use &,, otherwise l3
Synthetic 3: æ11 < 0 we use lg. otherwise Eg
The logits have been adapted from the originals in (Yoon et al. 2019) to produce probabilities closer to 0 or 1. This is so all the models have stronger purely predictive performance. The train set is size 60,000, and the validation and test sets are both size 10,000. AUROC is used as the evaluation metric.
Cube. The Cube dataset is synthetic dataset that is regularly used to evaluate Active Feature Acquisition methods (Rückstieß et al. 20133 (Shim et al. 2018 Zannone et al.) 2019). We specifically use the version where all features are normally distributed (Zannone et al. 2019). There are twenty continuous features, where different features are relevant for different classes. All features are drawn from normal distribution with mean 0.5 and standard deviation 0.3, except for the following cases:
Class I: Features 1, 2, 3 have mean [0,0, 0] and diagonal standard deviation 0011, 0.1).
0,1,
Class 2: Features 2, 3, 4 have mean [1, 0, OJand diagonal standard deviation [0.1, 0.1, 0.1].
Class 3: Features 3, 4, 5 have mean [0, 1, 0100aaan and diagonal standard deviation 0.11, 0.1).
0.1,
Class 4: Features 4, 5, 6 have mean [1, 1, OJand diagonal standard deviation [0.1, 0.1, 0.1].
Class 5: Features 5, 6, 7 have mean [0, 0, 1]and diagonal standard deviation [0.1, 0.1, 0.1).
Class 6: Features 6, 7, 8 have mean [1, 0, and diagonal standard deviation [0.1, 0.1, 0.1).
Class 7: Features 7, 8, 9 have mean [0, 1] and diagonal standard deviation (0.1,
Class 8: Features 8, 9, 10 have mean [1, 1, 1] and diagonal standard deviation [0.1, 0.1,
We use a train set with size 60,000, and the validation and test sets are both size 10,000. Accuracy is the evaluation metric. Bank Marketing. The Bank Marketing dataset (Moro et al. 2014) can be found at: htttss://aacctt Setu/dataset/222 11/22//ankttt taannkkarkktinng.. The data is taken from marketing campaign conducted by Portuguese bank. The task is binary classification, where the label indicates whether a client subscribed to term deposit at the bank. The features are both the client's 's information and information about the calls. There are 15 features in total (after combining the month and day of the call into one feature). 7 are continuous and are categorical. A full list of features can be found at the dataset origin. We use an 80:10: split, giving train, validation, and test sizes of 16,168, 4,521, and 4,522. The evaluation metric AUROC.
ics uci
The California Housing dataset is obtained through Scikit-Learn PPerreooaa datasets. fetch call 1 fornia sastaa_husiinggg html The labels are median house prices in California districts expressed in 100,000 dollars. There are 8 continuous features that can be found at the above URL. To convert this to a classification task, we discretize the labels into 4 equally sized bins. We use an 80:10:10 split, giving train, validation, and test sizes of 16,512, 2,064, and 2,064. The evaluation metric is accuracy.
California Housing.
et al. 2011)
learn. www.//gggaaaee.o.oooo wettable/mee modu lemodules//
24

---

Stochastic Encodings for Active Feature Acquisition
MiniBooNE. MiniBooNE is experiment Fermilab designed to detect neutrino oscillations, namely muon-neutrinos
into electron-neutrinos (Roe al. 2005 Roe 2010). The data was obtained from https: The task is binary classification, distinguishing eectron-neutrrng events from background events. The dataset does not have balanced classes. In early experiments, we saw that this meant all models were highly predictive with very few features, and results were indistinguishable. To make the task harder for the models, we enforced balance by reducing the number of background events at random to match the number of signal events. We also used a subset of the features selected using STG (Yamada et al., 2020), as preprocessing step. The selected features were [2. 6, 14, 15. 17, 20, 21, 22, 23, 25, 26, 29, 34. 39, 40, 41, 42, 43, 44]. The train set size 56,499, and the validation and test sets are both size 10,000. The evaluation metric is AUROC.
i ics vci
MNIST and Fashion MNIST. MNIST and Fashion MNIST are image classification datasets with 10 classes, consisting of images of handwritten digits and items of clothing, respectively. We preprocess by selecting a subset of 20 pixels each for computational reasons an acquisition trajectory with 784 features, where the majority are redundant, will be very slow, especially for methods such as EDDI and VAE, where the whole acquisition is (dd). To do this we use STG (Ymaada et al. 2020) deep learning method for feature selection. After flattening the images to vectors, the features found by STG were:
MNIST: [153, 154, 210, 211, 243, 269, 271, 295, 327, 348, 350, 375, 405, 409, 427, 430, 461, 514, 543, 655]
Fashion MNIST: [10, 38, 121, 146, 202, 246, 248, 341, 343, 362, 406, 434, 454, 490, 546, 574, 580, 602, 742, 770]
For both datasets, we split the provided train set into train set with size 50,000 and validation set with size 10,000. We use the provided test sets, each with size 10,000. The evaluation metric is accuracy.
METABRIC.
The Molecular Taxonomy of Breast Cancer International Consortium (METABRIC) database consists of clinical and genetic data for 1,980 breast cancer subjects (Curtis al., 2012 Pereira
jet al. 2016). The data was accessed https::/ kwggle.ggg.... com/
breast-
We construct a classification task, predict-
ing the PAM50 status using gene expressions as features. There are six classes:
1. Luminal
Luminal
3. HER2 Enriched
4. Claudin Low
5. Basal Low
6. Normal
As with the other high-dimensional datasets, used STG to select subset of twelve continuous gene expressions given by:
1. CCNB1
2. CDK1
3. E2F2
4. E2F7
5. STATSB
6. Notch
RBPJ
8. Bcl-2
9. eGFR
10. ERBB2
11. ERBB3
12. ABCB1
We use 80:10:10 split, resulting in train, validation, and test sizes of 1,518, 189, and 191. The evaluation metric is
TCGA. The Cancer Genome Atlas (TCGA) consists of genetic data for over 11,000 cancer patients (Weinstein et al., accuracy.
2013). The data was accessed at Phttps://wwet cancer. wttpe//oooocc//geeeeeee rewp/rccrreeerccccccc/ca.... walelaarcccgggeeeeeeee gww.me- genonee-seqenccngg//csee construct the classification task of predicting the location of the tumor based on DNA methylation data. We use 17
locations as the classes:
1. Breast
2. Lung
3. Kidney
4. Brain
5. Ovary
6. Endometrium
7. Head and Neck
8. Central Nervous System
9. Thyroid
10. Prostate
11. Colon
12. Stomach
13. Bladder
14. Liver
15. Cervix
16. Bone Marrow
17. Pancreas
As the first step of dimensionality reduction, we removed features with more than 15% missingness. Following this, we
used STG to select 21 features:
1. C7orf51
2. DEF6
DNASEIL3
EFS
25

---

Stochastic Encodings for Active Feature Acquisition
5. FOXE1
6. GPR81
7. GRIA2
8. GSDMC
9. HOXA9
10. KAAGI
11. KLF5
12. LOC283392
13. LTBR
14. LYPLALI
15. PON3
16. POU3F3
17. SERPINBI
18. ST6GALI
19. TMEMI06A
20. ZNF583
21. ZNF790
We then removed subjects with more than 10% missing features and used an 80:10:10 split. This gave train, validation, and test sizes of 6,327, 790, and 792. The evaluation metric is accuracy.
I. Model Details and Implementations
All models were implemented using PyTorch (Paszke et al. 2017); code is available at https::/ com/ a a-norcliffe/SEn
I.I. General Model Details
Here we provide details that tend to be shared across models. We explicitly state if a model does not follow the descriptions below and provide model-specific details in the next subsection.
Input Representation. Continuous features are represented as [x @ m, m). Categorical features use a one-hot encoding, where we include an additional class to indicate missing feature. Continuous and categorical representations are then concatenated input to the main model.
Deep Networks. All deep networks are given by MLPs with ReLU activation followed by Batch Normalization (loffe & Szegedy 2015). All hidden layers in a given network are the same width, which is hyperparameter that can be tuned as well as the number of hidden layers. The exception to this is the Opportunistic RL model, where we replace Batch Normalization with dropout with 0.5 probability in accordance with the method's implementation (Kachuee et al. 2019a).
Acquiring Features. After a method calculates its acquisition objective RMetod (Xo, i), this is multiplied by (1 m2) so we do not acquire features we have already observed. We also multiply by mrrest Data, do not acquire features that not available in the test data. This would not apply deployment where we have the ability measure all features.
1.2. Model Specific Details
Here we include any key details that are specific to given models, such hyperparameter names and roles. We highly recommend seeing each method's 's paper for full details of each model. Unless otherwise stated, each method follows the general rules stated previously.
Fixed MLP. The Fixed MLP uses a simple MLP structure as described above. It is trained for 120 epochs. We prevent overfitting during training by choosing the iteration with the best validation accuracy/AUROC. The greedy fixed order found after training by masking out all features and calculating the evaluation metric on the train set for each feature individually. The best feature is chosen and is unmasked for the model. The procedure is repeated with the best feature being kept to find the second best feature. This is repeated until all features have been placed in fixed greedy order.
GDFS. GDFS (Covert et al., 2023) has two separate networks, one for prediction, one for scoring features. Both have softmax final activation to give a probability distribution over the label and a positive score for each feature. Our implementation follows the original. We use the same hidden width and number of hidden layers for both networks. The Boolean "Share Parameters" hyperparameter says whether to share half the hidden layers between the two networks, this is presented in the paper as a possible way to increase performance, we treat it as a hyperparameter. We carry out pretraining on the predictor network for 80 epochs, we then carry out main training on both networks. This done using a geometric temperature progression of [Tx[1000] 0.56, 0.32, 0.18, 0.1], where the initial temperature T is hyperparameter. Main training is carried out for 15 epochs for each temperature in the progression, please see the original paper for full details. Main training consists of sampling feature acquisitions and training the scoring network to choose features with the best greedy prediction from the predictor network.
DIME. DIME (Gadgil et al. 2024) uses two separate networks, one for prediction and one for predicting the CMI of features with the label. The information network used score each feature. The information network limits the output
26

---

Stochastic Encodings for Active Feature Acquisition
to a minimum of zero and maximum of the entropy of the current predictions H(Y (x)) The majority of the DIME implementation follows the GDFS implementation above. Instead of temperature progression, use progression during main training following the original paper. This is given by InitialX[I.O, 0.25, 0.05, 0.005]. This gives the probability of choosing a feature uniformly at random compared to the best feature predicted by the information network. Main training is also done for 15 epochs for each e value, the information network is trained to predict the change in loss when a given feature is acquired.
Opportunistic RL. Opportunistic RL (Kachuee et al. 2019a) is a Deep Q learning method, where the reward is given by the 4 norm of the change in prediction distribution after an acquisition. The target network is updated compared the main network with rate of 0.001 suggested. Predictions are made by using drepout provide different network parameters at test time with 50 samples taken and averaged as suggested. The P and Q networks share representations as described in the paper. The model is trained for 20000 episodes with evaluation every 100 episodes. For the first 2000 episodes only the predictor network is trained, using uniformly random actions. Following this, the probability of a random action decays by 0.1 zndon every episode to a minimum of 0.1. After 10000 episodes and for every 2000 episodes after that, the learning rate decays by a factor of 0.2. This is all in line with the original implementation. The only change is that in each episode we do not consider individual samples from the dataset (it is not an online stream of data); instead, we train using batch of samples each episode. This improves the training, improving Opportunistic RL compared to its original online setting.
VAE. The VAE method is vanilla generative modeling approach to the AFA problem to analyze the viability of generative models. We use a Variational Auto-Encoder (Kingma, Diederik P and Welling, Max 2014) to model the distribution of the features. We train with the standard ELBO. The encoder and decoder are separate networks with separate widths and numbers of hidden layers. We then train a separate predictor that uses a standard MLP. Features are scored by taking samples of the unknown features conditioned on the observed ones (we use 50 samples). These samples go through the predictor to give an estimated label distribution. The mutual information is then estimated with [X;;; Y (xs) Ep(n,/xs)) [DKL (p(Y| (X,]x, (lp)) [xs)]]] We train for 120 epochs. We prevent overfitting during training by choosing the iteration with the best validation ELBO.
EDDI. EDDI Ma et al. 2019) is an advanced generative modeling method for AFA. The encoder is a Partial VAE. We encode the label in the same way as a categorical feature. We do not include a separate predictor, instead we follow the original paper to make predictions: features are encoded to a latent distribution and samples are decoded to y. The absence of a dedicated predictor negatively impacts the results for EDDI. The decoder follows the same structure as for the VAE. We train for 400 epochs, to prevent overfitting we choose the iteration with the best validation ELBO. Features scored based on sampled KL divergence calculated in the latent space as described in the original paper, we use 50 samples.
ACFlow. ACFlow is another advanced generative modeling method that is designed to handle arbitrary conditional likelihoods. It works by using a flow based model to maximize the conditional likelihood of p((x, y)u|(x v)o), ). such that arbitrary subsets of the features and label can be used as the condition for other arbitrary sets. We follow the publicly available implementation, training for 120 epochs.
GSMRL. GSMRL is an RL method that uses information from a generative surrogate model as additional information to a policy network. As suggested in the original paper, the policy is trained with PPO (Schulman al. 2017p. and trained ACFlow model is used as the surrogate model. We train for 300 episodes, each episode consisting of 8 epochs.
SEFA. Our method, as described in the main paper, encodes each feature separately to a normal distribution. The structure of the individual encoders depends on whether the feature is continuous or categorical:
Continuous Features. Continuous features first undergo a copula transformation ã, = -- (F,(z))), where F; is the empirical CDF of the feature and + is the standard normal CDF. This transformation is specific to stochastic encoders that regularize (Xs; 2). since it enforces a symmetry associated with the mutual information and also encourages sparse disentangled representations (Wieczorek et al.) 2018). We give [m, m.] to that feature's encoder, which is given by the MI.P described previously. The MLP ends with Batch Normalization which we found sped up training.
Categorical Features. Categorical features replace the MLP and copula transform by simply having a learnable embedding matrix (where we include an additional category representing a missing feature).
To train, we subsample the features first and use the predictive loss as described in Section 5.1 We use 100 latent samples for training and train for 120 epochs. We also provide pseudo-code for the loss calculation for batch size in Algorithm |
27

---

Stochastic Encodings for Active Feature Acquisition
To calculate the acquisition objective, we 200 latent samples and the objective given by We provide pseudo-code for scoring single feature in Algorithm (the actual implementation is able to score features in parallel on a batch of data).
Algorithm 1 Loss calculation on batch size 1.
Algorithm 2 Calculating the acquisition score for feature i.
Input: features x, data mask m, label
Input: features x, observation mask
PRemoval ~ 2(0 1)
for f = to do
for f = to d do
Pf,05 - S8, (F)
u ~ (0,1)
# == [o,0,]
[,
P5.0z= = Jā, (ar; m x I(u > pRemoval)
/ = p, a = a
end for
p(Y xo) = 0
end for
for sample to N₁ do
LR- - DRI DKI (N (NN p(((xo) 0
z ~ N~,,
p(||xo) p(((x))>x)))))))) (2(2)
for sample- = 1 to N do
end for
2~ ~~N(H, 0²)
Score = 0
p(( = p (xo) + (z)))
for y = 1 to C do
end for
for sample- = to N₂ do
INLI log(p(( = glxo))
z ~ (H,
Return: LNLL + BLR
g Vz (fg(z))))
Score = Score + p(Y = y/xo) 1 1R0 ll2
No E, 1100,
end for end for Return: Score
J. Model Runtimes
There are two places to consider runtime: training and acquisition. In Table 7 we provide the computational complexities of each method with respect to the number of features d.
As RL, DIME, and GDFS train by simulating acquisition, training time scales linearly with the number of features. Generative models (and SEFA) are constant to train since they only train to predict well. However, during inference, RL, DIME, and GDFS only require one forward pass of their policy/CMI network, whereas EDDI and VAE must individually score every feature. SEFA takes gradients of the predicted class outputs, so the runtime is linear in the number of classes, which is typically far fewer than the number of features. The main takeaway is that SEFA scales better than half the methods at training time, better than the other half during acquisition (assuming fewer labels than features), and never the worst.
Table Computational complexities for the models that actively acquire features, treating forward pass of neural network and taking the maximum among d values as approximately constant compared to calculating l(xo. i).
Computational Complexities
Single Acquisition Times (8)
Training Time (s)
Model
Acquisition Step
Training Step
Syn
MiniBooNE
MNIST
Syn 1
ACFlow
((d)
O(1)
1.483 0.003
2.208 4 0.002
21.318 1 0.191
376.885 +: 1.871
DIME
O(1)
((d)
0.017 0.000
0.021 1 0.000
0.024 0.000
575.283 1 1.332
EDDI
O(d)
O(1)
2.871 + 0.059
9.039 + 0.245
12.360 + 0.029
273.358 + 1.060
GDFS
O(1)
((d)
0.018 0.000
0.020 + 0.000
0.024 +.000
277.651 + 0.716
GSMRL
O(1)
O(d)
0.047 0.000
0.045 0.000
0.097 0.000
1010.650 + 3.186
Opportunistic RL
O(1)
O(d)
0.029 + 0.001
0.034 1: 0.001
0.001
0.030 0:000000 0.000
716.517 + 2.962
VAE
O(d)
0(1)
0.153 0.000
0.291 1 0.001
0.452 0.007
126.774 1 1.403
SEFA (ours)
O(vl))
0(1)
0.246 + 0.001
0.314 + 0.002
1.498 0.002
297.061 + 1.701
In Table 7 we also include the time per acquisition step on a subset of the datasets. As expected, the acquisition time for
28

---

Stochastic Encodings for Active Feature Acquisition
SEFA does increase significantly with the number of features (Syn 1 MiniBooNE), but does the number of classes increases (MiniBooNE to MNIST). The generative models scale the worst, and models with policy networks are the fastest to make acquisitions. For completeness, we also include the total training times for Syn 1. These results should be treated carefully, since this depends on the number of epochs, different methods converge at different rates. However, the pattern is as expected, the models that train policy network by simulating acquisition are slower to train than the generative models and SEFA. The main takeaway from the recorded wall-clock times is the same as for the computational complexities, SEFA scales better than half the methods at training time, better than the other half at acquisition time, and never the worst.
K. Experimental Details
experiments were run on an Nvidia Quadro RTX 8000 GPU. The specifications can be found at
All
attps://vis
pdf] All experi-
ments were repeated five times over parameter initializations to obtaín means and standard error estimates.
Training. We train all models using the Adam optimizer (Kingma Ba 2015), the learning rate and batch size are hyperparameters that are tuned using validation set. All methods (except for Opportunistic RL which its original implementation) learning rate scheduler that multiplies the learning rate by 0.2 when there have been set number of epochs without validation metric improvement the patience, which is also tuned.
We prevent overfitting during training by tracking a validation metric every epoch and using the model parameters that produce the best value. The validation metric we choose (unless explicitly stated for a given model) is the area under the acquisition curve. Starting from zero features, we acquire features individually, calculating the accuracy/AUROCC each acquisition, and then the validation metric the area under the acquisition curve divided by the total number of features.
Hy Hypeeparaaeeerr Tuning. For every model, initial hyperparameter tuning was conducted by finding ranges for each hyperparameter that produced strong acquisition performance on the synthetic datasets. Following this, for each model, we generated 9 hyperparameter configurations using the ranges. For each method, we test each configuration three times producing a mean value for the area under the acquisition curve. The configuration with the highest mean value is separately trained five times the main experiments. The nine configurations for each method are provided in Tables
10
and 16 We give the selected hyperparameter configurations for each dataset in Table
Table Hyperparameter configurations for ACFlow.
Hyperparameter
2
Prior Hidden Width
100
200
200
150
100
120
No. Hidden Prior Layers
2
2
No. Flow Modules
3
5
5
Flow Module Hidden Width
100
100
100
150
100
100
120
No. Hidden Flow Module Layers
1
1
1
1
1
ANLL
0.1
0.05
0.05
0.01
Learning Rale
0.001
0.001
0.001
0.001
0,0005
0.001
0.0005
0,0005
0.001
Batchsize
128
128
256
256
128
256
128
128
256
Patience
5
5
Table Hyperparameter configurations for DIME.
Hyperparameler
1
3
5
7
Hidden Width
200
200
200
200
100
200
100
100
100
No. Hidden Layers
2
2
2
Share Parameters
False
False
False
True
Truc
Truc
False
False
Truc
Pretraining Learning Rate
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
Main Training Leurning Rate
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.0001
Batch Size
128
128
128
128
128
128
512
256
512
Patience
5
5
2
5
Initial
0.4
0.2
0.1
0.4
0.2
0.1
0.4
0.2
0.1
29

---

Stochastic Encodings for Active Feature Acquisition
Table 10. Hyperparameter configurations for EDDL
Hyperparameter
3
5
7
9
Dim
200
200
50
100
20
250
100
Latent Width
200
200
100
50
20
80
250
40
60
No. Hidden Devoder Layers
2
2
2
2
1
Decoder Hidden Width
200
200
200
200
200
100
100
200
No. Hidden Encoder Layers
2
2
2
2
2
3
3
Encoder Hidden Width
200
200
200
200
200
100
100
200
Learning Rate
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
Batch
128
512
128
256
512
128
128
256
512
Decoder
0.2
1.0
0.2
0.2
0.2
0.2
0.2
0.2
0.2
Patience
5
5
5
5
5
Table 11. Hyperparameter configurations for Fixed MLP.
Hyperparameter
3
5
7
Hidden Width
200
100
200
100
300
100
250
50
120
No. Hidden I Layers
2
2
Learning Rate
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.0005
Batch Size
128
128
128
128
256
128
256
64
128
Patience
5
5
10
Table Hyperparameter configurations for GDFS.
Hyperparameter
2
Hidden Width
200
200
200
200
200
200
100
200
200
No. Hidden Layers
2
2
2
2
Share Parameters
False
False
Falsc
True
True
True
Truc
Falsc
Truc
Pretraining Learning Rate
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
Main Training Learning Rale
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
0.001
Batch Size
128
128
128
128
128
128
512
512
512
Parience
2
2
2
2
Temp Initial
2.0
1.0
0.1
2.0
1.0
0.1
2.0
1.0
0.1
Table 13. Hyperparameter configurations for GSMRL.
Hyperparameter
3
5
7
8
9
Hidden Width
100
100
200
100
100
150
75
200
100
Hidden Actor Layers
No. Hidden Critic Layers
2
2
Discount Factor
0.95
0.95
0.9
0.99
0.99
0.95
0.9
0.99
0.95
AGAE
0.95
0.95
0.9
0.95
0.9
0.8
0.95
Auntooy
0.01
0.01
0.0
0.005
0.0
0.0
0.01
0.01
0.1
Gradient Clip Norm
1.0
1,0
0.5
1.0
1.0
No. Epochs per Episode
5
8
5
10
10
8
Rollout Batch
256
256
256
512
512
300
256
256
Optimization Batch Size
128
128
128
128
256
256
100
50
50
Use Surrogate Reward
True
Falsc
False
True
True
Falsc
True
Falsc
True
Auxiliary Samples
Learning Rate
0.001
0.001
0.001
0.001
0.0005
0.0001
0.0001
0.0005
0.001
Patience
2
Table 14. Hyperparameter configurations Opportunistic
Hyperparameter
Hidden Width
200
200
200
100
200
200
100
200
100
No. Hidden Layers
Discount Factor
0.5
0.75
0.25
0.5
0.75
0.25
0.5
0.75
0.25
Learning Rate
0.001
0.001
0.001
0.001
0.001
0.0001
0.001
0.0001
0.001
Batch Size
128
128
128
256
256
128
256
256
128
30

---

Stochastic Encodings for Active Feature Acquisition
Table 15. Hyperparameter configurations for VAE,
Hyperparameter
3
5
9
Latent Width
30
10
50
30
50
10
30
40
20
No. Hidden Decoder Layers
2
2
2
2
3
Decoder Hidden Width
100
200
150
200
200
100
200
250
No. Hidden Encnder Layers
2
2
2
2
3
Encoder Hidden Width
100
200
150
150
250
No. Hidden Predictor Layers
2
2
2
2
2
Predictor Hidden Width
100
200
200
200
200
100
Leurning Rate
0.001
0.001
0.001
0.001
0,001
0.001
0.0005
0.001
0.001
Batch Size
128
128
128
128
256
512
64
128
512
Decoder
0.2
0.2
0.2
0,2
0.2
Patience
5
5
5
5
5
5
5
Table 16. Hyperparameter configurations for SEFA.
Hyperpurameter
2
3
5
6
Latent Components per Feature
4
No. Hidden Predictor Layers
2
2
2
2
I
2
Predictor Hidden Width
250
100
150
250
250
180
250
250
No. Hidden Encoder Layers
2
2
1
2
Encoder Hidden Width
20
150
20
50
150
100
40
100
100
0.0005
0.001
0.001
0.0005
0.005
0.0008
0.005
Learning Rate
0.001
0.0005
0.001
0.001
0.0003
0,001
0,0005
0,0005
0.0005
Batch Size
128
128
256
128
128
Patience
5
5
5
5
5
5
Table Selected hyperparameter configurations dataset.
Dataset
ACFlow
DIME
EDDI
GDTS
GSMRL
Opportunistic RL
SEFA
Syn
9
4
Syn 2
Syn 3
6
8
6
Bank Marketing
8
9
4
Califoria Housing
MiniBooNE
7
MNIST
B
Fashion MNIST
METABRIC
5
4
2
S
31


Reviews:
[{"review_id": "pgaZ0AsjSQ", "reviewer": "Reviewer_ABGZ", "summary": "The paper proposes a Stochastic Encodings for Feature Acquisition (SEFA) framework addressing the limitations of CMI and RL methods. SEFA uses an encoder-predictor architecture with intermediate stochastic latent variables. The architecture makes predictions and calculates the acquisition objective.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review, we are grateful for the feedback, we answer your questions below. All changes below will be added.\n\n# Scalability\nWe provide the computational complexities for a single training step and single acquisition step for each model in Table 4 (Appendix H.3). The main takeaway is that SEFA scales better than half the methods at training time, better than the other half during acquisition and never the worst.\n\nAs suggested by Reviewer Pdip, to add to this, we have timed each model to determine the wall-clock times for inference and training. The average times in seconds for a **single** acquisition are given below:\n\n|Model|Syn 1|Syn 2|Syn 3|\n|---|---|---|---|\n|DIME|$0.017\\pm0.000$|$0.018\\pm0.000$|$0.024\\pm0.000$|\n|EDDI|$2.871\\pm0.059$|$2.613\\pm0.002$|$2.854\\pm0.056$|\n|Fixed MLP|$0.014\\pm0.000$|$0.011\\pm0.000$|$0.017\\pm0.000$|\n|GDFS|$0.018\\pm0.000$|$0.015\\pm0.001$|$0.022\\pm0.000$|\n|Opportunistic RL|$0.029\\pm0.001$|$0.021\\pm0.001$|$0.032\\pm0.000$|\n|VAE|$0.153\\pm0.000$|$0.154\\pm0.000$|$0.267\\pm0.005$|\n|GSMRL|$0.047\\pm0.000$|$0.040\\pm0.000$|$0.055\\pm0.000$|\n|Random|$0.014\\pm0.000$|$0.011\\pm0.000$|$0.018\\pm0.001$|\n|SEFA (ours)|$0.246\\pm0.001$|$0.165\\pm0.001$|$0.240\\pm0.002$|\n\n|Model|Cube|Bank Marketing|California Housing|MiniBooNE|\n|---|---|---|---|---|\n|DIME|$0.021\\pm0.000$|$0.016\\pm0.000$|$0.010\\pm0.001$|$0.021\\pm0.000$|\n|EDDI|$8.872\\pm0.213$|$1.456\\pm0.002$|$0.583\\pm0.006$|$9.039\\pm0.245$|\n|Fixed MLP|$0.020\\pm0.000$|$0.012\\pm0.000$|$0.006\\pm0.000$|$0.019\\pm0.000$|\n|GDFS|$0.021\\pm0.000$|$0.017\\pm0.002$|$0.008\\pm0.000$|$0.020\\pm0.000$|\n|Opportunistic RL|$0.031\\pm0.000$|$0.022\\pm0.001$|$0.010\\pm0.000$|$0.034\\pm0.001$|\n|VAE|$0.278\\pm0.002$|$0.136\\pm0.000$|$0.048\\pm0.001$|$0.291\\pm0.001$|\n|GSMRL|$0.061\\pm0.000$|$0.035\\pm0.001$|$0.021\\pm0.000$|$0.045\\pm0.000$|\n|Random|$0.019\\pm0.000$|$0.012\\pm0.000$|$0.006\\pm0.000$|$0.017\\pm0.000$|\n|SEFA (ours)|$0.861\\pm0.004$|$0.130\\pm0.001$|$0.077\\pm0.000$|$0.314\\pm0.002$|\n||||||\n|Model|MNIST|Fashion MNIST|METABRIC|TCGA|\n|DIME|$0.024\\pm0.000$|$0.023\\pm0.000$|$0.003\\pm0.000$|$0.010\\pm0.000$|\n|EDDI|$12.360\\pm0.029$|$8.543\\pm0.200$|$0.108\\pm0.001$|$0.765\\pm0.011$|\n|Fixed MLP|$0.020\\pm0.000$|$0.020\\pm0.000$|$0.002\\pm0.000$|$0.008\\pm0.000$|\n|GDFS|$0.024\\pm0.000$|$0.025\\pm0.000$|$0.003\\pm0.000$|$0.009\\pm0.000$|\n|Opportunistic RL|$0.030\\pm0.000$|$0.031\\pm0.000$|$0.003\\pm0.000$|$0.012\\pm0.000$|\n|VAE|$0.452\\pm0.007$|$0.452\\pm0.007$|$0.009\\pm0.000$|$0.031\\pm0.001$|\n|GSMRL|$0.097\\pm0.001$|$0.069\\pm0.001$|$0.016\\pm0.000$|$0.021\\pm0.000$|\n|Random|$0.019\\pm0.000$|$0.019\\pm0.000$|$0.002\\pm0.000$|$0.008\\pm0.000$|\n|SEFA (ours)|$1.498\\pm0.002$|$1.497\\pm0.003$|$0.019\\pm0.000$|$0.098\\pm0.000$|\n\nAs the number of classes increases (MiniBooNe to Cube and Cube to MNIST/Fashion MNIST), SEFA's acquisition time increases, despite the number of features staying the same, this matches Table 4. The time for Generative Models increases significantly with the number of features in line with Table 4. As expected, models with policy networks are fastest, and scale well with the number of features, but SEFA is never the slowest (faster than EDDI, often by a significant amount).\n\nThe average time in milliseconds for a **single** training step (one gradient descent step) on Syn 1 is given below:\n\n|Model|Iteration Time (mS)|\n|---|---|\n|DIME|$35.156\\pm0.081$|\n|EDDI|$7.290\\pm0.028$|\n|Fixed MLP|$5.152\\pm0.027$|\n|GDFS|$21.661\\pm0.056$|\n|Opportunistic RL|$89.565\\pm0.370$|\n|VAE|$6.761\\pm0.075$|\n|GSMRL|$90.383\\pm0.285$|\n|SEFA (ours)|$15.843\\pm0.091$|\n\nModels that train by simulating acquisition are slowest to train despite being faster at inference time. Supervised models, including SEFA, are faster because they do not need to generate rollouts.\n\n# Additional Baselines\nThank you for the suggestion, we have run the experiments on the following new baselines:\n- GSMRL [Li and Oliva 2021], a more recent and popular RL baseline for AFA. \n- A random ordering with an MLP predictor as requested by Reviewer V2wP (to show feature ordering is necessary).\n\nThe results are given below:\n\n|Model|Synthetic 1|Synthetic 2|Synthetic 3|\n|---|---|---|---|\n|GSMRL|$5.570\\pm0.127$|$6.227\\pm0.185$|$8.199\\pm0.067$|\n|Random|$9.484\\pm0.006$|$9.499\\pm0.005$|$9.987\\pm0.008$|\n|SEFA (ours)|$4.017\\pm0.003$|$4.098\\pm0.007$|$5.081\\pm0.021$|\n\n|Model|Cube|Bank Marketing|California Housing|MiniBooNE|\n|---|---|---|---|---|\n|GSMRL|$0.891\\pm0.001$|$0.879\\pm0.006$|$0.638\\pm0.003$|$0.946\\pm0.001$|\n|Random|$0.699\\pm0.001$|$0.816\\pm0.003$|$0.569\\pm0.003$|$0.912\\pm0.001$|\n|SEFA (ours)|$0.904\\pm0.000$|$0.919\\pm0.001$|$0.676\\pm0.005$|$0.957\\pm0.000$|\n||||||\n|Model|MNIST|Fashion MNIST|METABRIC|TCGA|\n|GSMRL|$0.701\\pm0.002$|$0.683\\pm0.001$|$0.665\\pm0.002$|$0.781\\pm0.003$|\n|Random|$0.661\\pm0.001$|$0.648\\pm0.001$|$0.647\\pm0.005$|$0.753\\pm0.003$|\n|SEFA (ours)|$0.761\\pm0.001$|$0.721\\pm0.000$|$0.708\\pm0.002$|$0.845\\pm0.002$|\n\nSEFA is better than both new baselines. Random performs poorly in all cases.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "Yes.", "experimental_designs_or_analyses": "Yes.", "relation_to_broader_scientific_literature": "The contributions are novel addressing the limitations of CMI and RL methods.", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "Strengths -\n\nThe theoretical proofs to establish the limitations of CMI and RL and the proposed SEFA framework, look solid. The performance look pretty good. \n\nWeaknesses -\n\nThere's a high probability of encountering scalability issues with the proposed method and may turn out to be difficult to extend this to various other applications.\n\nSome recent RL baselines can be used as baselines.", "other_comments_or_suggestions": "NA", "questions_for_authors": "Please check the weaknesses section.", "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": "mQk8iFeqLX", "reviewer": "Reviewer_V2wP", "summary": "This paper considers active feature acquisition, which is a dynamic test-time selection (usually sequential) of features to make predictions on each test instances. The selected features are chosen independently for each considered test instance. The authors argue (with some theory) how prior methods based on conditional mutual information can fail in pathological cases, and how they might fail more generally. Instead the authors produce an information bottleneck inspired approach that utilizes a latent space to encode surrogate behavior for how unobserved variables might couple with selected features, in predicting the label. This relies on a decoupling of latent variables, one grouping for each candidate, along with a gradient-based acquisition function. There are extensive synthetic and real-world experiments to support the use of this method in practice.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review, we are grateful for the feedback, we answer your questions below. All changes below will be added.\n\n# Renyi Entropy\nThank you for bringing this to our attention, we agree that Renyi entropies are relevant. In particular how the min-entropy relates to Shannon Entropy focussing on making low likelihoods lower (Section 4). We shall add these citations.\n\nThat being said, Renyi entropies would still suffer from making myopic decisions. They would still fail on the Indicator for the same reasons as Shannon Entropy due to Proposition 4.1. Additionally it is highly non-trivial to incorporate these metrics into the SEFA architecture.\n\n# Random Selections\nWe originally tested random orderings to determine whether datasets would benefit from a non-random ordering. Random orderings performed poorly so we did not present the results to save space. We have provided the original results in our response to Reviewer ABGZ in the Additional Baselines sections. The random ordering always performs worse than the Fixed Global MLP Ordering - the datasets benefit from non-random orders.\n\n# Using the Latent Space\nMarginalization over $p(\\mathbf{x}\\_{U} | \\mathbf{x}\\_{O})$ is possible, we do it in our Feature Space ablations (Table 1 and Table 3). Instead the choice to use the latent space was made to improve performance, due to the following reasons:\n- Acquisitions are not made using gradients calculated with feature level noise. By using latent space regularization we make acquisitions with a feature representation with noise removed containing only label relevant information.\n- We avoid the need to train a generative model, avoiding the complex feature distributions (e.g. multi-modal) with complex interdependencies. In latent space we can model the distribution as independent Gaussians and the predictor learns to predict using these, and handles the interdependencies. See more in the Feature Independence response for a further explanation.\n- Tabular data may contain features that are categorical (gradients are not meaningful), or features at different scales (the gradients are different scales), so gradients with respect to the raw features is not a reliable way to score them. We are able to train the encoders so that the latent components are all approximately the same scale (by regularizing the latent space) and real numbers, so that the gradients give reliable feature scores.\n\nIn summary, by calculating scores using **representations** of the features, we avoid both the need to train a generative model, and the associated pitfalls (feature level noise, complex distributions, complex interdependencies, less meaningful gradients).\n\n# Feature Independence\nAs suggested by the ICML Program Chairs we have addressed this shared point in our response to Reviewer N1oA. Briefly, the complex interdependencies are able to be accounted for by the predictor network. The encoders do not need to learn the couplings between the features, because the predictor is able to treat latent components **differently** based on observed features, and therefore the gradients change based on this.\n\n# Scoring Metric\n**Overcoming Shannon Entropy Drawbacks**: The two main properties that Shannon Entropy lacks are (1) considering unobserved values in the current decision and (2) focussing on identifying the most likely class. SEFA addresses the first property by taking an expectation of $r(c, \\mathbf{z}, i)$ over $\\mathbf{z}$ considering the possible latent realizations of the unobserved features. SEFA addresses the second property by using probability weighting (the outer sum over $p\\_{\\theta, \\phi}(y| \\mathbf{x}\\_{O})$ in Equation 1). \n\n**Intuition behind  $r(c, \\mathbf{z}, i)$**: As pointed out by Reviewer Pdip, a good perspective on $r(c, \\mathbf{z}, i)$ is via Feature Attribution. Gradients are a common technique for explaining which features are most important for predicting a particular instance [Baehrens et al. 2009, Simonyan et al. 2014]. And LIME [Ribeiro et al. 2016] suggests attributions using locally linear approximations of the predicted probability for each class separately. So we take many samples from the latent space, representing possible future latent realizations. For each of these, $r(c, \\mathbf{z}, i)$ uses gradients as local linear approximations, to calculate feature attributions via the latent components based on how important they are for predicting class $c$. After the mean over the latent samples is taken, we have feature attributions for each class, we take a sum over classes, weighted by the current predicted probability, $p\\_{\\theta, \\phi}(y| \\mathbf{x}\\_{O})$, to give more focus to features that are relevant to more likely classes.    \n\n# Figure 4\nThank you for the suggestion, we agree that the figure can be improved to highlight the class relevant features. As suggested by the ICML Program Chairs, we have addressed this shared point in more detail in our response to Reviewer N1oA.", "claims_and_evidence": "The empirical results are strong, with sufficient comparisons against baselines and ablations, and results that seem significant over these baselines. I also really liked the explanation of why conditional mutual information might fail in proposition 4.1 and the surrounding discussion. This was well done.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "I did not look carefully at the proofs", "experimental_designs_or_analyses": "Yes. There are a nice range of analyzes including both number of features needed to reach the optimal set (synthetic data) as well as accuracy vs number of acquired feature curves, and qualitative visualizations showing the \"trajectories\" of selected features.", "relation_to_broader_scientific_literature": "While I appreciate the discussion around Shannon entropy not necessarily encouraging a maximal probability most likely class, I think more discussion is needed here around alternate metrics. Sometimes Renyi entropies are used instead of Shannon entropy to encourage different properties. For instance the min-entropy exactly encourages maximum probability on a single class. Extrinsic Jensen-Shannon divergence is another example, as an alternative to mutual information that better distinguishes between likely hypotheses (see e.g., Naghshvar & Javidi 2012).\n\nNaghshvar, M., & Javidi, T. (2012, July). Extrinsic Jensen-Shannon divergence with application in active hypothesis testing. In 2012 IEEE International Symposium on Information Theory Proceedings (pp. 2191-2195). IEEE.", "essential_references_not_discussed": "See above. These alternate metrics instead of Shannon entropy are not in the critical path of the paper, but would strengthen the discussion", "other_strengths_and_weaknesses": "- one additional baseline that I'm wondering about is, why didn't you compare against the baseline of selecting a *random* order of features, for each dataset?\n\nI think Section 5 needs more exposition. It just gets right into the details of the design choices made, but this paper would be strengthened with additional discussion around each of these choices. In particular:\n- What is the core intuition for z? Does it essentially serve as a surrogate for a marginalization over p(xU | xO), which would otherwise be too computationally expensive to consider?\n- As discussed in the limitations, the factorization of the latent encoding is a huge assumption. This entirely decouples any relationship between the selected features. It is not clear to me at all why this is a good idea, intuitively. In fact, I'm a bit surprised that the empirical results are as good as they are, since the features do not seem to be coupled in any way. Am I missing something?\n- The discussion around Shannon entropy set things up really nicely to argue how a different scoring metric is needed to better encourage maximal class probabilities on a single class. Yet, to me the choice of the gradient-based r(c,z,i) is still very much a heuristic. I think more discussion is needed around this scoring function. Why is this a good idea? Sensitivity of p(Y|z) is listed as the reason, but how does this in any way fulfill the property that Shannon entropy is lacking?", "other_comments_or_suggestions": "I like the visualizations showing the \"trajectories\" of the selected features. However, without the context of the scientific literature discussion in Section 6.3, I don't know how much Figure 4 is really adding. For someone without a medical background, the trajectories look random. This figure could be strengthened by annotating which features are actually scientifically important for each tumor location, so we can see that it is actually grabbing the important features, early on.", "questions_for_authors": "Please see the comments above\n\n## update after rebuttal\nThank you for your response. I maintain my score at accept", "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": "FtnIomYzcO", "reviewer": "Reviewer_Pdip", "summary": "This paper addresses Active Feature Acquisition (AFA), the task of sequentially selecting which features to measure for a specific test instance to improve prediction accuracy while minimizing the number of features acquired. The authors identify limitations in existing approaches: Reinforcement Learning (RL) methods face training difficulties, while Conditional Mutual Information (CMI) maximization makes myopic decisions that don't consider future acquisitions.\n\nThe method was evaluated on synthetic datasets with known optimal feature orderings and on real-world datasets including image classification and cancer classification tasks. The authors demonstrate that SEFA consistently outperforms baseline methods, including both RL-based and CMI-based approaches. For the cancer classification task, they also validate that the features selected by SEFA align with biomarkers identified in scientific literature.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review, we are grateful for the feedback, we answer your questions below. All changes below will be added.\n\n# Scalability\nWe have addressed this shared point in our response to Reviewer ABGZ.\n\n# Generalization to Regression\nSEFA can be modified for regression so that the predictor network has two heads, one that predicts the label directly and one that predicts the label as a classification problem (after it has been discretized into equally sized bins). We can use the regression head to predict the label, and the classification head to carry out the feature acquisitions.\n\n# Feature Costs and Budget\n**Feature Costs**: There are two common ways to include cost: (1) additively, e.g. minusing the cost from the reward in RL [Li and Oliva 2021], and (2) to divide by the cost either in the reward [Kachuee et al. 2019] or in the CMI estimate [Gadgil et al. 2024]. We believe dividing by the cost is better. Intuitively, the adjusted scores represent the advantage of acquiring a feature per unit cost. And it enforces the desired property that features with zero cost are acquired. Additionally, like CMI, the scores and training of SEFA are independent of cost, scores are adjusted based on cost **after** calculation. So if feature costs vary over time, SEFA would not need to be retrained like RL methods that learn a policy.\n\n**Acquisition Budget**: Figure 3 shows the evaluation metrics throughout the acquisition - how the models would perform if the acquisition was ended after $n$ features. SEFA consistently performs best throughout the acquisition.\n\n# Theoretical Claims\nWe have addressed this shared point in our response to Reviewer N1oA.\n\n# Hyperparameter Ranges\n**$\\beta$**: The value of $\\beta$ typically needs to be small, similar to a weight decay parameter. This is because it limits how much information from the features can be found in the latent space, papers such as Alemi et al. 2017 typically use values smaller than 0.01. In the sensitivity analysis in Figure 10 we see that if $\\beta$ is too large, SEFA fails because it cannot predict effectively, and if $\\beta$ is too low, then the latent space is not regularized enough. Based on this, we recommend 0.0001 - 0.01 as an effective range.\n\n**Number of Samples**: SEFA improves with the number of acquisition and training samples, shown by the sensitivity analyses in Figures 11 and 12. This is because during training it encourages a diverse latent space, and during acquisition the diversity can be fully sampled. We recommend 50-200 as a good tradeoff between acquisition performance and computation.\n\n**Number of Latent Components**: We use more than one latent component per feature so that there is a large capacity in the latent space for a rich representation of each feature. However, using too many increases the possibility of overfitting. Since each feature is only one value we recommend a value between 5-10. We tested 4, 6 and 8 as possible values and these were effective.\n\n**Other Hyperparameters**: The initial ranges were determined by manually finding good values on the synthetic datasets and using ranges around those. We started with broadly accepted values, e.g. 0.001 as the default PyTorch Adam learning rate.\n\n# References\nWe agree with all suggestions and will add them. We will also cite Baehrens et al. 2009 and Simonyan et al. 2014 as gradient based feature attribution methods.\n\n# Feature Independence\nWe have addressed this shared point in our response to Reviewer N1oA.\n\n# Variational Upper Bound\nThe regularization term to minimize is $I(Z; X)$. This requires access to the marginal $p(z)=\\int p(z|x)p(x)dx$, which is intractable. Instead we use the standard normal as a variational approximation since it produces an analytical upper bound. The derivation can be found in equations 13 and 14 of Alemi et al. 2017.\n\n# Noisy Features\nWe have run Syn 1 (with two additional baselines) with added Gaussian noise with three standard deviations: 0.1, 0.2 and 0.4. The number of acquisitions to acquire the correct features are:\n\n|Model|No Noise|0.1|0.2|0.4|\n|---|---|---|---|---|\n|DIME|$4.079\\pm0.057$|$4.688\\pm0.211$|$4.703\\pm0.276$|$5.368\\pm0.267$|\n|EDDI|$9.183\\pm0.187$|$8.988\\pm0.285$|$9.216\\pm0.136$|$9.382\\pm0.245$|\n|Fixed MLP|$6.009\\pm0.000$|$6.312\\pm0.202$|$7.321\\pm0.378$|$6.009\\pm0.000$|\n|GDFS|$4.568\\pm0.195$|$4.566\\pm0.187$|$4.543\\pm0.164$|$5.583\\pm0.283$|\n|Opportunistic RL|$4.203\\pm0.034$|$4.347\\pm0.091$|$4.720\\pm0.142$|$5.488\\pm0.084$|\n|VAE|$6.593\\pm0.085$|$6.866\\pm0.041$|$7.005\\pm0.086$|$7.095\\pm0.078$|\n|GSMRL|$5.570\\pm0.127$|$5.495\\pm0.126$|$5.778\\pm0.120$|$6.858\\pm0.269$|\n|Random|$9.484\\pm0.006$|$9.495\\pm0.010$|$9.495\\pm0.010$|$9.495\\pm0.010$|\n|SEFA (ours)|$4.017\\pm0.003$|$4.100\\pm0.004$|$4.207\\pm0.008$|$4.406\\pm0.004$|\n\nThe models get worse with more noise. SEFA remains the best model at all levels of noise. This is likely due to the latent space regularization removing noise between the feature space and latent space.", "claims_and_evidence": "Overall, the paper's claims are generally well-supported by evidence, though with some limitations:\n\n## Well-supported claims:\n\n1. **Theoretical limitations of CMI maximization**: The authors provide both theoretical arguments and concrete examples showing why greedy CMI maximization can be sub-optimal for AFA (Section 4).\n\n2. **Performance of SEFA vs. baselines**: The extensive empirical evaluation across multiple datasets consistently shows that SEFA outperforms baseline methods. The authors report means with standard errors across 5 runs, providing statistical confidence in their results.\n\n3. **Ablation studies**: The impact of each SEFA component is demonstrated through comprehensive ablations on synthetic and real datasets, showing that each design element contributes to performance.\n\n4. **Feature relevance on medical datasets**: For TCGA data, the authors cite scientific literature supporting that their model selects biologically relevant features for different cancer types.\n\n## Claims with more limited evidence:\n\n1. **Overcoming RL difficulties**: While the authors cite RL challenges as motivation for SEFA, they don't directly demonstrate how these specific challenges affect performance on their tasks. SEFA does outperform the RL baseline, but the exact connection to the cited RL difficulties isn't experimentally verified.\n\n2. **Non-greedy acquisitions**: The synthetic experiments show SEFA makes better decisions than greedy methods, but it's not entirely clear if this is due to the specific \"non-greedy\" design or other aspects of the method. The indicator example helps, but more direct experimental validation of this specific claim would strengthen it.\n\n3. **Scalability claims**: While the paper discusses computational complexity in Table 4, there's no empirical evaluation of actual runtime comparisons between methods, which would help verify the practical implications of the theoretical complexities.\n\n4. **Generalization beyond classification**: The method is currently limited to classification tasks, as noted in the limitations section. Claims about the general superiority of the approach should be considered in this context.\n\nThe paper's claims are largely substantiated by the evidence presented, with the experimental methodology being thorough and the conclusions reasonably drawn from the results.", "methods_and_evaluation_criteria": "## Proposed Methods\n\nThe proposed SEFA method makes sense for the Active Feature Acquisition (AFA) problem and addresses known limitations of existing approaches:\n\n1. **Latent space reasoning**: Using a regularized latent space to simplify decision-making is appropriate for complex feature relationships, particularly since the Information Bottleneck approach helps focus on label-relevant information.\n\n2. **Stochastic encoding**: The approach of using multiple latent samples to consider diverse feature realizations addresses the key issue of myopic decision-making in CMI-based methods.\n\n3. **Probability weighting**: Weighting by predicted class probabilities is a sensible way to focus on distinguishing between likely classes rather than ruling out unlikely ones.\n\n4. **Supervised training**: Avoiding reinforcement learning complexities by using supervised training with a predictive loss is a reasonable design choice given the known difficulties with RL.\n\nThe factorization of the latent distribution (each feature responsible for specific latent components) is a pragmatic simplification, though as the authors acknowledge, it limits modeling conditional dependencies between features.\n\n## Evaluation Criteria\n\nThe evaluation approach is comprehensive and appropriate:\n\n1. **Dataset selection**: The mix of synthetic datasets (with known optimal strategies), image datasets, tabular datasets, and medical datasets provides good coverage of different application scenarios.\n\n2. **Synthetic datasets**: Using synthetic data with known optimal feature ordering provides clear validation of the method's ability to learn non-myopic acquisition strategies.\n\n3. **Metrics**: Using AUROC for binary classification and accuracy for multi-class classification are standard and appropriate choices. The \"average evaluation metric during acquisition\" appropriately captures performance across the entire acquisition process.\n\n4. **Baselines**: The comparison against diverse baselines (RL-based, CMI-based, and fixed ordering) provides a thorough evaluation landscape.\n\n5. **Ablation studies**: The detailed ablation experiments effectively isolate the contribution of each component of the method.\n\n6. **Qualitative analysis**: For the TCGA dataset, connecting selected features to literature-supported biomarkers adds biological plausibility to the results.\n\nOne potential limitation is the lack of evaluation under different acquisition budget constraints, which would be relevant for resource-constrained applications. Also, while the authors pre-select features for high-dimensional datasets (MNIST, TCGA, etc.) using STG for computational efficiency, this pre-selection might impact the true feature acquisition performance in those domains.\n\nOverall, the methods and evaluation criteria are well-aligned with the AFA problem and provide convincing evidence for the effectiveness of SEFA.", "theoretical_claims": "1. The proof of Proposition 4.2 doesn't fully explain why considering possible unobserved feature values in the acquisition objective enables optimal feature acquisition more generally beyond the indicator example.\n\n2. The paper doesn't provide formal guarantees about SEFA's optimality, only demonstrating that considering unobserved feature values is necessary for optimality in the indicator example.\n\n3. The connection between the theoretical insights from Propositions 4.1 and 4.2 and the design of SEFA isn't made fully explicit—particularly how the stochastic latent space approach relates to the integral in Proposition 4.2.", "experimental_designs_or_analyses": "The synthetic dataset experiments (Section 6.1) are well-designed to test whether models can learn optimal feature orderings. The use of known optimal paths, multiple runs with error reporting, and clear visualizations provides strong validation of acquisition strategies. This controlled setting effectively demonstrates SEFA's ability to prioritize features optimally.\n\nThe real dataset experiments (Section 6.2) appropriately evaluate performance across varied acquisition steps. The diverse dataset selection, consistent protocols, and multiple runs strengthen validity. One limitation is the pre-selection of features for high-dimensional datasets using STG, which might introduce bias if different acquisition methods would naturally prefer different feature subsets.\n\nThe cancer classification experiments (Section 6.3) effectively combine quantitative performance with qualitative analysis. The supporting literature citations for selected features and analysis across cancer types validate that the model makes biologically plausible decisions. While not exhaustive, this literature validation adds credibility to the feature selections.\n\nThe ablation studies are comprehensive, systematically removing each proposed component while maintaining consistent evaluation protocols across synthetic and real datasets. The accompanying visualizations of acquisition behavior clearly demonstrate each component's impact, strongly supporting the claim that all components are necessary for optimal performance.\n\nThe sensitivity analyses and hyperparameter selection procedures follow sound practices, with appropriate exploration of parameter ranges and validation metrics. The reporting of final configurations enhances reproducibility, though more explicit justification for initial hyperparameter ranges would strengthen methodological transparency.\n\nOverall, the experimental designs employ appropriate controls, metrics, and validation procedures. The analyses thoroughly support the paper's claims about SEFA's effectiveness for active feature acquisition.", "relation_to_broader_scientific_literature": "The paper's critique of CMI-based acquisition relates to broader discussions about greedy vs. non-greedy information acquisition. Chen et al. (2015a) established theoretical bounds for when greedy information maximization is near-optimal, but SEFA demonstrates scenarios where this breaks down. SEFA's factorized latent space approach conceptually relates to Ma et al.'s EDDI (2019), but fundamentally differs by focusing on gradients in latent space rather than estimating CMI.\n\nSEFA addresses the exploration-exploitation tradeoff in active learning without using reinforcement learning. This connects to broader challenges in RL, offering an alternative supervised approach.\n\nSEFA's probability weighting mechanism addresses a known issue in information-theoretic approaches where minimizing entropy can focus on low-probability events. This relates to research on decision-theoretic active learning  and optimal experimental design, where the goal is to directly minimize expected error rather than maximize information gain.\nSEFA's approach of encoding features individually before integration relates to research on partial variable-wise encoders in missing data contexts. While prior works like ACFlow (Li et al., 2020) used normalizing flows for this purpose, SEFA provides a simpler yet effective alternative specifically designed for the acquisition task.", "essential_references_not_discussed": "While the paper provides a thorough literature review, there are a few notable omissions that would help contextualize their contributions:\n\n- Monte Carlo Tree Search (MCTS): The non-greedy acquisition problem shares similarities with planning problems where MCTS has been successful (Lim et al., 2012). MCTS explicitly balances immediate rewards with long-term payoffs through rollouts, conceptually similar to how SEFA samples multiple latent realizations to evaluate acquisition decisions.\n\n- Feature attribution methods: SEFA's gradient-based approach for calculating feature importance shares similarities with feature attribution methods like Integrated Gradients (Sundararajan et al., 2017) and LIME (Ribeiro et al., 2016). These connections could strengthen the theoretical foundation of the gradient-based scoring function.\n\n- Partial observability in sequential decision making: The paper lacks references to Partially Observable Markov Decision Processes (POMDPs) literature", "other_strengths_and_weaknesses": "## Strengths\n\n**Novel Integration of Techniques**: The paper creatively combines stochastic encoders, latent space regularization, and gradient-based feature scoring into a coherent framework. While each individual component builds on existing ideas, their integration into a unified acquisition approach is original and effective.\n\n**Principled Design Choices**: Each aspect of SEFA addresses specific limitations of prior methods. The stochastic latent space addresses myopic decision-making, probability weighting focuses on distinguishing between likely classes, and the factorized architecture enables tractable feature scoring. These design choices are well-motivated by theoretical insights.\n\n**Strong Empirical Results**: The consistent outperformance across diverse datasets is impressive. Particularly noteworthy is the performance on medical datasets where domain knowledge validates that selected features align with biological mechanisms.\n\n**Pragmatic Balance**: The method strikes a good balance between theoretical principles and practical implementation. By using supervised training instead of reinforcement learning, the authors provide a more accessible approach to AFA that performs better and is easier to implement.\n\n**Clear Writing and Visualizations**: The paper is generally well-written with excellent visualizations. The heat maps showing acquisition trajectories (Figures 2, 4) are particularly effective at conveying complex patterns in feature selection across different scenarios.\n\n## Weaknesses\n\n**Limited Theoretical Guarantees**: While the paper provides examples where SEFA works well, it lacks formal guarantees about its optimality or convergence properties. The theoretical foundation primarily motivates the approach rather than providing rigorous performance bounds.\n\n**Parameter Sensitivity**: The method introduces several hyperparameters (β, number of latent components per feature, number of samples). While ablations show their importance, the paper doesn't provide strong guidance on how to select these parameters for new datasets beyond empirical tuning.\n\n**Computational Overhead**: At inference time, SEFA requires sampling from latent distributions and computing gradients, which may be computationally demanding compared to simpler methods. The paper acknowledges this limitation but doesn't provide empirical runtime comparisons.\n\n**Restricted to Classification**: As noted in the limitations, the method is currently designed for classification tasks. This restricts its applicability to regression problems, which are common in many domains where feature acquisition is relevant (e.g., predictive maintenance, climate modeling).\n\n**Feature Independence Assumption**: The factorized latent space, while computationally convenient, assumes independent encoding of features. This may limit the model's ability to capture complex interactions between features, particularly in datasets with strong feature correlations.\n\n**Limited Discussion of Scalability**: While the paper demonstrates effectiveness on datasets with moderate dimensionality, it's unclear how the approach would scale to problems with thousands or millions of features, which appear in genomics and other high-dimensional domains.\n\nOverall, the paper presents a significant contribution to the AFA literature with strong empirical validation. Its originality lies more in the creative combination of techniques and their application to an important problem than in fundamental theoretical breakthroughs.", "other_comments_or_suggestions": "1. When describing the latent space regularization (Section 5.3), the paper refers to using a variational upper bound but doesn't explicitly state the form of this bound. For clarity, it would be helpful to provide the specific form.", "questions_for_authors": "1. **Computational Overhead Analysis**: How does the inference time of SEFA compare to the baseline methods across different datasets? Since SEFA requires multiple latent samples and gradient calculations during acquisition, it would be helpful to understand the practical trade-off between performance gains and computational costs. A response showing reasonable computational requirements would strengthen the practical applicability of the method.\n\n2. **Adaptation to Varying Acquisition Costs**: Many real-world scenarios have different costs associated with acquiring different features (e.g., some medical tests are more expensive or invasive than others). How would SEFA be adapted to account for variable acquisition costs? A clear adaptation strategy would significantly enhance the method's practical utility in cost-sensitive domains.\n\n3. **Extension to Regression Tasks**: Given that the current probability weighting mechanism is specific to classification, what modifications would be required to extend SEFA to regression tasks? Understanding this extension would clarify the method's generalizability beyond classification problems, which is currently listed as a limitation.\n\n4. **Robustness to Noisy Features**: How does SEFA perform when features contain noise or measurement errors? Since the method relies on gradient information to score features, it's important to understand its sensitivity to noise. Evidence of robustness would strengthen confidence in SEFA's real-world applicability.\n\n5. **Selection of Latent Dimensionality**: How should practitioners determine the number of latent components per feature for new datasets? This hyperparameter seems crucial for balancing model capacity with regularization, and guidance beyond empirical tuning would make the method more accessible to new users.", "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": "Bx6S37J359", "reviewer": "Reviewer_N1oA", "summary": "Considering the training challenges of reinforcement-learning approaches and the myopic shortcomings of conditional mutual information (CMI) strategies, this paper proposed a method called Stochastic Encodings for Active Feature Acquisition (SEFA). By encoding features into a regularized, stochastic latent space, it can better predict future observations; it then scores each unobserved feature using a gradient-based measure of how it would disambiguate among likely classes. This paper tested their method on standard tabular/image data, and cancer gene-expression classification.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the positive review, we are grateful for the feedback, we answer your questions below. All changes below will be added.\n\n# Theory\nWe agree that the theory provided is specific to the indicator and does not provide general statements about bounds on SEFA's performance. The purpose was not to provide fully general proofs, but insight into why CMI maximization can fail (it does not take into account possible future feature observations).\n\n**Purpose of Proposition 4.2**: Proposition 4.1 gives a necessary condition for optimality, proposition 4.2 shows that it can be sufficient, we include 4.2 because only necessity without showing sufficiency is a weaker result. The connection to SEFA is to use the **idea** of considering possible future observations in the model design (the expectation in the latent space), since this meets a necessary condition for optimality. We verify this empirically in our ablations (Table 1 and Table 3 Deterministic Encoder). We do not use the objective from 4.2 because it is intractable and does not fit into our latent set up.\n\n**Generalizing Beyond the Indicator**: Proposition 4.1 and 4.2 can be generalized beyond the indicator: CMI will be suboptimal where there are features that are jointly informative but individually non-informative. By definition, the CMI of the jointly informative features is zero, so any other features that are individually informative (even if they are very noisy) will be selected first by CMI. In contrast, the objective in 4.2 can capture the effect of jointly informative features, since it considers the possible values of the unobserved features.\n\n# Feature Independence\nThere are two reasons why we believe our conclusion overstates this as a limitation:\n\n1. Whilst the encoders do not model the feature interdependence, this can be accounted for by the predictor network $p\\_{\\phi}(y|\\mathbf{z})$. The complex interdependencies are captured by the layers after the latent space. Existing work uses this principle, for example, $\\beta$-VAE [Higgins et al. 2017] has disentangled latent spaces where the complex relationships are handled by the decoder. The key point is that we are not modelling distributions of the features themselves, but of their representation - their effects on the predictor network. The predictor can learn to act differently based on other latent values, for example, using samples of feature 2's latent components differently depending on if latent component 1 is positive or not. Whilst the distribution of feature 2's latent components cannot be affected by observing feature 1, how the latent samples are used by the predictor can be affected. As a result the gradients are affected, and therefore the acquisition scores. This is another reason to run the acquisition in the latent space. The predictor network can learn the complex relationships (provided it trains on enough latent samples), allowing the encoder to have a structure that does not learn the complex feature distributions. Whereas if we ran the acquisition in feature space, the predictor would only be trained on real feature values, and so the generative model would need to be able to learn the complex multivariate feature distributions.\n\n2. Empirically we see this. The Feature Space Ablation (Table 1 and Table 3), uses a generative model for $p(\\mathbf{x}\\_{U}|\\mathbf{x}\\_{O})$, to run the acquisition in feature space and does not perform as well as SEFA that runs in the latent space, despite being able to model the dependencies. Additionally, our datasets have high interdependencies (image and genomics datasets), and SEFA still performs well on these.\n\nWe appreciate this is still a modelling restriction, however, we believe that due to the predictor being able to model the complex interdependencies and the ablation results, we have overstated this as a limitation. We will adjust our conclusion to accurately reflect this.\n\n# Figure 4\nWhen we were making Figure 4, we initially had just the feature numbers on the y-axis. However, the trade-off we made was to put the feature names, providing more immediate information, but making the text smaller. The figure is a vector image so can be zoomed in. We shall add that the figure is best viewed zoomed in to the caption.\n\nAs suggested by Reviewer V2wP, without the supporting text it is not immediately clear which features are relevant for certain tumor locations. Therefore, we will edit Figure 4, such that class-wise selections with high frequencies (darker rectangles on the heat map) are highlighted using green bounding boxes, as was done with Syn 1-3. And notable low frequency selections, (where the selection frequency is low, but high for other classes) will be highlighted with black bounding boxes. This way the instance-wise differences are more noticeable, the selections are easy to identify when reading about them in the text, and the figure should be clearer by guiding attention to relevant parts of the figure.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "The paper’s benchmark datasets and evaluation setup from synthetic datasets designed to showcase the pitfalls of myopic conditional mutual information to tabular, image, and real-world medical tasks to demonstrate both practical applicability.", "theoretical_claims": "These formal statements (Propositions 4.1 and 4.2) are demonstrated on that indicator toy example. Seems that the authors use it as a focused scenario to prove where purely myopic strategies fall short and how considering possible unobserved values can address that limitation. They do not provide general proofs.", "experimental_designs_or_analyses": "The experimental settings, including datasets, baselines and evaluation metrics are standard.", "relation_to_broader_scientific_literature": "SEFA draws inspiration from several established concepts and methods, like RL, CMI. SEFA goes beyond these existing methods by incorporating a stochastic, latent-space perspective that is regularized via an information bottleneck. It introduces a gradient-based sensitivity strategy in which each feature owns a block of the latent representation.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "Strengths: The method has been tested across synthetic data, real tabular, image, and cancer genomics datasets.\n\nWeaknesses: The method’s factorized latent encoder implicitly assumes feature independence in the unobserved blocks, however, considering some strong feature interdependencies situations, this method might not work well.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "Figures like Figure 4 are really hard to read.", "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:

Meta-review: The reviewers thinks that this is a solid paper: it highlights the limitation of conditional mutual information-based methods, which motivated their novel approach of using latent space. Extensive synthetic and real-world experiments supported the use of this method in practice. Additional baselines, noisy feature experiments, as well as scalability evaluation in the rebuttal stage helped confirm the utility of the proposed methods.
Decision: Accept (poster)