We are grateful for the positive feedback and for the time reviewing. Following your review we have updated our manuscript to include an additional baseline - GSMRL. Please see our answers to all questions below.

Sensitivity of Classes with Respect to Features

To explain the use of gradients, it is first useful to consider our full proposed acquisition objective (Equation 1):

$R(\mathbf{x}\_{O}, i) = \mathbb{E}\_{p\_{\theta}(\mathbf{z} | \mathbf{x}\_{O})} \mathbb{E}\_{p\_{\theta, \phi}(y | \mathbf{x}\_{O})} r(y, \mathbf{z}, i)$

The gradient-based subscoring $r(y, \mathbf{z}, i)$, is one component of this objective. The main contribution of our method is to perform the scoring using the expectations over $\mathbf{z}$ and $y$, since this way we are able to place more weight on the competing classes (average over $y$). Additionally, we are able to consider many latent samples (average over $\mathbf{z}$), effectively searching through possible latent realizations that match the current observations. These are the key contributions that address the weaknesses of CMI maximization we identified in Section 4. We construct $r$ so that it tells us the following:

If the true value of the label is $c$, how important would feature $i$ be for confirming that prediction in the context of an entire sampled latent vector $\mathbf{z}$.

We can think of our method as considering possible instances that match the current observations, finding which measurement would be the most useful for identifying the label for that instance, and then averaging over them.

This is under the assumption that $r$ is able to tell us this information. We chose to use a gradient-based sensitivity, since this is scalable and tells us what local changes lead to the largest change in prediction, that is, which latent variables have the largest effect on the prediction. These gradients are not taken with respect to the features, since missing features being replaced by observations is not a continuous change, and thus the gradient will not be an accurate representation of the label’s sensitivity with respect to the feature. Instead, the gradient is taken with respect to the latent encodings. We designed the encoding architecture so that each latent component is only encoded by one feature, allowing us to trivially accumulate the gradients for the latent components to score each feature.

How Integrals are Calculated

The integrals over the latent variable in our work do not have an analytical solution, and are high dimensional. Therefore, we solve using Monte-Carlo estimates, which is standard practice for integrals of this nature.

We have also conducted sensitivity analyses to examine how important the number of samples are in these estimates, these can be found in Figures 11 and 12 in the Appendix. They show that both at training and evaluation time the performance improves with more samples, and plateaus at around 100.

Comparison to GSMRL

We appreciate the feedback about comparing to a more recent RL baseline. Whilst we compare to state-of-the-art AFA methods (DIME from ICLR 2024 and GDFS from ICML 2023), we chose to compare to Opportunistic RL, since it is regularly used as an RL baseline (see DIME and GDFS) and is not a hybrid method like GSMRL (Li et al. 2021), so we can compare directly to the overall RL approach. Following the feedback we have tested GSMRL on the Cube dataset. We have used the public implementation of GSMRL (https://github.com/lupalab/GSMRL), using the exact hyperparameters suggested in the public repository for the task. We provide the average accuracies during acquisition below, we also include the same numbers for IBFA and Opportunistic RL (from Table 3) for easy comparison:

IBFA and Opportunistic RL outperform GSMRL. We have also added this Table and the acquisition curves to the updated manuscript to show this result.

References

1. Li, Y. and Oliva, J., 2021, July. Active feature acquisition with generative surrogate models. In International conference on machine learning (pp. 6450-6459). PMLR.

• Copula: We have rerun a subset of experiments on IBFA without the copula transform. Note that since the features are normally distributed for the Synthetic and Cube experiments, the copula transform does not affect those datasets. We have included the average AUROC during acquisition below, we have also included the results for DIME and Opportunistic RL for convenience:

We see the change is small, but the copula does lead to improved performance for IBFA. For bank marketing, IBFA without the copula would still have been the best model, and for MiniBooNE it is still able to outperform DIME. Thus, the copula transform does improve results, but is not the main source of IBFA performance. We use this transform as a method for IB models introduced by Wieczorek et al. 2018 because the method encourages sparse, disentangled latent representations, it also enforces an invariance for IB (the method is invariant to monotone transformations) which is not in place by default. We do not include this transform in the baselines because: 1) we stayed as close as possible to the original methods’ implementations, tuning the hyperparameters that they tune, and not making changes outside of their original implementations; and 2) this is a transform designed for methods that use IB, to enforce the invariance and encourage sparse disentangled latent representations. There is theoretical justification for including this in IBFA but not the other models.

• Sensitivity to Gradient Measure: Below we give the number of acquisitions to acquire the correct Synthetic features when we don’t normalize the $r$ scores. We also provide the results from DIME and Opportunistic RL for convenience:

The performance is slightly worse than when we normalize. However, IBFA is still the best model on these datasets. We can also consider the performance on the real world datasets. Below we provide the mean accuracy/AUROC during acquisition:

The results here are approximately unchanged, only for Fashion MNIST the performance changes within error. Therefore, the performance is not particularly sensitive to this normalization, but we do see on the synthetic datasets, where we know the optimal behavior, it leads to a small improvement. This supports our claim that the exact form of the subscoring $r$ is not as important as the overall architecture and scoring features by taking expectations over both the latent space and the current predictions.

• Stopping Criteria and Feature Costs: It is correct that our model currently does not learn stopping criteria, and the objective does not consider feature costs. It is important to note our contribution in this work is a novel overall approach to AFA, which steps away from both RL and CMI maximization. We have demonstrated that this approach (and this is the first iteration of an approach like this) can perform better than these other two approaches. Ultimately it is not possible to solve every problem in a single paper, but we have set the foundation for further research into methods that use our approach. Which includes stopping criteria and feature costs. An example of stopping criteria could be to stop acquisition when a certain cutoff prediction uncertainty is reached. Or alternatively a user can decide when to stop, we believe this is quite an important option, we do not make the claim that AFA needs to be fully automated, or that IBFA has to work entirely in isolation, and often a human in the loop is very important when it comes to high stakes applications. One approach to feature costs could be to use the Opportunistic RL approach, there the reward is given by $|| \Delta p ||\_1/ \text{cost}$, where $\Delta p$ is the change in probability distribution. Here we can adjust our score by dividing each feature score by the feature cost. Note: Our scoring has no units, since we take a gradient divided by a sum of gradients, and then take expected values of this. Other approaches could take the score minus the cost, or can incorporate higher order information, such as the cost already spent acquiring the features. Or as suggested, an expected increase in the confidence, perhaps by sampling reduced regions of the latent space for each feature and seeing how the predictions change. There are various ways to include feature costs and stopping criteria which may be domain dependent, for now this moves away from the main message of this paper, but we agree these are valuable areas for future research.

