Adaptive Labeling for Efficient Out-of-distribution Model Evaluation
Authors: Daksh Mittal
Yuanzhe Ma
Shalmali Joshi
Hongseok Namkoong
Venue: NeurIPS 2024 poster (2024)
Keywords: Model Evaluation
Uncertainty Quantification
Markov Decision Process
Policy Gradient
Auto-differentiation
Primary area: evaluation
TL;DR: Supervised data suffers severe selection bias when labels are expensive. We formulate a MDP over posterior beliefs on model performance and solve it with pathwise policy gradients computed through an auto-differentiable pipeline.

Abstract: Datasets often suffer severe selection bias; clinical labels are only available on patients for whom doctors ordered medical exams. To assess model performance outside the support of available data, we present a computational framework for adaptive labeling, providing cost-efficient model evaluations under severe distribution shifts. We formulate the problem as a Markov Decision Process over states defined by posterior beliefs on model performance. Each batch of new labels incurs a “state transition” to sharper beliefs, and we choose batches to minimize uncertainty on model performance at the end of the label collection process. Instead of relying on high-variance REINFORCE policy gradient estimators that do not scale, our adaptive labeling policy is optimized using path-wise policy gradients computed by auto-differentiating through simulated roll-outs. Our framework is agnostic to different uncertainty quantification approaches and highlights the virtue of planning in adaptive labeling. On synthetic and real datasets, we empirically demonstrate even a one-step lookahead policy substantially outperforms active learning-inspired heuristics.

