Partially Observable Cost-Aware Active-Learning with Large Language Models

Authors: nicolás  astorga
tennison  liu
nabeel  seedat
mihaela  van der schaar
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 10:44:27

Abstract: Conducting experiments and gathering data for machine learning models is a complex and expensive endeavor, particularly when confronted with limited information. Typically, extensive _experiments_ to obtain features and labels come with a significant acquisition cost, making it impractical to carry out all of them. Therefore, it becomes crucial to strategically determine what to acquire to maximize the predictive performance while minimizing costs. To perform this task, existing data acquisition methods assume the availability of an initial dataset that is both fully-observed and labeled, crucially overlooking the **partial observability** of features characteristic of many real-world scenarios. In response to this challenge, we present Partially Observable Cost-Aware Active-Learning (POCA), a new learning approach aimed at improving model generalization in data-scarce and data-costly scenarios through label and/or feature acquisition. Introducing $\mu$POCA as an instantiation, we maximise the uncertainty reduction in the predictive model when obtaining labels and features, considering associated costs. $\mu$POCA enhance traditional Active Learning metrics based solely on the observed features by generating the unobserved features through Generative Surrogate Models, particularly Large Language Models (LLMs). We empirically validate $\mu$POCA across diverse tabular datasets, varying data availability, acquisition costs, and LLMs.

Main paper: 
 introduction: 
 In real-world machine learning (ML) applications, fully-observed, pristine training data is an exception rather than the norm. This challenge is especially evident during the initial stages of model development when training data is limited and varies in its informativeness across samples [1] [2] [3] . At this stage, obtaining additional data is crucial for improving model generalization but is fraught with challenges [4] [5] [6] . In particular, acquiring new data can be costly, often resulting in only essential features and labels being collected, leading to partially observed features in training data. Therefore, it's vital that acquisition is efficient, yet it remains unclear which features and labels from each instance will ultimately prove essential. Furthermore, data sources themselves can also be partially observed, with different features available across samples, further complicating the acquisition process. These challenges emphasize the importance of a new problem we call Partially Observable Cost-Aware Active-Learning (POCA) illustrated in Figure 1 . Before its formalization in Section 2, we provide an intuitive overview:"In situations with limited labeled data and partial feature observations, our objective is to enhance the generalization capabilities of a predictive model by strategically collecting features and/or labels. This goal should account the cost associated with data collection, as well as the varying levels of informativeness of labels and features across different instances"Addressing the POCA problem is vital when building systems with partial observation or relevant features are yet to be defined, particularly in fields like customer churn, monitoring, healthcare, and finance (see Appendix A). For example, developing a churn customer prediction system might start with some basic client information, such as demographics and income. However, to build such a system, additional features may be needed, which could be gathered through further customer interactions or surveys. At the outset, it's uncertain which specific features will prove essential, and acquiring additional information and relevant labels (e.g., churn events) necessary to refine the ML system involves costs related to money, time, or risks limiting the data acquisition in practice. From a practical perspective, we envision POCA to be useful in applications or fields where missing features exist, and also data acquisition techniques like Active Learning (AL) are necessary. Applications from different fields dealing with missing features and/or applying AL can be found in Table 2 . Related work. The most related data acquisition technique is AL [7] [8] [9] [10] . AL centers around enhancing model generalization through the acquisition of only additional labels. It operates under the assumption of having access to an initial small, fully observed training set (referred to as the historical labeled set) and seeks to acquire additional labels for samples from an unlabeled dataset (referred to as the pool set). This set is assumed to be fully observed in features, missing only labels. The distinctions between POCA and AL are illustrated in Figure 1 . Tangentially, Active Feature Acquisition (AFA) methods [11] [12] [13] have been proposed to enhance the prediction of individual samples at test time-where the sample is partially observed. Like AL, it assumes a fully-observed historical labeled set, on which a model has already been trained. Given the trained model, the task then becomes identifying the most relevant unobserved features to acquire for partially observed instances at test time. We emphasize that AFA's primary focus is on optimizing feature acquisition for individual test samples, differing from our broader goal of data collection to enhance model training.Towards an Instantiation of POCA. Given this problem definition, it is natural to wonder whether traditional AL metrics can be employed straightforwardly in the POCA setting. These metrics are usually derived from a predictive model that typically operates with fixed-size inputs. Consequently, predictions on partially observed instances can adversely affect the accuracy of AL metric estimations, leading to acquiring poor quality samples [14] . To overcome this challenge, we incorporate Generative Surrogate Models (GSM) to impute missing features in partially observed inputs, facilitating a more precise estimation of AL metrics. The effectiveness of GSM hinges on its ability to discern feature interrelations from available but unlabeled data. This task is particularly challenging due to the varying degrees of missingness in the instances and the constraints of limited sample sizes. To address these complexities, we employ Large Language Models (LLMs) to instantiate GSMs, utilizing their generation ability based on arbitrary conditioning and strong sample efficiency, allowing robust imputations to support the estimation of AL metrics under partial observability [15] [16] [17] Uncertainty POCA. We term this instantiation uncertainty POCA (µPOCA), due to its connection with Bayesian Experimental Design [18] [19] [20] [21] [22] , and its application in Bayesian Active Learning (BAL) [7, 23] and Bayesian Optimization (BO) [24] [25] [26] . From a Bayesian perspective, µPOCA maximizes the expected information gain or also known as expected uncertainty reduction, in the model's hypothesis resulting from an experiment. More specifically, µPOCA extends the concepts of expected information gain in the model's parameters (EIG) and expected predictive information gain (EPIG) to partially observed scenarios, introducing PO-EIG and PO-EPIG, respectively [7, 27, 28] . Here, these methodologies maximize the expected uncertainty reduction when acquiring labels and a subset of features. Since the impact of unacquired features cannot be directly assessed, GSMs facilitate the computation of these metrics.In summary, we make the following contributions: 1 ⃝ We address the unexplored challenge of costly data acquisition to enhance model generalization in partially observed scenarios. This leads us to introduce and formalize POCA, a novel ML paradigm for the acquisition of features and/or labels in the partially observed setting. 2 ⃝ We propose µPOCA, a cost-aware Bayesian instantiation of POCA that maximizes the uncertainty reduction when acquiring data. µPOCA extends traditional AL metrics by imputing partially observed instances using GSMs. We theoretically show that the uncertainty reduction is larger than using vanilla AL metrics.⃝ We propose the use of LLMs as a specific instance of GSMs, designed to address challenges in partially observed scenarios, including data efficiency, arbitrary information conditioning, and handling both categorical and numerical feature values.⃝ We empirically demonstrate µPOCA outperforms standard active learning on a variety of partially observability scenarios spanning datasets, sample availability, and acquisition metricshighlighting the usefulness and applicability of µPOCA.2 POCA: Partially Observable Cost-Aware Active-Learning Preliminaries. Partially Observable Cost-Aware Active-Learning is a data acquisition problem that focuses on improving the predictive performance of p ϕ (y|x) in the supervised setting, with ϕ the models we can employ. We denote x ∈ X and y ∈ Y as instances of observed features and target, alongside the respective random variables (RV) X and Y . Bold variables, expressed as x = {x j } J j=1 , represent a set of variables, in this case, features indexed by j ∈ [J] = {1, . . . , J}, where the bold form of j indicates a set of sub-indices j. The sample index i ∈ [I] = {1, . . . , I}, representing possible indexes in the pool set, is omitted when unnecessary, i.e., x i,j ≡ x j . We denote x o as the observed features with o ⊆ [J]. 1 In the general case, we assume that each feature x i,j considered for acquisition and the output of interest y i have associated acquisition costs c i,j and c i,J+1 . Here, c i,j represents the total cost of acquiring the variables indexed by j for instance i.In the context of partially observed data, our focus is on efficiently gathering features and/or labels to optimize a utility function, U t (•) subject to an acquisition constraint r t (•) at iteration t. U t (•) quantifies the trade-off between the costs of data acquisition and the increased generalization capabilities of the model ϕ, estimated from the available information x o and the hypothetical acquisition of a specific set of features and/or labels. We formulate the optimization of this utility as follows:(i, j) * = arg maxEQUATIONU t (•) is broadly defined, potentially estimated as result of using Bayesian techniques [18, 23, 27] , frequentist techniques [29] [30] [31] , RL techniques [32, 33] , or can even be subjectively defined through human desires. Note, optimizing U t (•) involves an iterative process of 1 selecting the instance and variables (i, j) * to acquire (features and/or labels); 2 adding these variables into the training set; 3 updating the model ϕ using the updated training set. In a more general case, this could also encompass batch acquisition [34, 35] by using i instead of i. Note that Eq. ( 1) represents the most general form of POCA, supporting model generalization when only features are acquired, as in semi-supervised or self-supervised learning. Our specific µPOCA instantiation (Section 2.1) focuses on the supervised case, where selected features and labels are acquired.Common modalities for POCA. We anticipate that most applications of POCA will center on the tabular domain (see Table 2 ). However, it could also find valuable uses in fields like medical and satellite imaging, where noise-induced occlusion is common. In these cases, determining when a sample requires additional information (features) is essential for enhancing prediction accuracy and model training. Likewise, interactive robots that learn through vision may benefit from this approach, as they need to discern which scenarios (samples) merit interaction to effectively learn the relationship between features (objects) and labels (task to solve). 