• Factorized Latent Space: When discussing our Indicator example, in the first bullet point, there was a brief statement about the unobserved latent encodings not conditioning on the observed features, which we address now. The factorized latent space means that each latent component is only encoded by one feature. This allows us (by choice) to design an acquisition objective so that we can score latent components, and then this allows us to trivially score features. Note: the acquisition objective still accounts for the observed features, since our objective uses the whole latent vector in $r(c, \mathbf{z}, i)$ so IBFA still makes acquisitions based on current observations. Ultimately, we had to make a tradeoff between the trivial linking of latent components to features and using a more general encoding architecture. Empirically IBFA performed better than the baselines, justifying this decision.

References

1. Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep Variational Information Bottleneck. In International Conference on Learning Representations, 2017.
2. Chao Ma, Sebastian Tschiatschek, Konstantina Palla, José Miguel Hernández-Lobato, Sebastian Nowozin, and Cheng Zhang. EDDI: Efficient Dynamic Discovery of High-Value Information with Partial VAE. International Conference on Machine Learning, 2019.
3. Tishby, N., Pereira, F.C. and Bialek, W., 2000. The information bottleneck method. arXiv preprint physics/0004057.
4. Li, Y. and Oliva, J., 2021, July. Active feature acquisition with generative surrogate models. In International conference on machine learning (pp. 6450-6459). PMLR.
5. Li, Y., Shan, S., Liu, Q. and Oliva, J.B., 2021. Towards robust active feature acquisition. arXiv preprint arXiv:2107.04163.

Additional Questions

• STG: In our experiments we saw that IBFA carries out inference in a reasonable amount of time (see our approximate wall-clock times). Note: increasing the number of features also has a significant effect on the training times for Opportunistic RL, DIME and GDFS, not just the inference times for EDDI. We chose to include a feature selection preprocessing so that within our computational resources (a single Nvidia Quadro RTX 8000, modest by modern standards) we were able to run a large variety of models and datasets rigorously. That is, hyperparameter tuning with multiple seeds, using multiple runs to have uncertainty estimates etc. It is not a new experimental choice to reduce the number of features in AFA, for example Li et al. 2021 also reduce the size of their MNIST experiment “to accommodate baselines such as EDDI that have trouble scaling”. The number of features is also within the range in the majority of current AFA papers. For example Figure 2 from GDFS or Figure 2 from DIME, the standard ranges tested for real tabular datasets are on the order of 20. We stress that we have used the same experimental setting for all models equally. Additionally, this is the first iteration of a method that solves AFA using our new approach. Scalability is one potential area of research, and is in fact an active area of research in AFA (see Towards Robust Active Feature Acquisition by Li et al. 2021 B for example), it is not the case that all other models are scalable. Finally, we argue that in many real-world applications there are no limitations of using a hybrid approach, a single model making up the entire pipeline is not a requirement. In this case the practical approach of a combined scalable feature selection method followed by an advanced AFA method might be preferable.

• MNIST: The discrepancy in performance is likely due to our feature preprocessing as mentioned above. Since we have limited the number of features, there are some features that are not available to the models, and so the prediction performance is slightly worse. Note: All models were evaluated equally within the same experimental setup, which was used so that within our limited resources we could conduct a large variety of experiments on multiple baselines and datasets. We also respectfully point out there is a larger discrepancy between our Opportunistic RL results on MNIST and those from DIME and GDFS, our results are almost 50% greater in accuracy, and match the result from Opportunistic RL more closely (see Figure 2 in Opportunistic RL).

• Wall-clock Times: We have looked over the training and inference times and provided a table of wall-clock times to augment the computational complexities. The table below gives the aggregated times over the datasets as orders of magnitude. Note the caveat is that many methods were running in parallel on the same GPU, which will affect the exact times, hence reporting the times as orders of magnitude.

These empirical observations match the computational complexities provided in Appendix 5.

• MiniBooNE: We found that using the full class balance was too “easy”, all models were able to achieve very high AUROC with very few features and the results were indistinguishable between models. By enforcing class balance we have made the task harder (not easier), with less data and more balance. In this setting where all models have been evaluated equally, we see that IBFA performs best.

Loss Function

You are completely correct that one interpretation of our loss is learning to predict $y$ via $\mathbf{z}$ and imposing a KL penalty on the latent space. The same argument can be made about VAEs, saying that they are learning to reconstruct $\mathbf{x}$ via $\mathbf{z}$ while imposing a KL penalty on the latent space. However, it does not mean that the Evidence Lower Bound interpretation is incorrect, and in fact using the ELBO interpretation provides more theoretical motivation for the loss, the use of a KL divergence rather than Jensen-Shannon divergence for example. The same argument applies to our loss function. One correct, simpler interpretation is the one you have given, another correct interpretation is the one we have given, using the IB framework. However, the IB framework provides more theoretical grounding for our loss, the terms involved and more importantly how we can adapt it for our situation. For example, the IB interpretation explains the use of a KL divergence over any other regularizer. Additionally, the purpose of Theorems 1 and 2 is to explain why we have moved away from the Deep Variational Information Bottleneck loss (Alemi et al. 2017). The VIB loss is:

$L\_{\text{VIB}} = \mathbb{E}\_{p\_{\theta}(\mathbf{z} | \mathbf{x}\_S)} \bigg[-\log(p\_{\phi}(y | \mathbf{z}))\bigg]+ \beta D\_{\text{KL}}(p\_{\theta}(Z|\mathbf{x}\_S) || p(Z))$

Where one sample is taken in the expectation. This is the standard loss for deep IB methods given on line 312. We instead use our custom loss:

$L = -\log \big( \mathbb{E}\_{p\_{\theta}(\mathbf{z} | \mathbf{x}\_S)}\big[p\_{\phi}(y | \mathbf{z})\big] \big) + \beta D\_{\text{KL}}(p\_{\theta}(Z|\mathbf{x}\_S) || p(Z))$

Where we have moved the expectation inside the logarithm and take multiple samples. It is important to explain why we have made this change and using the interpretation of a predictive loss with KL regularization does not provide any motivation for moving away from the established VIB loss. However, within the IB framework we do motivate this, since Theorem 1 says that further regularization of the latent space is possible to make it more conducive to acquisition, and Theorem 2 says that our loss is the IB loss with this added regularization term. The change in the loss does not improve predictive performance (hence why the simpler interpretation cannot explain why we make the change), however it does improve acquisition performance. This is seen in ablations when the VIB loss is used, see the "1 Train Sample" row in Tables 2 and 4.

For completeness in our response, an alternative, more intuitive reason for this change is as follows: We desire to construct a latent space where there are separated regions in the latent space associated with each class. Such that initially we have a large latent variance, covering all of these regions. As more features are acquired the latent variance is reduced. However, if we train with only one sample of the latent space with subsampled feature subsets (VIB loss), then for the case of no features the model should predict high uncertainty for every sample (rather than the integral over samples). And so the label uncertainty is encoded into the position of the latent space which will create regions of the latent space that predict no class. The latent encoding with no features is around one region with a low latent variance where every sample predicts no class. Crucially taking samples in the latent space to make an acquisition will not explore a diverse set of possible latent realizations anymore and acquisitions suffer. Theorems 1 and 2 formalize this intuitive reasoning.

