Active Learning with LLMs for Partially Observed and Cost-Aware Scenarios

We are grateful to the reviewer for their insightful feedback and constructive comments that have improved the paper.

[P1] Inclusion of label cost in Equation 1

Thank you for your concerns. We appreciate the opportunity to clarify the inclusion of label cost in Equation 1.

In Equation 1, we include label costs because this expression allows for the acquisition of a subset of features without the label. This separation of label cost is intentional.

Note that one of our primary goals is to formalize the POCA problem, and for this, we opted for the most generalized problem definition. Notably, the POCA problem, i.e., acquisition of a subset of features and/or labels, along with its associated cost, is not explicitly addressed in the existing experimental design literature [C3], which is a contribution in itself as remarked by other reviewers (CwS1).

The utility function in Equation 1 measures the level of generalization based on acquisition. Semi-supervised and self-supervised learning are examples of how generalization can be achieved using only unlabeled data.

[P2] Training when some features are available

Thank you for highlighting this concern. We can interpret this issue in two ways:

• Training a downstream model with only a subset of features. In the context of $\mu$POCA, we work with predictive models that require fixed-size inputs. However, not all features may be observed due to the acquisition of a subset of features. To handle this situation, we employ a GSM to impute any missing features. This approach ensures that the downstream model receives the fixed-size inputs it requires. This imputation is indicated in L208.

• Training a GSM with only a subset of features. This is a common case when the GSM is trained on the pool set (results in Figure 11). Note, that our training of the GSM for any available unlabeled data. As described in L246, the process involves generating random masks to create an input $m \odot \boldsymbol{x}_o$, and an output, $(1 - m) \odot \boldsymbol{x}_o$. However, we ensure the input always has at least two features and the output at least one feature. Therefore, we only consider training samples with at least three observed features. This detail is now included in Appendix E.2.

As illustrated in the $\mu$POCA Algorithm (see attached PDF), the predictive model is trained after each data acquisition. In contrast, the GSM is trained with the available unlabeled data before the acquisition process begins.

[P3] GSM quality

We wish to clarify the impact of GSM quality on downstream performance. The ability of GSMs to accurately approximate the underlying data distribution is crucial for effective generative feature imputation, which in turn affects both predictive accuracy and acquisition performance in downstream tasks. In [G1, G2], we identify the specific conditions under which GSMs can successfully approximate the data distribution. We will flag this more clearly in the revised manuscript.

[P4] Additional baselines

Thank you for highlighting this concern. We direct you to Fig 13 (App I) where we have included a subset of uncertainty-based AL methods. These results feature EPIG (from the information family in Table 1) and other uncertainty-based metrics (marginal entropy, mean-std) [C2].

Actions taken: We have made the links to these additional results more prominent in the main paper.

Thank you for your feedback. We have addressed your individual points in this response and in the revised manuscript

[P1.1] GSM for metric estimation and training

Thank you for your suggestion; we acknowledge that this important detail should be clear.

Update. In response, before line 165 (after the new inclusion of [P1.3]), we added: "For models with fixed-size input, like Random Forest, GSMs can be used to impute missing features (see Section [3.2] for details)."

[P1.2] Clarity on indexes

Index omissions Bold variables are defined as a set of variables in line 94, which we agree can be easily overlooked. To clarify, we have now highlighted "Bold variables" for emphasis.

Additionally, we corrected line 149 by bolding j, and we found no other omissions. The omission of $i$ is noted in line 96; to make this clearer, we state ".... is omitted when unnecessary, i.e., $x_{i,j}\equiv x_{j}$".

Index J+1 is an important detail. We thank the reviewer for spotting it. Consequently, for improved clarity, now line 99 states: "...indexed by $\boldsymbol{j}$, for instance, $i$, with index $j = J+1$ indicating the label acquisition".

[P1.3] Stochastic GSMs

As the reviewer correctly mentions, the stochasticity of the GSM is a key in uPOCA. To be more direct about this, we made the following updates:

• Update 1, after the presenting the functionality of GSMs, we include in line 165 the following: "Why generative imputation can help Active Learning? Note, that various (unobserved) features can be compatible with the observed features $x_{\boldsymbol{o}}$, leading to potentially different outcomes. In this case, a traditional AL metric would select this instance because $x_{\boldsymbol{o}}$ alone is insufficient to predict $Y$. However, generative imputation incorporates information about unobserved features into the predictive model, allowing it to select instances that remain uncertain even after this information is considered. Additionally, generative imputation can identify features that are not worth acquiring, as changes to these features do not impact $Y$.