Adaptive Labeling for Efficient Out-of-distribution Model Evaluation
Daksh Mittal, Yuanzhe Ma; Shalmali Joshi, Hongseok Namkoong
Columbia University
(dm3766, ym2865, sj3261, hn2369)0colluu edu
Abstract
Datasets often suffer severe selection bias; clinical labels are only available on patients for whom doctors ordered medical exams. To assess model performance outside the support of available data, we present a computational framework for adaptive labeling, providing cost-efficient model evaluations under severe distri- bution shifts. We formulate the problem as a Markov Decision Process over states defined by posterior belief's model performance. Each batch of new labels in- curs a "state transition" to sharper beliefs, and we choose batches to minimize uncertainty model performance at the end of the label collection process. In- stead of relying on high-variance REINFORCE policy gradient estimators that do not scale, our adaptive labeling policy is optimized using path-wise policy gra- dients computed by auto-differentiating through simulated roll-outs. Our frame- work agnostic different uncertainty quantification approaches and highlights the virtue of planning in adaptive labeling. On synthetic and real datasets, we em- pirically demonstrate even one-step lookahead policy substantially outperforms active learning-inspea inspired heuristics.
Introduction
Engineering progress AI is predicated on rigorous empirical evaluation. Evaluating pervised ML models is easy when there is distribution shift. However, naive offline Y evaluations on observational data cannot reli- ably assess safety due to selection bias. When ground -tutt outcomes are expensive to col- lect, existing supervised datasets are often heavily affected by selection bias. For ex- ample, Al-based medical diagnosis systems are typically evaluated using past clinical la- bels, which are only available for patients who were initially screened by doctors. This creates a fundamental limitation: we cannot even assess the model's performance on pa- tients who were never screened in the first place. Moreover, due to the absence of positive pathological patterns in the available data, pre- viously unseen patient types may never get diagnosed by the AI diagnosis system. Evaluating the model's performance on previously unscreened patients is first step toward mitigating this negative feedback loop.
Labels N/A
Severe selection bias underestimate MSEI
Labels N/A
+
Outcome
Extrapoineed
Interpolated
X
X
x
X
Figure 1: Adaptive labeling to reduce epistemic uncertainty over model performance. Among the two clusters of unlabeled examples (left vs. right), we must learn to prioritize labeling inputs from the left cluster better evaluate mean squared error.
"Equal contribution
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Observe labels and updale posterior
Action Pick (X :c 1
Act
Action Pick (
Pick (2² = X pool.2)
C X pook
12 :- p1( Dz
pT
Dr)
1A0
f State
f
State
State
State
Reward
E (g(f))
Figure 2: Overview of our adaptive sampling framework. At each period, we select batch of inputs X to be labeled, and obtain a new labeled data De. We view posterior beliefs (+(-) on (Y( as "states", and update it as additional labeled data is collected. Our goal is to minimize uncertainty on the performance of the prediction model (-) at the end of ' periods.
In this paper, we study model evaluations under severe distribution shifts, where the support of the available supervised data is different from examples encountered during deployment. Randomly selecting inputs X to collect additional outcomes/labels Y incurs prohibitive cost, especially when the outcome distribution varies over nigh-dimensional feature space. To address this, we must focus our sampling efforts on outtof-distribution (OOD) examples where the model's performance is uncertain. Going beyond the OOD detection literature that focuses the distribution of input X [41], we further connect the rarity of X with how smooth we believe Y|X to be. For example, in the extreme case where we believe Y|X is linear, it is easy to extrapolate model performance to regions of X where no data is available; such extrapolation is more challenging for nonlinear models.
We propose an adaptive sampling framework that focuses sampling effort on inputs X where there is large epistemic (actionable) uncertainty the shape of Y|X. We differentiate this with aleatoric (irreducible) uncertainty, which is due to idiosyncratic noise in the outcome measurement process. Since updating the sampling policy requires engineering organizational effort in practice consider a few-horizon setting where at each period, a batch of inputs is selected for labeling.
Our main contribution is the formulation of a planning problem for selecting batches of inputs from a large pool of unlabeled data (Section 3). We develop adaptive policies that incorporate how beliefs on model performance get sharper as more labels are collected. We model our current belief on the data generation process Y|X as the current state. After observing a new batch of labeled data, we update our belief on the epistemic uncertainty via a posterior update ("state transition").. Our goal to minimize the uncertainty over model performance the end of label collection, as measured by the variance under the posterior after the final batch is observed. Modeling each batch time period, we arrive at a Markov Decision Process (MDP) where states are posterior beliefs, and actions correspond to selecting subsets (batches) from the pool of examples to be labeled (Figure 2). Notably, our problem is less ambitious than an active learning problem [34] where the goal is to adaptively label data to improve model performance, over large number of batches. Narrowing our focus to model evaluation allows us to address a broad class of evaluation problems including regression models, and solve the problem more effectively.
Solving our planning problem is computationally challenging due to the combinatorial action space and continuous states. To address this, we develop tractable approach by considering a smooth relaxation of the problem, where we optimize over smoothly parameterized policies that select an entire batch of inputs [39]. Although the well-known "score trick" enables the computation of policy gradient estimates based on function evaluations alone ("REINFORCE") [36, 5], this method suffers from high variance and may lead to poor performance in practice.
To compute reliable policy gradients, we develop an dito-differentiable planning framework in Sec- tion 4. Since the dynamics of the MDP (posterior updates) are known, we perform "roll-outs" using the current posterior beliefs: we simulate potential outcomes that we may obtain if we label par- ticular batch, and evaluate the updated uncertainty model performance based on the imagined pseudo labels of this new batch. By implementing all the operations within an auto-differnniabble framework, we can backpropagate the observed reward to the policy parameters and calculate the

---

pathwise policy gradient. While our original problem is nonsmooth and discrete, auto- auto-diffeerntiating over this smoothed MDP provides reliable policy optimization methods.
Our approach is agnostic to the uncertainty quantification approach, and allows incorporating latest advances in Bayesian modeling (e.g., [5, 28)). Our "model-based" approach rests encoding state transitions (posterior updates) in an uuto-differentiation framework. We stress that we do not require the posterior updates have a closed form (conjugacy), and allow approximate posterior inference using gradient-based optimizers. We demonstrate our planning framework over beliefs formed by Gaussian processes, and Deep Ensembles [18. 26] that enable leveraging the inductive biases of neural networks in high dimensions and gradient descent methods to update probabilistic beliefs over Y X. Empirically, we observe that even single planning step can yield significant practical benefits (Section 5). On both real and simulated datasets with severe selection bias, we demonstrate that one-step lookahead policies based on our diteeeiifernniniooonn framework achieve state-of-the performance, outperforming REINFORCE policy gradients and active learning-based heuristics.
Limitations Our empirical validation highlights several open research directions required to scale our framework. First, although recent advances in Bayesian modeling have achieved substantial progress in one-shot uncertainty quantification, we empirically observe that it is often difficult to maintain accurate beliefs as more data are collected. Second, we detail engineering challenges in implementing our auto diffffffferentiiionn framework over multi purlti- roll-outs, highlighting the need for efficient implementations of high-order auto- titto-differentation.
2 Related work
Our work is closely related to active learning [1, 34), where the modeler adaptively collects labels to improve model performance. Active learning methods typically focus on classification models and select inputs to label based on uncertainty sampling techniques such
sargin-sampling or entropy, or by leveraging disagreements among different models, exemplified by query-by-commitee Bayesian Active Learning by Disagreement (BALD) [14]. Although there is notion of planning, some query strategies account for the expected reduction in error or model change [34]. Since is important to collect labels over a diverse set of inputs. heuristic criteria based on diversity and density are sometimes incorporated to account for the feature space distribution [34]. Prior work addressing distribution shifts in active learning [42] uses importance weighting, which is not feasible in our setting, where the distribution shift is out of support. Batched settings are a long-standing challenge in active learning [34]. In Section 5, we show that adapting greedy-based heuristics to batched scenarios yields poor performance. Recent extensions of active learning to batched settings (c.g., BatchBALD [17], BatchMIG [37]) continue rely on myopic greedy algorithms for elassification. On the other hand, we leverage our limited scope on evaluation to formalize planning problem and derive a unified computational framework that can handle regression problems. We provide lookahead policies [2, 6, 7] that plan for the future and flexibly handle different objectives and uncertainty quantification methodologies.
and
Our problem can also be viewed as a version of pure exploration bandits [23] where we want to minimize some given function of the posterior. In contrast to conventional formulations, our central focus on batching induces combinatorial action spaces and few horizons. Bayesian optimization methods optimize black-box functions that are computationally challenging to evaluate by modeling the function as a draw from a Gaussian process [8]. Notably, the knowledge gradient algorithm [9] maximizes the single period expected increase in the black-box function value and can be viewed as a one-step lookahead policy. Using auto auto-differee iitt-ddiffreettattion methods, we extend these classical ideas to model evaluation by incorporating batching (combinatorial action spaces).
While Gaussian processes works well in low dimensions, quantifying epistemic uncertainty on (Y|]) for high dimensional inputs is an active area of research, including popular techniques such as dropout [10], Bayes by Backprop [3], Ensembles/ Ensemble+ [18, 26), and Epistemic Neu- ral Networks [28]. Our framework is compatible with deep learning-based uncertainty quantification models, such as Ensembles, where we can only do approximate posterior updates.
3 Markoy decision process over posterior beliefs
Our goal to evaluate the performance of the model V. : X' ++ y over the input distribution Px that we expect to see during deployment. Given features X E X, outcomes/ labels are generated from
3

---

some unknown function y f where E E~ N2(0022) When ground truth outcomes are costly to obtain, previously collected labeled data D typically suffers selection bias and covers only a subset of the support of input distribution Px over which we aim to evaluate the model performance. Assuming we have a pool of data Xooo, we design adaptive sampling algorithms that iteratively select inputs in Xooo to be labeled. Since labeling inputs takes time in practice, we model real-world instances by considering batched settings, Our goal is to sequentially label batches of data to accurately estimate model performance over Px and therefore we assume we have access to a set of inputs Xeval ~ Px. We use the squared loss to illustrate our framework, where our goal is to evaluate Ex~Px [YY [- Under the "likelihood" function p(ylf, g(f) be the performance of the Al model (-) under the data generating function f. which we refer to as estimand of interest. When we consider the mean squared loss, g(f) is given by
(y - f(x)), let
s()) - Ex-Px [EY~P(ISX)] (Y- - -(X)))
(1)
Our framework is general and can be extended to other settings. For example, a clinically useful metric is Recall, defined as the fraction of individuals that the model (-) correctly labels as positive among all the individuals who actually have the positive label g(J) ExxPx (X 1 (V(X)= > 00]Y=1]]
Since the true function I' is unknown, we model it from Bayesian perspective by formulating a posterior given the available supervised data. We refer to uncertainty over the data generating function J as epistemic uncertainty- measurement noise & as aleatoric uncertainty. Assuming independence given features X, we model the marginal likelihood of the data via the product p()'lf,'' = [I(xYY) V(YIf. X)) where D
we can resolve it with more data-and that over the
(X' x ))) Our prior belief ju over functions f reflects our uncertainty about how labels are generated given features.
To adaptively label inputs from Xpool, we quantify epistemic uncertainty over the labels reliably using an uncertainty quantification (UQ) method. Roughly speaking, a UQ method provides us the ability to formulate posterior belief μuf DD given any supervised data D. As we detail in Section 4.1, our framework can leverage both classical Bayesian models like Gaussian processes and recent advancements in deep learning-basse
UQ methods. As new batches are labeled, we update our posterior beliefs about / over time, which we view as "state transitions" of a dynamical system. Recalling the Markov decision process depicted in Figure 2, we sequentially label batch of inputs from Xpoot (actions), which lead to state transitions (posterior updates). Specifically, our initial state is given by (o)()= = 1a(- DD) and at each period f, we label a batch of K2 inputs x+++1 C Xpooo resulting in labeled data D=+1 = (x++1 x )xtt)). After acquiring the labels at each step t, we update the posterior state to (+++(-)) pe(. DI+1). Modeling practical instances, we consider a small horizon problem with limited adaptivity T. Formulating an MDP over posterior states has long conceptual roots, dating back to the Gittin's 's index for multi-armed bandits [12].
We denote by Fe the adaptive labeling policy at period t. We account for randomized policies x1+ ~ Te (T( with flexible batch size | Kg. We assume Te is F; meaaeauuee for all t I T, where Fe is the filtration generated by the observations up to the end of step t. Observe that M4+1 contains randomness in the policy Te as well as randomness in yx-1 (x*+11 .4t). Letting
π (To, ..... -
we minimize the uncertainty over gff) at the end of data collection
H (TT) A EDIT [G(pr)] e EDIT [G(u(- DD) EDIT... [Drr,]] (2)
where G(HT)) 5ars~urg(f1 gf)) In the above objective (2) we assumed that the modeler pays fixed and equal cost for each outcome. Our framework can seamlessly accommodate variable labeling cost as well. Specifically, we can define a cost function c(-) applied on the selected subsets and update the objective accordingly to have a term Ac(D1:') where is the penalty factor that controls the off between minimizing variance and cost of acquiring samples.
In the plamning problem (2), states are continuous and the action space is combinatorial. To ad- dress these issues, we propose a continuous policy parameterization To (with parameter 0) in the next section. We solve this MDP using policy gradients. Let the gradient be defined as (o)( & VoEANTI [oEA)] Here, we slightly abuse notations by letting A a D1:T and G(A) & G(u(- DO, DI.T)), where the distribution of depends on up- To approximately solve this MDP using policy gradients, the score function estimator (REINFORCE [38]) uses the fact

---

that VoEA~Ta [G(A)] = Ea-aa [G(A)Va log Ta (T)) for any function G(-). Despite being unbi- ased, Monte-Carlo estimation of the above expectation typically suffers from prohibitively large variance [33] especially when no (A) is small. Although stream of work strives to provide variance reduction techniques for REINFORCE [19, 29], they require value function estimates that are also challenging to compute in our planning problem (2).
Planning using pathwise policy gradients
Our main algorithmic insight is that for systems with known dynamics (posterior updates) or where the dynamics can be approximated sufficiently well, we can leverage auto-differentiation directly estimate approximate pathwise policy gradients (4) instead of relying on high-variance score- based gradients. The policies derived using our efficient differentiable simulator exploits the system struc- ture and achieves significantly improved performance compared to policies that do not rely on gra- dient information, detailed in Section 5. Intuitively, this akin to finding random variable Z ~ pz distributed independent of policy To, such that A = h( 0) and
VoEANTH|G(AN = VoEzapz [G(M(Z,0)) Bz~pz VoG(h(Z,0))
(3)
However, the equality (a) above holds only if G(h(Z. 8)) is differentiable w.r.t. A. In our case, as we will see later, h(Z, 0) is non-differentiable w.r.t. 0 because the actions are discrete and combi- natorial. To address this, we will use smoothed approximation, h (Z, 0), such that G(h, (Z. 0)) becomes differentiable. Here, I is a temperature parameter that controls the degree of smoothing. Consequently, we can approximate the gradient as follows
V AeEAA [[EA)]]G(A]]] ~ (gE|G((,
- (2,(h,
(4)
is important to note that sometimes even G(-) is non-differentiable, and that case we must also consider smooth approximation of G(-) (e.g., refer to Section D.1 for details on smoothing the Recall objective discussed earlier in Section 3). Our approach is inspired by the line of work on differentiable simulators [4, 15, 21, 40, 35]. Theoretical analysis in stochastic optimization [11. 20] also highlights the benefit of these estimators over zeroth-order gradient estimates of a stochastic objective such REINFORCE. Suh et al. [35] notes that first-order gradient estimators typically perform well when the objective is sufficiently smooth and continuous. In what follows, we em- pirically demonstrate that through a careful smoothing of the planning objective, it is possible to trade-off bias and variance through nuto-differentiation.
Algorithm Autodiff 1-lookahead
1: Inputs: Labeled batch data DO, horizon T, UQ module, pool data Xpoot, batch size K
Returns: Selected batches (2** 3r*) for 1 S t < T and updated estimate of the objective G(ur)
2: t= 0:: Compute initial posterior state 140 based on UQ module and batch of labeled data D
3: for <t < T 1: do
4: .ee Algorithm 2 (Inputs: Posterior μ4, pool (At(t)) batch size K)
pool Te) for S و K} 5: x {X, X, SortDescending( (xtttiinnnnnnt)nn)nn
6: Obtain labels y to create D4++1
7: Update posterior state: Pt+1 (-(- [D(+1
8: Estimate the objective G(
9:
end for
We showcase the power of planning using pathwise policy gradients by considering the simplest possible planning algorithm: one-step lookaheads. Empirically, we demonstrate that even one-step look ahead achieves significant improvement in sampling efficiency over other heuristic baselines. Additionally, we highlight open engineering directions that can enable efficient multi-step planning methods. At each time step f, our algorithm selects a batch of inputs * to minimize the uncer- tainty over the one-step target estimand G(H441) based on current posterior belief f4, After selecting the batch 1 we observe the corresponding labels y and update posterior belief over the data- generating function f to 1 (see Algorithm l). The samples x to query are determined using policy te,0 e Te. which we optimize through pathwise policy gradients (Algorithm 2). To optimize
5

---

State Pe Policy Tt.0 pool Sample ) pool
K-subset sampling
Select batch S x++1
Pseudo Posterior update
Estimate Var (g(f))
Soft posterior update (differentiable)
Soft K-subset
a(O)
a()) M++
Differentiable estimate
sampling
Figure 3: Differentiable one-step look ahead pipeline for efficient adaptive sampling
the policy T, simulate "roll-outs" of posterior belief's over using imagined pseudo labels Pooo generated from current belief μ4. Algorithm summarizes each of these steps of the overall procedure. Our conceptual framework general and can leverage multi-step look-aheads [2, 6, 7].
Our algorithm consists of three major components which we will describe in detail shortly. First, we parameterize single-batch policy Ty and use K-subset sampling [39] to choose K samples. Sec- ond, we use a UQ module that characterizes posterior beliefs over f. Finally, to reliably optimize 0, we adopt an cite-differentiable "roll-out" pipeline through which we approximate policy gradients by smoothing the pipeline. The overall pipeline (and its differentiable analogue) is briefly described in Figure We also explain our algorithm graphically in Figure 8 of Section A. Our codebase is available at https: /tt/t/gtttt/ttt// com/ namkoong lab/ adaatiiaeeeeeeee Seaative-labellle labeling.
4.1 Uncertainty Quantification (UQ) Module
The UQ module maintains estimate b4 over posterior beliefs of over time, enabling planning framework. Our method is agnostic to the UQ module, and we illustrate this using two instantiations: i) Gaussian Processes (GP) that extensively used in the Bayesian Optimization literature and are known to work well in low dimensional data and regression setting, and ii) recently developed neural network-based UQ methods such Ensembles / Ensemble+ [27].
Gaussian Processes GPs f(-) ((,)))))(()))) are defined by a mean function m(-) and ker- nel (C)))), where for any inputs X, ((X) is Gaussian with mean m(X) and Cov (ox) ((X)) KCX,, Xz). Additionally, the observation consists of (X, Y) with Y (*( + and Æ
(O, of for some noise level a > 0. Posterior updates have closed form and are differentiable, as we review in Section B.
Ensembles Ensembles [18] learn an ensemble of neural networks from given data, with each network having independently initialized weights. Ensemblet [26] extends this approach by com- bining ensembles of neural networks with randomized prior functions and bootstrap sampling [25]. The prior function is added to each network in the ensemble, trained using L2 regularization. Re- cently, Epistemic neural networks (ENNs) [28] have also been shown to be an effective way of quantifying uncertainty. All these deep learning models are parametrized by some parameter For a given sample {(X,, Yy)),er the model weights 1j are update through gradient descent under a loss function É(), with the update rule expressed as: neew = 17 Σeez V, Yy.
4.2 Sampling Policy
At each time step, a batch of K samples is queried improve the objective. The samples selected from a pool Xoool of size n, based on weights w()). Specifically, given weights w 2 0 and a batch size K, the Kssubset sampling mechanism generates a random vector S {0, 1)",, whose distribution depends on w such that 2 Bi=1 S K. Let e/ denote the j-th unit vector, and consider
sequence of unit vectors je .e...... e* e Using parametrization known as weighted reservoir sampling (wrs) [39], the probability selecting a sequence is given by:
Porr((e" e* /w)
Wii
Wia
WiN
(5)
Σ Wy
Win
E wy j==1 W/
The probability of selecting batch S can then be defined as: p(S)w) ExET(s) Pwrs (s/ss)))) where I(S) = { 8::S 5 =ビi= }bi]] denote all sequences of unit vectors with sum equal to S.