We appreciate the feedback on this section, and as a result have improved the clarity and motivation for these theorems in the context of our model and its training.

CMI Shortcomings - Entropy Minimization

Regarding the point about the weakness of entropy, the purpose of classification is identifying a single class. And so we argue that entropy is not a useful measure of uncertainty in the classification setting, because it is measured across all classes. Instead, a better measure of certainty is the predicted probability of the most likely class $\max\_{c} p\_{\theta, \phi}(Y=c | \mathbf{x}\_{O})$. We apply the same concept of certainty to our formalism in Section 3: We want to continue to acquire features that make the model more confident in its prediction, but we care only about the prediction for what the label is, and not the prediction of what the label is not. The difference between $[0.7, 0.15, 0.15]$ and $[0.5, 0.5, 0.0]$ is that in the first distribution there is a relatively confident prediction for what the class is (class 1), whereas in the second there is only a confident prediction for what the class is not (class 3). As we see, minimizing entropy can be achieved by further reducing already low probabilities rather than distinguishing between likely classes. Of course, in real-world settings, it is important to rule out options, but the analogy in the medical setting would be suspecting a fracture in the leg and requesting an arm X-Ray to confirm the problem is not in the arm. As more classes are added, more extreme examples are possible, for example if we have 101 possible classes with two distributions:

• 0.5 for two classes, and 0 for the rest, entropy = 0.693
• 0.9 for one class and 0.001 for the rest, entropy = 0.786

We used this as motivation for weighting our criterion by the current probabilities, to place more focus on disambiguating the most likely classes. We see in Table 4 that empirically this improves the results.

You are correct it might be possible to formalize this as an alternative uncertainty measure, for example a selection criterion for a general model could be

$\text{argmax}\_i \mathbb{E}\_{p\_{\theta}(x\_i | \mathbf{x}\_{O})}\bigg[\max\_{c} p\_{\theta}(Y=c | \mathbf{x}\_{O}, x\_i)\bigg]$,

instead of minmizing the entropy. However, greedily maximizing this will still make myopic acquisitions because it does not consider the values of unobserved features. And since the whole purpose of our approach is to be a non-greedy alternative to CMI and RL, this is not desirable. Incorporating this into other methods like DIME is outside the scope of this work, which concerns our IBFA method (especially given this uncertainty measure would still make myopic acquisitions).

CMI Shortcomings: Indicator Example

We discuss the indicator example in the paper because it concretely provides insights into why CMI maximization can fail (this has been identified in the review as a strength). The purpose of proposition 1 is to show an example where CMI fails and more importantly what the failure is (in this case the failure is not selecting the indicator first, a feature that provides no immediate value but significant long-term value). The purpose of proposition 2 is to show that with the knowledge from proposition 1, it is possible to construct a criterion for this problem, that when maximized will be optimal. It is true that we do not use this criterion (since as we stated, it is intractable). However, the purpose was not to give an exact objective that would generalize to any data, but to explore the CMI failure cases in a fully analytical setting that is relatively straightforward to understand; providing insight into the reasons for failure and a possible solution. We disagree that the example is unrelated to IBFA because we take the ideas that we discuss in the indicator example, using them as motivation for our solution, where there is a requirement to be more robust. As we have shown previously, considering the values of unobserved features in the criterion is carried out in IBFA by sampling the regularized latent space. And similarly placing more weight on higher probability classes is addressed in our criterion by taking the average over $p\_{\theta, \phi}(y | \mathbf{x}\_{O})$.

A phenomenon from the indicator example that does generalize is the importance of considering interactions between features, and not features in isolation. Individually a feature may not have much predictive power, however if it is jointly known with another feature (or features) the predictive power can increase significantly. This is a common phenomenon, in genetics for example it is known as epistasis, where the effect of a gene mutation is dependent on the presence or absence of a mutation in another gene. Another recent paper (Kim et al. 2024) calls these interactions synergistic, where variable $X\_1$ gives more information about $Y$ when conditioned on $X\_2$ than it does unconditionally. These interactions would be missed by CMI maximization, whereas a criterion that considers values of unobserved features has the ability to capture these interactions.

We can construct another example, let $X\_1 \in \\{-1, 1 \\}$, $X\_2 \in \\{-1, 1 \\}$, and let $y = x\_1 x\_2$ (this is effectively XOR). If for all other features we have $x\_i = 0.0001 y + \varepsilon$, $\varepsilon \sim \mathcal{N}(0, 1)$, then features 1 and 2 can only ever give information about $y$ if the other is known. And so CMI maximization would focus on acquiring the noisy features before the useful features, because they are jointly predictive but not individually. The purpose here was to show that despite the exact details of the indicator example not necessarily generalizing, the high-level ideas do.

Finally, on the indicator example, you are quite right that the objective in proposition 2 would not generalize outside of the indicator example due to highly correlated features giving zero value. Firstly, we reiterate that we do not use this criterion, and that it was never intended to be a fully general criterion, the main purpose was to demonstrate that in this indicator example it is possible to construct a criterion that considers the unobserved features and is optimal (i.e. it is sufficient), giving motivation for our solution. That being said, there are ways it can be fixed. For example, we could instead extend it to take a sum over all possible subsets of $U$ including the empty set. This criterion would tell us how much information a feature will give us if we had any subset of the unobserved features (including no additional features), in addition to the observed ones. This would be even more intractable, but the main takeaway again is that it is not enough to consider the mutual information with no conditioning on the unobserved features, and this insight does generalize to situations like those described above where we have interactions between features. Another practical solution for the criterion in proposition 2 could be to include a preprocessing feature selection step to remove highly correlated features. There is no reason why an AFA method has to act in isolation, and using feature selection prior to AFA would both reduce the search space and remove correlated features. But we stress again, the main purpose of propositions 1 and 2 was not to create a criterion that generalizes. But to provide valuable insight into the CMI failure cases, why it is suboptimal and what a possible solution could look like, motivating the actual design of our method.

Problem Formulation

We address each point separately:

1. The core purpose of AFA is that we do not select the same $S$ for each test instance. Instead a model is able to adapt to each test subject based on the evolving understanding of that specific instance. Therefore the $S$ optimization must not be carried out over a distribution of many data points, since we do not desire a fixed global ordering (the fixed MLP in our experiments uses a fixed ordering and performs poorly as a result). Instead, we desire a separate optimization for each test point (giving a separate $S$ for each instance). The long-term goal provided in Section 3 is for one test point, and it generalizes to a test set by running separate optimizations for each instance.