µpoca: a bayesian implementation of poca: 
 Although several techniques can be used to implement POCA, we opt for a Bayesian approach due to its widespread success in data acquisition literature. Building on the foundational principles of Bayesian Experimental Design [18, 23, 27] , which provides a comprehensive framework for integrating various sources of information [36, 37] , we introduce an instantiation of POCA within a Bayesian framework. This new approach, termed uncertainty POCA or µPOCA, leverages information theory [38] to recast Eq. ( 1) as a cost-aware uncertainty reduction problem [27, 39] . The core of µPOCA is centered on reducing uncertainty through a class of models that are exclusively trained using supervised learning, focusing on feature and label acquisition in partially observed scenarios. Here, we denote μi,j as the uncertainty reduction for acquiring the label and features j for sample i, which varies based on the approximation or method used.µPOCA We reformulate the optimization problem (1) by substituting U t with a utility function Ũ . This function Ũ is designed to capture the trade-off between uncertainty reduction, μi,j , and the acquisition costs associated with features and labels, represented by ci,j . The new objective can be expressed as:EQUATIONhere ci,j = c i,j + c i,J+1 . In our research, we explore one specific instantiation, among potentially infinite options, denoted by Ũi,j = μi,j and r(i, j) = ci,j < c, respectively, with c indicating the iteration's budget. 2How to obtain this uncertainty? We aim to minimize epistemic uncertainty [40, 41] by acquiring data, decreasing the predictive uncertainty produced by the possible hypothesis explaining the data. We work within the supervised model framework, hence we represent hypotheses as distributions over parameters. Our approach assumes the predictive model p ϕ (y|x ′ ) can be expressed as:EQUATIONwhere ω ∈ W is an instance of the parameter space and Ω its associated RV. Here, ϕ specifies the model choice, defining the functional form of p ϕ (ω) = p(ω|D), the posterior given the observed training set D, and the posterior predictive distribution p ϕ (y|x ′ ), marginalized over ω. Here, x ′ represents a partially observed input, so estimating p ϕ (y|x ′ ) must be adaptable to varying lengths of x ′ .To achieve this flexibility, models capable of handling variable-length inputs (such as Transformers) or, more broadly, marginalization techniques introduced in Section 3.2 can be employed.This formulation is general, encompassing Bayesian models, neural networks with certain stochastic parameters [42, 43] , and ensemble models [44, 45] . It also applies to Gaussian processes [46] when the posterior p(ω|D) is interpreted as a distribution over functions.3 Method: Optimizing µPOCAThe challenge in optimizing µPOCA is in developing an uncertainty reduction metric, μi,j , that accurately represents the decrease in uncertainty when acquiring a subset of features j for instance i, which has not been thoroughly investigated in the Bayesian literature. To address this, let's first provide some key background information. For data acquisition in ML, the primary focus has been on maximizing the expected uncertainty reduction, also known as expected information gain, when acquiring data [27] . This concept can be mathematically defined as:EQUATIONwhere H(A) quantifies the uncertainty (entropy) about A, and H(A|B) represents the uncertainty of A after observing B (in expectation).Existing AL approaches that utilize the expected reduction of uncertainty are summarized in Table 1 . These methods maximize the uncertainty reduction of I(G, Y |•) when Y is observed. Here, we use G to represent any random variable aligned with the generalization capabilities of the model and • any arbitrary conditioning. In the Appendix, for completeness, we derive the estimation for these acquisition metrics.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Section 5 discusses possible limitations with mitigation strategies.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