• Update 2, We include the canonical example (see attached PDF) complementing this explanation in Appendix K.

[P2] Independence assumption

Thank you for highlighting this concern, which has helped us to be more explicit in illustrating why this is not a strong assumption.

Lines L176-181 state that the assumption is valid for the supervised learning models we consider. Appendix B explains this further, and Appendix H provides empirical results confirming the validity of this assumption through Eq. (7). We made two modifications to clarify this:

• Update 1. Now sentence L176 reads: "Note that the independence assumption of Proposition 1 is not a strong assumption and is valid for the supervised models we consider.". Also, in L181, we corrected the link to Appendix H, which was previously incorrect.
• Update 2. Now sentence L180 reads: "... detailed explanation of this independence assumption's validity. In the same Appendix, we include a more informal but intuitive and practical explanation". The intuitive and practical explanation is the following: Intuition. Let the RV of observed and unobserved features be $X_1$ and $X_2$, respectively. Consider $\mathcal{G} = \Omega$, representing the EIG case, where the model parameter is related to generalization. Our downstream model is a Bayesian neural network (BNN) that samples $w_1$ half the time and $w_2$ the other half. This BNN is trained on dataset $\mathcal{D}$, using samples from $[X_1, X_2]$ as input. Given an instance with observation $x_1$, we use GSM to sample $x^\*_2$. The prediction $y$ is obtained by sampling $w^\*$ from the BNN and $x^\*_2$ from GSM, resulting in $y^\* = w^\*([x_1, x^\*_2])$. It is clear to notice that knowing $w^\*$ gives no information about the possible values of $x^\*_2$; this is because how we sample $w^\*$ depends on $\mathcal{D}$ and how the model is trained. This is the independence assumption. It is only possible to know something about $x^\*_2$ when we know $y^\*$, thanks to $y^\* = w^\*([x_1, x^\*_2])$. In the observation, we refer to this as the immorality.

[Q1] Study on different missingness patterns

We believe that researching pattern missingness and its impact on GSM and acquisition performance is an important area of study that deserves a comprehensive study on its own. That said, we agree that a better understanding of how partial observability affects the acquisition process is necessary. In our study, we already consider real-world datasets that were engineered to construct partial observability. However, to be able to fully control the level of correlation of unobserved and observed feature, we extend this study using synthetic datasets in [G2].

[Q2] when $\mu$POCA outperform traditional AL under limited GSMs

Active Learning's main limitation is that its acquisition metrics do not consider uncertainty reduction in feature acquisition. Our work addresses this by using GSMs to estimate the distribution of unobserved features and samples from this distribution to estimate the uncertainty reduction, though, as noted by reviewers, the GSM can be limited in some applications. uPOCA-based methods offer the advantage of approximating unobserved features, and even if imperfect, it can help identify non-useful features. To determine which unobserved features are relevant, we only need to observe some effect in the outcome when varying these features. In contrast, irrelevant features will have minimal impact on the outcome despite variability.

We are grateful to the reviewer for their insightful feedback that has improved the paper.

[P1] Computational costs of MC sampling

Thank you for highlighting this concern. We agree that it is important to discuss $\mu$POCA's computational costs. We approach this discussion by (1) providing a detailed algorithm for complete acquisition in Algorithm 1 (in the attached PDF), (2) detailing the complexity below, and (3) discussing how this cost can be alleviated in future works.

Training process. Before we discuss inference complexity, we note that GSMs will be trained using available unlabeled data, which may be either fully observed (if a bank of fully observed unlabeled data is available) or partially observed (using the pool set itself). After the GSM is trained, it is used for inference during the acquisition process.

Inference complexity. Using notation consistent with the paper, $I$ is the number of instances, $J$ is the number of features, and $S$ is the number of Monte-Carlo samples. A direct acquisition would consider selecting subsets of features over all instances, resulting in a complexity of $\mathcal{O}(I\*J^2\*S)$. To reduce this overhead, we introduce an approximate algorithm (detailed in S4.3), where we (a) first select the most informative instance (assuming all features are acquired), then (b) select the subset of features of that instance to acquire:

• (a) As we now make more explicit in Alg 1 (in attached PDF), the first step involves identifying the instance to acquire, which has a sampling cost of $\mathcal{O}(I\*S)$ (note, all unobserved features are sampled jointly).
• (b) Once the instance is selected, the cost of acquiring a subset of features is bounded by $\mathcal{O}(J^2\*S)$.

Our approximate procedure thus has a total complexity of $\mathcal{O}(I\*S + J^2 \*S)$, compared to $\mathcal{O}(I\*J^2\*S)$ of the exact procedure. This analysis indicates that the dominant factor is the number of features $J$ as it affects sampling costs quadratically.

Future work. In our work, we estimate uncertainty using Monte Carlo samples. However, a promising area for further research is improving efficiency by avoiding sampling with tractable analytical expressions. For example, assuming that features and labels are jointly Gaussian allows us to compute Eq. 3 in closed form. Other conjugate distributions can also be explored. Another approach to enhancing the efficiency of sampling-based approaches is in-context learning, as generative imputation can be done in parallel for multiple instances.

Actions taken: The above analysis and Alg 1 (attached PDF) have been included in App J, with appropriate references in the main text.

[P2] Focus on tabular data

We agree that the primary applications of POCA are in tabular data, and we focused our empirical analysis on this domain. However, we emphasize that tabular data is ubiquitous in the real world, particularly in fields such as finance, healthcare, industrial engineering, and retail, where the POCA problem is most likely to be applicable.   However, our formalism and method can be readily extended to other modalities, such as images, by treating pixels or pixel patches as features. Although this is less common, practical applications could exist in fields like medical and satellite imaging, where occlusion caused by noise is prevalent. In these cases, it's crucial to determine when a sample requires additional information (features) for accurate prediction and model training. Similarly, interactive robots that learn through vision can also benefit from this approach.

Actions taken: We have (1) added discussions on extensions to other modalities in the Future Works section (S5) and (2) extended Table 1 to include example applications involving other modalities, demonstrating the versatility of POCA.

[P3] Conditions that affect POCA performance

Thank you for highlighting this concern. We acknowledge that the performance of $\mu$POCA-based methods can vary in certain scenarios. The key question we want to answer is: in what settings do $\mu$POCA methods perform poorly? We hypothesize that this can occur when the GSM does not accurately learn the underlying data distribution, which is affected by (1) the approximation power of the GSM, and (2) the intrinsic properties of the partial observability.

To investigate this question and our hypothesis, we conducted a comprehensive investigation in [G1, G2] of the global response, highlighting how GSM approximation performance and intrinsic correlations in the dataset affect $\mu$POCA performance.

Actions taken: Thank you for this suggestion, which has improved our analysis of $\mu$POCA performance in different settings. We have introduced a dedicated section in the appendix for [G1, G2].

[P4] Additional baselines

Thank you for highlighting this concern. We direct you to Fig 13 (App I) where we have included a subset of uncertainty-based AL methods. These results feature EPIG (from the information family in Table 1) and other uncertainty-based metrics (marginal entropy, mean-std) [C2].

Actions taken: We have made the links to these additional results more prominent in the main paper.

We hope that most of the reviewer’s concerns have been addressed. We’d be happy to engage in further discussions.

We are grateful to the reviewer for their insightful feedback that has improved the paper.

[P1] Emphasizing novelty

Thank you for highlighting this concern. Allow us to further clarify our novelty and theoretical contributions for partially observable data acquisition.

Theoretical analysis. We start by focusing on $\mu$POCA, which we theoretically analyze in Proposition 1, applying to a family of Bayesian methods. Although the theorem may initially seem intuitive, it crucially implies that the new acquisition term requires estimating the distribution of the unobserved features. This requirement led us to develop GSMs to estimate this distribution. In other words, this theoretical rationale, which begins by analyzing uncertainty reduction in a partially observable setting, highlights the necessity of GSMs in this context, representing a novel contribution.

Methodological novelty. To demonstrate the methodological novelty of $\mu$POCA, we compare it with Vanilla Active Learning (VAL) and Active Learning with Imputation (ALI) using EIG as an illustrative example:

• PO-EIG. The acquisition metric is represented by $I(\Omega, Y|x_{\boldsymbol{o}}, \mathcal{D}, X_j)$, which measures predictive uncertainty considering the variability of $X_j$. This helps identify which values of $X_j$ are worth acquiring, as its variability allows measuring the impact of uncertainty reduction on predictive performance. Fig 3a in the attached PDF illustrates this, showing that exploring the unobserved regions $X_2$ and $X_3$ helps understand areas of uncertainty in the predictive model. This capability enables us to measure what features are relevant and in what regions.
• VAL, ALI. We use $\tilde{x}$ to denote an imputation, the acquisition metrics for VAL and ALI are then represented by $I(\Omega, Y|x_{\boldsymbol{o}}, \mathcal{D})$ and $I(\Omega, Y|x_{\boldsymbol{o}}, \mathcal{D},\bar{x})$, respectively. These metrics cannot determine which features to acquire: VAL does not consider unobserved features and their impact on predictive performance; ALI provides only a point estimate in the unobserved region (as shown in Fig 3a from the PDF), which is insufficient for understanding uncertainty in the unobserved space.

Actions taken: This discussion and the illustrative example in Fig 3 have been added to App B.

[P2] Effect of partial observability on performance

We appreciate your feedback and agree that partial observability (PO) is a critical factor that can influence the ability of GSMs to accurately approximate the underlying data distribution. This approximation becomes vital as it directly impacts both predictive accuracy and acquisition performance in downstream tasks.

Actions taken: We hypothesize that PO can affect acquisition performance when the GSM does not accurately learn the underlying data distribution, which is affected by (1) the approximation power of the GSM, and (2) the intrinsic properties of the partial observability. To investigate our hypothesis, we conducted a comprehensive investigation in [G1, G2] of the global response, highlighting that GSM approximation performance and intrinsic correlations in the dataset affect $\mu$POCA performance.

[P3] Acquisition of only features in POCA

Thank you for raising this point. As you correctly identified, Eq 1 describes that acquiring labels or features can improve generalization. We intended for this formalism to be as general as possible, as our work is the first to formalize the POCA problem.

We clarify that in a purely supervised learning setting, acquiring only features does not improve generalization performance. However, this does not preclude this application for other learning problems. Two prominent examples are semi- and self-supervised learning, where models use additional unlabeled data to enhance generalization (by learning underlying structures).

Actions taken: We have clarified our statement in L100 and explicitly referenced applications of Eq 1 in semi and self-supervised learning.

[P4] Clarifying Eq 3

Thank you for spotting this. We agree with the reviewer that in Eq 3, the input of the predictive model can be partially observed data.

Actions taken. We have updated Eq 3 to describe the predictive model as $p_{\phi}(y|\mathbf{x}')$, where $\mathbf{x}'$ represents partially observed data. We further clarify that while models that accept variable-length inputs (eg Transformer, random forests) naturally handle this, models that expect fixed-size inputs would require imputation.

[P5] Clarifying $\phi$

Effectively, $p_{\phi}(y|\boldsymbol{x})$ represents a likelihood over $Y$ when conditioned on the input $\boldsymbol{x}$. $\phi$ represents the model, which is specified by the functional form, the parameters $\omega$, and its distributions. For example, in a Bayesian neural network:

$p_{\phi}(y|\mathbf{x})=\int_{\Omega}p(\omega|\mathcal{D})p_{\phi}(y|\mathbf{x}, \omega) \, d\omega$

Where $p(\omega|\mathcal{D})$ is the posterior distribution over weights, $p(\omega|\mathcal{D}) \propto p(\mathcal{D}|\omega)p(\omega)$. This model is specified with an architecture and priors over weights, $p(\omega)$. As such, $p_\phi(y|\mathbf{x})$ is the posterior predictive distribution marginalized over $\omega$.

We note that the posterior $p(\omega|\mathcal{D})$ does not necessarily have the same distribution as the likelihood $p_\phi(y|\mathbf{x})$ (eg in BNNs, $p(\omega|\mathcal{D})$ are typically Gaussian, but the likelihood function could be categorical for classification). As the reviewer has pointed out, obtaining an analytical expression might require the use of conjugate distributions. However, in many practical cases, this is not necessary (eg by using variational inference).

Actions taken: We have clarified the meaning of $\phi$ in Eq 3 using a paraphrased version of this discussion.