2. It is important to recognize here that the $S$ optimization given is what an AFA model is solving at test time. It is therefore impossible to know what the true class is, and so it cannot be maximized by the model. Instead it has to increase the confidence in its own predictions, as measured by its predictive component $p(Y=c | \mathbf{x}\_{O \cup S})$. This is what happens in the real world. We cannot know the true label, there is not an oracle giving it to us, so we have to make measurements until we have confidence in our predictions, if we knew the true label we would not need to acquire any features. Note: this assumes the predictive power of the model is well-calibrated in the first place, i.e. it has good test predictive performance, but also that when there are few features it does not make overly confident incorrect predictions.

3. By including both the budget and the penalty we are remaining more general, we can remove the budget constraint by setting $B=\infty$, and remove the penalty by setting $\lambda=0$. We have indeed set up our experiments to show how each model would perform if we were to stop acquisition at different points. We see in Figure 3 that if we were to stop the acquisition at certain points then the overall accuracy/AUROC of IBFA is higher. However, the inverse is also true, if the aim is to each a certain level of accuracy/AUROC we see that IBFA can achieve this in fewer acquisitions, so the penalty is not entirely redundant.

Acquisition Objective and Model Architecture

There is an important distinction to be made here between greedily maximizing an objective and greedy acquisitions. Nearly all methods will greedily maximize some objective at inference time, even Deep RL methods at inference time will choose the action that maximizes their value or Q function. The important question is what does the objective measure, and is maximizing this leading to myopic behavior. RL methods train such that their value function takes long-term reward into account, so even though it is being maximized the selections made are not greedy. On the other hand, CMI maximization leads to greedy acquisitions, since it will never select a feature that gives no immediate information about the label but might give us information about which features to measure next. Whilst IBFA greedily maximizes our new criterion, this does not mean the acquisition itself is greedy, and in fact the purpose of our work was to design this new architecture and criterion that fundamentally does not make greedy acquisitions.

To see why this is the case we first consider the architecture. Our encoding architecture is factorized so that each latent component can only be associated with a single feature. Therefore, if we can identify which latent components we most want to reduce the uncertainty for, then we know which feature needs to be measured. Now we consider our criterion, where the order of expectations has been flipped for the purpose of this discussion:

$R(\mathbf{x}\_{O}, i) = \int p\_{\theta}(\mathbf{z} | \mathbf{x}\_{O}) \bigg( \sum\_{c \in [C]} p\_{\theta, \phi}(Y=c | x\_{O}) r(c, \mathbf{z}, i) \bigg) d\mathbf{z}$

We construct $r$ so that it tells us the following:

$r(c, \mathbf{z}, i)$ tell us that if the true label is $c$, how important is feature $i$ for confirming that prediction in the context of an entire sampled latent vector $\mathbf{z}$.

Under a subscoring that matches this, our criterion works in the following way:

For a given latent sample (which represents one possible instance with the current observations in the latent space), if the true label were $c$, score the latent components based on how important they are for confirming this. Sum over all possible labels, weighted by current predicted likelihoods so that we give more weight to more likely classes. And then we average over many samples from the latent space, effectively searching through possible instances that match the current observations. This allows us to score each group of latent components and our architecture makes it trivial to link the highest scoring group to a feature. The key components are the two expectations, as well as the architectural and training decisions that allow the objective to work effectively, rather than the exact form of $r$.

This criterion and architecture addresses the identified issues of CMI as follows:

1. Myopic acquisitions are avoided since many latent samples are used which represent latent realizations of unobserved features, which we identify as a necessary condition for non-myopic behavior. Note: there is no technical fault in sampling the latents associated with observed features, they are still sampled to make full predictions, and uncertainty in the latent variable is still non-zero but reduced. Our synthetic results confirm that IBFA does not make myopic acquisitions as the only model to consistently select feature 11 first in the synthetic datasets.

2. The issue of not distinguishing effectively between high likelihood classes is addressed by using the current predictions as weights in the criterion. This is also empirically shown since removing the weighting produces worse performance (Table 4).

Intuitively, we can think of our method as considering the possible instances that match the current observations, finding which measurement would be the most useful for identifying the label for that instance, and then averaging over these. So whilst we do not believe a probabilistic interpretation is necessary for a model to be performant, we disagree that there is no interpretation of our criterion.

Crucially, this is under the assumption that $r$ is able to give us the information we required. $r$ cannot be any function, we chose to use a gradient-based sensitivity, since this is scalable, and tells us locally which latent variables have the largest effect on the prediction, so it is not entirely heuristic. This builds on the requested intuitive explanation using a linear predictor (the second to last question in the review) since for each latent sample we can consider a locally linear approximation of the predictor. We stress again that whilst there were design choices constructing $r$, its gradient based scoring is still not the main contribution of our method, but a solution that worked well with the overarching method.

Importance of Information Bottleneck

The $\beta$ hyperparameter in Information Bottleneck is typically quite a small value, as correctly identified we tried values between 0.0001 and 0.001 when hyperparameter tuning. This is in line with the expected order of magnitude of $\beta$, for example see Deep Variational Information Bottleneck (Alemi et al. 2017), where the values of $\beta$ are just as small. From that paper for instance Table 1 $\beta = 0.001$, Figure 1 values of $\beta$ are investigated from $10^{-9}$ up to $1.0$, with any values higher than $10^{-2}$ showing poor performance, Figure 2 also visually shows this. That paper also demonstrates there are critical values that are too high, e.g. in their MNIST results if $\beta > 0.01$ the model fails because not enough information is present in the latent variable to distinguish between 10 classes. Thus our values of $\beta$ are within an expected range. It is similar to weight decay, the weight decay parameter is also typically very small in an optimizer, the reason for both values being small is that regularization is only useful if the model is able to learn in the first place, therefore the predictive loss has a larger weight in the total loss.

Our results demonstrate that when $\beta$ is set to 0.0 (no IB regularization), IBFA performs worse. As mentioned, in Table 2 for each synthetic example the number of acquisitions required increases by approximately 0.5, which is more than a 10% increase. Without IB in these examples the performance is worse than DIME, as well as GDFS and Opportunistic RL in some cases. This is further shown in Figure 10, our sensitivity analysis, where the number of acquisitions is large if $\beta$ is too high and if it is too low. There is an optimal value which is non-zero and moving away from this value in either direction makes the performance worse. The same tests on real data in Table 4 show that in relatively noiseless regimes such as MNIST and Fashion MNIST IB is less beneficial, but in noisy regimes like the TCGA Cancer task the performance improves with IB.

Regarding the reasoning behind using IB to regularize the latent space, it is well established that IB is designed to encode $\mathbf{x}$ to $\mathbf{z}$ to only contain information about the label (see the introduction of the original IB paper Tishby et al. 2000 for example). This means that the features cannot be reconstructed from the latent variable and any noise at the feature level has been removed during the encoding process. You have correctly pointed out that the network cannot learn anything unrelated to the prediction task, however, the IB principle is more about removing the irrelevant information that is present in the features. And this is our motivation for using IB, since our acquisition objective is calculated using samples from the latent space we wish to use information that is only relevant to the label. On top of this, within the IB framework we are able to derive a further regularization term in our loss to make the latent space more conducive to effective acquisitions as given by Theorems 1 and 2 (we discuss this more later). Therefore we have given both theoretical motivation for the use of IB and empirically shown it is a necessary component of the model.