metrics for uncertainty reduction in partially observed scenarios: 
 Challenges in designing μi,j . In real-world scenarios, the challenge is estimating uncertainty reduction based solely on accessible data x o . Traditional AL acquisition metrics, denoted as µ ϕ (x o ), estimate uncertainty scores assuming x o ≡ x. However, in partially observed scenarios where only a subset of inputs, x o ⊆ x, is available, the observed features may lack sufficient informativeness for precise y estimates and reliable uncertainty scores µ ϕ (x o ). Generative Surrogate Model (GSM) to estimate metrics. A more accurate estimate of current metrics can be achieved using the aforementioned AL metrics by imputing the potential missing features in expectation:EQUATIONHere, the samples xj are obtained with a GSM denoted as p θ (x j |x o ), which sample possible unobserved features x j based on the observed x o . It's worth noting that training p θ (x j |x o ) could be done leveraging unlabeled data. In Figure 2 , we illustrate the acquisition process of µPOCA using GSMs.Learning? In Bayesian active learning, acquisition is closely linked to the concept of uncertainty reduction. To identify which features need to be acquired, it is essential to estimate the possible unobserved values. If these values lie in areas of high uncertainty within the hypothesis space, acquiring these features is beneficial, as it will help reduce this uncertainty. Conversely, if the possible values for certain unobserved features show little or no impact on uncertainty, then acquiring these features may not be necessary. Notably, deterministic imputation cannot achieve this, as the lack of variability prevents assessment of its effect on uncertainty within the hypothesis space. This concept is illustrated in Figure 11 from Appendix H.Are we doing better? We demonstrate the theoretical value of this approach for a family of acquisition metrics presented in Table 1 , delving into their impact on the optimization process. These propositions convey the intuitive idea that acquiring more information, in this case, features, leads to a higher reduction in uncertainty for the predictive model (proofs can be found in Appendix B). Proposition 1. Let µ(x o ) be an acquisition metric that can be written as I(G, Y |•), with G and • representing the same variables observed in traditional AL (Table 1 ), and with Y , X j as previously defined. If G ⊥ ⊥ X j |•, the following equality holds:EQUATIONCorollary 1. Under the assumptions of Proposition 1, the subsequent inequality is established:EQUATIONEquality is attained when G ⊥ ⊥ X j |Y, •.Proposition 1 states that the uncertainty reduction of G (e.g., the random variable of the parameters, Ω) by knowing Y and X j is equivalent to the expected uncertainty reduction achieved by knowing Y while conditioning on unobserved variables x j . This is convenient as the conditioning on x j can be computed using Monte-Carlo approximation [48] .Corollary 1 implies that acquiring both labels and features results in greater uncertainty reduction compared to acquiring only labels, the objective maximized in traditional AL (Table 1 ). The uncertainty reduction is equivalent when, given • and Y , the unobserved features X j don't have any impact in generalization G.Note that the independence assumption of Proposition 1 is valid in the supervised models we consider. In essence, this is because acquiring features without labels do not aid parameter updates and in consequence generalization improvements. The foundation of this assumption lies in the predictive mapping process from X → Y ← Ω, rather than in the data itself. Appendix B provides a more detailed explanation of this independence assumption's validity. Additionally, empirical evidence supporting the validity of Corollary 1 and, by extension, Proposition 1, is shown in Appendix K.Equations ( 6) and ( 7) always apply to the true random variable of unobserved features or any of its approximations. However, the terms in Eq. ( 6) reflect the uncertainty reduction of obtaining the actual features when the approximated distribution of the GSM accurately reflects the distribution of the true random variable. We empirically investigate this approximation and its practical utility.PO Active learning metrics.Building on Proposition 1 and Corollary 1, we extend BALD and EPIG as ▶ Partially Observable Expected Information Gain (PO-EIG): E x j I(Ω, Y |x j , x o , D) and ▶ Partially Observable Expected Predictive Information Gain (PO-EPIG): E x j I(Y eval , Y |x j , x o , X eval , D). Corollary 1 states that these metrics provide a higher uncertainty reduction than their vanilla counterparts. We use Monte-Carlo for estimation (see Appendix C).Our derivations are based on a distribution perspective, considering different numbers of conditioned variables. For instance, when calculating PO-EIG, expressed asE x j I(Ω, Y |x j , x o , D), it is necessary to compute the distribution p ϕ (y|x o , x j ).Here, x o could vary in length from one instance to another and x j varies based on the number of features considered for computing the uncertainty reduction metric. In practical terms, this means that the predictive model, attempting to approximate this distribution, must effectively handle inputs with varying variables and lengths.To address this challenge, we employ GSMs to impute the missing information to enable predictive models that expect fixed-size inputs. This imputation is separated in two different steps (1) conditioning and (2) marginalization. Essentially, when evaluating the uncertainty reduction of an unobserved subset of features x j considered for acquisition, we condition on this subset x j and x o (the observed features), marginalizing over the remaining subset of unobserved features x j ′ (where x j ∪ x j ′ is the set of all unobserved features). This approximation process is mathematically formalized as follows, with supplementary visual aids provided in Figure 8 of Appendix C.3:EQUATIONHere the predictive model simulates the behavior, wherein the predictive model only has access to x o and x j but it is computed using a model as it would have all the features. The marginalization step is essential for accurate metric estimation in the pool set and can also be applied during training. However, to reduce costs, we use GSM to impute features not acquired in the training set.Algorithm 1 Acquisition process1: P = [ ], F = [ ] 2: for i ∈ [I] do 3: j * = [J] 4:while r(i, j * ) do5: v * = arg max v∈j * Ũi,j * \v s.t. r(i, j * \v) 6: j * = j * \v * 7:end while 8:P.add(μi,j * ), F.add(j * ) 9: end for 10:i * = arg max i∈[I] P [i], j * = F [i * ] 11: Return: (i * , j * )Our goal is to maximize Ũi,j (•) ≡ Ũ (μ i,j , ci,j ), which incorporates the uncertainty reduction μi,j . It is crucial to recognize that μi,j could encompass all possible combinations of unobserved features. However, computing μi,j for every possible combination of (i, j) is impractical, since it is of order O(2 J ).To overcome this challenge, we propose estimating the uncertainty reduction for all unobserved features and subsequently excluding the less relevant ones, i.e. those contributing minimally to uncertainty reduction. This ensures that we always retain the most relevant features until the constraint r in Eq. ( 2) is satisfied, in order O(J 2 ). The acquisition process is summarized in Algorithm 1, with feature selection steps highlighted in teal. Appendix C.3 provides details on an efficient approach to computing the marginalization step necessary for estimating μi,j . This efficiency can be further improved by selecting the most informative samples, followed by the application of Algorithm 1 (see Appendix D). For a comprehensive overview, including cost analyses, and details on GSM training and sampling, refer to Appendix D.LLMs as GSMs. For the scenarios outlined in POCA, we specify the following desiderata for GSMs: (P1) generative capability, (P2) ability to learn from partially observed data, (P3) sample efficiency, and (P4) seamless integration of mixed-type variables. We argue that LLMs are well-suited to meet these criteria due to their ▶ generative capabilities and flexibility in training under ▶ arbitrary conditioning contexts [49] [50] [51] . Moreover, recent research highlights their exceptional performance in ▶ few-shot settings [52, 53] and their generative capabilities applied to ▶ tabular data comprising mixed-type attributes [51] . These strengths provide strong justification for focusing our research on LLMs as GSMs. However, any imputation method that fulfills these criteria may also serve as a suitable GSM, as further discussed in Appendix G.We use LLMs as GSMs leveraging the unlabelled information via Supervised Fine-Tuning (SFT). When working with tabular data, we serialize rows of the data, thereby converting it to natural language. For example, a set of features is serialized as "Age is 25, Gender is Female, . . . , Blood pressure is 0.57". The LLM is then used to predict unobserved features based on available information. To achieve this, we utilize SFT on the LLM with the available observed features. The training data can encompass all unlabeled data, including historical and pool set data. The process entails generating random masks to form an input, m ⊙ x o , and an output, (1m) ⊙ x o , for SFT across all available data. This empowers the LLM to predict missing information by leveraging various combinations of observed features. 3 For more details, refer to Appendix F.2.The effectiveness of µPOCA in partially observed settings is closely tied to the GSM's ability to approximate the distribution of unobserved features. Two primary factors influence the accuracy of this approximation: (1) the approximation capacity of the GSM and (2) intrinsic characteristics of the dataset. A detailed examination of these factors is provided in Appendix J.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 We evaluate µPOCA across three dimensions4 : First, in the case that all features are acquired, we demonstrate that µPOCA acquisition metrics are more informative in selecting instances with informative features than AL metrics. Second, we present a synthetic experiment accompanied by theoretical insights. Finally, we explore scenarios with budget constraints demonstrating that µPOCA on more challenging scenarios.Comparing µPOCA with the current AL models is complex, as the latter are designed for fixed-size inputs. To address this challenge, we developed Scenario 1 (see visual aid in Appendix H). This scenario involves dividing each instance in the pool set into the same observed and unobserved feature sets. We specifically select half of the features to remain unobserved, chosen by their high relevance to the predictive task as identified by a preliminary RF. It is important to note that while µPOCA methods can handle any form of missing data, Scenario 1 ensures a fair comparison by allowing AL models to operate without any modification, which could bias our evaluation. This scenario presumes the availability of a historical unlabeled dataset for training the GSM, using instances that include data on unobserved features. In practical applications, the pool set can often serve as the training set itself, representing a more realistic scenario we may encounter. We refer to this setting as Scenario 2.Results for this scenario are presented in Appendix I, where GSM is trained on partially observed data, while vanilla AL employs deterministic imputation to manage this case.We select Magic, Adult, Housing, Cardio, and Banking tabular datasets based on their use in AL [28] , tabular generation [49, 51] , LLM-based classification [54] , and relation with potential real-world A ). These datasets have diverse characteristics: sample size, number of features, number of categorical, and numerical variables. We prioritize datasets with over 1000 samples to guarantee sufficient samples for the pool set. We showcase results using a RF trained with 100 estimators. We start training with two fully observed samples per class, conduct 150 acquisition cycles, repeat each experiment over 60 seeds, and display a 95% confidence interval. We train Mistral7B-Instruct-v0.3 using 8 Monte-Carlo samples for generative imputation.Objective. To assess the need for more generalized methodologies such as µPOCA, we analyze the performance of PO-EIG, a partially observed extension of BALD (EIG)-the most widely used metric in active learning literature. Additionally, we incorporate EPIG into the study, a recently developed active learning metric within the 1 family. According to our theoretical framework (see Corollary 1), PO-metrics outperforms their vanilla counterparts in terms of uncertainty reduction.Our goal is to examine whether this uncertainty reduction leads to improved downstream performance when all features are acquired based on the same information, x o , or, in other words, if the selected instances possess features that are more relevant.Setup. To ensure a fair comparison, we evaluated PO-metrics and Vanilla-metrics under Scenario 1, using Random and Oracle as reference baselines. Here, Oracle represents the Vanilla-metrics acquisition metric, but with access to all features. Ideally, when GSM functions optimally, the performance of PO-EIG should align with that of Oracle.Analysis. The first thing to note is that EIG metrics computed with partially observed features can be significantly worse than simple baselines like random as shown in Magic dataset from Figure 3 ) (top). Figure 3 (top) demonstrates that PO-EIG generally either outperforms or worst case matches their fully observable counterparts BALD across all datasets. A similar behavior is observed for PO-EPIG, which generally outperforms their vanilla metric counterpart. This suggests that an increase in uncertainty reduction translates into an increase in downstream performance. While PO-EIG and PO-EPIG metrics consistently outperform baselines, they occasionally fall short of oracle performance, notably in the Housing datasets. This may stem from two factors: Firstly, the GSM has poor prediction performance on the unobserved data due to insufficient data or model capacity. Secondly, even with adequate capacity and data, weak correlation between unobserved data and the target hinders the acquisition process. We study these factors in Appendix J. We note that it is non-trivial to quantify the GSM's capability or correlations of unobserved data to the target. Thus, the practical implication is that both PO-metrics should be preferred in PO settings, providing a performance boost or at least matching their vanilla counterparts. Additionally, in Appendix I.1, we include other relevant Active Learning metrics that, while not fitting into the family of studied metrics, also demonstrate performance gains with the proposed framework.First, AL metrics computed on partially observed features can dramatically fail for selecting relevant instances. Second, PO-EIG and PO-EPIG generally outperform or match fully observed counterparts.Objective. We investigate the implications of our theoretical findings (Eq. ( 7)) on the acquisition process; determining whether a weak correlation between unobserved features and the target, results in a small gap between PO-EIG and BALD. We also explore how correlation affects performance.Setup: We create an intuitive synthetic 2D experimental setup (Figure 4 ) with a variable target. The target is determined linearly with varying slopes, leading to different correlations with the features. Our chosen features-X 1 , X 2 , and X 3 -represent data along the x-axis, y-axis, and a Gaussian category, respectively. Introducing the Gaussian category injects stochasticity into the marginalization process, ensuring non-trivial solutions. The observed feature is X 1 , with possible acquisition of X 2 , X 3 . We examine three scenarios: 1) Low Corr(X 2 ,Y ), where the class depends solely on X 1 due to vertical slope; 2) High Corr(X 2 ,Y ), where the class depends solely on X 2 , rendering X 1 irrelevant; and 3) Mid. Corr(X 2 ,Y ), where X 1 has some impact. Note, we evaluate acquisition metrics and performance until the convergence of the oracle (BALD with all features) Analysis. Figure 5A empirically validates that PO-EIG is always equal to or greater than BALD, consistent with our theoretical insights (Eq. ( 7)). Figure 5B illustrates the evolution of the metric gap between PO-EIG and BALD under varying correlations between X 2,3 and Y . In low correlation scenarios (orange line), the gap diminishes towards the acquisition's end, aligning with Corollary 1 where both metrics should converge when G ⊥ ⊥ X 2,3 |x 1 , •, i.e., when the unobserved features don't impact generalization. Initially, the gap exists as the model learns from data the redundancy of unobserved features. The same figure shows larger correlation leads to a wider gap, observed most notably in the large correlation scenario (purple) and moderately in the medium correlation scenario (teal). Figure 5 shows that, generally, the degree of problem correlation provides a proxy correlation with acquisition performance. For example, the purple line exhibits the largest difference between BALD and EIG. Particularly in low correlation scenarios, the performance difference between PO-EIG and BALD is negligible across the acquisition (orange line).Low Corr X2 YObjective: We evaluate the performance of µPOCA (specifically PO-EIG) under budget-constrained feature acquisition, aiming to determine if acquiring only a subset of features, denoted as j, offers an advantageous trade-off in performance. This selective acquisition approach enables acquiring a larger number of instances within the same budget. Furthermore, we aim to show that imputation alone cannot fully replace the need for direct data acquisition.Setup. We use the Magic dataset as a case study to examine the impact of cost constraints on predictive performance and the feature acquisition process. To facilitate this assessment, we introduce costs associated with both features and labels. For simplicity and visualization clarity, we assume the cost of an instance to be 1, representing the sum of the costs for all features and the label, with each feature assigned an equal cost. This setup allows us to analyze four distinct approaches: (1) the Vanilla acquisition metric (EIG), (2) PO-EIG, (3) PO-EIG with a maximum feature acquisition limit of 60%, and (4) PO-EIG with unrestricted feature acquisition. We evaluate the performance of these approaches in three ways: by accuracy based on acquired instances (Figure 6 , left), by performance relative to the budget utilized (assuming no label costs) (Figure 6 , middle), and by performance with varying label costs under a fixed total budget of 50 (Figure 6 , right).Analysis. Figure 6 (left) illustrates that acquiring fewer features generally results in decreased performance; however, it still outperforms the EIG baseline. While limited feature acquisition impacts performance, it allows for a more efficient budget allocation across instances, enabling the acquisition of a larger instance pool. This trend is visible in Figure 6 (middle), where performance is plotted against the total budget spent, assuming no label cost. Here, methods focusing on selective feature acquisition excel, as they gather more overall information through increased instance count and key features. Figure 6 (right) demonstrates that the optimal PO-EIG method depends on feature acquisition cost: when cost is heavily weighted toward feature acquisition (right histogram), the best method is PO-EIG with 20% of features acquired, whereas a 50% label cost favors PO-EIG with 60% feature acquisition. While these findings might suggest that acquiring fewer features and imputing the rest is optimal for maximizing instances, this approach may introduce noise into the training set, potentially biasing the model. To explore this, we analyze model performance at different levels of feature acquisition in Figure 7 , with varying levels of pool data, using the full pool set for training (excluding non-acquired features). As shown on the y-axis, acquiring more features enhances performance. When the budget is unlimited, acquiring all available data is preferable; however, in practice, this may not be feasible, making POCA approaches advantageous.First, µPOCA metrics (PO-EIG) can be more cost-effective than common active learning metrics. Second, imputation is useful for missing data but shouldn't replace data acquisition. 