---

Algorithm 2 One-step look-ahead policy gradient
1: Inputs: Posterior belief μr, pool data (Aoool), batch size K, training epochs: M
Returns: Policy me
2: Initialize the policy Te (parametrized by 0)
3: for 1 1 m I M: do
4: Generate peusdo labels pool for all the Xoool, where pool is drawn from the current poste-
rior state μ.
5: Use Algorithm 3 (Input: no == w(f) to generate random soft vector a()) corresponding
to policy x
6: Update (pseudo) the posterior state to μ4 a()) using the pool data and pseudo labels P
7: Estimate the objective G(1++ (00)
8: SGD update: < 0 - c VoG(
9: end for
Return: Policy x
4.3 Differentiable pipeline
We propose a fully differentiable pipeline by ensuring that each component of the pipeline: i) K- subset sampling, ii) posterior updates (UQ module), and iii) the objective are differentiable. We use parametrizations 0 for the policy function π which yields sampling probabilities w(O). Posterior beliefs μ are parametrized by (e.g., Ensemblet weights). In the following, we describe how smooth the respective components. The one-step objective parametrized by 0 is given by
H (0) (p IG (uf)
(6)
where G was defined in (2). Here, p(-|w((0) is the distribution over subsets governed by as defined by (5). Further, 4 is the current posterior state and μ4+ is the updated posterior state after incorporating the the additional batch S selected from Xpool and the corresponding pseudo-outcomes from Pooo (drawn based current posterior state μ)). We smooth the objective (0) to estimate VaH(O) using a soft K -ssubeet sampling.
Soft K K-subeet sampling We use the smoothed K K-subset sampling procedure proposed in Xie and Ermon [39]. Briefly, the Algorithm 3 introduced in [39] produces a random vector a such that Ei= a. = K (soft version of S), such that (0)) ~ p(-1w(()) = (+))) (5), and is differentiable. We defer the details to Section C.
Smoothing the UQ module To enable differentiable posterior updates, we approximate the pos- terior updates using the soft K-subset samples a0). In the case of GPs, the soft K K-subset selection variable can be interpreted as weighting of samples in a GP. Thus, closed form GP updates can be analogously used with appropriate weighting, described in the Appendix D. To perform weighted updates for deep learning based UQ modules (Ensembles / Ensemble+ ), we approximate previous gradient updates by incorporating soft sample weights (0):
m (a)
- n70 Σ ((X,,
where é the loss function of the concerned model, with gradients evaluated at 1). We interpret 7) as an approximation of the optimal UQ module parameter with one gradient step. Similarly, we define
a higher-order approximation 7/ (a) as
(a) A my (0()) Σ a,,00, Vy.n) (n=ma(a())
(8)
Note here that effectively, the gradient is now of the form VGG (argmin,, e (n (a())))
V (7+ (a(0))) To compute this, higher [13] and torchopt [32] which allows us approx- imately differentiate through the argmin. In practice, there a trade-off between the computational time and the accuracy of the gradient approximation depending on the number of training steps do to estimate the argmin.
7

---

0.04
REINFORCE
Error of estimated model l2 loss
0.6
REINFORCE
Autodiff 1-kokkshead (ours)
Autodéf 1-laokahead (nurs)
Uncortainty sampling sequential
Uncertainty sampling sequeetiall
Random sampling
0.5
Random sampling
0.03
Uncortainty sa mpling static
Uncertainty sampling static
Var of t2 loss
0.4
0.02
0.3
0.01
0.2
0.1
0.00
0
4 6
10
0
10
Time horizon
Time horizon
Figure 4: (Synthetic data) Variance of mean squared loss evaluated through the posterior be- lief 14 at each horizon t. This is the objective that policy gradient methods like REINFORCE and Autodiff 1-lookahead optimizes. 1 1-step lookaheads are surprisingly effective even in long horizons.
Figure (Synthetic data) Error between MSE calculated based on collected data D0:7 vs. pop- ulation oracle MSE over 'Deval ~ Px. Reduc- ing uncertainty over posteriors directly leads to better OOD evaluations. 1-step lookaheads sig- nificantly outperform active learning heuristics in small horizons.
5 Experiments
We empirically demonstrate effectiveness of our planning framework on both synthetic and real datasets by focusing on the simplest planning algorithm: 1-step lookaheads. Using two uncertainty quantification modules-GPa and Ensembles Ensemble+ we show that planning with path- wise gradients is promising approach to adaptive labeling. Throughout this section, we focus on evaluating the mean squared error of regression model V and develop adaptive policies that minimize uncertainty on this quantity (g(f)). When GPs provide valid model of uncertainty, our experiments show that our dite-eifferettialle planning framework significantly outperforms other baselines. We further demonstrate that our conceptual framework extends to deep learning-based uncertainty quantification methods like Ensemble+ while highlighting computational challenges that need to be resolved in order to scale our ideas. For simplicity, we assume naive predictor, ie., eh(+) O. However, we emphasize that this problem is just as complex as if we were using a sophisticated model (.). The performance gap between the algorithms primarily depends on the level of uncertainty in our prior beliefs.
To evaluate the performance of our algorithm, we benchmark it against several baselines. Our first set of baselines are from active learning [1]:
Active Learning Heuristics: (1) UNCERTAINTY SAMPLING (STATIC): In this approach, we query the samples for which the model is least certain about. Specifically, we estimate the variance of the latent output (X) for each X E Xpool using the UQ module and select the top-K points with the highest uncertainty. (2) UNCERTAINTY SAMPLING (SEQUENTIAL): This is greedy heuristic that sequentially selects the points with the highest uncertainty within a batch, while updating the poste- rior beliefs using pseudo labels from the current posterior state. Unlike UNCERTAINTY SAMPLING (STATIC), this method takes into account the information gained from each point within batch, and hence tries to diversify the selected points within a batch.
We also compare our approach to solving the planning problem using (3) REINFORCE- pol- icy gradients, which implements one-step look ahead policy gradient using the score trick. Finally, we study (4) RANDOM SAMPLING, which selects each batch uniformly at random from the pool. We repeat all experiments with 10 random seeds.
5.1 Planning with Gaussian processes
We now briefly describe the data generation process for the GP experiments, deferring a more de- tailed discussion of the dataset generation to Section E. We use both the synthetic data and the real data to test our methodology. For the simulated data, we construct a setting where the general pop- ulation is distributed across 51 non-overlapping clusters while the initial labeled data D just comes from one cluster. In contrast, both Dpool e (Xpool, Ypoal), Deval A (Xevat, Veval) are generated from all the clusters. We begin with low-dimensional scenario, generating one-dimensional regression setting using a Gaussian Process (GP). Although the data tata-generating process is not known to the
8

---

Error of estimated model t, loss
0.010
0.875-
0.850
Var of 12 loss
0.009
0.825
0.008
0.800
0.775
REIFFORCE
REINFORCE
0.007
it tours!
Autoo lookahead (ours)
Uncertaiity sampling - sequental
00.750
Uncertainty sampling sequential
Random sampirg
Random samplig
0.006
Uncertaint, sampling static
0.725
Uncertainty sampling static
0
4 Time horizon
8
10
4 Time horizon
8
Figure 6: (Real-world eICU data) Variance of mean squared loss evaluated through the poste- rior belief ft at each horizon t. Even 1-step lookaheads are extremely effective planners, and auto-diffcrentiation- based pathwise policy gra- dients provide a reliable optimization algorithm based on low-variance gradients.
Figure 7: (Real-world eICU data) Error between MSE calculated based on collected data Do:7 vs. population oracle MSE over Deval ~ Px. Re- ducing uncertainty over posteriors directly leads to better OOD evaluations. Our method signifi- cantly outperforms active learning heuristics, and 1-lookahead with REINFORCE policy gradients.
algorithms, we assume that the GP hyperparameters are known to all the algorithms to ensure fair comparisons. This can be viewed as a setting where our prior is well-specified, allowing us to isolate the effects of different policy optimization approaches without any concerns about the misspecified priors or in-accurate posteriors. We select 10 batches, each of size K = 5 across T 10 time horizons.
To examine the robustness of our method against the distributional assumptions made in the sim- ulated case, we then move to real dataset where the correct prior not known. We simulate selection bias from the eICU dataset [30], which contains real-world patient data with in-hospital mortality outcomes. We conduct a k-means clustering to generate 51 clusters and then select data from those clusters. We view this to be credible replication of practice, as severe distribution shifts are common due to selection bias in clinical labels. To convert the binary mortality labels into a regression setting, we train random forest classifier and fit a GP on predicted scores, which serves as the UQ module for all the algorithms. As before, the task is to select 10 batches, each consisting of samples, across 10 time horizons.
In Figures 4 and we present results for the simulated data. Figure 4 shows the variance of fg loss, while Figure 5 presents the error in the estimated & loss using Me (relative to true lz loss, that is un- known the algorithm). As we can see from these plots, our method Autodiff 1-lookahead gives substantial improvements over active learning baselines. Compared to the other one-step lookahead planning approach using REINFORCE based policy gradients, we observe that pathwise policy gradients provide significantly more robust performance over all horizons.
In Figures 6 and 7, we observe similar findings on the eICU data. We see that planning policies (REINFORCE and ours) consistently outperform other heuristics by large margin. Active learning baselines perform poorly in these small-horizon batched problems and can sometimes be even worse than the random search baselines. Overall, our results show the importance of careful planning in adaptive labeling for reliable model evaluation.
We offer some intuition as to why Autodiff 1-lookahead may outperform other heuristic algo- rithms. First, UNCERTAINTY SAMPLING (STATIC) while myopically selects the top-K inputs with the highest uncertainty, it fails to consider the overlap in information content among the "best" in- stances; see [1] for more details. In other words, it might acquire points from the same region with high uncertainty while failing to induce diversity among the batch. Although UNCERTAINTY SAM- PLING (SEQUENTIAL) somewhat addresses the issue of information overlap, significant drawback of this algorithm is the disconnect between the objective we aim to optimize and the algorithm. For example, it might sample from region with high uncertainty but very low density.
5.2 Planning with neural network-based uncertainty quantification methods (Ensemble*)
We now provide a proof-of-concept that shows the generalizability of our conceptual framework to the deep learning based UQ modules, specifically focusing on Ensemble+ due to their previously

---

observed superior performance [28]. Recall that implementing our framework with deep learning based UQ modules requires us to retrain the model across multiple possible random actions (@) sampled from the current policy no. This requires significant computational resources, in sharp con- trast to the GPs where the posteriors are in closed form and can be readily updated and differentiated.
Due to the computational constraints, we test Ensemblet on toy setting to demonstrate the gen- eralizabilitty of our framework. We consider a setting where the general population consists of four clusters, while the initial labeled data only comes from one cluster. Again we generate data using GPs. The task to select a batch of points in one horizon. We detail the Ensemble+ architec- ture in the Appendix, and we assume prior uncertainty to be large (depend on the scaling of prior generating functions). The results are summarized in the Table 1.
Table 1: Performance under Ensemble+ + as UQ modulc
Algorithm
Variance of &g loss estimate 7129.8 + 1027.0
Error of & loss estimate
RANDOM SAMPLING
136.2 + 8,28
UNCERTAINTY SAMPLING (STATIC)
10852 + 0.0
162.156 + 0.0
UNCERTAINTY SAMPLING (SEQUENTIAL)
8585.5 + 898.9
144 + 6.93
REINFORCE
1697.1 = 0.0
45.27 = 0.0
Autodiff 1-lookahead
1697.1 + 0.0
45.27 = 0,0
Gradient Estimation TTheoretical Insight
In this section, we present theoretical analysis of the statistical properties (bias-variance trade- off) of the REINFORCE and AUTODIFF semoothed-pathwise (4)) gradient estimators. To com- pare the two estimators, we consider simplified setting with two actions A = {-1. and the policy is parametrized by E [0, with To (-(1)) = and (1) = 1 - 0, Additionally, we are interested in the gradient (o)) where (0) Eawn, G(A) with G(A)
A for sim- plicity. Given N i.i.d. samples of A; ne, the REINFORCE estimator is defined as VRF V log(m The pathwise gradient estimator Vgrad is the N-sample approx- imation of Euu-uno, (oG (h, (U,))). where h-(U, 0) = (exp (5=²) - 1) / (exp ( 1)
+ DE N
(smoothing of h(U,0) := 2U > 0) - _ A). Our result (proof in Section F) highlights the conditions under which 7.N achieves a lower mean squared error (mse) compared to VRF.
Theorem 1. For 0 E [0, and N < -ा-ー 1, there exists 7 depending (N,0) 0) such that
MSE grad F.N < 4 < MSE((
Additionally, for any N. to and + xc, we have the following MSE(( = ((k)), F.N) < 4. The same statement holds for 0 = 1 &/2 well. This implies that the mse RE is unbounded, while the mse of gradient estimator is bounded. of N
MSE( (
7 Conclusion and Future Work
Supervised data often suffers severe selection bias when labels are expensive. We propose new framework for adaptive labeling for model evaluation under out-of-support distribution shifts. Fol- lowing a Bayesian framework, we formulate an MDP over posterior beliefs on model performance. We show that these planning problems can efficiently solved with pathwise policy gradients, computed through carefully designed auto-differentiable pipeline. We also demonstrate that even one-step lookahead policies outperform heuristic active learning algorithms. However, this improve- ment comes at the cost of additional computational resources. Our formulation is thus appropriate in high-stake settings such as clinical trials where labeling costs far exceed computational costs.
We further highlight some important nuances which we learned from our experiments. There are some important properties that a UQ module should have for it to perform well in our framework. The most important feature is that posteriors should be consistent, that is, early stopping of the train- ing should not cause a significant change across the ranking of the subsets of the pool. Second, the posterior update should be done efficiently, which can be achieved either by doing parallelization or having UQ modules that provide readily available posterior updates, such as GPs. Recent advances in Bayesian transformers [24, 22] will be an interesting direction to explore in this regard.
10

---

References
[T] C. C. Aggarwal, X. Kong, Q. Gu, J. Han, and S. Y. Philip. Active learning: A survey. In Data classification, pages 599-634. Chapman and Hall/CC, 2014.
[2] D. Bertsekas and N. Tsitsiklis, Neuro-dynamic programming. Athena Scientific, 1996.
[3] C. Blundell, J. Cornebise, K. Kavukeuoglu, and D. Wierstra. Weight uncertainty in neural network. In Proceedings of the 32nd International Conference on Machine Learning, pages
1613-1622. PMLR, 2015.
[4] F. de Avila Belbute-Peres, K. Smith, K. Allen, J. Tenenbaum, and J. Z. Kolter. End-to-end differentiable physics for learning and control. In Advances in Neural Information Processing
Systems. Curran Associates, Inc., 2018.
[5] T. Du, Li, Xu. A. Spielberg. K. Wu, D. Rus, and W. Matusik. [pg): Deep differentiable
deterministic policy gradients, 2020.
[6] Y. Efroni, G. Dalal, B. Scherrer, and S. Mannor. Multiple-step greedy policies in approximate
and online reinforcement learning. In Advances in Neural Information Processing Systems,
2018.
[7] Y. Efroni, M. Ghavamzadeh, and S. Mannor. Online planning with lookahead policies. In
Advances in Neural Information Processing Systems, 2020.
[8] P. I. Frazier. Bayesian Optimization, pages 255-278. 2018.
[9] P. I. Frazier, W. B. Powell, and S. Dayanik. A mnowledge-gradien policy for sequential infor- mation collection. SIAM Journal on Control and Optimization, 47(5):2410-2433 2008. doi:
10. 10.1137/070693424.
[10] Y. Gal and Z. Ghahramani. Dropout as Bayesian approximation: Representing model un- certainty in deep learning. In Proceedings of The 33rd International Conference on Machine Learning, pages 1050-1059, 2016.
[11] S. Ghadimi and G. Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic programming. SIAM Journal on Optimization, 23(4):231-2238. 2013.
[12] J. C. Gittins. Bandit processes and dynamic allocation indices. Journal of the Royal Statistical Society, Series B, 4177):148-177, 1979.
[13] E. Grefenstette, B. Amos. D. Yarats, P. M. Htut, A. Molchanov, F. Meier, D. Kiela, K. Cho, and S. Chintala. Generalized inner loop meta-learning. arXiv preprint arXiv:191..01777. 2019. [14] N. Houlsby, F. Huszár, Z. Ghahramani, and M. Lengyel. Bayesian active learning for classifi- cation and preference learning. arXiv: / 12.55755 [cs.CV],20) 2011.
[15] Z. Huang, Y. Hu, T. Du, S. Zhou, H. Su, J. B. Tenenbaum, and C. Gan. Plasticinelab: A soft-body manipulation benchmark with differentiable physics. In International Conference on Learning Representations, 2021.
[16] E. Jang, S. Gu, and B. Poole. Categorical reparameterization with Gumbel-softmax. In Inter- national Conference on Learning Representations, 2017.
[17] A. Kirsch, van Amersfoort, and Y. Gal. BatchBALD: Efficient and diverse batch acquisition for deep bayesian active learning. In Advances in Neural Information Processing Systems, volume 32, 2019.
[18] B. Lakshminarayanan. A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In Proceedings of the 31 International Conference on Neural Information Processing Systems, page 6405-6416, 2017.
[19] A. Mnih and K. Gregor. Neural variational inference and learning in belief networks. In Pro- ceedings of the s1 International Conference on Machine Learning, volume 32 of Proceedings
of Machine Learning Research, pages 1791-1799, 2014.
[20] S. Mohamed, M. Rosca, M. Figurnov, and A. Mnih. Monte Carlo gradient estimation in ma- chine learning. Journal of Machine Learning Research, 21(1), 2020.
[21] M. A. Z. Mora, M. Peychev, S. Ha, M. Vechev, and S. Coros. Pods: Policy optimization via differentiable simulation. In Proceedings of the 38th International Conference on Machine Learning. 2021.
[22] S. Müller, N. Hollmann, S. P. Arango, J. Grabocka, and F. Hutter. Transformers can do bayesian inference. In Proceedings of the Tenth International Conference on Learning Repre- sentations, 2022.
[23] R. MUNOS, S. BUBECK, and G. STOLTZ. Pure exploration for multi-armed bandit problems. Lecture Notes in Computer Science, 2009.
[24] T. Nguyen and A. Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In Proceedings of the 39th International Conference on Machine
11

---

Learning, 2022.
[25]
Osband and Van Roy. Bootstrapped Thompson sampling and deep exploration.
arXiv: 1507.000000000 [stat.ML.], 2015.
[26]
Osband, Aslanides, and A. Cassirer. Randomized prior functions for deep reinforcement
learning. In Advances in Neural Information Processing Systems 31, volume 31, 2018.
[27] I. Osband, Z. Wen, M. Asghari, V. Dwaracherla, X. Lu, M. Ibrahimi, D. Lawson, B. Hao, B. 1..O'DDmm O'Donoghue, and B. V. Roy. The Neural Testbed: Evaluating Joint Predictions. In Advances
in Neural Information Processing Systems, 2022.
[28] I. Osband, Z. Wen, S. M. Asghari, V. Dwaracherla, M. Ibrahimi, X. Lu, and B. V. Roy. Epis- temic neural networks. In Thirty-seventh Conference on Neural Information Processing Sys- tems, 2023.
[29] M. Papini, D. Binaghi, G. Canonaco, M. Pirotta, and M. Restelli, Stochastic variance-reduced policy gradient. In Proceedings of the 35th International Conference on Machine Learning, pages 4026-4035, 2018.
[30] T. Pollard, A. E. Johnson, J. D. Raffa, L. A. Celi, R. G. Mark, and O. Badawi. The eICU Collaborative Research Database, a freely available multi-center database for critical care re- search. Scientific data, 5(1):1-13, 2018.
[31] C. E. Rasmussen. Gaussian processes for machine learning. pages 63-71, 2006.
[32] Ren, X. Feng, Liu, X. Pan* Fu, L. Mai, and Y. Yang. TorchOpt: An Efficient Library for Differentiable Optimization. Journal of Machine Learning Research, 24(367):1-14, 2023. [33] D. Rezende, S. Mohamed, and Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 3 International Conference on Machine Learning, pages 1278-1286. 2014.
[34] B. Settles. Active learning literature survey. 2009.
[35] H. Suh, M. Simchowitz, K. Zhang, and R. Tedrake. Do differentiable simulators give better policy gradients? In Proceedings of the 39th International Conference Machine Learning,
volume 162, 2022.
[36] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[37] C. Wang, Sun, and R. Grosse. Beyond marginal uncertainty: How accurately can Bayesian regression models estimate posterior predictive correlations? In International Conference
Artificial Intelligence and Statistics, pages 2476-2484, 2021.
[38] R. J. Williams. Simple statistical gradient-followigg algorithms for connectionist reinforce-
ment learning. Machine Learning, 8:229-256, 1992.
[39] S. M. Xie and S. Ermon. Reparameterizable subset sampling via continuous relaxations. In Proceedings of the 28th International Joint Conference on Artificial Intelligence, IJCAI' page 3919-3925, 2019.
[40] J. Xu, V. Makoviychuk, Y. Narang, E. Ramos, W. Matusik, A. Garg, and M. Macklin. Accel- erated policy learning with parallel differentiable simulation. In International Conference on Learning Representations, 2021.
[41] Yang, K. Zhou, Y. Li, and Z. Liu. Generalized out-of-distribution detection: A survey. 17X.1:2110.11334 [cs. CVJ, 2021.
[42] E. Zhao, Liu, Anandkumar, and Yue. Active learning under label shift. In Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, pages 3412- 3420, 2021.
12

---

A Graphical Representation of Our Algorithm
Autodiff 1-Lookahead
Intial Data
Train module
on Initial data
(state Fu)
UQ module (state Fs)
UQ module (state pt)
uq module (state pt)
Repeat t times
Update UO module H using data (X,, Y,) to
1-step lookaheed (Algorithm
Relurn policy
Select batch X; be labeled (using and get true labels Y'
1-step lookahead (Algorithm 2) at time step t
Initialize policy ne
Model we want ovallaato
Generate pseudo labels (8paa))
Pseudo posterior update (1. using pseudo labels and soft K-subset sampling [n(9)
Estimate variance MSE under updated posterior [c (2209))
(8psa)) using module (ax)
f- True generating function Posterior state (distribution over at time
SGD update policy x, using end-to-ent autb-differentiation (Repeat N times)
[UQ trained on data (X.a.1..)) Vn.01
Figure 8: Description of our algorithm Autodiff 1- 1-lookahead
B Gaussian Process posterior updates
For any inputs X, we assume f(X) is Gaussian with mean m(X) and Cov (xx, = K(X,, X₃). In addition, the observation consists of (X, Y) with Y (*(X) and (00, I² for some noise level a 0. Given the training data (X, Y) and test points X", closed form posterior estimates over f* can be computed
X"Y,X" N(fI", Cov(f")),
where fl" A X(X*,
+
Cov(f") K(X", X') - K(X", X)IK(XXX)) -
Using the above expression, we can update our belief about f* after observing new data.
C Details of K-subset sampling algorithm
Algorithm 3 for soft K K-subset sampling introduced in [39]. We present this algorithm here and refer the reader to [39] for further details.
D Details of weighted GP
We present weighted GP algorithm (Algorithm 4) adapted from Algorithm 2.1 in [31]. Recall that GPs f(-) ~ (P(m()) are defined by mean function m(-) and kernel IC(:)))
Now assume and we aim to compute posterior estimates over f* for test points X* considering the weights w. Let us define the following terms: K A K(X,, X), K, A X(X*,, X) and K.,, - K(X* X Using these definitions, GP Algorithm 4 provides posterior estimates of f. for the test samples X given that the input data (X, YY has weights w.
given input data (X, Y),
also given some weights w for these inputs,
13

---

Algorithm Soft K -ssubeet sampling algorithm
1: Inputs: Weight vector w E R_
2: Sample independent standard Gumbel random variables 91, 92: ... In
3: Compute keys f; = 9 + log(w;; for all
4: Initialize x! fr for all I....,
5: For j = I,.., K, set
df exp(12/TT for all -, ..... n
Σk=1 exp(Nie/T)
1+ = xi + log(1 - af) for all i 1,... nt
Ki
6: Return: The soft K-hot vector, a a² +... a
Algorithm Weighted Gaussian process regression
1: Inputs: K, K., Ka,ay Y,
2: K K((1- Tyww" + I)
3: Kw,- w @ K.
4: L Cholesky ( + 1)
5: a L (L Y)
6: f KT .
7: I. Kw
8: V K.,- U TT
9: Return: mean and covariance: V)
D.1 Smoothing the recall objective
Performance metric, such as the Recall of a model (.). not differentiable. Recall that H(0) tp (G' (mL) where G(uE) Var gf),, and g(f) -
Ex~Px Ly-x() (X) > 0} Y =1 with for depending ne. To estimate VyH(0) in differentiable manner, we can smooth approximation of g(f) using the softmax trick [16] and draw 2n i.idd Gumbel(0, 1) samples R with / e {1,2}. Define:
R-(f:R) & EL Y? () Σ (((FR
exp (llg fi) + R!)/T)
where (f; = exp (flog fi) + R2)/T) + exp ((log( fi)) + R2)/T)'
here >> 0 is the temperature hyperparameter that controls the smoothing of g-(J; R). easy to show that BY (J:: R)] z fi. making (J R) natural approximation of the previously defined Recall g()).
E Experimental details
In this section, we provide detailed information about the experiments discussed in Section 5.
E.1 Planning with Gaussian Processes synthetic data experiments
As mentioned earlier, we generate our data using a Gaussian Process (GP). Specifically, with RBF kernel: fi GP(m, where m(X) 0 and IC(X, X') oj exp
we use a GP
27
further Gaussian noise N(0, x²) is added to the outputs. We set f 1, of 0.69, and x² 0.01. Recall that the marginal distribution Px consists of 51 non-overlapping clusters. To achieve this, we create a polyadic sampler, which first samples 51 anchor points spaced at linearly increasing distances from the center to avoid overlap. These anchor points serve as cluster centers. We
14

---

sample the points around these anchor points by adding a small Gaussian noise [0, 0.25). The training points are drawn from a single cluster, while the pool and evaluation points are drawn from all the 51 clusters. The setup includes 100 initial labeled data points, 500 pool points and 285 evaluation points used to estimate the objective.
Autodiff 1-lookahead training and hyperparameters: As mentioned earlier, the underly- ing uncertainty quantification (UQ) module for these experiments the GP. Although our algo- rithm does not know the true data generating function, but it has access to the GP hyperparam- eters. Therefore we use a GP with an RBF kemel:: fs gP(m, KC), where m(X) = and
with Gaussian noise NO, o2) added to the outputs. We set <(X, = exp (- ||X 1
26a
1, 0.69 and x² = 0.01. For soft K-subset sampling (see Algorithm 3), we set to 0.1. Further for evaluating the objective function Var(g( take 100 samples of (eaeaa)) from the posterior state 14=", see Algorithm 2. For policy optimization in each horizon, we use the Adam optimizer with a learning rate of 0.1 perform policy gradient steps over 100 epochs. The results presented are averaged over 10 different training seeds.
Policy gradient (REINFORCE) training and hyperparameters: As mentioned earlier, all the algorithms have access to the true to GP hyperparameters. For evaluating the objective Var(g(f)) under given action, take 100 samples of (Xeeal) To optimize the policy in each horizon, we use the Adam optimizer with a learning rate of 0.1 to perform policy gradient updates over 100 epochs. The results presented are averaged over 10 different training seeds.
E.2 Planning with Gaussian Processes - real data (eICU) experiments
The eICU dataset is healtheare dataset that contains data from various critical care units across the United States. To create a supervised learning setup from this dataset, we first extracted the 10 most important features using a Random Forest classifier (number of trees 100, criterion "gini"). The outcome variable was in-hospital mortality. Some examples of the extracted features include: Hospital length of stay, Number of days on the ventilator and the Last recorded temperature Day 1 of ICU admission. We transformed the classification outcome variable (in-hospital mortality) into a regression task using the probability of in-hospital mortality predicted by the classifier. To adapt the extracted data to our setting, we introduced selection bias in the data. We generated 51 clusters through the standard -- means algorithm. Our initial labeled data comes from only cluster, while the pool and evaluation data comes from all the 51 clusters. Our dataset consists 100 initial labeled data points, 500 pool points, and 285 evaluation points used estimate the objective.
We then fitted a GP with an RBF kernel to above data using GP-regression. The resulting GP hyperparameters were: m(xx) 0.255, E = 0.50, of = 1.0, o² 0.0001. This GP model serves as uncertainty quantification (UQ) module.
Autodiff 1-lookahead training and hyperparameters: For soft K-subset sampling (see Algo- rithm 3), we set 7 to 0.1. Further to evaluate the objective function Var(g(f)), we take 100 samples For policy optimization in each horizon,
of (Xeva)) from the posterior state /4 see Algorithm we use the Adam optimizer with a learning rate of 0.1 to perform policy gradient steps over 1000 epochs. The results presented are averaged over 10 different training seeds.
Policy gradient (REINFORCE) training and hyperparameters: To evaluate the objective Var(g())) under an action, we take 100 samples of (Xeva). For optimizing the policy in each horizon, we use an Adam optimizer with learning rate of 0.1, performing policy gradient updates over 1000 epochs. The results presented are averaged over 10 different training seeds.
E.3 Planning with Ensemblet experiments
Once again, we use GPs as our data generating process. Specifically, we use GP with an RBF kernel: fa GP(m, K.), where m(X) 0 and K(X, =03 exp (- || Additionally, Gaussian noise o2) is added to the outputs. We set f = 1, of 0.69, and 0² 0.01. The marginal distribution Px consists of 4 non-overlapping clusters. We follow the same process as in the Gaussian process experiments to form the cluster centers and sample points around them. While the training points are drawn from single cluster, both the pool and evaluation points are drawn
2
15

---

from all the 4 clusters. Our setup includes 20 initial labeled data points, 10 pool points, and 252 evaluation points used to evaluate the objective.
Ensemble+ Architecture: We use an ensemble of 10 models. Each model being a 2-hidden layer with 50 units per hidden layer. In addition, each models includes additive prior. The additive prior function for each model is fixed to be 2-hidden layer MLP with 50 units in each hidden layer. Specifically, the m-th model for S m _ 10 of the Ensemble+ takes the form
fm... (X) In.. (X) + cxm
where Ia... is the trainable part of the network while oxpm (-) is the additive prior function. Here or controls our prior belief about the uncertainty -higher the a, the greater the uncertainty. The prior functions (.) differs across models m due to different initializations. In our setup, we seta 100 reflect high level of uncertainty in the prior beliefs.
Ensemble+ training: For a given dataset DD = (X; VO) we train the m-th model so as to minimize the following loss function
810m: D)
Σ 9ym (X,) opm (X,) Y2) A 11
We tune to 0.1 and use the Adam optimizer with a tuned learning rate of 0.1. Each model is trained for 50 iterations. Although in standard Ensemble+ each model is trained different bootstrapped subset from the dataset D, in our case we train all the models within Ensemble+ on the entire dataset available.
Autodiff 1-lookahead training and hyperparameters: Recall that in soft K-subset sampling there is a hyperparameter (see Algorithm 3), which we set to 0.1. To differentiate through the argmin operation, we employ the differentiable optimizer from the torchopt package, specifically the MetAdam Optimizer with learning rate of 0.1. For optimizing the sampling policy, the Adam optimizer with learning rate 0.05, performing policy gradient updates over 500 epochs.
F
Proof of Theorem
We begin by analyzing the REINFORCE based gradient estimator. With N samples, we have:
MSE( = Var( ( = N
Var Var AV, log(To(A)))
Using the definition of A, we derive the following:
Var (AVo log (A))) = (1 0)0 4.
As result, we have:
1 N (1 0-0 1)
MSE(VRF)
Next, analyze the pathwise gradient estimator grad variable, where U Uni(0, 1), so that:
7. Let h (U,
exp expl
be random
D-0
MSE(
(U,,( + N Var(Vo
(10)
To compute the bias term in (10), we note that
1 1 ++ exp (-U=0)
eep( 1-0
+1
Eh,, (U, E
E
2r log
1.
exp (U-0) +
exp( -
+
VoE(h,(U,8) -
1 exp (-0)
1 1 exp (=0)
16

---

Therefore,
2
Bias (oas( =4 1+
(11)
+ exp 1-0
1 +
exp
Note that the above term goes to 4 as xc. The other part contributing to the mse (10) is the variance term and it satisfies
exx -0) +
3 exx (1-2) +
Var(Val,, (U,0) S <VVV((+(U3))))
6 (exp
+ +1)
(exp (1-9) + 1)
Combining the above equation with (11), we arrive at MSE(V grad
grad, F,N)
S uN((T) (T). where
(N(T) a 1 1 exp (-=))
1 4
3 exp ( +1
3xx (7ا7) +
+ exp -0)
N
6 (exp ( (- + 1)
6 (exp
(--) + 1)"
For all N 2 1, verify that lim UN(T) = 4 and for large enough (T is increasing in Thus, there exists some F > such that MSE(V F.N) < 4. Comparing this with (9) for MSE(VV) completes the proof.
17

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
Sections 3, 4, and 5
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create 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 au- thors should reflect 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the ap- proach. For example, a facial recognition algorithm may perform poorly when image resolution is low images taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Ys]]
18

---

Justification:
Sections 6 and F
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yss]
Justification: Sections 5 and E Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary either make it possible for others replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or a way to re- produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
19

---

Answer: [Yes]
Justification: Section 4
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips cc/
Policy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: cc/ wtc//pubicc// Sewudee//oeeeiiieeeeeeeeeee
ssi onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Sections 5 and E Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
Sections 5 and E
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
20

---

is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: //terr///n....//eeeeee wc/ppblicc/thiiccuuieeeeeeeee
Answer: [Yes]
Justification:
Our research conform, in every respect, with the NeurIPS Code of Ethics
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
Sections 1 and
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact 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 spe- cific groups), privacy considerations, and security considerations.
21

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility MI).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
[NA]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere usage guidelines or restrictions 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 require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [Yes]
Justification: Sections 4 and 5
Guidelines:
The answer NA means that the paper does not 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 URL.
The name the license (c.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of use the pack- age should be provided. For popular datasets, paperswi whwode com/ /ataases has curated licenses for some datasets. Their licensing guide can help determine the li- cense of 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.
22

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer:
[No]
Justification:
We do not have any assets of now.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the per include the full text instructions given to participants and sereenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does involve crowdsoureing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) 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 NeurlPS 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.
23


Reviews:
[{"review_id": "xEeKclHlT0", "reviewer": "Reviewer_GtV8", "summary": "This paper addresses the issue of selection bias in previously collected data when ground truth labels are expensive. It introduces a computational framework for adaptive labeling, aiming to provide cost-efficient model evaluations under severe distribution shifts. The problem is formulated as a Markov Decision Process (MDP), where states are defined by posterior beliefs about model performance. New batches of labels lead to \"state transitions,\" resulting in sharper beliefs, with the goal of minimizing overall uncertainty on model performance. The proposed method optimizes the adaptive labeling policy using pathwise policy gradients via auto-differentiation through simulated rollouts, avoiding high-variance REINFORCE policy gradient estimators. The framework is adaptable to various uncertainty quantification methods and emphasizes the importance of planning in adaptive labeling. Empirical results on both synthetic and real datasets show that even a one-step lookahead policy significantly outperforms active learning-inspired heuristics.", "questions": "- Line 102: Why do you mean by \"Batched settings\"?\n- Line 107: What are lookahead policies?\n- Line 134: What's the purpose of rewriting the likelihood function $p(y|f, x)$ as $p_{\\epsilon}(y - f(x))$? What does this likelihood function mean?\n-  Line 160: What's \"F_t-measureable\"?\n- Eq 2: Is uncertainty equal to Var? In the intro, you defined epistemic and aleatoric uncertainty, but what's the uncertainty you're trying to minimize in Eq2?\n- Line 168: What's the definition of A?\n- Line 179: Missing subject.\n- Line 180: What's the definition of g(W_i, \\theta)?\n- Figure 4: Missing y-axis label.\n- Line 285: Providing intuition of empirical results before presenting the experimental results seems weird to me. I suggest stating the intution after presenting the results.\n- Figure 4 and 5: What is the difference in terms of purpose between them? What should readers takes away from the comparison on variance of mean squared loss and error of MSE on collected data? Line 314 refers to them but doesn't explain the details.\n- Figure 6 and 7: Same issues with Figure 4 and 5.\n\n\nMinor comments:\n- The related works are a lot. Consider making it a section. It's too dense.", "limitations": "Yes, the limitation is discussed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "To my knowledge, formulating model evaluation as RL is a novel and potentially useful idea.", "weaknesses": "The presentation is extremely confusing. See my questions.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised by the reviewer. Additionally, please refer to the global response for the changes we intend to make in the camera-ready version based on all the reviews.\n\n**Line 102: Why do you mean \"Batched settings\"?**\n\nIn a batched setting, we select a group of inputs to be labeled in each step, rather than just a single input. This approach is particularly important because, in practical scenarios, labels are typically requested in batches rather than individually. Additionally, it makes the problem more challenging by inducing a combinatorial action space. We discuss this in Section 1 (Lines 47-48) and will elaborate further on the problem setup in the camera-ready version.\n\n**Line 107: What are lookahead policies?**\n\nLookahead policies are a class of policies for approximately solving the MDPs (Markov Decision Processes). These policies make decisions by explicitly optimizing over a shorter horizon of the MDP rather than the entire horizon. Depending on the length of this shorter horizon, they are referred to as 1-step, 2-step lookahead policies. We will include the following relevant references in the main paper.\n\n[1] D. Bertsekas and J. Tsitsiklis. Neuro-dynamic programming. Athena Scientific, 1996\n\n[2] Y. Efroni, M. Ghavamzadeh, and S. Mannor. Online planning with lookahead policies. Advances in Neural Information Processing Systems,  2020\n\n[3] Yonathan Efroni, Gal Dalal, Bruno Scherrer, and Shie Mannor. Multiple-step greedy policies in approximate and online reinforcement learning. In Advances in Neural Information Processing Systems, 2018.\n\n**Line 134: What's the purpose of rewriting the likelihood function $p(y|f, x)$ as $p_{\\epsilon}(y - f(x))$? What does this likelihood function mean?**\n\nWe want to highlight that the only source of randomness in the label $y$, given $f$ and $x$, is the noise $\\epsilon$. Furthermore, this randomness depends solely on $y-f(x)$ and the distribution of the noise $\\epsilon$. \n\n**Line 160: What's \"$F_t$-measureable\"?**\n\n$\\mathcal{F}_t$-measureable states that the posterior state $\\mu_t$ is known at time $t$. This is a formal way of stating that the posterior state depends only on the prior and data acquired up to time $t$.\n\n**Eq 2: Is uncertainty equal to Var? In the intro, you defined epistemic and aleatoric uncertainty, but what's the uncertainty you're trying to minimize in Eq2?**\n\nWe aim to minimize the epistemic uncertainty, specifically by considering  the variance of $g(f)$, where $f \\sim \\mu_T$. This variance represents the uncertainty due to a lack of knowledge about the underlying $f$, which is the epistemic uncertainty. \n\n**Line 168: What's the definition of A?**\n\n$A$ is a random variable distributed according to $\\pi(\\theta)$. It is introduced to highlight the difference between the REINFORCE estimator and the pathwise gradient estimator.\n\n**Line 179: Missing subject.**\n\nThanks for pointing this out, we will address it in the main paper.\n\n**Line 180: What's the definition of $g(W_i, \\theta)$?**\n\nA random variable $A$ distributed according to $\\pi_\\theta$ can always be expressed as $h(W, \\theta)$ for some function $h(.)$, where $W$ is distributed independent of $\\theta$. We use  a smoother approximation $g(W, \\theta)$ of $h(W, \\theta)$ to enable differentiability with respect to $\\theta$. For example, in Algorithm 2, we discuss how the standard Gumbel random variable $W$ can be used to generate a softer version $a(\\theta)$ of K-subset samples. \n\n**Figure 4: Missing y-axis label.**\n\nThank you for bringing this to our attention. We will add the y-axis labels in the paper. For your reference, we have mentioned the labels in the figure description.\n\n**Line 285: Providing intuition of empirical results before presenting the experimental results seems weird to me. I suggest stating the intuition after presenting the results.**\n\nThank you for the suggestion. We will move the intuition after presenting results in Section 4.\n\n**Figure 4 and 5: What is the difference in terms of purpose between them? What should readers take away from the comparison on variance of mean squared loss and error of MSE on collected data? Line 314 refers to them but doesn't explain the details.**\n\nFigure 4 demonstrates that the variance of the model’s L2 loss is reduced, which is the objective of the algorithm.\nIn Figure 5, besides comparing the variance of the L2 loss, we also compare the accuracy of the updated L2 loss estimation after acquiring new points. This illustrates that the objective is appropriate: by minimizing the variance, we also achieve a better evaluation of the model.\n\n\n**Figure 6 and 7: Same issues with Figure 4 and 5.**\n\nThis has been addressed above.\n\n**There are many related works. Consider making it a separate section, as the current content is too dense.**\n\nThanks for pointing this out. We will make it a separate section.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IAXTxuEAql", "reviewer": "Reviewer_CyPE", "summary": "This paper focuses on sampling efforts on out-of-distribution (OOD) examples to minimize uncertainty in the prediction model's performance over T periods. The main idea is based on a novel view of MDP, which considers beliefs as states and samples as actions. The authors also demonstrate how to make the entire pipeline differentiable for optimization. Experiments show that the proposed framework achieves better sample efficiency in reducing variance.", "questions": "1. in Line 270, what do mean by \"We just consider a naive predictor, i.e., ψ(·) ≡ 0 for simplicity.\"", "limitations": "The paper includes a section on limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The approach of considering beliefs as states and samples as actions to minimize uncertainty in the prediction model's performance over T periods is both interesting and novel. Efficiently addressing OOD samples is also a very important problem.\n- Experimental results provide significant evidence that Autodiff can efficiently reduce variance.", "weaknesses": "- I do not fully understand some parts:\n-- Can the authors provide a more detailed formulation on: (1) how the G(u) is calculated? (2) Algorithm 2 returns \\pi_{\\theta}.  Does the pseudo-label-based part also update the policy? How is this updating connected to the updating with ground truth labels acquired in line 6 of Algorithm 1?\n\n- In lines 102-108 of the related work, the authors mention several active learning methods in batched settings, highlighting that they focus on classification settings while this paper mainly focuses on regression settings. Given that these methods are not used as baselines, I wonder how different the active learning settings in classification and regression are. Can methods in classification settings be easily adapted to regression settings?\n\n-What is the size of the training dataset?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised by the reviewer. Additionally, please refer to the global response for the changes we intend to make in the camera-ready version based on all the reviews.\n\n\n**I do not fully understand some parts: -- Can the authors provide a more detailed formulation on: (1) how the $G(\\mu)$ is calculated? (2) Algorithm 2 returns $\\pi_{\\theta}$. Does the pseudo-label-based part also update the policy? How is this updating connected to the updating with ground truth labels acquired in line 6 of Algorithm 1?**\n\n**(1)** As shown in line 164, given any posterior belief $\\mu$, $G(\\mu)$ is defined as $\\mathrm{Var}_{f \\sim \\mu}  g(f) $. Depending on the form of $g(f)$ and the belief $\\mu$, $G(\\mu)$ can either be derived analytically or estimated by drawing samples $f \\sim \\mu$ and then calculating the empirical variance of $g(f)$ across these samples.\n\n\n**(2)** We refer the reviewer to the diagram explaining our algorithm in the attached PDF (alongside the global response). \nBroadly, at each time step $t$, the input to Algorithm 2 is the current posterior state $\\mu_t$, which is used to generate pseudo labels. These pseudo labels are then used to optimize the policy $\\pi_\\theta$ in Algorithm 2. Any update of the UQ module/posterior state within Algorithm 2 using these pseudo labels is temporary and the posterior state is reverted back to $\\mu_t$ at the end of Algorithm 2. Next, we select batch $X_{t+1}$ according to the optimized policy $\\pi_\\theta$ from the Algorithm 2 and obtain the true labels $Y_{t+1}$. Finally, we update the UQ module/posterior belief $\\mu_t$ using $(X_{t+1},Y_{t+1})$ resulting in the new posterior state $\\mu_{t+1}$ at time $t+1$.\n\n**In lines 102-108 of the related work, the authors mention several active learning methods in batched settings, highlighting that they focus on classification settings while this paper mainly focuses on regression settings. Given that these methods are not used as baselines, I wonder how different the active learning settings in classification and regression are. Can methods in classification settings be easily adapted to regression settings?**\n\nAs mentioned in the global response, we refer the reviewer to  our active learning-based baselines in Section 4.1. In this section, we adapt common active learning heuristics in classification settings such as uncertainty sampling (static), uncertainty sampling (sequential), and random sampling, to regression problems we focus on.\n\n**What is the size of the training dataset?**\n\n\nWe direct the reviewer to Section D for the experimental details, particularly noting the training dataset sizes mentioned on lines 485-486 and 511-512.\n\n**in Line 270, what do mean by \"We just consider a naive predictor, i.e., ψ(·) ≡ 0 for simplicity.\"**\n\n\nRecall our objective is to evaluate the performance of the model $\\psi(.)$. For simplicity, we assume a naive predictor/model which outputs $0$ for all the inputs. We want to emphasize that this problem is just as complex as having a sophisticated model $\\psi(.)$, as it merely serves as an input to our algorithm through objective $G(\\mu)$.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jrkjjKVsON", "reviewer": "Reviewer_4yoX", "summary": "This work formulates the process of data labeling as a Markov Decision Process (MDP). It introduces three key components: K-subset sampling, an Uncertainty Quantification (UQ) module, and a differentiable simulator. These components enable the application of the REINFORCE policy gradient algorithm to solve the MDP. The proposed method demonstrates significant performance improvements on both synthetic and real-world data.", "questions": "I have no further questions. Please refer to the weaknesses mentioned above.", "limitations": "The authors have discussed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The problem of data labeling is crucial in the field of machine learning. This work formulates the data labeling challenge as a Markov Decision Process (MDP) and addresses it using a Reinforcement Learning (RL) algorithm, providing valuable insights.\n\n2. This study adapts three key components to make the REINFORCE algorithm suitable for this problem.\n\n3. The proposed method achieves significant performance improvements on both synthetic and real-world data.", "weaknesses": "The organization of Section 4 could be improved. I suggest introducing the overall description of the proposed method and the pseudo-code for Algorithm 1 after the introduction of the three key components.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough review and valuable feedback on our manuscript. Below, we address the specific questions raised by the reviewer. Additionally, please refer to the global response for the changes we intend to make in the camera-ready version based on all the reviews.\n\n**The organization of Section 4 could be improved. I suggest introducing the overall description of the proposed method and the pseudo-code for Algorithm 1 after the introduction of the three key components.**\n\nWe thank the reviewer for their advice on improving the presentation of the paper. As mentioned in our global response, we will incorporate all their suggestions in the final version of the paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FT5tdTNZln", "reviewer": "Reviewer_PJGF", "summary": "This paper studies model evaluation under distribution shifts. It formulates a planning problem to select inputs for labeling from a large pool of unlabeled data. The paper then proposes an auto-differentiable planning framework to compute reliable policy gradients and update the sampling policy accordingly. To demonstrate the utility of this framework, the paper considers one-step lookahead policy based on their framework and shows that it outperforms active learning-based heuristics and REINFORCE policy gradients.", "questions": "- Why not evaluate how close $\\mu_T$ is to the ground truth $f^*$? Is $\\mu_T$ always guaranteed to be a distribution concentrated around $f^*$?\n- How should $n$ be chosen in Algorithm 3? How should the sampling size $K\\$ be selected, and how do different choices affect the results?", "limitations": "See Weaknesses", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- **Clear presentation**: The paper is well-structured, with figures (e.g., Figure 2 and Figure 3) illustrating the adaptive sampling framework, which helps in understanding the proposed method.\n- **Detailed explanations**: The authors provide detailed explanations of the UQ module in Section 3.1, as well as Soft K-subset sampling and Smoothing UQ module in Section 3.3. Even for readers unfamiliar with these terms, the explanations are clear enough to convey their meanings.", "weaknesses": "- **Computational inefficiency**: Their method requires significant computational resources, especially when the posteriors aren’t in closed form. Due to these computational constraints, Section 4.2 only considers a simple toy setting and does not compare the results with active learning and REINFORCE policy gradients. Additionally, the paper only implements one-step lookahead based on their framework, leaving the question of how to efficiently conduct multi-step planning methods unresolved.\n- **Fixed and equal labeling cost assumption**: The paper assumes that the modeler pays a fixed and equal cost for each label. It would be more interesting to consider variable labeling costs, which are more common in practice.\n- **Lack of theoretical guarantees**: The paper does not provide any theoretical guarantees for their method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the thorough review and valuable feedback on our manuscript. Please refer to the global response for our detailed answers to the concerns raised and the changes we intend to make in the camera-ready version based on all the reviews. We also address the concerns raised point by point here.\n\n\n**Computational inefficiency: Their method requires significant computational resources, especially when the posteriors aren’t in closed form. Due to these computational constraints, Section 4.2 only considers a simple toy setting and does not compare the results with active learning and REINFORCE policy gradients. Additionally, the paper only implements one-step lookahead based on their framework, leaving the question of how to efficiently conduct multi-step planning methods unresolved.**\n\nOur work introduces a novel adaptive labeling formulation for evaluating model performance when existing labels suffer severe selection bias. As we highlight in Section 2, the adaptive labeling MDP is hard to solve because of continuous states and combinatorial actions. We make initial methodological progress by designing a one-step lookahead planning algorithm and demonstrate that even this approach can outperform existing active learning heuristics. However, this improvement comes at the cost of additional computational resources. Our formulation is thus appropriate in high-stake settings such as clinical trials where labeling costs far exceed computational costs, and we hope our new formulation spurs subsequent methodological innovations for multi-step lookaheads.\n\nWe conduct an additional experiment for planning with Ensemble+ as the UQ module to compare the performance of our algorithm with other active learning based heuristics. This comparison is conducted in a more complex setting than previously discussed in the paper. The results can be found in the attached PDF (included with global response). As shown, Autodiff 1-lookahead and REINFORCE outperforms the active learning heuristic, while the performance of the Autodiff 1-lookahead and REINFORCE is similar. Based on our extensive GP experiments in Section 4.1, we anticipate that our algorithm will perform better than REINFORCE in more complex scenarios for Ensemble $+$. This presents an interesting research question for future exploration.\n\n\n**Fixed and equal labeling cost assumption: The paper assumes that the modeler pays a fixed and equal cost for each label. It would be more interesting to consider variable labeling costs, which are more common in practice.**\n\nOur framework can seamlessly accommodate variable labeling costs. Specifically, we define a cost function $c$ applied on the selected subsets $\\mathcal{D}^{1:T}$ and update the objective accordingly to have a term $\\lambda c(\\mathcal{D}^{1:T})$ where $\\lambda$ is the penalty factor that controls the trade-off between minimizing variance and cost of acquiring samples. \n\n**Lack of theoretical guarantees: The paper does not provide any theoretical guarantees for their method.**\n\nWe provide a theoretical explanation of why pathwise gradients outperform REINFORCE. In a tractable example we detail in the global response, we prove pathwise gradients can achieve a lower MSE than REINFORCE.\n\n**Why not evaluate how close $\\mu_T$ is to the ground truth ${f^{*}}$?**\n\nNote that the modeler never has access to the ground truth $f^*$. Similar to supervised learning losses that use $Y$ as a proxy of $f^*(X)$, our objectives considers the modeler’s belief over $f^*$ based on $Y$’s observed so far. \n\nEven in our experiments (Section 4) the $f^*$ is not known in the closed form—it is only known through an evaluation dataset generated by it and it is difficult to compare $\\mu_T$ (a distribution over functions $f$) to  $f^*$. In addition to comparing the variance of the L2 loss, we also show the difference between L2 loss estimated under $\\mu_T$ to the true L2 loss estimated under $f^*$ (through the available evaluation set)  in Figures 5 and 7. \n\n**Is $\\mu_T$ always guaranteed to be a distribution concentrated around  ${f^{*}}$?**\n\nIf our model is well-specified, that is, the prior $\\mu_0$ contains $f^*$, then we are asymptotically guaranteed to concentrate around $f^*$ as we collect more samples. This is due to the well known results from Doob [1]. Empirically, this means that if the ML uncertainty quantification methodology has enough capacity, $\\mu_T$ will converge to $f^*$ asymptotically. We will highlight this point in our revision.\n\n\n**How should $n$ be chosen in Algorithm 3? How should the sampling size $K$ be selected, and how do different choices affect the results?**\n\nHere $n$ is simply the number of samples in the pool set we want to query labels from. The size $K$ is assumed to be a given that is input from the analyst based on acquisition cost and labeling budget.\n\nReference:\n[1]  Doob, J. L. (1949). Application of the theory of martingales. In Actes du Colloque\nInternational Le Calcul des Probabilités et ses applications (Lyon, 28 Juin – 3\nJuillet, 1948), pages 23–27", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for the valuable feedback. We appreciate the constructive comments and will incorporate all discussions below in the camera-ready version of the paper.

**Summary of contributions:** Our work introduces a *novel problem formulation* for evaluating model performance when existing data suffer severe selection bias, where state-of-the-art methods such as importance sampling are insufficient. Specifically, we formulate the adaptive labeling problem as a Markov Decision Process (MDP) over states defined by posterior beliefs on model performance. The rest of the paper highlights the potential of our novel formulation using simple methods derived from it. As we highlight in Section 2, the adaptive labeling MDP is hard to solve because of continuous states and combinatorial actions. We design a tractable *single-step lookahead* planning algorithm and demonstrate even small lookahead *outperforms existing active learning heuristics*. 

Our results underscore the benefits of formulating the adaptive labeling problem as an MDP and designing specific RL algorithms to solve it. As Reviewer PJGF correctly points out, this comes at the cost of additional computational resources. Our formulation is thus appropriate in *high-stake settings such as clinical trials* where labeling costs far exceed computational costs, and we hope our new formulation spurs subsequent methodological innovations. We thank Reviewers GtV8, CyPE, 4yoX, for recognizing the unusual yet exciting dimension of having the formulation as our main contribution.

To implement one-step lookahead, we move away from standard REINFORCE policy gradients that suffer high variance and propose a novel pathwise gradient based on auto-differentiation. As we summarize in the attached PDF, our approach smoothes the non-differentiable operations to enable back propagation. Despite the bias incurred by smoothing, we observe pathwise policy gradients outperform REINFORCE buoyed by low variance (Section 4). While a substantive methodological contribution, the computational efficiency of our algorithm is bottlenecked by that of the auto-differentiation framework and we highlight promising future directions of research in the Section 4.

**Comparison with active learning:**  Reviewer CyPE asks for comparison with active learning heuristics, and we would like to point to Section 4.1 where we compared our algorithm to common active learning heuristics adapted from classification to regression settings, such as uncertainty sampling (static), uncertainty sampling (sequential), and random sampling.

**Variable labeling cost:** Our framework can seamlessly accommodate variable labeling cost. Specifically, we define a cost function $c$ applied on the selected subsets $\mathcal{D}^{1:T}$ and update the objective accordingly to have a term $\lambda c(\mathcal{D}^{1:T})$ where $\lambda$ is the penalty factor that controls the trade-off between minimizing variance and cost of acquiring samples. 

**Experiments:** We have conducted an additional experiment as suggested by Reviewer PJGF. Specifically, for planning with Ensemble+ as the UQ module, we compare the performance of our algorithm with other active learning based heuristics. This comparison is conducted in a more complex setting than previously discussed in the paper. Please refer to the attached PDF for the results. Based on the extensive GP experiments in Section 4, we expect that our results will extend to even more complex scenarios for Ensemble $+$.

**Theoretical insights:** Incorporating Reviewer PJGF’s suggestion, we provide a theoretical explanation of why pathwise gradients can outperform REINFORCE, despite being biased due to smoothing. Recall that in our adaptive labeling framework, we deal with a combinatorial action space where actions $A\in\mathcal{A} = \lbrace a_1,...,a_m \rbrace$ are drawn according to a policy $\pi_\theta$. The objective is to minimize $G(\theta)=E_{A\sim\pi_\theta}f(A)$, which requires estimating $\nabla_\theta G(\theta)$. We consider a tractable yet insightful setting where $\mathcal{A} =  \lbrace -1,1  \rbrace$ with $\pi_\theta(-1)=\theta$ and $f(A)=A$.


The REINFORCE estimator $\nabla_N^{RF}$ is a $N$-sample approximation of $E_{A\sim\pi_\theta}[f(A)\nabla_\theta\log(\pi_\theta(A))]$. Note that $G(\theta)$ can be rewritten as $E_Uf(h(U,\theta))$, where $h(U,\theta):=2I(U > \theta)-1\stackrel{d}{=}A$ and $U\sim[0,1]$. To enable the pathwise policy gradient, we consider a smoothed $h^{\tau}(U,\theta)$ using sigmoid approximation, where $\tau$ is the temperature. Pathwise gradient estimator $\nabla_{\tau,N}^{grad}$ is the $N$-sample approximation of $E_U[\nabla_\theta f(h^{\tau}(U,\theta))]$.

**Theorem:** For $\theta\in (0,1)$ and $N<\frac{1}{4\theta(1-\theta)}$, there exists $\tau$, depending on $N,\theta$ such that $E[mse(\nabla_{\tau,N}^{grad})]<4<E[mse(\nabla_N^{RF})]$. Further, for any $N$, as $\theta=\frac{1}{k}$ and $k\to\infty$, $E[mse(\nabla_N^{RF})]=\Omega(k)$, while $E[mse(\nabla_{\tau,N}^{grad})]<4.$ The same statement holds for $\theta=1 - \frac{1}{k}$.


**Presentation:** We have taken to heart reviewer comments on improving our presentation. We outline a roadmap for enhancing the presentation of the paper below:

(1) Section 1: As suggested by reviewer GtV8, we will create a separate section for related works.

(2) Section 3: We will first explain all the major components of our algorithm before introducing Algorithm 1 and 2. Additionally, we plan to include the figure (from the attached PDF) in the main paper.

(3) Section 4: We will reorganize this section and will provide intuition after we present our experimental results, as suggested by Reviewer GtV8. In addition, we will make active learning based heuristics a separate subsection to convey the results better.

Meta-review: The paper studies the problem of model evaluation under distribution shift. It formulates adaptive labeling as a planning problem to select inputs for labeling from a large pool of unlabeled samples. To show the efficacy of this approach, it considers one-step look-ahead policy and shows that it outperforms several heuristics based on active learning and REINFORCE policy gradient method. Although the authors found the problem studied in the paper interesting and the approach novel, there are serious concerns about the clarity of presentation, the computational inefficiency of the proposed approach, and the lack of theoretical guarantees. Despite all the above, they still found the paper (slightly) above the bar. I would recommend the authors to revise the paper, significantly improve its presentation and properly address the comments/questions raised by the reviewers (to include some of the responses during the rebuttals in the paper) for the next revision.
Decision: Accept (poster)