Regarding renaming the method (and subsequently the paper), we are not against this. However, we disagree that the gradient-based component is the main focus or contribution of our work. We discuss our objective in more detail later, briefly, the gradient-based subscoring in $r$ is only one part of the overall objective. The more fundamental part of the objective is using both an expectation over the latent space to capture effects of non-observed features, and an expectation over the labels to distinguish between competing classes. Therefore a more appropriate name might be “Active Feature Acquisition via the Latent Space”. Our preference is to IB in our method name, since for people familiar with IB it makes a clear implication that we use a latent space and that it is a stochastic encoding. Additionally, IB was an important part of our model (performance is worse without it), and we are able to motivate the loss function we use to train within the IB framework.

We are incredibly grateful for the very detailed feedback and the time spent on the review. Following the review we have updated our manuscript by improving clarity in the notation throughout our paper to make the proofs clearer; we have also updated Section 5.3 to make the motivation behind our loss clearer. We answer your queries below.

Clarifying the Notation in the Proofs

We appreciate this point about our proofs, this is a case of overly-simplified notation. In various places we have dropped subscripting probability densities with the model parameters. In these cases the intended meaning was the likelihood as calculated by the model, and not the true underlying density (since this is not available). For example:

$p(y|\mathbf{x}\_O) = \int p\_{\phi}(y | \mathbf{z}) p\_{\theta}(\mathbf{z} | \mathbf{x}\_O) d\mathbf{z}$

Should instead be written as:

$p\_{\theta, \phi}(y|\mathbf{x}\_O) = \int p\_{\phi}(y | \mathbf{z}) p\_{\theta}(\mathbf{z} | \mathbf{x}\_O) d\mathbf{z}$

To indicate that this is the likelihood calculated by the model. The same applies to any mutual information terms $I(X; Y)$ should be written as $I\_{\theta, \phi}(X; Y)$.

We apologise for any inconvenience our choice of notation had, and will make our notation more precise in our updated manuscript to reflect this.

The use of the model to calculate these terms is justified since in real applications we do not have access to the true distributions, and therefore need to approximate them with a model. For example, classification is framed as maximizing the log-likelihood of the true class calculated by the model. Similarly, Deep Variational Information Bottleneck (Alemi et al. 2017) uses deep networks, and the IB loss is calculated using the model’s distributions, here the loss is written as:

$I(Z ; Y | \theta) - \beta I(X ; Z | \theta)$

In other works, the mutual information is written as $I\_{\theta}(Z; Y)$. In both cases it is clear that this is calculated with respect to the model by including $\theta$ in the term, we will opt for using a subscript to make this clear. Since these terms are calculated using the model’s distribution, introducing the modelling assumption to simplify the terms is a valid operation, which we do in both proofs. An example of this is EDDI (Ma et al. 2019), which uses the VAE approximation in their Appendix proofs to yield a tractable acquisition objective which approximates the true mutual information.

Thanks to the authors for their detailed response, I appreciate the effort put into these clarifications and believe they'll improve the work. I still see some important issues, as I'll describe below, and in the interest of efficient communication I'll keep my responses brief.

Clarifying notation in the proofs. I don't think the issues in your proofs are a matter of unclear notation. In your revised submission, you now show parametric versions of distributions that you don't actually estimate, like $p_{\theta, \phi}(x_i, y \mid x_O)$. Furthermore, I've never heard of parameterized mutual information terms like $I_{\theta, \phi}(X_i; Y \mid x_O)$: when you show this in Appendix E it's in the context of the greedy CMI technique, which seeks to maximize the true CMI $I(X_i; Y \mid x_O)$, not some parameterized version $I_{\theta, \phi}(X_i; Y \mid x_O)$. I also cannot understand the purpose of substituting every distribution in your analysis like $p(y \mid x_O)$ for its parametric version $p_{\theta, \phi}(y \mid x_O)$, besides quickly resolving the issues mentioned in my review. These seem like hasty revisions that need to be thought through more carefully. You may be able to address the problems by revisiting your proofs assuming that $p(y \mid x_O) = \mathbb{E}\_{p_\theta(z \mid x_O)} [ p_\phi(y \mid z) ]$, then considering the disconnect introduced by these not being equal, then revisiting the point about your loss function basically learning to predict $y$ given $z$, and showing how this minimizes some KL divergence between the aforementioned distributions. For now, I believe the current results are unsound and I'm going to recommend rejection.

Importance of information bottleneck. It sounds like you agree there are settings where IB doesn't help, and that there are other more central aspects of your method's design like the choice of $r$, the choice to sample in a latent space, and the use of a gradient-based sensitivity measure. It's up to you what to do for the paper/method name, but leaning into one of these ideas seems to make more sense than the IB component, which looks like more of an add-on.

Acquisition objective and architecture. There are aspects of your method that make sense, like the idea of focusing on classes in proportion to their predicted probability. The parts that are less intuitive are 1) how to reason about unconditionally sampling latent dimensions $z_i$ vs conditionally sampling unobserved features $x_i \mid x_O$ (which is how you would ideally make optimal selections given access to the data distribution), 2) what value there is in sampling latents $z_i$ for observed features, 3) whether the gradient-based sensitivity is what we actually want or just a heuristic that works empirically. For example, based on your description, I find it difficult to reason analytically about how your method would handle the situation outlined in Section 4. I won't belabor the point because it's not the biggest issue in the paper, but this all remains somewhat unclear to me.

Problem formulation. Thanks for your clarifications, this mostly makes sense. What still seems odd about your problem formulation is it doesn't show the correct object we should be maximizing over: it's a policy whose rollouts should result in a low loss in expectation over a data distribution, perhaps subject to constraints or penalties on the resulting subsets. Again, not the biggest issue in the paper.

CMI shortcomings: indicator example. I think including this example is fine, but the proposal you give to resolve it doesn't make sense and is barely related to IBFA. I would recommend removing it, or ideally replacing it with something more closely related to IBFA.

CMI shortcomings: entropy minimization. Great. The uncertainty criterion you're talking about, $\max_c [p(y = c \mid x_O)]$ can be justified from the perspective of expected 0-1 loss, whereas entropy is related to expected log-loss. It seems straightforward to design a procedure to do maximize this at each selection step, and I would be curious to see how it performs, but agree it's beyond the scope of this work.

Loss function. I would argue the reverse of what you're saying: minimizing the predictive error is clearly the right thing to do (you're trying to predict $y$, after all), a regularizer on the posterior to ensure non-zero variance might be helpful (otherwise there might be no ability to sample latents), and the harder bit is to explain is the connection with the standard VIB formulation. Up to you what you'd like to do here, your clarifications in the paper are probably an improvement, but this argument about starting from VIB and moving the integral inside the log comes off as a bit odd.