discussion: 
 We introduce and formalize POCA a data acquisition framework, addressing the vital but underexplored challenge of partially observed settings. Through µPOCA, a practical implementation of this framework, we demonstrate the feasibility of acquiring unobserved features and labels based on those partially observed features, using more generalized AL utility metrics -computed by estimating features generated using an LLM-based GSM. Our results over various scenarios are substantially more effective than alternatives -of substantial value for data acquisition in cost-restrictive environments. We hope the POCA framework and our subsequent findings will spur additional work to advance data acquisition in partially observed settings.Limitations. Our work focuses on the values of features, providing a general framework where restrictions are the main source of constraints in terms of acquisition. However, we do not assess how these restrictions are selected, which could be a promising area for future research. We also note that we use LLMs in the context of data acquisition. Like any GSM, LLMs can indeed exhibit biases that affect the acquisition process. In this study, we did not consider this issue, and it represents an interesting avenue for future work. If necessary, current debiasing techniques can be applied.In the PO setting with data "missingness," GSM imputation is essential for acquisition. Future work could quantify uncertainty [43, 55] to assess GSM efficacy. (2) LLM capability also impacts acquisition; while we use a 7B-parameter model, larger models could further enhance performance, though this is beyond our current scope.Appendix: POCA: Partially Observable Cost-Aware Active-Learning with Large Language Models.A Appendix A: Real-World Use Cases Table 2 : Real-world use-cases of POCA. We outline real-world scenarios where the POCA framework can have an impact. For each problem domain, we describe partially observable features, labels, and the underlying predictive task. We categorize references into three types: A) where active learning is employed, B) where predictive modelling is performed in the presence of partially observed features, and C) active learning is applied to partially observed settings (with data pre-processing to handle missing features). The symbol ▶ stands for acquisition costs. Results from specific sensor data (soil moisture, specific pollutant levels) ▶ Operational costs of measuring data Environmental condition classifications ▶ field surveys, lab analysis of samples A: [95] , B [96] Table 2 illustrates that active learning is extensively utilized in a variety of real-world application scenarios. Furthermore, it is not uncommon in these contexts to encounter situations with incomplete data, which can harm generalization [2, 6, 97, 98] capabilities of downstream models. The breadth of related work covers diverse sectors including customer churn prediction, marketing research, healthcare diagnostics, and predictive maintenance. While active learning is adept at selectively querying labels in scenarios where data is fully observed, its application in the context of missing data is less clear. The challenge is compounded by the fact that, in similar problem settings, it is not always guaranteed that features will be fully observed. This reality underscores the need for alternative machine learning techniques to address such challenges. Our proposed approach, POCA, offers a novel solution for applying active learning in scenarios with partially observed data, taking into account realistic cost constraints.B Appendix B: Acquisition metrics for partially observed scenarios.Proposition 1. Let µ(x o ) be an acquisition metric that can be written as I(G, Y |•), with G and • representing the same variables observed in traditional AL (Table 1 ), and with Y , X j as previously defined. If G ⊥ ⊥ X j |•, the following equality holds:EQUATIONProof: We can decompose the left part of Eq. ( 9) as:EQUATIONUsing G ⊥ ⊥ X j |• =⇒ I(G, X j |•) = 0, the first term on the right of Eq. ( 10) cancels, obtainingI(G, (Y, X j )|•) = I(G, Y |X j , •) = E xj I(G, Y |x j , •) concluding the proof.Corollary 1. Under the assumptions of Proposition 1, the subsequent equivalent inequality is established:EQUATIONEquality is attained when G ⊥ ⊥ X j |Y, •.Proof: Symmetrically as before we can decompose the left part of Eq. ( 9) as:EQUATION) Using proposition (1), we obtain:EQUATIONClearly the equality is obtained when G ⊥ ⊥ X j |Y, • since the mutual information is zero. When taking I(G, X j |Y , •) ≥ 0 we obtain:EQUATIONconcluding the proof.Observation: Note that, the independence assumption G ⊥ ⊥ X j |• always hold for the class of supervised learning models we consider. G is a random variable representing the generalization capabilities of the model ω ∈ W, with random variable Ω. This random variable is subject to model training A, which can be written as the mapping between the training set D and hyperparameters h ∈ H to the output ω, i.e., A : D × H → W. Additionally, the model prediction is a mapping P between the model ω and the input x ∈ X to the output Y , i.e., P : W × X → Y.It is crucial to acknowledge that the "world generator" influences D, X, and Y , but does not directly affect Ω. Given that D is observable, any connection through this path is cut. The sole connection of Ω to X and Y is through the mapping P, which establishes a causal structure:Ω → Y ← X.According to this structure, Ω and X are generally independent unless Y is observed, leading to a dependence due to P creating a configuration known as an "immorality" among these variables. This explains why G ⊥ ⊥ X j |• holds; however, this independence may not persist in scenarios whereG ⊥ ⊥ X j |•, Y .In the context of our work the input X can be decomposed in X = x o , X j the observed part of the random variable and the unobserved part of the random variable. Consequently, X j follows the same independency assumptions of X.In our we evaluate BALD and EPIG, along with their partially observed counterparts, PO-EIG and PO-EPIG. For completness, we show the estimatation the vanilla metrics, and later their corresponding partially observed extensions.We follow a similar notation to [28] . For categorical variables, BALD can be decomposed as:EQUATION≈ -EQUATIONwhere K represent the total number of parameter samples ω k ∼ p ϕ (ω) from the posterior distribution given D. Here,EQUATIONPO-EIG extends this formulation to include missing features as conditioning samples xj from the GSM, accommodating partially observed settings. To illustrate, let's first consider the case where the metric is conditioned over all unobserved features, i.e., j ′ = ∅, which correspond to results of Figure 3 . Following a similar decomposition, we can approximate E xj I(Ω, Y |x o , D, xj ) as:EQUATIONhere L represent the total number of Monte-Carlo samples from the GSM, with xl j one possible sample. Here,EQUATION)What happen when j doesn't consider all unobserved features?. This scenario is useful when we want to asses the impact acquiring of subset of the unobserved features. Utilizing a smart notation, this approximation can be stated identically as Equation ( 19), but the estimation of the predictive distribution changes slightly:EQUATIONhere, P are a total of new MC samples from the GSM. This marginalization trick is necessary to deal with model that expect fixed size inputs.For PO-EPIG and we replace the sub-index eval by * . Similar to before, for categorical variables, EPIG can be decomposed as:EQUATIONEQUATIONhere,EQUATIONEQUATIONEQUATIONSimilarly as before, PO-EPIG extends this formulation to include missing features when conditioning on samples xj from the GSM. We now derive the estimation of this metric when estimating PO-EPIG when considering all the unobserved features j. We can decomposeEQUATION)EQUATIONEQUATIONEQUATIONNote that the marginalization step of the "evaluation set" is also necessary since we also assume it to be partially observed. Here, x h j are Monte-Carlo samples necessary make the predictive model receive fixed inputs. While expensive, PO-EPIG can also compute metrics when only a subset of features is estimated to be acquired. The trick is similar to PO-EIG, being necessary a marginalization across the j ′ for the the predictive mode receive fixed-size inputs (as ( 22)).As in the earlier sections, we approximate the PO-metrics through the techniques of conditioning and marginalization. Figure 8 provides a visual representation of this approximation.In Eq. ( 8), for efficient computation in the marginalization of j ′ , we substitute p θ (x j ′ |x o , x j ) with p θ (x j ′ |x o ), using the same samples used to compute the uncertainty reduction across all unobserved features. This approximation does not significantly compromise precision, as we exclude the least relevant feature sequentially. Consequently, the samples used to marginalize an "irrelevant" feature remain minimally impacted by the overall sampling strategy. Dl ← impute_missing_data(Dl, θ) # Imputing when using a predictive model that receive fixed-size input 5:φ ← maximize_likehood(ϕ, Dl) 6:if use_heuristic then 7:I * ← compute_top_uncertain_instances( φ, PS, [I], [J])) 8:end if 9:P = [ ], F = [ ] 10: for i ∈ [I] do 11:if use_heuristic and i / ∈ I * then 12:P.add(-inf), F.add([]) 13: continue 14:end if 15:j * = [J] 16:while r(i, j * ) do 17:U = [ ] 18:for v ∈ j * do 19:if r(i, j * \v) then 20:μi,j * ← compute_uncertainty( φ, PS, i, j * \v)) 21:U.add( Ũ (μi,j, ci,j)) 22: else 23:U.add(-inf) 24:end if 25:end for 26:v * = arg max U 27: j * = j * \v * 28:end while 29:P.add(μi,j * ), F.add(j * ) 30:end for 31:i * = arg max i∈[I] P [i], j * * = F [i * ] 32:Dl ← Dl ∪ P[i * ][j * * ] 33:P ← P \ P[i * ] 34: PS ← PS \ PS[i * ] 35: I ← I -1 36: end whileWe assume 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), respectively. This assumption enables us to train the GSM and generate imputed values for the pool set before the acquisition process begins. The complete acquisition is detailed in Algorithm 2.We define the generation cost as C g and the downstream cost as C d . Algorithm 2 indicates that the sampling cost for GSMs is O(I • J • S • C g ), where I is the number of instances, J is the number of features, and S represents the number of Monte Carlo samples. The inference cost of the downstream model isO(I • J • S • C d ).The cost of acquiring a subset of features depends on the restriction r(i, j), which is bounded by J (the case where J features are discarded for acquisition). The cost for acquiring a subset of features for a single instance is O(J 2 • S • C d ) using Algorithm 1. To reduce this overhead, we first select the most informative instance (assuming all features are acquired) inO(I • J • S • C d ) (L7 in Algorithm 2), and then select the subset of features to acquire using Algorithm 1 (10 in our case) in orderO(J 2 • R • S • C d ).This analysis shows that the most critical factor is the number of features J, as it affects both sampling and the downstream model (quadratically in this case).Another consideration is when the available data is insufficient for a reliable GSM. In this case, any additional features acquired during the acquisition process can be used to update the GSM weights periodically, which increases the costs.Datasets were constructed such the number of pool samples were numerous enough to determine the impact of the acquisition performance any confounding effect. The distribution of pool set were selected maintaining the distribution of the original dataset similar to previous work [28] . The "evaluation distribution" in EPIG is follows the same distribution of pool set. As mentioned in the text, we maintain a small unlabeled historical set for GSM training allowing fair comparison with Active Learning. For large datasets like Adult and Housing we limit their maximum original data size avoiding unnecessary costs in Monte-Carlo sample generation. The test set is defined as the 30 % of the intial dataset.• Magic [99] : Original data size of 19020 samples. Historical set of 1000 samples. Pool set distribution; Class0: 4980 samples. Class1: 2700• Adult [100] . Original data size of 19020 samples (after cut). Historical set of 1000 samples. Pool set distribution; Class0: 5760 samples. Class1: 1920.• Housing. Original data size of 19020 samples (after cut). Historical dataset of 1000 samples. Pool set distribution; Class0: 3840, Class1: 3840.• Cardio [101] . Original data size of 100k samples. Historical dataset of 1000 samples. Pool set distribution; Class0: 3000, Class1: 3000.• Banking [102] Original data size of 45211 samples. Historical dataset of 400 samples. Pool set distribution; Class0: 2000, Class1: 500.We did slight modification in datasets allowing LLM better comprehension. For Magic, we didn't include fConc1 since its name similarity with fConc. For CMC dataset, we change categorical values 0 and 1 to the categories that were represented on metadata. For example, in column "Standard-of-living_index" instead of using 0,1, 2, 3 we use Low, Medium-Low, Medium-High, and High.With this adjustment the final columns of these datasets are:• Magic. 9 numerical columns: fLength, fWidth, fSize, fConc, fAsym, fM3Long', fM3Trans, fAlpha, fDist.• Adult. 8 categorical columns: workclass, education, marital-status, occupation, relationship, race, sex, native-country. 6 numerical columns: age, fnlwgt, education-num, capital-gain, capital-loss, hours-per-week.• Housing. 8 numerical columns: MedInc, HouseAge, AveRooms, AveBedrms, Population, AveOccup, Latitude, Longitude.• Cardio. 10 numerical columns: ID, age, age_years, height, weight, ap_hi, ap_lo, bmi. 7 categorical variables: gender, cholesterol, gluc, smoke, alco, active, bp_category_encoded.• Bank. 2 numerical columns: age, balance. 6 categorical variables: job, marital, education, contact, day_of_week, default.For model training, we use one hot encoding for categorical variables.We analyzed primarily one open-source Mistral7B-Instruct-v0.3 with more in Section J. Preliminary findings show that instruct LLMs generally outperformed their non-instruct counterparts. In what follows, we give component used and a general overview of the prompts used:• Input: The input is the tabular information provided to the LLM. Here, it is in the form of "{feature A: value A, feature B: value B, feature C: value C}", for possible features and values.• Query: The query represents the information requested by the LLM to generate a response. It should only include the attribute names, formatted as follows: "{feature D: [Mask], feature E: [Mask], feature F: [Mask]".• Response: The input is the tabular information provided to the LLM. Here, it is in the form of only values: "{feature D: value D, feature E: value E, feature F: value F}".The specific prompts were adjusted to specific LLM format.F. We tested two masking strategies. First, when we had access to all the historical data, particularly when comparing against traditional AL in the main text. Second, when certain features might be missing, as encountered in Appendix I. In the first scenario, we randomly masked some observed information, always ensuring at least two features as input to prevent overly complex tasks for the LLM. In the second scenario, due to the partial observability of data, some features could be less observed than others, leading to varying degrees of missingness.We train the models using QLoRA using 4 bit quantization, r = 32, lora_alpha= 64, for 10000 steps, with 2 samples per batch size, and a learning rate of 7.5e-5.In this section, we outline the core desiderata for Generative Surrogate Models to be used within the µPOCA framework and compares related imputation methods against the desiderata.• [P1] Generative capability: The model must model a non-deterministic distribution over the unobserved features to effectively identify and prioritize the most relevant features. To exemplify, we consider x j * as the subset of current features under review for acquisition, and x j ′ as the features already excluded from acquisition. PO-EIG's acquisition metric strategically selects the feature,v * , that maximizes uncertainty reduction among all possible features, v, considered for exclusion from the set j * . This effectively minimizes information loss. Formally, this is expressed as:EQUATIONIn contrast, employing a deterministic GSM modifies the acquisition strategy for imputed features x, simplifying to:EQUATIONThe second equality arises from the fact that we are limited to a single value due to estimated conditioning and marginalization. Consequently, any permutation of conditioning sets j and marginalization sets j ′ will lead to the same metric for a deterministic GSM. (34) explicitly demonstrates that employing deterministic GSM results in an acquisition metric independent of v, the feature under consideration for exclusion through a greedy approach. In conclusion, the GSM should model a stochastic distribution for acquisition purposes. [P1] ✓ ✓ ✗ ✗ [P2] ✓ ✓ ✓ ✓ [P3] ✓ ✗ ✓ ✗ [P4] ✓ ✗ ✗ ✗Flexibility of Large Language Models. We believe that Large Language Models hold significant potential as sample-efficient methods for feature estimation. Although LLMs are not specifically pretrained for tabular data estimation, their general capabilities have demonstrated effectiveness across various domains. These general abilities have been successfully applied in optimization [24, 117] , reasoning [118, 119] , planning [120] [121] [122] , concepts [123, 124] , autonomous adaptation [125] [126] [127] , digital twin construction [128, 129] , and as tabular data generators [49] . They have also shown promise as few-shot tabular learners in supervised settings [52, 53] . The application of LLMs in the intersection of tabular data tasks, particularly in the last two contexts, provides direct evidence of their viability in this area and to be used as GSMs in the context of data acquisition.[G2] Performance across varying data characteristics Setup. Next, we turn our attention to investigating how the data distribution affects acquisition performance. We are particularly interested in analyzing the effect of the correlation between unobserved features X unobs and observed features X obs on acquisition performance.To demonstrate this, we examine a scenario where (1) X unobs correlates with the outcome, while (2) observed features do not. In this context, the GSM becomes crucial for downstream performance, as X obs alone provides insufficient information to predict outcomes accurately. As such, the GSM must effectively model the relationship between X obs and X unobs to acquire missing features critical for predicting the outcome.We model both X obs and X unobs as two-dimensional random Gaussian variables centered at zero and establish a specific controllable correlation between them through ρ:Σ = I 2 ρ X obs X unobs I 2 ρ X obs X unobs I 2 I 2The label Y is then constructed to be independent of X obs using the orthogonalization:X orthogonal = X unobs -X obs (X T obs X obs ) -1 X T obs X unobswhich we use to construct the label using logits = 1 1 + e -Xorthogonal , C = 1 logits>0 Analysis. Figure 15a illustrates how varying ρ between X obs and X unobs empirically affects variable correlation, c, validating our synthetic experiment design. Figure 15b analyzes EIG (traditional active learning without GSM) and PO-EIG with varying ρ. We note that when the correlation between X obs and X unobs is low, GSMs provide no performance benefits. However, as correlation increases, the performance gains of PO-EIG over EIG expand significantly, confirming our hypothesis. Justification: The abstract and introduction accurately reflects the claims made in the paper. We formalize the POCA problem and provide an instantiation of POCA which we validate empirically.Guidelines:• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: Not applicable -our paper presents a process for data acquisition -which does not fall into this category.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: Appendix D provides details and/or citations for all the open source data assets used in the paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.In contrast with AL, POCA assumes xo ⊆ x instead of the fully observed assumption of AL xo ≡ x. In addition, POCA considers the acquisition of features and/or labels, in this case, represented as j.Alternative utility functions may balance uncertainty against costs as Ũi,j = μi,j /c i,j . Other constraints could consider c as the overall experimental budget.Without loss of generality, in-context learning is viable for an LLM-based GSMCode can be found at: https://github.com/jumpynitro/POCA or https://github.com/ vanderschaarlab/POCA 
 
 Appendix:  