As mentioned, it is important to us that we are able to align on these proofs. So to help discussion we re-present the proof below, where we have improved the clarity. Our intention by introducing the parameter subscripts was to avoid the supposed conflation of true and approximate distributions, however, it is clear this has unfortunately doen the opposite. Therefore, as you have suggested, we reason about the true mutual information, using underlying distributions for as long as possible before introducing assumptions. Once we bring in modelling assumptions or distributions it becomes the estimated MI. We hope this removes the ambiguity around notation, we shall say explicitly where assumptions are brought in. Consider the mutual information:

$I(X\_i ; Y | x\_O) = \int p(x\_i, y | x\_O) \log \frac{p(y |x\_i, x\_O)}{p(y | x\_O)} dx\_i dy$

Since we only need to consider maximizing this with respect to $i$ we remove terms that do not depend on $i$ (the $p(y | x\_O)$ in the logarithm). We include a marginalization over any new random variable.

$I(X\_i ; Y | x\_O) = \int p(x\_i, y, z | x\_O) \log p(y |x\_i, x\_O) dx\_i dydz$

Bayes' theorem gives us $p(y | x\_i, x\_O) = \frac{p(y| z, x\_i, x\_O)p(z| x\_i, x\_O)}{p(z | y, x\_i, x\_O)}$

Replacing this in the logarithm, flipping the fraction inside the logarithm and subsequently multiplying by -1 gives

$I(X\_i ; Y | x\_O) = -\int p(x\_i, y, z | x\_O) \log \frac{p(z | y, x\_i, x\_O)}{p(z| x\_i, x\_O)p(y| z, x\_i, x\_O)} dx\_i dydz$

Separating out the logarithm gives

$I(X\_i ; Y | x\_O) = -\int p(x\_i, y, z | x\_O) \log \frac{p(z | y, x\_i, x\_O)}{p(z| x\_i, x\_O)} dx\_i dydz + \int p(x\_i, y, z | x\_O) \log p(y| z, x\_i, x\_O) dx\_i dy dz$

The first term by definition is $-\mathbb{E}\_{p(x\_i | x\_O)}I(Z; Y | x\_i, x\_O)$. In the second term we can multiply the tem in the logarithm by $\frac{q(y |z)}{q(y|z)}$. This becomes:

$\int p(x\_i, y ,z | x\_O)\log \frac{p(y| z, x\_i, x\_O)}{q(y |z)}dx\_i dy dz - \int p(x\_i, y ,z | x\_O)\log q(y |z)dx\_i dy dz$

The second part of this is not affected by the optimization over $i$ so can be disregarded. This gives:

$I(X\_i ; Y | x\_O) = -\mathbb{E}\_{p(x\_i | x\_O)}I(Z; Y | x\_i, x\_O) + \mathbb{E}\_{p(z, x\_i | x\_O)}D\_{KL}( p(y| z, x\_i, x\_O)|| q(y |z))$

With additional terms not affected by the regularization. Note: we have not specified the additional variable $z$ or $q(y |z)$ yet, everything has been achieved by simple manipulation of the expressions. What we see is that minimizing $\mathbb{E}\_{p(x\_i | x\_O)}I(Z; Y | x\_i, x\_O)$ is maximizing a lower bound of the mutual information. We introduce our modelling choices by letting $z$ be the latent variable of our model and $q(y|z)$ be our predictor network. Our modelling assumption is used to reason that the second term is zero, using the approximating $p(y| z, x\_i, x\_O) \approx q(y |z)$. Assuming agreement with the previous lines, the only place where there can be an error is in this assumption step. And as the provided citations have shown, it is widely accepted to include the modelling assumptions. The modelling assumptions we have used are common in deep information bottleneck models (https://arxiv.org/pdf/1912.13480 provides theoretical reasoning why this modelling assumption is an effective one).

How does this affect Theorem 1? Theorem 1 can be restated as, minimizing $\mathbb{E}\_{p(x\_i | x\_O)}I(Z; Y | x\_i, x\_O)$ maximizes a lower bound on $I(X\_i ; Y | x\_O)$, and under our modelling assumption $p(y, z |x) = p(y|z)p(z|x)$ it maximizes the MI. This can be viewed as an estimate using the model, rather than the exact quantity, recall we don't have access to the underlying distributions, and this involves distributions that are parts of our learnt model as well as our assumption. The result of the theorem still stands, minimizing $I(Z; Y | X_S)$ during training will help regularize the loss to improve the acquisitions, even if it is not used as the acquisition objective. Without explicitly calculating these MI terms, we are still able to reason about them, and how they can guide our loss function. Note: we made no assumption on the modelling architecture, we only include our modelling assumption in the last step to make the lower bound tight.

I've taken some time to go through your results to explain exactly what's wrong with them. Let's start with Theorem 1 as this is enough to illustrate the problem. As setup, let's see if we can agree on this:

• There exists a data distribution $p(x, y)$, and we can do arbitrary marginalization and conditioning to get $p(x_i \mid x_o)$ and $p(y \mid x_o)$. At your suggestion we will not use these terms, especially $p(y \mid x_o)$, and instead opt for versions given by your models.
• Your model specifies two distributions $p_\theta(z \mid x_o)$ and $p_\phi(y \mid z)$. An issue that occurred to me upon first reading your paper, which I forgot to mention in the review, is that the notation is imprecise: these distributions should have separate subscripts/superscripts specifying the variables they're conditioned on. For example, the latent you sample when conditioning on $x_o$ is conditionally independent of any other $x_i$'s, and yet if your network conditioned on both $(x_o, x_i)$ you would have another latent variable that clearly does depend on $x_i$; these cannot both be denoted as $z$, it's ambiguous because they're two different random variables. To be more precise and ensure we handle conditional independence correctly and consistently, you should call these something like $z_o$ and $y_o$ for each $x_o$. This also helps distinguish your network's output from the real $y$, because you want your theory to be based entirely on the model's approximations (there should be no reference to $y$ whatsoever). Now, perhaps we can agree that what your networks actually specify is two distributions $p_\theta(z_o \mid x_o)$ and $p_\phi(y_o \mid z_o)$; because these variables only exist in the context of your model, there's no risk of ambiguity so you could even drop the $\theta$/$\phi$ subscripts. Returning quickly to the $x_o$/$x_i$ conditional independence bit mentioned previously, we now have $p(z_o \mid x_o, x_i) = p(z_o \mid x_o)$, but no similar conditional independence for $p(z\_{o \cup i} \mid x_o, x_i)$.

As a first check, I would ask if you agree up until this point.

Now returning to Theorem 1, this starts by considering a mutual information term you've written as $I\_{\theta, \phi}(x_i, y \mid x_o)$. A couple remarks:

• I can't tell which $y$ you're trying to characterize here. In the greedy CMI method, we would be trying to characterize the real $y$, but you claimed your analysis is all about the outputs from your model. One interpretation is that you're thinking of $y_o$, the random variable that gets sampled using the Markov chain $x_o \to z_o \to y_o$. If there's another interpretation please let me know, but I'll proceed with this for now.
• If we interpret the mutual information as described above, we come to an interesting conclusion: that $I(x_i; y_o \mid x_o) = 0$ for all $x_i$. You can follow your proof's mechanics to see this, or you can notice that $y_o$ is simply independent of any $x_i$ conditional on $x_o$. So this result doesn't mean anything; by making it about the "$y$" implied by your model, you ensured by construction that no $x_i$ can provide any information about it.

Now, you might be wondering why this wasn't apparent previously. I can track it down to a couple inconsistencies in how you applied conditional independence in your proof. I'll revert to your ambiguous notation for this part:

• To get to line 1216, you used the fact that $p(y \mid z, x_i, x_o) = p(y \mid z)$. If we replace $y \to y_o$ and $z \to z_o$, this is actually true: $p(y_o \mid z_o, x_i, x_o) = p(y_o \mid z_o)$ by construction.
• However, in line 1216 you also preserved the terms $p(z \mid x_i, x_o)$ and $p(z \mid y, x_i, x_o)$. If we were consistent with interpreting $y \to y_o$ and $z \to z_o$ then these should be simplified to $p(z_o \mid x_o)$ and $p(z_o \mid y_o, x_o)$. When combined with the remaining terms inside the log, this would lead you to conclude that the integral is equal to 0! Crucially, you did not make this simplification, presumably because you now began to view $z \to z\_{o \cup i}$ which actually does depend on $x_i$. This seems to be the main mistake that makes this result wrong.

As a second check, I'll ask if you agree up until this point.

Separately, another issue is that you integrate over $p\_{\theta, \phi}(x_i, y \mid x_o)$. Like I mentioned previously, your models don't estimate this distribution, it's never defined. At some point this becomes an integral over $p\_{\theta, \phi}(x_i, y, z \mid x_o) = p\_{\theta, \phi}(x_i \mid x_o) p\_{\theta, \phi}(y, z \mid x_i, x_o)$. I don't know what the difference is between $p(x_i \mid x_o)$ and $p\_{\theta, \phi}(x_i \mid x_o)$, since the latter isn't defined by your model.

Overall, I think the remedy is 1) improve the ambiguous notation, 2) reconsider whether you truly want your analysis to be about distributions implied by the model vs real ones under the data distribution. Let me know what you think.

Apologies for the delayed reply. Thanks for your last message, let me respond to your latest proposed changes. A couple preliminary points before I get into the details of your new derivation:

• It appears my suggestion to distinguish between latent variables depending on the conditioning subset was misunderstood. The notation $z_o$ was intended to indicate which subsets of $x$ are conditioned on, not a subset of latent dimensions. So a few of my points were lost in translation, for example the conditional independence bit $p(z_o \mid x_o, x_i) = p(z_o \mid x_o)$. To avoid ambiguity, perhaps better notation would have been $z^o$ so you can still subscript into each dimension, e.g., $z^o_i$. I'll return to this later in my reply.
• One of your main changes was to switch back to analyzing the mutual information with the true random variable $y$ rather than the output distribution induced by your model. I believe that’s the right decision, because the alternative leads to vacuous results like variables $x_i$ contributing zero CMI.

I’ve gone through your new proof, and your algebraic transformations are correct. However, there are two flaws in your derivation:

1. Your assumption $q(y | z) = p(y | z, x_i, x_o)$ is not justified. You can’t simply state that your network matches this distribution, you must show that some component of your learning algorithm encourages them to be close. Please show where this appears in your objective, perhaps in the form of a KL divergence between the two distributions. Because we've discussed this previously, I believe you can show that your objective minimizes an upper bound on $D\_{KL}(p(y \mid x_o) \mid \mathbb{E}\_{p_\theta(z \mid x_o)}[p_\phi(y \mid z)])$, but you're assuming something stronger than that.

2. More fundamentally, and as mentioned in my last response, you don't have a single random variable $z$ that's valid to analyze in this fashion. Your algebraic transformations are correct for any random variable, but $z$ is not a random variable that exists like $x$ and $y$. It's induced by your network separately for each conditioning subset; this is a subtle point, but that means there are different random variables $z^o$ for each subset $x_o$. Consider this: if there were a single random variable $z$, then we should be able to get its distribution $p(z \mid x_o)$ directly from your network $p_\theta(z \mid x_o)$ and identically by marginalizing over another variable, say $\int p(x_i \mid x_o) p_\theta(z \mid x_o, x_i) dx_i$. You cannot guarantee that these distributions are identical, therefore there isn't a random variable $z$. Fixating on a single dimension in your factorized setup, you can see that $p_\theta(z_i \mid *)$ is not guaranteed to equal $\int p_\theta(z_i \mid x_i) p(x_i \mid x_o)$. Surely you can agree that this point is true and acknowledge that it's odd; it's also inconsistent with the notion of a random variable $z$. Anyway, as a result, besides your modeling assumption being unjustified, the entire derivation is invalid because you haven't correctly specified the random variable you're analyzing.

I've tried to think about how to fix your derivation. Since your algorithm generates the latent variable conditioned on $x_o$, it could make sense to replace all instances of $z$ with $z^o$ in your derivation. This will unfortunately lead you to trying to show $p(y \mid z^o, x_o, x_i) = p_\phi(y \mid z^o)$, which of course is not true, both because 1) single sampled latent values do not contain of all $x_o$'s information, and 2) $x_i$ contains new information beyond the latent variable and $x_o$. For that reason, I cannot see how to fix your derivation.

To summarize, I believe your theoretical results are still incorrect. The AC can be a judge of that, and my concerns with the different versions of your theory are all laid out here. I tend to agree with you that the theory is not a central part of your paper, and the empirics are promising on their own, but I can't recommend acceptance for a paper with incorrect theorems.

Factorized Latent Space

The factorized latent space means that each latent component is only encoded by one feature. This allows us (by choice) to design an acquisition objective so that we can score latent components, and then this allows us to trivially score features.

The implication of this modelling choice is that measuring one feature has no effect on the latent encodings of any other features. When in reality if there are correlated features they should affect the latent encoding of the other. For example, if features 1 and 2 are highly correlated and feature 1 is measured, the model’s uncertainty of feature 2 is unchanged (even though a-priori this should be reduced). Note: the acquisition objective still accounts for the observed features, since our objective uses the whole latent vector in $r(c, \mathbf{z}, i)$ so IBFA still makes acquisitions based on current observations.

Ultimately, we had to make a tradeoff between the trivial linking of latent components to features and using a more general encoding architecture. Empirically IBFA performed better than the baselines, justifying this decision.