Decision: The paper introduces 𝜇POCA, a novel framework for optimizing data acquisition in scenarios with partial observability and high acquisition costs. 𝜇POCA, the proposed instantiation, leverages Generative Surrogate Models (GSMs), particularly Large Language Models (LLMs), to impute missing features and maximize uncertainty reduction during acquisition. The reviewers agree that the problem addressed is timely and that integrating GSMs for feature imputation is a novel contribution with practical relevance.

However, the reviewers also highlight several areas for improvement. The reliance on Monte-Carlo sampling introduces significant computational costs, and the paper's clarity could be enhanced. Additionally, the empirical validation is primarily focused on tabular data (in the original submission), limiting the generalizability of the findings. The authors have provided a rebuttal addressing most of these concerns. If accepted, the authors are encouraged to refine the theoretical analysis and broaden the scope of their experiments.

Reviews: review: paper_summary: paper_summary: This paper introduces a novel active learning framework for optimizing data acquisition in scenarios with partially observed features and high costs. The proposed method, POCA utilizes GSM, specifically LLMs, to impute missing features and improve active learning metrics. By integrating Bayesian methods to maximize uncertainty reduction, the proposed active learning method can select features and labels to enhance predictive performance while minimizing acquisition costs.
strength_weakness: strength_weakness: 1. The Monte-Carlo-based approach in POCA would be slow and with high computational cost due to its reliance on repeated random sampling to estimate quantities of interest.
2. This model is limited to tabular data.
3. Sometimes the performance of POCA is not as good as random sampling, POCA may vary depending on the specific characteristics of the dataset and the underlying correlations​.
4. Please consider more active learning methods as baselines except for BALD and random sampling.
questions: questions: There are many typos, the author should carefully check the full manuscript. For example, Line 721 "?." and line 723 "Figure ??".
limitations: limitations: The authors have adequately addressed the limitations and potential negative societal impact of their work.
paper_summary: paper_summary: This paper  addresses the challenge of efficiently gathering features and labels in partially observed settings to enhance model generalization while considering data acquisition costs. It introduces POCA and its Bayesian instantiation, leveraging Generative Surrogate Models (GSMs) to impute missing features and compute uncertainty-based active learning metrics. The paper demonstrates the practical usage of µPOCA through empirical validation across different datasets.
strength_weakness: strength_weakness: 1.While the practical approach is well-formalized, the theoretical contributions and novelty are somewhat limited. The problem addressed by POCA has similarities with experimental design, and the extension of existing active learning methods to partially observed data is a relatively straightforward adaptation.
2.The paper does not adequately address how partial observability of features impacts the overall performance and reliability of the predictive model. There is a need for a deeper exploration of the theoretical implications and limitations of using GSMs for feature imputation.
3.In Equation 1, it does not seem to require that the data samples' labels must be annotated. If a sample only reveals a few features, how does such annotation help supervised learning? The only benefit I can think of is for the generative model, but the authors did not explain this clearly.
4.The description of the predictive model in Equation 3 is inaccurate. According to the authors' definition and the footnote on page 3, bold bold x represents fully observed features, but the predictive model should generally accept partially observed data.
5.The meaning of p_{phi} is unclear to me. In line 92, the authors mention that phi is a model we can employ. Does this mean that p_{phi}() represents a specific density function (e.g., Gaussian)? If so, how can we ensure that both the likelihood and posterior follow the same distribution, whose density is denoted by p_{phi}? The authors did not impose any conjugate restrictions.
questions: questions: In Equation 1, it appears that the annotation strategy does not require that data samples' labels be annotated. If a sample only reveals a few features, how does such annotation benefit supervised learning? Is the primary benefit for the generative model? If so, could you provide a clearer explanation of this benefit?
limitations: limitations: NA
paper_summary: paper_summary: This paper introduces a novel problem setting in Active Learning (AL) called Partially Observable Cost-Aware Active-Learning (POCA). POCA deals with situations where acquiring both features and labels comes with a cost and where datasets are partially observed, meaning some features might be missing for certain data points. The authors propose µPOCA, an instantiation of POCA that uses Generative Surrogate Models (GSMs), specifically Large Language Models (LLMs), to impute missing features and estimate the uncertainty reduction from acquiring both features and labels. This uncertainty reduction then guides the acquisition process, aiming to maximize the model's generalization capability within a given budget. The authors theoretically show that acquiring both labels and features leads to greater uncertainty reduction compared to acquiring only labels (the traditional AL approach) and provide empirical validation of µPOCA's effectiveness on various synthetic tabular datasets.
strength_weakness: strength_weakness: - Clarity and flow: The paper's clarity and flow could be improved. For example, the use of GSMs for both training and metric estimation is not immediately clear and could be better explained upfront. The notation is sometimes confusing, for example, the use of 'j' to represent both a single feature index and a set of feature indices, as in 'ci,j' and 'xj', and J+1 being used for the label can be easily missed. Additionally, some crucial concepts, like the need for a stochastic GSM (P1 in Appendix F), are introduced late in the paper, making it difficult for the reader to grasp their importance. Finally, the difference between Scenario 1 and 2 could be better explained as it's not entirely clear what the different implications of the two are and how the unnatural pattern of Scenario 1 impacts the performance of µPOCA.
- Potentially strong assumption of independence: The theoretical justification seems to rely on the assumption of independence between generalization capability (G) and unobserved features (Xj) given observed data (•).  Intuitively it seems like this assumption would not hold in many real-world scenarios. The authors should discuss the practicality of the assumption, the implications of violating this assumption, and potential mitigation strategies.
questions: questions: - Have you considered evaluating µPOCA on real-world datasets with varying degrees and patterns of missingness? This would provide a more comprehensive understanding of its practical applicability and limitations.
- Are there specific situations where µPOCA significantly outperforms traditional AL, even when the GSM imputation accuracy is not perfect?
limitations: limitations: The authors briefly address the limitations of µPOCA.
paper_summary: paper_summary: This paper explores the problem of active learning where we can choose not just which data point to annotate, but also which features to obtain, when given an unlabeled pool of data points for which only some features are known. The paper proposes using generative models to sample the missing features and apply bayesian active learning techniques to estimate the expected information gain from annotating the data points.
strength_weakness: strength_weakness: There are some points of clarification in writing and the numerical analysis is missing some potential interesting ablations and comparisons with different baselines. See questions below
questions: questions: - What is the motivation for having the label cost be separate in problem (1)? If this problem is not really studied and if you are focusing only on (2) anyways where the label cost is baked in, then this seems unnecessary generalization.
- If you have only some features for some data, what is the training process that you use?
- How does the quality of the generative model affect the downstream performance of POCA? Have you tried different models, or at the minimum, injected artificial noise to estimates?
- The numerical analysis primarily focuses on comparisons against BALD. While this makes sense given that the method is an extension of BALD, it would be interesting to see how other benchmark active learning techniques fare on this problem. Specifically, could it be that other uncertainty or diversity-based active learning methods, beyond Bayesian ones, also perform well for example if the number of missing features is small?
limitations: limitations: Limitations are discussed in the paper

Comments: title: Thank you
reponse
Thank you
Thank you
Rebuttal response - POCA
comment: Thank you for your insightful feedback. We’re pleased that we could resolve your concerns and are grateful for your input, which has greatly improved the quality of our work.
I'm satisfied with your response and will raise my score to 6.
Thank you for your thoughtful feedback and for positively acknowledging our rebuttal. We appreciate your insights, which have been instrumental in improving the quality of our work.
Thank you for your valuable feedback. We are glad to have addressed your concerns and appreciate your insights, which have significantly enhanced the quality of our work.
Thank you for the detailed rebuttal. Most of my concerns are addressed. I want to adjust my rating to weak accept.
Thank you for the detailed clarification, especially regarding the independence assumption. After considering the other reviews and your response, I believe the combined impact and the applicability of this work warrant a higher rating of 7.