Regarding alternatives, it is non-trivial to design a method that relaxes this architectural choice and maintains the benefits of IBFA. For instance, a fully connected encoder accounts for correlations between features but it is not possible to link scores in the latent space to single features. For this reason we view this as interesting future research. One practical approach is to use a feature selection method before applying IBFA (since we do not make claims that one model has to constitute the entire AFA pipeline), this will filter out correlated features and has the added benefit of reducing the number of features in the AFA problem.

References

1. Alexander A. Alemi, Ian Fischer, Joshua V. Dillon, and Kevin Murphy. Deep Variational Information Bottleneck. In International Conference on Learning Representations, 2017.

2. Tishby, N., Pereira, F.C. and Bialek, W., 2000. The information bottleneck method. arXiv preprint physics/0004057.

3. Stewart, L., Bach, F., Berthet, Q. and Vert, J.P., 2023, April. Regression as classification: Influence of task formulation on neural network features. In International Conference on Artificial Intelligence and Statistics (pp. 11563-11582). PMLR.

Sensitivity to Number of Samples and Latent Components

Full sensitivity analyses on the number of samples for the synthetic tasks are given in the Appendix, Figures 11 and 12. The summarized results are also in the main text in Table 2. The results match the expectation that 1) more training samples improves acquisition at inference time by regularizing the latent space to be more conducive to acquisition (Theorems 1 and 2), and 2) that more acquisition samples are needed to sample the full diversity of the latent space at inference time. The performance plateaus at around 100 samples, therefore we used the optimal numbers that achieve good performance whilst remaining computationally efficient (100 train and 200 acquisition samples), we found that this produced strong results across all datasets. The recommendation would be to use as many samples as the computational resources allow, but as we see performance plateaus after 100 samples.

Regarding the number of latent factors per feature, this was one of the hyperparameters we tuned, testing three values 4, 6 and 8. To give an idea of the sensitivity we have provided the validation score (the area under the acquisition curve) for each hyperparameter configuration for a variety of the real world datasets below. Note that it is not only the number of latent components that changes between hyperparameter configurations.

The validation performance is not particularly sensitive to the number of latent components (for the values tested). In Table 13 we provide the selected hyperparameter configuration for each dataset, and we see there is no significant preference for the number of latent components. As with most hyperparameters, the best solution is to tune it with a suitable validation set and validation metric, we chose the area under the acquisition curve.

Extension to Regression

We view this as an interesting avenue for future research. The methods we are exploring are:

1. Setting up our predictor network to have two heads. One that predicts the regression target directly, which is used to make predictions. For the second prediction head we also discretize the label into bins with approximately equal frequency, so that this head is solving a classification problem. Therefore we can use the regression head to predict the regression label, and the classification head to carry out the feature acquisitions. Reframing regression as classification is not a new idea, it is often used to help stabilize training (see Stewart et al 2023. for example).

2. Using the same objective and extending the expectation over $y$ to continuous values. The acquisition objective in Equation 1 is given by

$R(\mathbf{x}\_O, i) = \sum\_{c \in [C]}p\_{\theta, \phi}(Y=c | \mathbf{x}\_{O})\int p\_{\theta}(\mathbf{z} | \mathbf{x}\_{O})r(c, \mathbf{z}, i) d\mathbf{z}$

We can change the sum over discrete labels to an integral over possible regression targets:

$R(\mathbf{x}\_O, i) = \int p\_{\theta, \phi}(y | \mathbf{x}\_{O}) p\_{\theta}(\mathbf{z} | \mathbf{x}\_{O})r(y, \mathbf{z}, i) d\mathbf{z}dy$

This integral can be carried out using Monte-Carlo estimation. The first method would be faster by not having to take large Monte-Carlo estimates, but approximation errors appear when discretizing the label. On the other hand, there are approximation errors that arise from Monte-Carlo estimates of this integral, so it is an interesting, non-trivial direction for further research.

We’re very grateful for the highly positive feedback on our work and the time spent reviewing. Following your review we have updated our manuscript to include memory requirements of IBFA in the discussion in our conclusion and to improve the structure of Section 5.3 where we provide theoretical justification for our loss. Please find our answers to your questions below.

Theoretical Analysis

Despite not primarily focusing on theory, we agree that the theoretical aspect of the paper is an important one. Apart from the theoretical exploration of CMI drawbacks, we have also provided theoretical reasoning for adapting our loss away from the standard Variational Information Bottleneck Loss (Alemi et al. 2017). These are given in Theorems 1 and 2 and demonstrate that our adapted loss adds another level of regularization that helps to shape the latent space to be more conducive to feature acquisition.

Specifically regarding the use of IB as a regularizer, a key point is that it is not IB in isolation that addresses the drawbacks of CMI, it is our method and acquisition objective as a whole. The two main ways we solve the issues associated with CMI are:

1. Taking samples from the latent space ($p\_{\theta}(\mathbf{z} | \mathbf{x}\_{O})$) to consider effects of unobserved features.
2. Weighting the scores by current predictions ($p\_{\theta, \phi}(y | \mathbf{x}\_{O})$) to distinguish the most likely classes.

However, these novelties alone are not enough; various decisions (such as factorizing the latent space) must be made to make IBFA work effectively. One of these decisions was to apply regularization specifically to the latent space (as opposed to regularization on the whole network such as dropout). Since we are using stochastic encoders, and we want the latent space to contain as much information about the label as possible whilst removing unwanted information about the features, Information Bottleneck is the natural choice of regularizer here; it is what IB was originally designed for (see Tishby et al. 2000). As well as this, as already discussed, by working within the Information Bottleneck framework we can add further regularization by moving away from the standard implementation of IB to our adapted loss (see Theorems 1 and 2). We appreciate the feedback and have made the reasoning behind using IB more clear in Section 5.3 as a result.

Computational Complexity

Please find a discussion about the computational complexities below (the discussion and table can also be found in the Appendix, Table 5).

There are two places to consider runtime: training and inference. The computational complexities of each method with respect to number of features $d$ are:

RL, DIME and GDFS train by simulating acquisition, so each step scales linearly with the number of features. Generative models (also IBFA) 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 network, whereas EDDI and VAE must individually score every feature. IBFA instead takes gradients with respect to the predicted class outputs, this is done with backpropagation so scales in the same way as the forward pass. Therefore, the runtime is linear in the number of classes, which is typically far fewer than the number of features.

Regarding memory, during training RL methods need to store past experiences, which is linear in the buffer size. IBFA uses multiple samples to train and carry out inference, so is linear in the number of samples in each case (we discuss how many samples are needed next). EDDI and VAE have a similar memory requirement at inference but are constant during training (since they only use one sample during training).

The main takeaway is that IBFA 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. IBFA uses more memory during inference than RL, DIME and GDFS depending on how many latent samples are used. We have added the memory requirement as a limitation in our conclusion.