ActiveCQ: Active Estimation of Causal Quantities
Authors: Erdun Gao
Dino Sejdinovic
Venue: ICLR 2026 Poster (2026)
Keywords: Causal Quantities
Active Learning
Uncertainty Quantification
Primary area: causal reasoning

Abstract: Estimating causal quantities (CQs) typically requires large datasets, which can be expensive to obtain, especially when measuring individual outcomes is costly. This challenge highlights the importance of sample-efficient active learning strategies. To address the narrow focus of prior work on the conditional average treatment effect, we formalize the broader task of Active Estimation of Causal Quantities (ActiveCQ) and propose a unified framework for this general problem. Built upon the insight that many CQs are integrals of regression functions, our framework models the regression function with a Gaussian process. For the distribution component, we explore both a baseline using explicit density estimators and a more integrated method using conditional mean embeddings in a reproducing kernel Hilbert space. This latter approach offers key advantages: it bypasses explicit density estimation, operates within the same function space as the GP, and adaptively refines the distributional model after each update. Our framework enables the principled derivation of acquisition strategies from the CQ's posterior uncertainty; we instantiate this principle with two utility functions based on information gain and total variance reduction. A range of simulated and semi-synthetic experiments demonstrate that our proposed framework significantly outperforms relevant baselines, achieving substantial gains in sample efficiency across a variety of CQs.

Published as a conference paper at ICLR 2026
ACTIVECQ: ACTIVE ESTIMATION OF CAUSAL QUANTI- TIES
Erdun Gao! & Dino Sejdinovic
"Responsible AI Research Centre, Australian Institute for Machine Learning School of Mathematical Sciences, Adelaide University (erdun, gao, dino. se jedinovi e
edu. au
ABSTRACT
Estimating causal quantities (CQs) typically requires large datasets, which can be expensive obtain, especially when measuring individual outcomes is costly. This challenge highlights the importance of sample-efficient active learning strategies. To address the narrow focus of prior work on the conditional average treatment effect, we formalize the broader task of Active Estimation of Causal Quantities (Ac- tiveCQ) and propose a unified framework for this general problem. Built upon the insight that many CQs are integrals of regression functions, our framework models the regression function with a Gaussian process. For the distribution component, we explore both a baseline using explicit density estimators and a more integrated method using conditional mean embeddings a reproducing kernel Hilbert space. This latter approach offers key advantages: it bypasses explicit density estimation, operates within the same function space as the GP, and adaptively refines the distri- butional model after each update. Our framework enables the principled derivation of acquisition strategies from the CQ's posterior uncertainty; we instantiate this principle with two utility functions based on information gain and total variance reduction. A range of simulated and semi-synthetic experiments demonstrate that our proposed framework significantly outperforms relevant baselines, achieving substantial gains in sample efficiency across variety of CQs.
INTRODUCTION
Causality (Pearl, 2009; Imbens Rubin, 2015; Hernan Robins, 2023; Ding, 2024) aims to un- derstand how interventions influence outcomes by modeling data lata-generating processes. Within this domain, causal inference (Hernan Robins, 2023; Ding, 2024) focuses on estimating treatment effects, which are vital for decision-making in fields like economics (Heckman, 2000) and health- care (Foster et al., 2011). While randomized controlled trials (RCTs) are the gold standard, they are often impractical (Benson & Hartz, 2017), prompting reliance observational data (Rubin, 2005). In particular, we are often interested in how causal effects differ across various subpopulations, and there several distinct causal quantities (CQs) such as the (conditional) average treatment effect ((C)ATE) and the average treatment effect on the treated (ATT) (Singh et al., 2024). Even observational studies, measuring individual outcomes for causal estimation can be costly and burdensome (Nwaimo et al., 2024; Kallus & Mao, 2024). In personalized medicine, this may involve invasive procedures or expensive tests (Bi et al., 2019; Turk, 2002; Nwankwo et al., 2025), while in economics, tracking outcomes like income changes often requires labor- intensiiee follow-ups (McKenzie, 2012). These challenges highlight the need for efficient methods to identify which instances (e.g.. patients or respondents) to prioritize for outcome measurement, enabling accurate estimation of CQs while minimizing data collection costs.
Active learning (AL) offers promising approach to addressing this challenge by improving accu- racy with fewer labeled samples, specifically by strategically selecting the most informative data points (Lindley, 1956; Chaloner & Verdinelli, 1995; Settles, 1994; Rainforth et al., 2024). Prior research on active causal inference has primarily focused on two fronts: developing active strategies all available covariates (Jesson
for learning generalizable CATE estimators, often conditioning et al., 2021; Qin et al., 2021; Wen et al., 2025a), and designing adaptable, loss- toss-targetin acquisition functions (Connolly et al., 2023). Beyond these, many scientific fields emphasize identifying causal

---

Published as a conference paper at ICLR 2026
Z
S
Aspirin
PSA
Cancer
Age
Statin
BMI
a
y
Figure 1: Causal graphs illustrating (Left) motivating health example (Aglietti et al., 2020) and (Right) the general data generation model used in our framework, where the shaded area indicates variables observable in the pool.
effects within specific population subgroups (Jaber et al., 2019). For example, in a healthcare context like the one illustrated in Fig. (left), researchers might investigate whether statin benefits older or younger patients, a question that is directly tied CATE, which examines how causal effects differ across subgroups defined by age (Abrevaya et al., 2015). Similarly, in social services, providers may need select which unstructured case notes annotate evaluate outreach interventions effectively lead to housing placement for specific vulnerable populations (Nwankwo et al., 2025). In business analytics, companies often seek to quantify whether sales campaigns yield higher long-term value for distinct customer segments, requiring targeted verification of outcomes (Wen et al., 2025b;;; Tran et al., 2024). Another common focus estimating the counterfactual outcomes for individuals currently receiving fixed dosage of statin, such what would happen had their dosage been increased, a scenario that pertains to the ATT. Furthermore, estimating treatment effects for target population using observations collected from a different source population presents a significant challenge, commonly referred to as average treatment effect under distribution shift (ATEDS) (Singh al., 2024; Gao et al., 2025). Clearly, for each distinct CQ, the choice of what constitutes the most informative data differs, and thus necessitates a tailored active learning strategy. This observation leads us to our key research question:
Key Question: What are the guiding principles for strategically selecting items from a pool dataset in order to build accurate estimators for given causal quantity of interest?
Challenges. Traditional information-theorrtie AL approaches typically aim to maximize the infor- mation gained about the model parameters minimize posterior uncertainty over the unlabeled pool of data. However, when estimating different CQs. the focus often shifts to the interventional distribution for some specific subpopulations. This shift naturally induces a distribution mismatch, where data samples are drawn from one distribution, but the outcome regression function needs to generalize to another. Consequently, conventional AL methods, such Bayesian Active Learning by Disagreement (BALD) (Houlsby et al., 2011) and total variance reduction (TVR) (Cohn et al., 1996), frequently fail to align with the goal of constructing accurate regressors for the intended target distribution. These limitations highlight the necessity for CQ-aware acquisition strategies that explicitly account for the target interventional objective.
Contributions. Our primary contribution is a unified framework that both formalizes the task of Active Estimation of Causal Quantities (ActiveCQ) (Sec. 3) and provides principled solution for it. This addresses a broad class of CQs often overlooked by the literature's 's focus on CATE. The cornerstone of our framework a novel integral representation that unifies these disparate CQs. We model the required regression function with a Gaussian Process (GP) (Sec. 4.1) and demonstrate that using Conditional Mean Embeddings (CMEs) for the target distribution provides a powerful synergy, sidestepping the difficult challenge of direct density estimation (Sec. 4.3). This unified representation subsequently enables the principled derivation of bespoke utility functions. We show how classic strategies like information gain and total variance reduction (Sec. 4.5) can be instantiated in elegant, dlosed-fom expressions, automatically and analytically tailored to the CQ of interest. We support framework with convergence guarantees for this class of CQs (Sec. 4.6). Finally, evaluate our method on the ActiveCQ task through several simulations and a semi-synthetic dataset, demonstrating that approach significantly outperforms baseline methods (Sec. 5).
2

---

Published as a conference paper at ICLR 2026
2 PRELIMINARIES
Causal DAG. We represent directed acyclic graph (DAG) as G = (v, E). where v is set of nodes corresponding to random variables, and is of directed edges. A DAG is termed causal DAG if each edge V; V indicates a direct causal effect. For a node a E v, an intervention, denoted by the do-operator do(a a a) or simply do(a). corresponds to an external action that sets to a fixed value a. This intervention modifies the data-genedt generating process, leading to inst-intervention distribution a), where pa (v; (1; denotes the parents of V, in G. Our analysis is based on the causal graph in Fig. (right), which considers a treatment é A, adjustment variables/confounders s S, optional effect modifiers z Z, and an outcome y E R. For simplicity, we restrict attention to scalar treatment and outcome variables; the framework extends straightforwardly to the multivariate setting. The corresponding observational and interventional distributions are:
P vido(a) with density y( (o)do(a)) (a)) Iv,v p( ())
(Observational)
p (=))p(()))))))) ())2)' (a)p(yla, 2, 8),
= (()=) (t)ppsss)))((()//)) 8) s)p(vla, 2, z.s).
(1)
(Interventional) p"
where p" (a) is the interventional treatment distribution. We occasionally denote the full set of input covariates as x (a. z.s). Further details on DAGs are provided in App. C.I.
Causal Quantities Estimation. In the causal inference literature, CQs typically refer to the expected effect of a do-operation, expressed as [[y]do(a]] over specific subpopulations. In what follows, we outline some commonly studied CQs, focusing on their estimation throughout this paper.
Definition 1. ATE: TATE ( a) == E[y]do(a)) represents the average effect over the whole population under the intervention (2) CATE: TCATE (0, == [[ a), z == =] represents the average effect over the subpopulation with z = z under the intervention a. (3) ATT: TATT(a, ā) == [[y]do(()) (a) a ā] represents the average effect over the subpopulation who received treatment had they instead received the intervention a. (4) ATE with distribution shift (DS): TDs = E[y do(a), P] denotes the average treatment effect under intervention evaluated over target population whose covariate distribution p(z, 8) differs from the observational distribution p(z,8) 8) used Eq. 1.
To estimate these CQs, certain assumptions necessary for identifiability. In this paper, we focus on scenarios where no unmeasured confounders exist between the cause and effect, and where the Stable Unit Treatment Value Assumption (SUTVA) and the positivity condition hold. More detailed discussions can be found in App. C.I.
Lemma 1. Under these assumptions of selection on observables and covariate shift, we have
(a)E(a)= ʃ( Elyla= = a,s 8 Pads)
TATT(a, SS Elyla= 0, F3F
TCATE(a,,) { EE [[yla= a- 8.2= 2) Po (ds/z)
ms((a, S EElyla a,s= s) (dss)
Remark 1. To unify the estimation various CQs, we can express all of them as TC= J3 E[y a= a,3 P) For CATE, where the effect modifier is fixed, the integration variable is and the integral is over the conditional distribution P.. For global quantities (ATE, ATT, and DS), the distinction between effect modifiers z and other confounders s is unnecessary for identification. Thus, for notational simplicity in these cases, overloud s to represent the full set of adjusiment variables (effectively merging z and s) and integaate over the entire joint distribution Pioo
3 THE ACTIVECQ PROBLEM FORMULATION
Datasets. In traditional CQ estimation tasks (Singh et al., 2024), it is typically assumed that an observational dataset ()) ()) 1 is accessible, where the observations are drawn from the joint distribution Pxy. In the AL setup, illustrated in Fig. I (right), we have access to two datasets: labeled training 'ataset Dr { (t), } and an unlabeled pool dataset Dp (x(") [n (Qin et al., 2021; Jesson et al., 2021). For convenience, we use symbols like X, and X p to represent
"I "abeledd indicates that the outcome is observable.
3

---

Published as a conference paper at ICLR 2026
the corresponding sets of input observations in Dr and Dp, respectively. Within this setup, budget is allocated to select ng observations XB (x(()) Ji-1 from Dp. The corresponding labels, yB ( ) Ree can then be requested and added Dr to update the training set. The updated D is subsequentyy used to estimate the desired CQs.
Adaptive and batch modes. The standard "adaptive" setting follows the Label-One-Theee Teee--hhe---rinn loop, where a single label is acquired per round, repeated for B rounds to acquire all labels (Holzmüller et al., 2023). In this process, the selection of xx()) at each step depends on the labels of previously selected data points, (xx(), ,,"))""))))) (Hübotter et al., 2024). Batch mode active learning involves requesting labels in smaller batches of size ni, over nn/nn, rounds, rather than requesting all na labels at once or one by one (Kirsch et al., 2019). Let D {æ J Ji=1 (where n NT + np) denote the union of the observations of x from both the labeled dataset Dr and the unlabeled dataset Dp. This union corresponds the gray shaded area in Fig. I (right) and remains fully accessible throughout the process, After each training round, we request labels for a batch of ny data points, (>(") J-11" from Dp and incorporate them into Dr. Notably, batch-mode active learning reduces to the adaptive setting when n4 In this paper, we focus on the batch-mode setup.
Note that our setting is purely observational. We can only query an individual's 's pre- pre-existing, factual outcome rather than intervene to assign a new treatment and observe a counterfactual. This fundamentally distinguishes our work from active experimental design, which requires the ability to perform interventions (Toth et al., 2022; Kato et al., 2024; Klein et al., 2025).
4 UNCERTAINTY QUANTIFICATION OF CQS
To efficiently estimate CQs with minimal labeling cost, we propose a Bayesian active learning framework (Chaloner Verdinelli, 1995). Our approach guided by an Eeormmation-thoorete principle: at each round, we select data points that are expected to maximally reduce the posterior uncertainty over the target causal quantity. We develop and illustrate this methodology for the CATE, fCATE(a, (a, z). abbreviated as i(a, 2) for clarity, as it represents the most general and complex case among the causal quantities identified in Lemma 1. While our main exposition focuses on CATE, the framework is general, and its detailed application to other quantities is provided in App. D.
4.1 REGRESSION FUNCTION MODELING
We assume that the outcome follows the model = [[[aa, sosss]++ + where lid. N'O, o() (Singh et al., 2024). The observations {(a,, 2,, 83, y:) [= are independently and identically sampled from P uuzsy. To quantify the epistemic uncertainty of the causal effect, we model the outcome using Gaussian Process (GP), where f(a, z, 8) B[yla, z, (Williams& Rasmussen, 2006; Kanagawa et al., 2025). The function f is assigned GP prior, f (O, k). with zero mean function and covariance kernel k, which constructed as a product kernel to handle multiple inputs. Given the training dataset Dr= {ar, Zr, Sr. yr), we compute the posterior GP. For new inputs and x', the posterior mean and variance of f can be calculated as follows:
m(x)= kexx (xxxx+ (Kx,))))xxXXXXxxxx++) 1) yx,; Kpoa( x') = kzx, Xx,x, +02 kx (2)
where kixx' = keakkes Kax, = kear @ KzZr @ kas,, and K Xr = Kayar @ K Z @ KS,S,. For brevity, we use k(.,-) as general kernel notation, where specific kernels follow from their arguments (e.g., 1(x, x') denotes the kernel on X).. Additionally, let d(x) denote the corresponding feature map. We define kiww k(x, x²), kax, [k(2 X1),..., k(x, X,, xn ] and 1T K Kx,x, = [k, x,),i
Remark 2 (On Scalability and Complexity). While the standard GP formulation in Eq. 2 has O complexity, framework's 's core contribution is orthogonal this specific implementation. The proposed framework is modular and readily accommodates various sealable approximations, such as sparse variational GPs, Random Fourier Features, and Nyström methods, handle larger datasets.
4

---

Published as a conference paper at ICLR 2026
4.2 CONDITIONAL DENSITY ESTIMATION
Next, we turn our attention to modeling the conditional distribution 2/ and explore two distinct approaches: a conditional density estimator (CDE) and conditional kernel mean embedding (CME).
Conditional Density Estimator. A natural approach to estimate the conditional distribution is to use CDE, such as a mixture density network (Bishop Nasrabadi, 2006), least squares density ratio estimator (Sugiyama et al., 2010), a conditional normalizing flow (Trippe & Turner, 2018), or squared neural family (Tsuchida et al., 2023; 2025). Importantly, estimating sa requires only paired observations (s. z). and does not rely on outcome labels, allowing it to be done using both training and pool datasets before label acquisition. Since (S, Z) is always available, incorporating the pool significantly enhances estimation accuracy with minimal overhead.
Conditional Embeddings in RKHS. We introduce an alternative representation of the conditional distribution in a reproducing kernel Hilbert space (RKHS). As we model the regression function using a GP, the input features z, and s are mapped into joint feature space. The product kernel used GP regression corresponds to the tensor product RKHS HAzs H. @ Hz @ Hs. We can then represent the conditional distribution P Ps/z= in His using the CME, which is defined as:
Mulz=z == = = Ss 6 s1s)
(3)
The CME serves the embedding of Ps/z- in H.s. As discussed in (Song al., 2013; Muandet et al., 2017), the CME can be associated with a Hilbert-Schmidt operator Csl2 : Hz - H.s. referred to as the conditional mean embedding operator (CMO). This operator satisfies Uı|z== = Czz0(2). and can be understood as Csz = Can with Cax == [ (s) @ Φ(z)] and Czu := [>( @ (()] representing the covariance operators. Similar to the CDE approach, we can use all paired
observations (2,S) in D to empirically estimate Cg as Caj2 & ZZzz + A where À > 0 a regularization parameter. The kernel for here need not be the same as the one employed in the GP, and the kernel for s will be discussed later.
4.3 INTEGRAL OVER CONDITIONAL DISTRIBUTION
Leveraging these two representations of P s= we then proceed to derive the final estimator for the CQ f(a, x). As a linear functional of the regression function f, f(a, x) follows distribution determined by the underlying GP, which is also a GP with mean and covariance functions derived
from the conditional structure as in (Chau et al., 2021a). For all a, a' e A and 2' e 2, we have:
v(a = Ep-Pate a, 2)),,
a ((a, 2), (d',)))) Pg-P s' [kpost ((a, (a',
To approximate the posterior mean and covariance, we employ sampling- based method. Specifically, we first draw n, samples (()),...,,s())) from Fa/z. and similarly (8/(1), ..., g((mm)) from PP,zz The posterior mean v(a, =) is then approximated by averaging the function m((a 8) over these samples, while the posterior covariance is estimated by averaging the kernel values over all pairwise combinations:
CDE-based method.
v(a,z) =)
1 - f1
a ((aa,)) E(),
Σ kpost (,.,,,""),,''',,,,,
(5)
The expressions for m and kipua are provided in Eq. 2. This sampling approach marginalizes out s, approximating the posterior distribution and thereby enabling CATE estimation with uncertainty. The estimator will be used for sample-efficient estimation in the next section.
Leveraging the CME (Lg/z we can evaluate the integrals defining the posterior mean v(a, 22) and covariance q (a. x). (d', x²)) in the previous equation analytically. By representing the conditional distribution in the RKHS via P. the integration of the GP posterior mean and covariance kposs can be reduced to closed-fomn kernel operations. This allows us to derive the CATE estimator directly from the conditional structure of the GP.
CME-based method.
5

---

Published as a conference paper at ICLR 2026
Proposition 1. Given the dataset Dr = {ar,, Zr, Sr,,S.yue yr) and D Z
Z, S). if the posterior
GP learned from Dr, then fCATE is a functional of f defined on (a. z) with the following mean and covariance estimated using @n == Ou @= Palz and Oz 02 02 Palz'
v(a (d mi)H kex, X Ax1 yxr
(6)
q g((a, 2), (a', = ker x (Kx,,, >x1))=
where the effective kernel terms incorporating the CME are defined as: kas Kig Kzzz ( kzz (Kzz
XI) Kss( Kzz XI) - ),z)), kax, = keay @ Zz O (k2z (Kzz+ + A1) Kss, K KxpXr K aaray K Zr Br K Sy, and x kaya' 'Zzzz'' (K K + AT) (2)))) X 0
is the regularization of the CME. X₂ > 0 is the noise term for
To ensure that the CME corresponds the RKHSs that appropriately match the kernels learned in the GP regression, the features ((a). (z), Ø(s) need to be updated after each batch acquisition round. By leveraging the CME, this approach effectively obviates the need for explicit conditional distribution estimation. It shifts the CATE estimation task from the computationally intensive numerical integration of a regression function to directly manipulating a distribution embedding, thereby streamlining computation and improving efficiency. Additionally, uncertainty from the CME itself can be introduced and propagated alongside the GP's inherent uncertainty, as explored in methods like BayesIMP (Chau et al., 2021b) and IMPspec (Dance al., 2024), although detailed exploration lies beyond the scope of the current paper. Further advantages that make CME particularly well-suited to our AL process detailed in App. G.2.
4.4 SPECIFYING THE SUBPOPULATION OF INTEREST
To evaluate the estimator f(a, =), it is crucial to specify the subpopulation and treatments of interest, as they determine where estimation accuracy matters most. We consider a set of treatment-effeet modifier pairs (a,, Zi) {(a,, 2;) >) over which performance is assessed. For instance, if the goal to understand the response of specific subpopulation defined by z = z to varying treatments, and draw a uniformly from finite treatment set A. i.e., a; Uniform((A) and (a,; 2,) (a,; 2). Alternatively, the interest lies in evaluating responses to a fixed treatment across different effect modifiers, each a; is fixed to and z, is sampled uniformly from set Z, ie., 2j ~ Uniform Z) and (o,, x;)= = (a, 2.). These scenarios reflect different inferential goals and ensure that the assessment of f [, =) is aligned with the intended application. The resulting sample {(a;, 2,)) forms the basis for measuring the estimator's 's accuracy in the region of interest.
one may fix each
4.5 UNCERTAINTY REDUCTION
We are now ready to address the Key Question. To this end, we propose the following guiding principle for selecting individuals from the pool dataset Dr whose outcomes should be acquired.
Key Principle: Select a subset of samples XB from the pool Dp manner that minimizes the posterior uncertainty of the estimator f(ar, Zz).
The distinction between active CQ estimation and traditional AL tasks is evident. Traditional AL methods, such as BALD and TVR, focus on minimizing the uncertainty of the regression function f over the distribution of the union dataset (or equivalently, the pool dataset) (Smith et al., 2023). In contrast, active estimation of CQ aims to reduce the uncertainty of the CQ of interest directly, with the regression function serving only as a means to an end. This principle guides data selection by quantifying uncertainty or label utility using methods like entropy and variance of the target estimator. We explore two strategies: information gain (IG) and TVR, and discuss their connection in App. D.5.
Data Acquisitions via IG. To evaluate the uncertainty of the target estimator F(a) Zr), we utilize its differential entropy, denoted as H[f(a,, ]. To obtain the labels for observations XB from Dp, we quantify posterior uncertainty using conditional entropy after observing yxp. This is expressed as Eyxp - (D) [H (f(a, Z,) yx,, ], which simplify H (f((a, Z,) yxx, Dr). Here, use y show its inherent randomness, which will be marginalized further. The IG, defined as I(f(a,, Zy):yx,, x (Dr) = H[f(a,, Z, Drr] H[f(a,, Zi)lyxn Dr captures the reduction in uncertainty about 7 f(ar, Zi) after observing y xg- For brevity, omit the constraint X D P in

---

Published as a conference paper at ICLR 2026
the following optimization expressions, assuming all selections are from the pool set unless specified otherwise. Based on the IG criterion, the acquisition rule is
XB arg mig H (f(as: Zi) yXn: Dr) arg max ((aonx((((aa).)))) x (D)).
(7)
Since employ a GP framework, this quantity can be expressed in closed form using the posterior covariance matrix of the predictive distribution. Namely, for is given by: (ON E)) log (2Tr) Σ] Thus, the mutual information simplifies to:
quantities, entropy
1 2 log det [Var[f(a,, Zi Lr)
det (Varl ((r. Z,) Pr])
X arg max
arg min (Var[f(ar, Zi) (8)
where det(.) represents the determinant of a matrix. Note that although the notation includes y x the denominator's 's covariance matrix does not actually depend on the specific values of outputs, only on the input locations XB.
B. This is standard property of the Gaussian process.
Data Acquisitions via TVR. The other measure of prediction uncertainty is the total variance, which is defined as the sum of the marginal variances over the target set Σ (a.=)e(a,.)) War[f(a, {)]. Therefore, we can have the TVR strategy as our data acquisition function as follows.
X B arg min't Tr (Var[f(a, Prr
(9)
Here, we define UIG (X) I(f((a, yx) (Drr) as the IG-based label utility function of XB- Similarly, the corresponding utility function for the TVR-basee approach is given by UTVR (TvR((B) (Xg)
Tr (Var[f Zr) Dr.yxn 1) representing the TVR-based label utility function of X B. Based on these utility functions, we can express the unified acquisition rule as X arg max U(R)).
To enable efficient data acquisition, we adopt batch-wise selection, choosing ni data points at time. A simple method involves using the utility function to rank all data points in DD and selecting the top n4. Recent studies (Gentile et al., 2024) emphasize the importance of ensuring diversity within batches for greater efficiency. This can be achieved through greedy approximation, where the selection of each data point æ; considers the previously selected points {z, {zz))=]] This can be formalized as maximizing a batch utility function U (X,), which values the joint contribution of points in the batch. As finding the optimal batch X3 arg x,/=n, U(x,) is computationally
Batch Selection.
intractable, we employ greedy approximation. This approach sequentially constructs the batch by iteratively adding the point that provides the highest marginal utility gain. Specifically, given the set of already selected points X;-1, the next point is chosen as:
F arg max DD
(10)
By employing this greedy selection strategy, we can enforce the diversity and informativeness of data points in one batch. The overall procedure for active CATE estimation is shown in Alg. 1. Moreover, softmax- -oftmax-BBLD,, introduced in (Kirsch et al., 2023) through importance-eiighedd sampling over the pool dataset, was also adopted in CausalBALD; more details and results are provided in App. F66.3
4.6 UNCERTAINTY DECAY ANALYSIS
In this subsection, we analyze the convergence of posterior uncertainty when using the proposed data acquisition function with the CQ estimator. We begin by introducing the following assumption.
Assumption 1. The utility function U is submodular over the pool dataset Xp..
Note that differential entropy often violates this assumption, while our utility function, information gain, satisfies it under mild conditions in the GP framework (Krause et al., 2008; Srinivas et al., 2012). Further justification of this assumption is provided in App. E.1. We then define two key quantities:
Definition 2. We define two key quantities: (1) the maximum information gain about the target set (a1, Zi) from nB observations Dr, and (2) the irreducible uncertainty, which represents the variance of fa, z) given full knowledge of Dp:
def
max T(f(ar, Zi):xx). 1b, def Var[f(a,
(11)
Xlsnn
Finally, we bound the marginal variance of the treatment effect estimator under a given acquisition strategy, showing it is controlled by both the irreducible uncertainty and the information gain:
7

---

Published as a conference paper at ICLR 2026
CATE
ATE
ATT
12 22 20 Number ddsarvataaaa
20 120 20 238 200 100
100 100
203
Number
1o 20 230 201
Figure Comparison of VAMSE on simulation datasets (shaded: standard error). Baselines: Random, 4-BALD, Coresets and TVR. Ours: "G" for greedy, others for top-b acquisition.
Theorem 2. Suppose Assumption 1 holds, and the data acquisition follows utility U (e.g., either the proposed IG or TVR strategy). Let ng represent the total number of individuals with observed outcomes that are acquired from Dp. Under these conditions, there exists constant C > 0 such that for any n.B 2 and for each pair (a, z) 6 (a1, Z1), the marginal variance is bounded as:
Var[f(a, 2)] s nB (a, + +C (7mal /BB)
(12)
The convergence analysis proof our proposed acquisition strategy is presented in App. E, mainly building the transductive active learning framework introduced by Hübotter et al. (2024).
EXPERIMENTAL RESULTS
We validate our approaches for active CQ estimation multiple simulations well the semi- synthetic IHDP (Louizos et al., 2017) and Lalonde Londe, 1986) datasets.
Baselines, Implementations, and Metrics. To evaluate the effectiveness of our proposed framework, we compare it with various baseline acquisition strategies, including random selection, total variance reduction, ,-BALD, and QHTE (which is based the core-set method). To ensure a fair comparison across all settings and methods, we use GP to approximate the regression function, and either Mixture Density Network (MDN) (Bishop & Nasrabadi, 2006) or CME for the conditional distribution estimation. Detailed implementations our proposed methods are provided in App. F.2. We evaluate the performance of methods using the Average Mean Squared Error (AMSE), which quantifies the accuracy of the estimated CQ compared to the true CQ. Note that the uncertainty reduction criterion not well-suited for batch active learning. A detailed explanation is provided in App. F. All results are reported as the mean # standard deviation, computed over 20 independent random test set runs for each configuration.
well as the hyper-paramettes of the baseline methods and
Synthetic Data Analysis. Limited access to counterfactual data often necessitates the use of synthetic or semi-syntheet datasets for evaluating treatment effect estimation methods (Bica et al., 2020; Gao et al., 2024). We design two simulation settings, each involving single conditioning variable. The first includes two adjustment variables to facilitate visualization, while the second uses four adjustment variables for numerical evaluation. For tasks beyond CATE estimation, the conditioning variable is treated as part of the adjustment set and is not explicitly illustrated. All simulation datasets are generated using predefined process, where treatment assignment is influenced by covariates, following a data generation process similar to that in (Abrevaya et al., 2015; Singh et al., 2024). We define two types of target treatments for all CQs: one with a fixed treatment value and another considering all possible treatment values, described in Sec. 4.4. For clarity, we present results for the second case in the main paper, while comprehensive results for all combinations of these settings, both for binary and continuous treatments, are available in App. F.6.3.
Results. Fig. 2 presents the results for CATE, ATE, ATT, and DS on the simulation dataset. Our proposed methods consistently achieve the best performance by prioritizing data that aligns with the target distribution of interest. The corresponding sampling results are shown in Fig. 3, demonstrating that our proposed method can acquire observations well-aligned with the target distribution. Moreover, for CATE, we can see that TVR with CME consistently outperforms MC sampling- baseed methods in both cases, CME directly operates on features relevant for the GP regression task, making it.
8

---

Published as a conference paper at ICLR 2026
0,0M
Teaa
Figure 3: Comparison of acquisition patterns for four methods (left to right): 2-BALD, TVR, IG- CME, and TVR-CME, Representative acquired points are shown for selected methods; complete visualizations are provided in App. F.
more prediction-oriented and efficient compared to estimating conditional densities. More detailed explanations on the advantages of CME in our AL setup are provided in App. G.2. For ATE, all methods, including the baselines, sample from the entire population, leading to similar performance among the uncertainty-awart methods, all of which outperform random acquisition. However, IG- based methods may suffer from numerical instability when computing the determinant of large covariance matrices, potentially leading to suboptimal performance. We also observe that in the ATE with DS case, all our proposed methods significantly outperform the baseline methods, due to the distribution shift between the target and sampling distributions. Additional ablation results on the stability of our methods, considering factors such different starting points, pool dataset sizes, batch sizes, and kernel choices, are provided in App. F.6.2.
Rumning Time Analysis. In Fig. 4, we present the running time for all methods on the CATE estimation task. While baseline methods that use naive top-n selection per acquisition round are slightly faster, our approaches incur additional cost due learning conditional distributions. The overall runtime is influenced by three main factors: (1) greedy acquisition, which selects points sequentially and triggers frequent posterior updates; (2) the pool size, which affects the scale of utility evaluation; and (3) the entropy computation in IG-based strategies. As shown in Fig. 4(b), smaller batch sizes intensify the cost of greedy selection due to repeated posterior recalculations. For GP-based methods, covariate dimensionality has limited impact, as runtime is mainly driven by distance computations. Despite these overheads, all methods remain computationally feasible in our experiments. Further details on computational complexity are provided in App. F.5.
case evaluation, we compare all algorithms on the widely used IHDP benchmark, which includes real covariates paired with simulated outcomes. This benchmark poses significant challenges due to its small size, imbalance, and limited overlap in covariate distributions. Notably, the two potential outcome functions, while supported the same covariates, have distinct functional forms, making the estimation of their difference particularly difficult. Our results show that for CATE and DS scenarios, where substantial distributional shifts exist between the target distribution and the pool data distribution, our proposed those leveraging the CME framework consistently outperform other approaches. A comprehensive description of the dataset, additional experimental results under various setups, and detailed discus- sions are provided in App. F.6.4. We additionally evaluate on the Lalonde dataset; further results and implementation details appear in App. F.6.5.
Semi-synthetic Data Analysis. For
6 RELATED WORK
Previous work on active causal learning has mainly emphasized two aspects: (1) patient recruit- ment (Deng et al., 2011; Song et al., 2023) and (2) selective acquisition of costly outcomes (Jesson et al., 2021; Wen et al., 2025b) to reduce uncertainty in CATE estimation given all covariates. Our work follows the latter, assuming treatments are assigned but outcomes are expensive to obtain. In contrast to prior studies, we aim to efficiently estimate a wider range of causal quantities beyond CATE, A related line of research is active causal inference, exemplified by ABCI (Toth et al., 2022), which leverages GPs and information-theonm theoretic criteria. However, it requires interventional data and is therefore not applicable to the observational setting we consider. This direction has since been extended in follow-up works (Annadani al., 2024; Zhou et al., 2024). Another relevant stream

---

Published as a conference paper at ICLR 2026
CATE
Rummr Eite
LAMAS
bumber Examgre
Number Acquied stamples
Figure 4: Running time comparisons.
Figure 5: Performance comparison on IHDP.
employs kernel-based approaches causal inference (Sejdinovic, 2025), where conditional distribu- tions are embedded an RKHS rather than explicitly estimated, yielding closed-form solutions for multiple causal quantities (Singh et al., 2024; Mastouri et al., 2021). Relatedly, BayesIMP (Chau et al., 2021b) addresses uncertainty quantification with a focus on data fusion. Further discussion of transductive active learning, active GP methods, and comparisons with ABCI provided in App. B.
7 CONCLUSION
In this paper, introduced and formalized the task of Active Estimation of Causal Quantities (ActiveCQ), broadening the field's focus beyond the commonly studied CATE. We proposed unified Bayesian framework for this new research direction that represents CQs integrals, pairing a Gaussian process for the regression function with conditional mean embeddings for the target distribution. This powerful combination allows for the model's adaptive refinement and enables the systematic derivation of uncertainty-awe aware utility functions. We demonstrated this by instantiating two strategies based on information gain and total variance reduction. Through simulations and semi- synthetic datasets, we showed that our principled approach significantly outperforms traditional active learning baselines, including BALD and vanilla TVR. Our work lays foundation for numerous future directions. While our GP-based instantiation inherits cubic complexity, the framework 's modularity key strength, inviting future work that integrates scalable approximations like sparse GPs or Nyström methods. Furthermore, proposed framework can be extended handle more complex causal settings involving hidden confounders or instrumental variables.
ACKNOWLEDGMENTS
This work is supported by the Responsible Research Centre (RAIR). BG also supported by the Centre for Augmented Reasoning (CAR).
ETHICS STATEMENT
This work does not involve human subjects, animal experiments, or sensitive personal data. The datasets used are publicly available and widely adopted in the research community. Our study does not raise concerns regarding privacy, security, discrimination, bias, or potential misuse. We have carefully followed the ICLR Code Ethics and confirm that research practices in this paper adhere to principles of integrity, transparency, and fairness.
REPRODUCIBILITY STATEMENT
We have taken several measures to ensure the reproducibility of our work. A detailed description of the datasets and preprocessing steps is provided in App. F.I. The implementation details of our method, including model architectures, training procedures, baselines, well comprehensive experimental settings and hyperparameter choices, are documented in App. E22 and App. D. An in-depth analysis of the experimental results is included in App.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Jason Abrevaya, Yu-Chin Hsu, and Robert P Lieli. Estimating conditional average treatment effects. Journal of Business & Economic Statistics, 33(4):485-50) 505, 2015.
Virginia Aglictti, Xiaoyu Lu, Andrei Paleyes, and Javier González. Causal bayesian optimization, In International Conference on Artificial Intelligence and Statistics, pp. 3155-3164, PMLR, 2020.
Ahmed M Alaa and Mihacla Van Der Schaar. Bayesian inference of individualized treatment effeets using multi-task gaussian processes. In Advances in neural information processing systems, volume 30, 2017.
Yashas Annadani, Panagiotis Tigas, Stefan Bauer, and Adam Foster. Amortized active causal induction with deep reinforcement learning. Advances in Neural Information Processing Systems, 37:44216-44299, 2024.
Kjell Benson and Arthur Hartz. A comparison of observational studies and randomized, controlled trials. In Research ethics, pp. 213-221. Routledge, 2017.
Wenya Linda Bi, Ahmed Hosny, Matthew B Schabath, Maryellen L Giger, Nicolai Birkbak, Alireza Mehrtash, Tavis Allison, Omar Arnaout, Christopher Abbosh, Ian F Dunn, et al. Artificial intelligence in cancer imaging: clinical challenges and applications. CA: a cancer journal for clinicians, 69(2):12--177, 2019.
Ioana Bica, James Jordon, and Mihaela van der Schaar. Estimating the effects of continuous valued interventions using generative adversarial networks. In Advances in Neural Information Processing Systems, volume 33, pp. 16434-16445, 2020.
Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.
Kathryn Chaloner and Isabella Verdinelli. Bayesian experimental design: review. Statistical science, pp. 273-304, 1995.
Siu Lun Chau, Shahine Bouabid, and Dino Sejdinovic. Deconditional downscaling with gaussian processes. In Advances in Neural Information Processing Systems, volume 34, pp. 17813-17825 2021a.
Siu Lun Chau, Jean-Francois Ton, Javier González, Yee Teh, and Dino Sejdinovic. Bayesimp: Uncertainty quantification for causal data fusion. In Advances in Neural Information Processing Systems, volume 34, pp. 3466-3477, 2021b.
Zonghao Chen, Masha Naslidnyk, Arthur Gretton, and Francois-Xate Xaviee Briol. Conditional bayesian quadrature. In Negar Kiyavash and Joris M. Mooij (eds.), Proceedings of the Fortieth Conference on Uncertainty in Artificial Intelligence, volume 244 of Proceedings of Machine Learning Research, pp. 648-684. PMLR, 15-19 Jul 2024.
David A Cohn, Zoubin Ghahramani, and Michael Jordan. Active learning with statistical models. Journal of Artificial Intelligence Research, 4:129-145. 1996.
Bethany Connolly, Kim Moore, Tobias Schwedes, Alexander Adam, Gary Willis, llya Feige, and Christopher Frye. Task-specific experimental design for treatment effect estimation. In Interna- lional Conference on Machine Learning, pp. 6384-6401. PMLR, 2023.
Hugh Dance, Peter Orbanz, and Gretton Arthur. Spectral representations for accurate causal uncer- tainty guantification with gaussian processes. arXiv preprint arXiv:2440. 2024.
Kun Deng, Joelle Pineau, and Susan Murphy. Active learning for personalizing treatment. In 2011 IEEE Symposium Adaptive Dynamic Programming and Reinforcement Learning (ADPRL), pp. 32-39. IEFE, 2011
Peng Ding. A first course in causal inference. CRC Press, 2024.
11

---

Published as a conference paper at ICLR 2026
Jake Fawkes, Lucile fer-Minassian, Desi R. Ivanova, Uri Shalit, and Chris Holmes. Is merging worth it? securely evaluating the information gain for causal dataset acquisition. In International Conference Artificial Intelligence and Statistics (AISTATS). PMLR, 2025.
Jared C Foster, Jeremy MG Taylor, and Stephen Ruberg. Subgroup identification from randomized clinical trial data. Statistics in medicine, 30(24))2867-2880 2011.
Erdun Gao, Howard Bondell, Wei Huang, and Mingming Gong. A variational framework for estimating continuous treatment effects with measurement error. In The Twelfth International Conference on Learning Representations, 2024.
Erdun Gao, Jake Fawkes, and Dino Sejdinovic. Causal-epig: prediction-oriented active learning framework for cate estimation. arXiv preprint arXiv:2509.21866 2025.
Claudio Gentile, Zhilei Wang, and Tong Zhang. Fast rates in pool-based batch active learning. Journal of Machine Learning Research, 25(262): 2024.
James Heckman. Causal parameters and policy analysis in economics: A twenticth century retrospective. The Quarterly Journal of Economics, 1 115(1):45-77, 2000.
M.A. Hernan and J.M. Robins. Causal Inference: What If. Chapman & Hall/CRC Monographs on Statistics & Applied Probab. CRC Press, 2023. ISBN 978 7781420076155.
Jennifer Hill. Bayesian nonparametric modeling for causal inference. Journal of Computational and Graphical Statistics, 20(1) 20(1)217-240, 2011.
David Holzmüller, Viktor Zaverkin, Johannes Kästner, and Ingo Steinwart. framework and benchmark for deep batch active learning for regression. Journal of Machine Learning Research, 24(164):1-81, 2023.
Neil Houlsby, Ferene Huszár, Zoubin Ghahramani, and Máté Lengyel. Bayesian active learning for classification and preference learning. arXiv preprint arXiv: 2011.
Robert Hu, Dino Sejdinovic, and Robin J Evans. A kernel test for causal association via noise contrastive backdoor adjustment. Journal of Machine Learning Research, 25(160):1-56, 2024.
Jonas Hübotter, Bhavya Sukhija, Lenart Treven, Yarden As, and Andreas Krause. Transductive active learning: Theory and applications. In The T'irttyeightt Annual Conference on Neural Information Processing Systems, 2024.
Guido W Imbens and Donald B Rubin. Causal inference in statistics, social, and biomedical sciences. Cambridge university press, 2015.
Amin Jaber, Jiji Zhang, and Elias Bareinboim. Identification of conditional causal effects under markov equivalence. Advances in Neural Information Processing Systems, 32, 2019.
Andrew Jesson, Panagiotis Tigas, Joost van Amersfoort, Andreas Kirsch, Uri Shalit, and Yarin Gal. Causal-bald: Deep bayesian active learning of outcomes to infer treatment-effects from observational data. In Advances in Neural Information Processing Systems, volume 34, pp.
30465-30478, 2021.
Fredrik D Johansson, Uri Shalit, Nathan Kallus, and David Sontag. Generalization bounds and representation learning for estimation of potential outcomes and causal effects. Journal of Machine Learning Research, 23(166):1-50, 2022.
Nathan Kallus and Xiaojie Mao. On the role surrogates in efficient estimation of treatment effects with limited outcome data. Journal of the Royal Statistical Society Series B: Statistical Methodology, pp. qkae099, 2024.
Motonobu Kanagawa, Philipp Hennig, Dino Sejdinovic, and Bharath K Sriperumbudur. Gaussian pro- cesses and reproducing kernels: Connections and equivalences. arXiv preprint arXiv:2506.17366, 2025.
12

---

Published as a conference paper at ICLR 2026
Masahiro Kato, Akihiro Oga, Wataru Komatsubara, and Ryo Inokuchi. Active adaptive experimental design for treatment effect estimation with covariate choice. In International Conference on Machine Learning. PMLR, 2024.
Andreas Kirsch. Black-bol box batch active learning for regression. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. Expert Certification.
Andreas Kirsch, Joost Van Amersfoort, and Yarin Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning. In Advances in neural information processing systems, volume 32, 2019.
Andreas Kirsch, Sebastian Farquhar, Parmida Atighehchian, Andrew Jesson, Frédéric Branchaud- Charron, and Yarin Gal. Stochastic batch acquisition: A simple baseline for deep active learning. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https://openreview net
jw. Expert Certification.
Omer Noy Klein, Alihan Hüyük, Ron Shamir, Uri Shalit, and Mihaela der Schaar. Towards regulatory-confiimee adaptive clinical trials: Machine learning opportunities and solutions. In International Conference Artificial Intelligence and Statistics, pp. 4969-4977. PMLR, 2025.
Suraj Kothawade, Nathan Beck, Krishnateja Killamsetty, and Rishabh Iyer. Similar: Submodular information measures based active learning in realistic scenarios. Advances in Neural Information Processing Systems, 18118888516697 2021.
Andreas Krause, Ajit Singh, and Carlos Guestrin. Near-optimal sensor placements in gaussian processes: Theory, efficient algorithms and empirical studies. Journal of Machine Learning Research, 9(2), 2008
Robert LaLonde. Evaluating the econometric evaluations of training programs with experimental data. The American economic review, pp.
1986.
Sara LaPlante and Emilija Perkovic. Conditional adjustment in markov equivalence class. In International Conference on Artificial Intelligence and Statisties, pp. 2782-2790. PMLR, 2024.
Dennis V Lindley. On measure of the information provided by an experiment. The Annals of Mathematical Statisties, 27(4):986-1005, 1956.
Christos Louizos, Uri Shalit, Joris M Mooij, David Sontag, Richard Zemel, and Max Welling. Causal effect inference with deep latent-variable models. In Advances Neural Information Processing Systems, volume 30, 2017.
David JC MacKay, Information -based objective functions for active data selection. Neural Computa- tion, 4(4):590-604, 1992.
Afsaneh Mastouri, Yuchen Zhu, Limor Gultchin, Anna Korba, Ricardo Silva, Matt Kusner, Arthur Gretton, and Krikamol Muandet. Proximal causal learning with kernels: Two-stage estimation and moment restriction. International conference on machine learning, pp. 7512-7523, PMLR, 2021.
David McKenzie, Beyond baseline and follow-up: The case for more in experiments. Journal of development Economics, 99(2):210-2)) 121211,, 2012.
Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, Bernhard Schölkopf, et al. Kernel mean embedding of distributions: A review and beyond. Foundations and Trends@ in Machine Learning, 10(1 1141,, 2017.
Chioma Susan Nwaimo, Ayodeji Enoch Adegbola, and Mayokun Daniel Adegbola. Transforming healthcare with data analytics: Predictive models for patient outcomes. GSC Biological and Pharmaceutical Sciences, 27(3):025-035, 2024.
Ezinne Nwankwo, Lauri Goldkind, and Angela Zhou. Batch-adaptive annotations for causal inference with complex-embedded outcomes. arXiv preprint arXiv:2502, /0605 2025.
13

---

Published as a conference paper at ICLR 2026
Junhyung Park, Uri Shalit, Bernhard Schölkopf, and Krikamol Muandet. Conditional distribu- tional treatment effect with kernel conditional mean embeddings and u-statistic regression. In International Conference Machine Learning, pp. 8401-8412. PMLR, 2021.
J Pearl. Causality. Cambridge university press, 2009.
Tian Qin, Tian-Zuo Wang, and Zhi-Hua Zhou. Budgeted heterogeneous treatment effect estimation. In International Conference on Machine Learning, pp. 8693-8702. PMLR, 2021.
Tom Rainforth, Adam Foster, Desi R Ivanova, and Freddie Bickford Smith. Modern bayesian
experimental design. Statistical Science, 39(1):100-114 2024
Tongzheng Ren, Haotian Sun, Antoine Moulin, Arthur Gretton, and Bo Dai. Spectral representation for causal estimation with hidden confounders, In NeurIPS 2024 Causal Representation Learning Workshop, 2024.
Donald B Rubin. Causal inference using potential outcomes: Design, modeling, decisions. Journal of the American Statistical Association, 100(469):322-3311 2005.
Dino Sejdinovic. An Overview of Causal Inference using Kernel Embeddings. Data Science for Econometrics and Related Topics, Studies in Systems, Decision and Control. Springer, 2025.
Burr Settles. Active learning literature survey. Machine Learning, 15(2):201-25 1994.
Rahul Singh, Mancesh Sahani, and Arthur Gretton. Kernel instrumental variable regression. In Advances in Neural Information Processing Systems, volume 32, 2019.
Rahul Singh, Liyuan Xu, and Arthur Gretton. Sequential kernel embedding for mediated and time-varying dose response curves. arXiv preprint arXiv:211 1.03900..0.,,, 2021.
Rahul Singh, Liyuan Xu, and Arthur Gretton. Kernel methods for causal functions: dose, heteroge- neous and incremental response curves. Biometrika, 111(2):497-51) 2024.
Freddie Bickford Smith, Andreas Kirsch, Sebastian Farquhar, Yarin Gal, Adam Foster, and Tom Rainforth. Prediction-oriented bayesian active learning. In International Conference Artificial Intelligence and Statistics, pp. 7331-7348. PMLR, 2023.
Difan Song, Simon Mak, and CF Wu. Ace: Active learning for causal inference with expensive experiments. In ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Workshop Causal Inference and Machine Learning in Practice, 2023.
Le Song. Kenji Fukumizu, and Arthur Gretton. Kernel embeddings of conditional distributions: A unified kernel framework for nonparametric inference in graphical models. IEEE Signal Processing Magazine, 30(4):98-111, 2013.
Niranjan Srinivas, Andreas Krause, Sham M Kakade, and Matthias W Seeger. Information-thhorret regret bounds for gaussian process optimization in the bandit setting. IEEE transactions on information theory, 58(5):3250-324) 2012.
Masashi Sugiyama, Ichiro Takeuchi, Taiji Suzuki, Takafumi Kanamori, Hirotaka Hachiya, and Daisuke Okanohara. Conditional density estimation via least-squares density ratio estimation. In International Conference on Artificial Intelligence and Statistics, pp. 781-788, 2010.
Christian Toth, Lars Lorch, Christian Knoll, Andreas Krause, Franz Pernkopf, Robert Peharz, and Julius Von Kügelgen. Active bayesian causal inference. Advances in Neural Information Processing Systems, 35:16261-16275, 2022.
Allen Tran, Aurelien Bibaut, and Nathan Kallus. Inferring the long-ten term causal effects of long-term treatments from short-term experiments. In International Conference on Machine Learning, pp- 48565-48577. PMLR, 2024.
Brian L Trippe and Richard E Turner. Conditional density estimation with bayesian normalising flows. In Advances in Neural Information Processing Systems, Workshop Bayesian Deep Learning, 2018.
14

---

Published as a conference paper at ICLR 2026
Russell Tsuchida, Cheng Soon Ong, and Dino Sejdinovic. Squared Neural Families: A New Class of Tractable Density Models. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
Russell Tsuchida, Jiawei Liu, Cheng Soon Ong, and Dino Sejdinovic. Squared families: Searching beyond regular probability models. arXiv preprint arXiv:2503.21128, 2025.
Dennis C Turk. Clinical effectiveness and costteffeectveeess of treatments for patients with chronic pain. The Clinical journal of pain, 18(6):355 -655, 2002.
Hechuan Wen, Tong Chen, Mingming Gong, Li Kheng Chai, Shazia Sadiq, and Hongzhi Yin. Enhancing treatment effect estimation via active learning: A counterfactual covering perspective. In International Conference on Machine Learning, pp. 66437-66466. PMLR, 2025a.
Hechuan Wen, Tong Chen, Guanhua Ye, Li Kheng Chai, Shazia Sadiq, and Hongzhi Yin. Progressive generalization risk reduction for data-efficient causal effect estimation. In Proceedings of the ACM SIGKDD Conference on Knowledge Discovery and Data Mining, 2025b.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for muchine learning, volume 2. MIT press Cambridge, MA, 2006.
Janine Witte, Leonard Henckel. Marloes H Maathuis, and Vanessa Didelez. On efficient adjustment in causal graphs. Journal of Machine Learning Research, 21(246):1 45,, 2020.
Liyuan Xu, Yutian Chen, Siddarth Srinivasan, Nando de Freitas, Arnaud Doucet, and Arthur Gretton. Learning deep features in instrumental variable regression. In International Conference on Learning Representations, 2021.
Qinyi Zhang, Sarah Filippi, Arthur Gretton, and Dino Sejdinovic. Large-scale kernel methods for independence testing, Statistics and Computing, 28:113-,30, 2018.
Zihan Zhou, Muhammad Qasim Elahi, and Murat Kocaoglu. Sample efficient bayesian learning of causal graphs from interventions. In The Thirty-eighth Anmual Conference on Neural Information Processing Systems, 2024.
Yuchen Zhu, Limor Gultchin, Arthur Gretton, Matt J Kusner, and Ricardo Silva. Causal inference with treatment measurement eтror: a nonparametric instrumental variable approach. In Uncertainty in Artificial Intelligence, pp. 2414 2424. PMLR, 2022.
15

---

Published as a conference paper at ICLR 2026
Appendix
Table of Contents
A Acknowledgment of LLM Usage
17
B More related works
17
B.I Related works (continuing)
17
B.2 Transductive Active Learning
17
B.3 In-depth Discussion of ABCI
17
Further Preliminaries and Definitions
DAG and Causal Graphs
18
C.2 Kernels and RKHS
20
C.3 Information-Theonee Mcasures
21
D Modeling for Various Causal Quantities
21
D.1 Derivation for the CATE
................................
22
Derivation for ATE
22
Derivation for ATT
23
D.4 Derivation for ATEDS
24
D.5 Connection between IG and TVR
25
E
Convergence Analysis
26
E.1
Justification of Assumption I.
26
E.2
Preliminary definitions and concepts
26
Approximate Markov Boundary
27
E.4
Proof of Thcorem
29
F
Experiments Details
29
F.1
Datasets
30
F2
Implementation Details
34
F3
Hardware... .ladwware.......................
................................
35
Metrics
................................
36
F5
Computational Complexity
................................
37
F6
Additional Experimental Results
38
G Discussions
47
G.I Summary of Contributions and Future Directions
47
G.2 The Advantages of CME in the Active Learning Setup
48
16

---

Published as a conference paper at ICLR 2026
A ACKNOWLEDGMENT OF LLM USAGE
In preparing this manuscript, LLMs were employed solely general- purposs writing aids. Their use was limited to word- and sentence-level polishing, including correcting typos, improving grammar, and refining phrasing. LLMs were not used for research ideation, scientific analysis, generation of results, or interpretation of findings. All scientific concepts, methods, analyses, and conclusions presented in this work are entirely the responsibility of the authors.
B MORE RELATED WORKS
B. RELATED WORKS (CONTINUING)
Building on the idea of learning causal quantities in RKHS, several works have extended this approach to more complex settings. For instance, Singh et al. (2019) and Mastouri et al. (2021), as well as Ren et al. (2024), address scenarios involving unmeasured confounders by leveraging instrumental variables and proxy variables, respectively. Additionally, Zhu et al. (2022) focus on settings with measurement error in treatments, while Singh et al. (2021) explore applications in sequential data. Other advancements include kernel-based methods for causal hypothesis testing (Hu et al., 2024) and modeling distributional treatment effects (Park et al., 2021). BayesIMP (Chau al., 2021b) introduced the mean process technique to incorporate uncertainty in causal estimation, particularly to handle dataset fusion problems. This work also investigated the uncertainty arising from learning the conditional distribution and proposed method to integrate uncertainties across multiple stages. Subsequently, IMPspec (Dance et al., 2024) introduced the spectral representation of Hilbert spaces to address the limitations restricted nuclear-dominant kernels identified in BayesIMP, Sejdinovic (2025) provided an overview of using kernel methods to address causal inference problems. Our framework has the potential to integrate these advanced techniques, enabling perhaps further improvements in performance. Technically, our proposed method for modeling different CQs using CME and GP closely related to the conditional mean process (Chau et al., 2021a) and conditional Bayesian quadrature (Chen et al., 2024).
B.2 TRANSDUCTIVE ACTIVE LEARNING
Traditional active learning (TAL) implicitly assumes that the observational data, including both the training dataset and the pool dataset, represent the overall target distribution of interest (Holzmüller et al., 2023; Kirsch, 2023). Under this assumption, information gain-based methods naturally select the most informative points from the pool dataset, aligning effectively with the learning objective. Recent research, however, has begun addressing scenarios where the target distribution differs from the observational distribution, framing this challenge as a transductive active learning problem (Hübotter et al., 2024; Smith et al., 2023; Kothawade et al., 2021). Our problem naturally aligns with TAL., enabling us draw upon similar insights to address the challenges we face. The key distinction lies in the target quantities: while TAL typically focuses optimizing the average performance across all data points within a distribution, our work emphasizes average metries over specific subpopulations, However, if we shift our focus to point-specific performance metrics, such as the conditional individual treatment effect or the individual treatment effect, our framework can directly benefit from TAL insights. Furthermore, unlike previous approaches, our framework requires the estimation of conditional distributions. By representing and learning these distributions in an RKHS, we provide more efficient and task-oriented solution. This feature-driven, task-specific approach to conditional distribution estimation not only enhances our framework but also introduces new perspectives for advancing transductive active learning research.
B.3 IN-DEPTH DISCUSSION OF ABCI
Here, we provide detailed comparison with the Active Bayesian Causal Inference (ABCI) method proposed by Toth et al. (2022), an active learning framework for integrated causal discovery and reasoning. ABCI is designed for settings where the causal structure is unknown and must be learned from interventional data, using information gain to select experiments and GPs to model causal mechanisms. While work shares high-level themes with ABCI, such the use of GPs and
17

---

Published as a conference paper at ICLR 2026
intommation-thoorttii acquisition functions, the two frameworks diverge fundamentally in their problem settings, methodological approaches, and core contributions.
Divergent Problem Settings. The most critical distinction lies in the problem formulation. ABCI operates in a setting where the causal graph is unknown: its primary goal is to actively select interventions jointly learn this structure and its underlying mechanisms. The uncertainty it seeks to resolve pertains to both the causal graph and the functional relationships, which it addresses by acquiring new interventional data. In contrast, our framework assumes the causal structure is known and concentrates on the efficient estimation of causal quantities from observational data. Our approach does not generate new data through interventions. Instead, it strategically selects the most informative samples from a fixed, pre- pre-existing pool of individuals with assigned treatments. This difference situates the two methods parallel yet distinct lines of research. ABCI is concerned with causal discovery via intervention, whereas our work is tailored for causal estimation from observation. Consequently, ABCI not directly applicable as a baseline in our setting. Even if adapted to assume known graph, its core mechanism of acquiring data through new interventions is incompatible with a fixed observational context.
Distinct Methodological Contributions. The methodological foundations of the two frameworks are also markedly different. Our work introduces a unified Bayesian framework for actively estimating broad spectrum of CQs, scope that extends beyond that of ABCI. A key technical innovation in our approach is the integration of Conditional Mean Embeddings with Gaussian Processes. This combination yields flexible and principled estimators capable of targeting diverse causal queries. This modeling strategy is entirely different from that of ABCI, which uses GPs to capture individual causal mechanisms within a DAG, Furthermore, we provide formal performance guarantee by deriving theoretical uncertainty decay rate for our acquisition function, contribution not present in ABCI. Therefore, although both methods leverage GPs and information- thhooretic ideas, these are superficial similarities. In our framework, these tools are applied to regression-basee estimation from observational data to answer specific causal queries. In ABCI, they are used to facilitate structural learning. Ultimately, the methods pursue different objectives under incompatible assumptions and are not methodologically comparable.
C FURTHER PRELIMINARIES AND DEFINITIONS
In this paper, upright Roman letters (e.g., v) represent graph nodes and their associated random variables: calligraphic letters (e.g.. V") denote measurable spaces; and italic letters (e.g., v v) indicate specific realizations. Bold letters (e.g., v) refer to sets of nodes or random vectors, while bold italic letters (e-g.. = v)) are used for their realizations.
C.I DAG AND CAUSAL GRAPHS
Directed Acyelic Graph (DAG). Let G = (v,c) represent a graph consisting of set of nodes (variables) v = (V1,'-, Vp}, and set of edges e. G is called a DAG if it contains only directed edges (-) and no directed cycles in G. Let p(v) v) an observational density over v, which is said to be Markov compatible with DAG G ev if it factorizes as p(v) = -v,vv p( p(v:lpa(:; G)), where pa(r,,9 9) includes the values of the parent nodes of v, g. We also assume the positivity condition, meaning we restrict our attention to distributions where p(v) 0 for all valid values of v.
Causal DAGs. Causality is designed to describe physical processes in the real world, and DAGs serve as a powerful framework for formally representing this concept (Pearl, 2009). When G is a DAG, it is considered a causal DAG if every directed edge V; + V, represents a direct causal effect of v, on Vj. Let a v be a node set in a causal DAG g. The intervention do(a = a) 2, or do(a) for short, denotes outside intervention that sets density p( (z)d( (a) is density resulting from such an intervention. Let P represent the set of all interventional densities pa)) We say that a graph G = (v,e) is compatible with P if and
fixed values a. An interventional
II the main paper, restrict the treatment variable a single dimension, whereas in the appendix, this constraint relaxed multiple variables.
18

---

Published as a conference paper at ICLR 2026
only if, for every p(oldo(a)) PP, the following condition holds:
p(oldo(a)) II (a
(13)
V.EVsa
We say that an interventional density p(oldo(a)) (a)) consistent with G if belongs set of inter- ventional densities Pt, where G is compatible with P (Pearl, 2009). Moreover, any observational density that is Markov compatible with G is also consistent with 9 (LaPlante & Perkovic, 2024).
We then present set of standard assumptions to establish the identifiability of the causal quantities discussed in the main paper, including the ATE, CATE, ATT, and ATEDS, It is important to note that the assumptions provided here are sufficient but not necessary conditions for identifying these causal quantities, with a primary focus on scenarios involving unmeasured confounders. Beyond these assumptions, alternative frameworks such as instrumental variable assumptions (Xu et al., 2021), front-door criteria (Chau et al., 2021b), and proxy variable approaches (Mastouri et al., 2021) can also enable identifiability in certain cases. However, these beyond the current scope of our work. Extending our framework to accommodate such conditions remains an interesting direction for future research. Throughout this paper, we make the following assumptions:
Assumption (Stable Unit Treatment Value Assumption (SUTVA)). The relationship between the observed outcome and the potential outcomes is stable. This entails two key conditions:
1. No Interference: The potential outcome of any individual unit depends only the treatment assigned to that unit and is unaffected by the treatment assignments of other units.
2. Consistency: For any unit, if the treatment a is set to a value equals the potential outcome y(a). Formally, y y(a) ifa a.
a, the observed outcome
This assumption ensures that the mapping from treatment to outcome well-defined and allows us to link observed data to theoretical causal quantities.
Assumption (Causal Markov Compatibility). The joint distribution of the variables is consistent with a causal DAG G, meaning:
p(vldo(a))= II v,Evsa
(14)
Here, pa( 9) denotes the set of parent variables v; in the graph G.
Assumption 4 (Backdoor Criterion). There exists a of variables s S (the backdoor adjustment set) such that s blocks all backaloor paths from the treatment a to the outcome in the graph G. Formally:
a a in the graph (aayyy) This ensures that adjusting for S removes confounding bias in the causal effect of a on y. Assumption 5 (Positivity). The treatment assignment is well-defined for all levels of the adjustment set s. Formally:
(15)
p(a= als= = s)2 >0 Va,s. 8,
(16)
where s is a realization of the adjustment set s. This assumption ensures that every treatment level has a non-zero probability of being observed for all values of s. Notice that for the CATE case, the adjustment set effectively encompusses both and z.
For the case of ATE with distribution shift, it is essential to impose constraints the shifted distribution IP to ensure the identifiability of this causal quantity. In addition to the assumptions outlined previously, we further introduce the following assumption:
Assumption 6 (Distribution shift (Singh et al., 2024)). Assume the following:
1. The joint distribution of y and under the shifted distribution is given by:
i(y, a,y,,s) a,s) IPyya sa,
(17)
2. The shifted marginal distribution (P(a,))) s) is absolutely continuous with respect to the original distribution sP.
Under the above assumption, populations TP and PS differ only in the distribution of the treatment variable and covariates. Furthermore, the support of encompasses the support of P, ensuring that the treatment effect well-defined across both populations. As a direct implication, the regression function remains invariant between IP and TP, regardless of the distributional shift.
19

---

Published as a conference paper at ICLR 2026
Z
S
e
2
S
S
a
y
a
y
a
Observational (CATE)
(CATE)
Simplified graph
Figure 6: Data generation processes depicted by DAGs, (a) and (b) illustrate the full causal graph for the CATE setting under observational and interventional distributions, respectively. (c) shows the simplified graph used for ATE, ATT, and ATEDS settings.
Difference from the PO Framework. The Potential Outcomes (PO) and causal graph frameworks provide distinct approaches for defining causal effects, differing in their foundational principles and particularly in their handling of continuous treatments. The PO framework defines causal effects via hypothetical counterfactuals, namely the potential outcomes under different treatment scenarios (Ding, 2024). contrast, the CG framework defines them directly from the structural relationships encoded in a graph, representing explicit causal assumptions (Pearl, 2009). While both frameworks align conceptually for binary treatments (e.g.. ATE as the expected difference between treated and untreated groups), their approaches diverge for continuous treatments. The CG framework generalizes more naturally. Causal quantities like the ATE and CATE are consistently defined through interventions the graph, corresponding directly to the Average Dose-Responss Function (ADRF) and its conditional version, respectively (Hernan & Robins, 2023). The PO framework. however, must extend its definition to a function of treatment doses, often requiring additional modeling assumptions like smoothness monotonicity to estimate the dose-response curve. In essence, the primary distinction is the representation of causality: PO models counterfactual outcomes explicitly, while CG models the causal mechanisms that generate them, offering more unified approach for both binary and continuous interventions.
ITE, CATE, and Our CATE. We distinguish our definition of the CATE from two related concepts in the literature: the Individual Treatment Effect (TTE) and the commonly used definition of CATE. Our entire analysis is grounded in the structural causal model framework, using the do-operator define causal quantities. key aspect of our approach that we focus on defining the conditional potential outcome under specific treatment To formalize these concepts, we consider a model (e.g.. Fig. 6) that includes observed covariates (s, z) and unobserved random vector e. This vector represents all remaining individual-level latent factors that are not confounders, ensuring that conditioning on (s, z, e) fully captures an individual's characteristics. The ITE is theoretical construct representing the causal effect for single individual i. It is defined by conditioning on all their characteristics, both observed and unobserved:
a, rather than the contrast between two different treatments.
TYEE (a; individual i) == Ey do(a = a),s 8,, Z,, e e
(18)
The ITE is impossible to estimate in practice because the latent vector e; is unknown. To achieve practical estimation, the literature typically focuses the CATE, which averages over the unobserved factors e. common definition conditions on the full set of observed covariates:
TCATE (a: (a:s, s, s, := E[y do(a = a), s s, 8,
(19)
This quantity represents the average treatment effect for the subpopulation of individuals sharing the same observed characteristics (s, 2). However, this definition can be restrictive, it compels the analysis to be conditioned on the entire set of observed covariates. In this work, we take a more flexible and practical definition of CATE. Our framework allows conditioning on any chosen subset of observed covariates, which we denote by z. The standard definition above thus becomes a special case of our framework where includes all observed covariates. This flexibility allows for the estimation of treatment effects for broader or more specific subpopulations based on the context of the analysis.
C.2 KERNELS AND RKHS
Reproducing Kernel Hilbert Spaces (RKHS). Consider any space A,B E (Z,S,X,) X.)). and let k A x A R be positive semi-definite kernel. The canonical feature map associated with
20

---

Published as a conference paper at ICLR 2026
this kernel is Pa == k(),)) for any A. The features matrix A is constructed by stacking the feature maps of each row of A as the columns, i.c. D == [ AAn.I "'A.... where A := [ai.".,,n,] We denote the Gram matrix as KA == 4 and the vector of evaluations A as (t(a,a)). "',,,,,,,, a, and Aa ka The RKHS spanned by this kernel is denoted by HA. consists of real-valued functions defined on A and is endowed with the inner product (.. We use @ and @ to represent tensor product and Hadamard product respectively and use HAS to represent the product space HA Hg. For any distribution I,, on A, we define the kernel mean embedding of I as Ma == J k(a, PP.(aa))
HA. which also corresponds to the Riesz representer of expectational functional f + [a)(a)]= = (a.La)HA due to the reproducing property that (f, ((, H) f(a) for all f E H.A Extending to the conditional distribution, for example P Palb-b, we define Pab == La k(a,) PPjB (P)bbd) E HA as the conditional mean embedding of Fa/b-b-
The tensor product RKHS expands the traditional RKHS framework to effectively manage functions with multiple arguments. Given two RKHSs, HA and H., associated with kemels k(a (a.)) and k(b,,). the tensor product RKHS H. H.A Hy creates space of functions defined over the Cartesian product A B, where A and B represent the input spaces of the respective RKHSs. The kernel of this tensor product space is the product of the individual kernels:
K ((a, b), (a', (a',b') b')) ka a') kb,
(20)
where (a. b). (a',b') b') A x B. The feature map for this space is the tensor product of the individual feature maps, (a, b) (a) (b))
b). The norm is multiplicative for elementary tensors: ( dus b) H = (()(())l HA. l (b) Ha. This framework is essential for modeling functions
f: A x R by capturing both the main effects of the variables and their interactions.
C.3 INFORMATION- THEORETIC MEASURES
This section provides concise definitions for the information-theone quantities used in work. Differential Entropy. For a continuous random variable a with probability density function (PDF) p(a) over its support A, the differential entropy quantifies its uncertainty:
H(a) := SA- -lo loggp(a pa)) dlP (a).
(21)
Unlike discrete entropy, differential entropy can negative and depends on the coordinate system. Conditional Entropy. The conditional entropy H(a b) measures the remaining uncertainty in given b:
H(a S JAMB log b) (aat(da, db).
(22)
Mutual Information. The mutual information T(a; b) quantifies the reduction in uncertainty about one variable from observing the other. It is defined as:
I(a: b) H(a) H(a b),
(23)
which is symmetric and can also be expressed using the joint entropy H(a, b):
I(a; = H(a) + H(b) - H(a,b).
(24)
These measures are fundamental to acquisition functions like BALD (Houlsby al., 2011), which seeks to maximize the mutual information between model parameters and unknown labels.
D MODELING FOR VARIOUS CAUSAL QUANTITIES
This section expands on the modeling of various CQs. As established in the main paper, our framework for the CATE integrates several key components. We began by modeling the outcome regression surface with a GP. The central challenge, evaluating the integral the CATE definition, was addressed by leveraging CMEs to represent the conditional distribution Pal.. This technique provides an analytic, closed-form posterior distribution for the CATE estimator. which is crucial for quantifying uncertainty and enabling effective active learning. Building on this foundation, now extend the methodology to other key causal quantities. We will detail the specific modeling adaptations and estimator formulations for the ATE, ATT, and ATEDS.
21

---

Published as a conference paper at ICLR 2026
D.1 DERIVATION FOR THE CATE
Here, we provide detailed, step-by-step derivation for the posterior mean v(a, x) and covariance ((1, 2), (a', 2 of the CATE estimator, as presented in Prop. The entire derivation is based on expressing the standard GP posterior formulas in terms of inner products in the RKHS. The key insight is the use an effective feature map O := &a @ Ø₂ @ As which allows the CATE estimation (an expectation over to be treated standard prediction problem.
The following equations show the expansion of these inner products.
The derivation for v(a, =) begins with the abstract inner product between the effective feature map and the GP posterior mean function, (da, mi) HA28
Similarly, the derivation for q( (a, 2), (a', =²)) starts with the abstract formula involving the prior kernel, (Φa (arr, and the update term, which depends on the cross-covariance inner products (+a; )-
The subsequent steps in the equations expand these abstract inner products into concrete matrix-vector operations. The notation (...)((...) is used to represent these operations in feature space. The underbrace annotations explicitly identify how these expanded expressions correspond to the final, compact terms like the effective cross-covariance vector kax,, the effective kernel kaw,, and the full training Gram matrix Xr,
(a, 2) (de, m
(oo @ oz @ P=|=)² (oar @ 2z @ 25 (Xxx) + +xd1) yxr
Kaay @ k o kzzz Kzz A1) K KSTT
RaX,
(25)
(Kaayar o K2,27 O Ks, + +
yxr
(KxpXx+h11)-
= kax, (Kx, x11) - yXr:
(((a,2) (d', = (aanda') - (dn, 1 xx (Xx,)(XXx) (Xxxx,
2wr)
Kin Kiz (k2z (Kzz(Kzzz + AI) - Kzz + AI) -kzx
(Waar O kzzz O (kizz (K -
1EXr
(26)
(Kayar @ K K2 K Ks Arl)
(Kx,Xx+>il)>
(Kaya' k2zz @ (Ks sKKzsKKzZ AI) k2z))
x
kas kax, Xxxxx Ar) kxpa'-
The parameter À 0 the regularization for the CME, and 0 is the noise variance for the GP
D.2 DERIVATION FOR ATE
Definition and Estimation Strategy. In this subsection. we detail the modeling and posterior derivation for the ATE The ATE for specific treatment is defined the expected outcome,
22

---

Published as a conference paper at ICLR 2026
marginalized over the entire population distribution of the covariates
(a) SS [ a=
(27)
We model the response surface E[y a a, s] with a GP f(a, s). The primary challenge is to evaluate the integral over the marginal distribution ... Unlike CATE, which requires conditional distribution, we can directly estimate Fs from all available samples in our dataset DD DDr u Dp. Our strategy is represent the distribution Is in the RKHS using its Mean Embedding (ME), Ms [.)) We form an empirical estimate of the ME from all individuals in D:
Px (s.))
(28)
This allows the ATE estimator be elegantly expressed inner product in the tensor product RKHS Ass
ÀT (a) = (f, ((a) jisH
(29)
This formulation is powerful because it converts the integration problem into a standard GP prediction problem at an "effective" input point.
Posterior Distribution. Based on the inner product formulation, the posterior distribution of the ATE estimator is also GP. The following proposition formally states its posterior mean and covariance.
Proposition 2 (ATE Estimator Posterior). Given the training dataset DD {ar, ST yr) and the full set of inputs D (a, let [D]. Iff(a,)) 8) is the posterior GP learned from DD. then the ATE estimator f follows a GP, ÂATE ~ (P((()),(((()) ((a,a))) Its posterior mean and covariance are derived using an effective feature map da ((a) @ Ps.
Detailed Derivation. formulas.
The following equations show the step-by-stee expansion of the GP posterior
v(a) 10 @ is: mr)
kear o 1 11,Kss, ((Karar @ ks, + yx,
(30)
(Kxpxy+Xd))
kax,
(ax, (Kx,Xr + Xx) yx.,
(aa,a)) ).=aa)))) kaar
kax, K (XxXx \11) -
Kea'
[
11,
n2
(31)
Keer
1 1 Kss 71
(( kKss,) Xx1=1 Kara' @
1 K sTsln n1
RaXT
1xya'
Here, Kx,x, Kayar ar @ K ST. In is a column vector of ones, and b 0 is the noise variance for the GP For simplicity in notation, we use 1" where the dimension is clear from context.
D.3 DERIVATION FOR ATT
Definition and Estimation Strategy. This subsection details the modeling and posterior derivation for the ATT. The ATT evaluates the effect of new treatment a the specific subpopulation that had previously received treatment à. For simplicity, we omit the context variable in this derivation. The ATT is formally defined as:
TATT a. 0) Ely a a, 8] PPla (dslā).
(32)
23

---

Published as a conference paper at ICLR 2026
We model the response surface E[y a, s] with GP f(a, s). The central challenge is to handle the integral over the conditional distribution Psla-a- Analogous to the CATE case, we represent this conditional distribution in the RKHS using its CME, May [
This allows the ATT estimator to be expressed an inner product in the tensor product RKHS HAs:
farT(a, a, a) (f, (()) @ Hsja)
(33)
where Paja is the empirical estimate of the CME. This formulation converts the integration problem into a standard GP prediction at an "effective" input point.
Posterior Distribution. Based on this formulation, the posterior distribution of the ATT estimator also a Gaussian Process. We summarize the result in the following proposition.
Proposition 3 (ATT Estimator Posterior). Given the training dataset DT {aT ST. yr) and the full input set 'D {a. S), if f(aa 8) is the posterior GP learned from Dr. then the ATT estimator Fxrr (a, â) follows GP, AATT GP (V(v(a, â), (( â), (d',d)) a)) Its posterior mean and covariance are derived using an effective feature map oa == (a) @ Psla-
Detailed Derivation.
The following equations show the ste-by--teep expansion of the GP posterior
formulas.
v a) (da @ Pula, Ms)
Kaar @ (kaa Kaa + +1) Kss,))
kaxr
(34)
x x(Karay @ K +xT) yxr
(Kxpx,+
kax, (Kxrx, + +x1) yxr:
((a,
(d',ā') (a
kaw - kax, (K (Kxxx + Ax) kx,
Kad' (kaa Kaa XI) K Kss( Kaa AI) kaa')
Kea!
(Kzar @ (Kaa (aa AI) K KSST)
(35)
kax,
x ((Karay @ K (TSST) + Xx1)-1
x (kaya' @ K Kaa AT) -aad
kxya'
Here, K XxyXr = K Karar K Sy. The parameter À is the regularization for the CME, and > is the noise variance for the GP f. Note that for the CME, the conditioning is on the treatment variable a, hence kernel matrices like Kaa are used for its estimation.
D.4 DERIVATION FOR ATEDS
Definition and Estimation Strategy. This subsection details the modeling and posterior derivation for the ATEDS. This quantity evaluates the effect of a treatment a over new target population, whose covariate distribution I.,, may differ from the source population IPx where the model trained. The ATEDS is formally defined as:
TATEDS (a) S. Ely a,s s] (dss).
(36)
The estimation challenge is that the regression function f(a, s) = E[y a= a,s= s] is learned from the source population data Dr, but the integration is performed over the target population distribution
24

---

Published as a conference paper at ICLR 2026
P.. We assume have access to set of ii.d. samples {8:) from the target population, which denote as D {5}.
Instead of learning the density p(s), we can directly use these samples to estimate the integral. We represent the target distribution P. via its empirical ME, estimated from the target samples:
1 Σ ((s)). Tiss
As
(37)
The ATEDS estimator can then be expressed as an inner product in the RKHS:
fa))Ds(a) (f, p(a) @ pL)
(38)
This formulation elegantly handles the distribution shift by converting the problem into a standard GP prediction an "effective" input point that encodes the target distribution.
Posterior Distribution. Based on this formulation, the posterior distribution of the ATEDS es- timator is Gaussian Process. The following proposition formally states its posterior mean and covariance.
Proposition 4 (ATEDS Estimator Posterior). Given the source training set Dr = (ar, Sr, yr) and a set of samples from the target population D = {S), let uı = | If ((a 8) is the posterior GP learned from DDr, then the ATEDS estimator fxTEDs( follows a GP, FATEDS GP(v(a), (((,aaaaa,,) Its posterior mean and covariance derived using effective feature map @a <(a) @ Ps.
Detailed Derivation. The following equations show the by-p-y expansion of the GP posterior formulas.
v(a)= = (0a @ iis; mM)
1 1T, K S5t
Kear @
(Karrrr k MI) yx,
74g
(39)
(x
hAXT
= kax (Kx, Xx + A11 yxr:
((a,a') kaa
kax, Kr X11) kxnæ
1
kaa'
-
K Egs I
n
Kaw
1 K S, s
Keap
11 K ssr
((Karar @ ks, + dx1)
kara'
glms)
T2g
kaXy
xxpaa
(40)
Here, Kx,x, Karar o Ksy Ing a column vector ng ones, and Xç is the noise variance for the GP f.
D.5 CONNECTION BETWEEN IG AND TVR
We have instantiated our key principle of uncertainty reduction using two prominent strategies derived from information theory and optimal design: IG and TVR. While these two approaches often presented as distinct heuristics, they share a deep connection, both aiming to shrink the posterior uncertainty of the CQ estimator. They differ, however, in how they quantify this uncertainty, focusing on different geometric properties of the posterior covariance matrix. The following remark clarifies this relationship by analyzing how each strategy operates on the eigenvalues of this matrix.
Remark 3. IG and TVR can be understood through the eigenvalues, {>:) of the posterior covariance matrix peaxt Warl f(( (1,
The TVR strategy aims to minimize the trace of the posterior covariance. Since the trace is the sum of the diagonal elements (the individual variances), and also the sum of the eigenvalues, the objective
X 3 arg min Tr Q Qpost) arg min Σ As.
(41)
XotDp
XBtDp
25

---

Published as a conference paper at ICLR 2026
This corresponds minimizing the arithmetic mean of the eigenvalues, effectively shrinking the average uncertainty across all dimensions. This is also known as A-optimality. The IG strategy is equivalent to minimizing the determinant of the posterior covariance. Since the determinant is the product of the eigenvalues, the objective is:
arg muin dei (Qooc) XB B X XgcDp
arg min I x
(42)
X XgcDp
This is equivalent to minimizing log( (edee(QQposl) E, log( which corresponds to minimizing the volume of the uncertainty ellipsoid, or the geometric mean of the eigenvalues. This also known as D-optimality.
E CONVERGENCE ANALYSIS
In this section, we analyze the uncertainty decay of the CQ estimator under our proposed method. Unlike inductive active learning, which reduces uncertainty predictions over the empirical distri- bution represented by the pool dataset Procl (2, t), our problem is inherently transductive; we aim evaluate the regression function over distribution that typically differs from Prusa (xr t), as is specific to the target CQ. Hence, our approach can be viewed as a form of transductive active learning (TAL). The key distinction is in the objective: TAL seeks to optimize function oover a target dataset, focusing on pointwise performance. In contrast, our method estimates an average over distribution, shifting the emphasis from individual points to distributional estimation. Nevertheless, our theoretical analysis leverages the TAL framework introduced by Hübotter et al. (2024).
E.1 JUSTIFICATION OF ASSUMPTION
To leverage the theoretical framework of Hübotter et al. (2024), we define the target set as the collection of all input points to the regressor f required for the numerical evaluation of the CQ integral, and the sample set S as the subset of the pool DD for which we could acquire outcomes. In the cases of ATE and DS, the structure of the integral ensures that the condition S s A holds. This is because estimating these quantities requires considering counterfactual outcomes for each individual. The positivity assumption, which maintain during acquísition, implies that any treatment could have been assigned to any individual, meaning the set of all potential evaluation points (LA) encompasses the set of observed factual points (S). Under this condition, the submodularity assumption naturally satisfied when using GPs for regression, as supported by Lemma C.9 in Hübotter et al. (2024).
For CATE and ATT, the situation is more nuanced because there can be overlap between and A without a strict subset relationship. In these cases, the assumption may not strictly hold. To address this, we introduce criterion called the Information Ratio later, which relaxes the assumption to weak submodular condition. The resulting bound is derived under the weak submodularity assumption and covers standard submodular case, affecting only a coefficient while maintaining the same convergence rate.
E.2 PRELIMINARY DEFINITIONS AND CONCEPTS
In the following, we focus on the analysis of CATE, which is the most complex case. We use the same notations in the main paper, where f is defined
function A x Z x S R. Before proceeding with the theoretical analysis, we clarify the notations and define key preliminary concepts used in this section. Instead of using ÊCANE (a, x) as defined in the main paper, simplify the notation by introducing X (a. z) and & E X A Z, with H = (a,,z)) and X, (a,, Z,). Note that x (a, 8) and should not be confused with æ. Additionally, o[f]== = Ss f(a, z, 8) )) (da|z). For brevity, we denote Ua = =u =). Since is a GP, Ua follows Gaussian distribution as is a linear functional of f. Our proposed method can be seen as acquiring outcomes to reduce the posterior uncertainty of the vector ux,, where the posterior mean and variance of any element Ua, are given in Prop. 1. We define xx, (ã,) = q ((a, z), (a, x)),
define operator
of del max x, og (zz), and o7 def of + o², where o² is the noise variance. We also connect to notation from the main text: Ua corresponds to the estimator f(a, x), and the irreducible uncertainty w/p (x) War (ar[v=l0 Ppr]
26

---

Published as a conference paper at ICLR 2026
Definition 3 (Marginal Gain and Maximal Marginal Gain). The marginal gain of Xp given X Xp is defined as
(x
def
Ux, (X (x)) -x,
(43)
which corresponds to the IG or TVR objective. The maximal marginal gain after i- greedy selections defined as
def
A (2(
(44)
X
Marginal gain quantifies the information contribution of a single point æ, while maximal marginal gain measures the peak possible contribution in a greedy sequence. Next, we define the information ratio.
Definition (Information Ratio). The information ratio of X Xp given s Xp, with [X] DI is defined
R (XID) def ExEx Ax, (x/ (x
E (0)
(45)
Remark 4. The insight behind these definitions is leverage the classic performance guarantee for greedy optimization of submodular functions. For a set of acquired points Eimn this is given by:
U(T1mA)
max U(X),
(46)
XEXP
[X|enA
where cy is a parameter. This implies the utility obtained by our strategy is bounded by a fraction of
the best achievable utility. Assuming Ass. holds, cg = for adaptive acquisition, while cy depends on the submodularity ratio for batch acquisition.
With these definitions in place, proceed to concept key to convergence analysis: the Approximate Markov Boundary.
E.3 APPROXIMATE MARKOV BOUNDARY
In inductive AL uncertainty decay analysis, the maximum information gain typically used to bound the information gain achieved by different methods. Acquiring all observation outcomes in Dp would naturally lead convergence this maximum information gain. However, our setting, even if all outcomes in Dp were acquired, the uncertainty over the target distribution of our estimator would not be fully eliminated. To account for this, we follow a similar approach to Hübotter et al. (2024) to define the irreducible uncertainty VDD (x) the lowest uncertainty attainable when all observation outcomes in Dp are acquired. However, a key challenge lies in distinguishing between the irreducible uncertainty and the reducible uncertainty, as analyzing the uncertainty decay rate requires isolating the reducible component. An idea inspired by the Markov Blanket can be used to derive a rough formulation: U2 fx fNB(B BB where B c Xp. This suggests that the uncertainty at i is conditionally independent of fx given the function values at the Markov Blanket MB( B), providing a potential way to separate the irreducible and reducible uncertainties.
Intuition 1.
2
Following the above Intuition, we introduce a new definition named Approximate Markov Boundary of æ in Dp, which is defined as Definition 5 (Approximate Markov Boundary (AMB)). For any > O, nr O, and i X,, we define B (Φ) as the smallest (multi-)set of DD satisfying
War
U2
Dr, 3 (
s (x)
(47)
where ur is the number of observations in Dr. We refer to Bur,(()) as the e-approximate Markov boundary of & in Dp.
The key idea behind constructing an AMB is that it acts as a bridge, linking the posterior uncertainty of Uh two components: the fundamental irreducible uncertainty and a controllable approximation term c. The rough idea, then, to analyze the uncertainty decay when observing the outcomes of points in B B ()). Since Eq. 47 provides upper bound on the remaining uncertainty. it can be leveraged to constrain the decay rate.
Before utilizing the AMB, we first establish its existence through formal proof.
27

---

Published as a conference paper at ICLR 2026
Lemma 3. e 0 and define the smallest integer satisfying
"r
CAmin Kx xx)
(48)
2mpofo}
where is definede as²
def
max i(fxp: X Xisr
(49)
For any Tur 20 and & E X,, there exists an e-approximate Markov boundary B Bug (+) for within Dp, with size bounded by
From this lemma, we conclude that for any i. given Wr and 6, there exists finite set B Buy (x) which satisfies Eq. 47.
Then, we move to examine key property of any set satisfying similar condition.
Lemma 4, Let > 0 and B ş Dp such that the following condition holds:
Var[ue yB] Cxmiu Kx (K),
(50)
Then for any e X1, have the following inequality:
Var[Us/ yBl & Var[Us
fx,
(51)
Proof.
Let the right-hand side of Eq. 50 be denoted by From the given condition, we can derive that
Varlus yB]
(1) E fxp unl War Uve fxp:y/ + Var farfx, pxplpx L [ve fxp: yn]]
Varea [vel fxr, yBl
Var fx lue E Uv fxr:8B]]
(52)
(10) Varval [Ua] [vm
fxp]
Var
xxp wa
Evi [v=
fx
irreducible uncertainty
roducible uncertainty
Step (i) follows from the law of total variance. Step (ii) utilizes the fact that the conditional variance of a GP depends only on observation locations, not their values. Step (iii results from the independence Vs IL YB fxp= as B Xp. The remaining task bound the reducible uncertainty.
We define a function hgg R R. fx, + B[Ua fx
Using the formula for the GP posterior
mean, we have
(53) where l = (Kx,x,) "r and Y = Cov U2). The explicit form of Y for the CME case is (#xp: du @ d2 @ Malx) Since ha is a linear function in fxp, we have for the reducible uncertainty that
he (fxx B[um] + 12 (fxp E[fxp]).
Varlha (fxp) [y]] L'Var[fxp lyB]
(54)
Amin K (Kx
This inequality can be derived in a similar manner to Lemma C.20 in Hübotter et al. (2024), using analogous steps. By the definition of e', this final term is equal to e.
Then, according to the results of Lemma C.19 in Hübotter et al. (2024), we let Var[falya [ s 28 Tz for all æ X,, we have for any X Var Vn (rr ya] s Var[valys]
k
(55)
Var[val [xpp +
The first inequality follows from the monotonicity of variance (i.e., more information does not increase variance), and the second inequality is the result of Lem. 4.
'Note that this 7, is defined over the regressor to leverage existing proofs, while the Yng in the main text is defined over the CQ estimator #. formal connection between them avenue for future theoretical work.
28

---

Published as a conference paper at ICLR 2026
E.4 PROOF OF THEOREM 2
To prove Thm. 2, we follow the analytical framework of Hübotter et al. (2024). The proof proceeds in three main steps: (1) we leverage the property of the AMB relate the current variance the information gain of the AMB set; (2) we bound this information gain using properties of submodular functions; and (3) we select a decaying approximation error e derive the final convergence rate.
First, from Lemma C.17 in Hübotter et al. (2024) and our AMB definition (Def. 5), we can bound the current variance of the estimator for any @ X as:
Var[ug [Dy.] 20 U2a (B PP) ())
(z) +E.
(56)
Next, we bound the mutual information term. By the submodularity of information gain (Assump- tion 1), the total gain from a set is bounded by the sum of marginal gains. This leads to bound involving the maximal marginal gain I' (similar to the logic in the proof of Thm 3.3 in the cited work):
I(Va: UB aer (Dr) s C
( ( nnys
(57)
where C1 is a constant related to the submodularityy/nforaaati ratio R. Let be -Bun,, (+)) be the size of the AMB. This gives:
Var[us/Dr] [rr] 20,CCbbI + ( +c.
(58)
Now, we select a specific value for e that decays with nT. Let E c
J.TT Vay
with c
2n 2npojol/>unn Amin (xx (x) From Lem. 48, this choice of e ensures that an AMB exists with size b, bounded by r /nr.
Substituting b. S nrr and the expression for into Eq. 58, we get:
Var Dr] s C my +H a
+
V /T
(59)
The problematic
The term an itself decays. Building on Theorem C.13 from Hübotter al. (2024), the maximal marginal gain is bounded by term that decays with nr. A simplified consequence for our analysis that we can find constant C2 such that:
Por S C2 Yer
(60)
TVT
Substituting this bound for I "z into Eq. 59 resolves the problematic term:
20 s 207CIV
C2 )mz)
(2070,
Tve /nT
(61)
nT
Combining all the pieces, the total variance is bounded by:
War U Dr] s (20)0,
Tor
+ Th (z) + 7
/ur
/nr
(62)
I n1 (x) C Ywr
where the final step combines all constants into a single constant C and uses the fact that for large nr, the Tiy term dominates the Y term. In the main paper, we use ng for the total number of acquired samples, which corresponds to nr here. The results for TVR can be derived similarly. It is worth noting that while our convergence analysis builds upon the framework for TAL, our contribution lies in successfully generalizing this framework from the simpler task of pointwise prediction to the more complex task of distributional integral estimation. Our analysis reveals that, despite the added complexity of our problem, the optimal convergence rate retains similar form, demonstrating the robustness and power principled active learning for this new and important class problems.
F EXPERIMENTS DETAILS
In this supplementary section, present a detailed account of the experimental results discussed in the main text. This includes a comprehensive description of the data generation process, as well as the implementation details for both the baseline methods and proposed methods. Specifically, provide information on the visualization dataset, the simulation dataset, and the the semi semi-synthetic dataset: IHDP (Hill, 201 1) and Lalonde onde,, 1986).
29

---

Published as a conference paper at ICLR 2026
F.1 DATASETS
General Setup For each experimental configuration, we report the mean and standard error over 20 independent trials, each initiated with unique random seed.
Datasets and Evaluation: All simulation datasets consist of 500 training, 200 validation, and 500 testing samples. We evaluate all methods on both the in-distribution training set and the testing set. The main paper presents the results on the testing dataset, while the Appendix here contains comprehensive results for both.
Across all experiments, the treatment a and outcome are one-dimensional. We denote single observation quadruplet (z, vy,s,a.y y) representing the conditioning covariate, adjustment covariates, treatment, and outcome, respectively. The dimensionality of the adjustment covariates varies by setting.
Variable Definitions:
Intervention Strategies: We evaluate types of interventions for all causal quantities:
The target treatment is single value, a, randomly selected for each trial. In settings with continuous treatments, we discretize the treatment space for evaluation on the training set, as continuous values may not be present in the data.
Hard Intervention:
Soff Intervention: The target treatment distribution, specified as p (a) Uniform[min maxin()))max(()))) This requires the estimator to average over range of treat- ments.
CATE-Specific Setups: For CATE estimation, which depends on the conditioning variable z, we test two distinct scenarios:
Fixed A single value of is randomly drawn from its training set range and held constant
across all 20 trials for a given configuration.
RRandom z: A new value of z is randomly and independently drawn for each of the 20 trials.
F.1.1 VISCALIZATION
In the visualization case, where the goal is to examine the match between the acquired data points and the target distribution, we set the dimensionality of the adjustment conditioning variable, and the other is a noisy variable that does not influence the treatment decision or the outcome variable. The process for generating single observation is follows:
One is controlled by the
1. Set the conditioning variable:
z Uniform(-2211
(63)
Define the adjustment variables $ and $2 as:
s Skew(z);
(64)
(2x 62:
where 61, (2.62 (0. For the skew function Skew(-), we define:
ys (xex(((x),xxx)xxxx))) where 2.5
(65)
with:
{(x) = 0.1 0.1 w(x) = 0.1 0.05, a(x) = -8 1 exp((xx)
(66)
Next, to generate the treatment variable, we focus on the continuous case (excluding the binary case). First, we concatenate the conditional variable with additional features form the random vector x, and the weight vector 3 is defined as:
x [z $1).
3
(67)
0.25
The intermediate variable aorg then generated as:
aong $ (x (3)) 1.5 0.5, c~,(0,1),
(68)
30

---

Published as a conference paper at ICLR 2026
Torter fiet Crrdere
Duttatten of Leater
Derrbutte of Trstrentees
8.
Skew function
(b) PDF of
PDF of Sg
Figure 7: Visualization of the visualization dataset. (a) Scatter plot of the conditioning variable and the adjustment variable $1: (b) the pdf of the first adjustment variable s for different treatments; (c) the pdf of the second adjustment variable for different treatments.
Then, we generate the corresponding treatment variable by
1.0
a
(69)
1.0 + exp(-aarg)
We observe that the values the treatment variable lie within the range (0, 1). discretization is required, use step size of 0.1 for the treatment variable. Finally, generate the outcome variable:
y 8 2. $1 + 6yy where Ey ~0,(00,).)))
(70)
F.1.2 SIMULATIONS
For the simulation, which mainly most of different scenarios, including the cate, ate, att and distribution shift. We mainly follow the similar data generation way from (Abrevaya et al., 2015; Singh et al., 2024). For the experimental results in Sec. 5, we make TR... A single observation is iid.
_N, 1), where = 1, 2, 3, 4,5. Then set generated as follow. Draw unobserved noise e ~
cos( (z) + +
$2 - 1+ 1 + E2
Uniform(-2,.
(71)
Sy sin z" cs
$1 exp(2e1) + ESt
Next, generate the treatment variable, this case, we have three different settings, including the binary treatment, the continuous treatment and the discrete treatment. First, we concatenate the conditional variable with the first three adjustment variable to form the random vector x, and the weight vector 3 is defined as:
x S2,
B j2
for j 1.,.44
(72)
The intermediate variable aorg then generated as:
Bor + (3 (x ()) 1.5 0.5, e~ (0,,
(73)
Then, the treatment is binary, we use the following decision rule to get the treatment
amng 0;
a=
(74)
0, else.
we need the continuous treatment, we make the following decision rule to get the treatment
a 1.0 + exp(-
(75)
We observe that the values the treatment variable lie within the range (0, 1). If discretization is required, step size of 0.1 for the treatment variable. Note that for the ATE and ATT cases,
31

---

Published as a conference paper at ICLR 2026
20
-20
La
-24
Figure 8: Scatter Plot with Conditional Distribution Overlay
retain only the first three adjustment variables and incorporate the conditioning variable directly into the adjustment variable set s. For the ATE under distribution shift, we define the target distribution as
81 Uniform(-2 (110mmmm(111)))))) 11,1) 1)
Uniform(-Le 1),,
S2 Uniform( (-0.5.
(76)
83 = Uniform(O. (,, 0.5)
F.1.3 SEMI-SYNTHETIC DATASETS
IHDP. The IHDP dataset, derived from randomized controlled trial, is widely used in causal inference to benchmark methods for estimating treatment effects. It contains 747 samples, with 139 in the treatment group and 608 in the control group, and 25 covariates: age, sex, race, married, education, income, maternal age, birth weight, gestation, mother smoking, mother drinking, mother health, father age, father education, father income, father smoking, father drinking, father health, num children, urban, prenatal care, previous pregnancy, low birth weight, premature birth, and health intervention. Among these, 6 covariates gestation, father age, and father income, while 19 covariates are binary. Treatment selection bias introduced by removing subset of treated individuals, and synthetic outcomes are generated to enable ground-truth calculations of ATE and CATE. The dataset's 's nigh-dimensional covariate space and realistic challenges, such as treatment heterogeneity, make it valuable resource for evaluating causal inference algorithms.
continuous: age, maternal age, birth weight,
Binary case. For the binary case, we do not need to generate the corresponding treatment and directly obtain them from the dataset. To generate the potential outcomes, we first extract the continuous covariates to form the random vector x, and the weight vector 3 is defined as:
x = [birth weight, birth head circumference, preterm birth, birth order, neonatal health, mother's 's age],
1 for j = 1,..,6
3 j
(77)
1.2 (x- B) +1+ co. 60 60~ (o~,(000.16) =ft=0: 0:
y 1.2 (x- B) + exp((x+ 0.5) 3 3- b.w + 1+ co. 6 ~~ (.,N00,011) else
(78)
where b.w represents the birth weight.
32

---

Published as a conference paper at ICLR 2026
Continuous and discrete case. For the continuous and discrete case, need to generate the corresponding treatment. To generate the treatment, we first extract the continuous covariates to form the random vector x, and the weight vector 3 is defined as: The intermediate variable aor is then generated as:
Borg (3 (x B)) 1. 0.5, <~ (0, Then, if the treatment is binary, the following decision rule to get the treatment
(79)
1, if aur >0; 0;
(80)
0, else.
If we need the continuous treatment, make the following decision rule to get the treatment
1.0
(81)
1.0 + exp(
We observe that the values of the treatment variable lie within the range (0, 1). If discretization is required, we use a step size of 0.1 for the treatment variable.
y y= 1.2 (x B) + 1.2 + b.w" + I b.w + 6u. 61 (1~N(000.16).
(82)
For the ATE with distribution shift case, we generate the six continuous covariates as follows:
Sj Uniform(O,0.5) 0.5) f=
(83)
When evaluating the performance of the estimator in the DS case, we generate new data by combining them with the original data. For both the training and testing sets, we sample the required number of points from the new distribution and replace the corresponding covariate observations of the original data with these new values. These modified data points, now incorporating the new covariates, are treated as being drawn from the shifted distribution. From the visualization of the IHDP dataset, we observe a relatively strong correlation among the first three variables: birth weight, birth head cireumference, and preterm birth. In all CATE setups, use birth weight as the conditioning variable to estimate the causal effect across different subpopulations, stratified by birth weight.
Lalonde. This is widely used benchmark in causal inference for evaluating methods that estimate treatment effects from observational data (LaLonde, 1986). It contains 2, 675 samples, with 151 in the treatment group and 2, 524 in the control group, and 8 covariates: age, education, black, hispanic, married, nodegree, re74 (earnings in 1974), and re75 (earnings in 1975). Among these, 4 covariates are continuous: age, education, re74, and re75, while covariates binary: black, hispanic, married, and nodegree. The dataset represents a real-world observational study where treatment assignment is not randomized, creating natural selection bias. The potential outcomes ceecoomp using established econometric methods, enabling ground-truth calculations of ATE and CATE. The dataset's 's realistic covariate structure and natural treatment heterogeneity make it valuable resource for evaluating causal inference algorithms in observational settings.
Binary case. For the binary case, do need to generate the corresponding treatment and directly obtain them from the dataset. The potential outcomes are pre-computed and stored in the dataset. To generate the potential outcomes, we first extract the continuous covariates to form the random vector x, and the weight vector 33 is defined as:
x Tage, education, re74, re75],
(84)
B, for 1,...,4
1.2 (x. () + co. y 1.2 (x- B) exppx++0.. B 3+3 3- age 1 co. G G~N(000.16) else where age represents the age covariate.
6o (0.00.6) if O:
(85)
For the ATE with distribution shift case, we generate the four continuous covariates as follows:
Sy UOiffrrm0.0.5) j 1..... 4.
(86)
When evaluating the performance of the estimator in the DS case, we generate new data by combining them with the original data. For both the training and testing sets, we sample the required number of points from the new distribution and replace the corresponding covariate observations of the original data with these new values. These modified data points, now incorporating the new covariates, are treated being drawn from the shifted distribution.
33

---

Published as a conference paper at ICLR 2026
F.2
IMPLEMENTATION DETAILS
F.2.1 GENERAL SETTING
Validation dataset. In the simulation dataset, we change the seed to sample a validation set from the same data generation process as the training dataset, ensuring that the observations in the validation set share the same distribution as the training data. In the real-world dataset IHDP, we directly randomly segment the training dataset into training set and a validation set. This validation set used to control early stopping during the training of the GP. However, this approach is not optimal because our goal is minimize the uncertainty of the target CQ estimator. Ideally, the GP should perform better the newly constructed distribution rather than directly on the observed data. This highlights a mismatch between the current validation setup and the optimal validation strategy. In principle, one could construct validation set that mimics the new target distribution. However, given the limited size of the validation set, such an approach could lead to insufficient data for validation, posing challenge for reliable model evaluation.
Trial Setup. For each trial, we initiated the process with a warm-start size, and the data for each experiment were randomly sampled from the pool dataset. This random sampling was done to ensure variability in the experiments and to evaluate the robustness of our methods under different initial conditions.
Simulations. In the binary treatment case, we used the Delta kernel, which is suitable for modeling binary outcomes. For future work, is possible explore the use of multi-task kernel (Alaa & Van Der Schaar, 2017), which could help model the inner correlation between different groups However, in our current implementation, we did not explore this possibility. In the CATE case, we employed an RBF kernel for the conditioning variable and used another RBF kernel for other adjustment variables. The RBF kernel chosen because of its flexibility in modeling smooth functions. For continuous treatment variables, we also used the RBF kernel. Additionally, we performed kernel ablation experiments by using alternative kernels, such as the Matern kernel and the Rational Quadratic (RQ) kernel, to explore the impact of different kernel choices model performance. However, the default setup uses the RBF kernel for all conditioning variables, and kernel bandwidths for the adjustment variables treated hyperparameters to be optimized during training.
Visualization. In the visualization experiment, the warm-start size was set to 20, and the acquisition step was set 5, meaning that data acquisition occurred in steps of 5 data points. This setup allowed us to monitor the effect of incremental data acquisition on model performance. The total number of training data points for this experiment used in this case. All other experimental settings for the visualization were identical to those used in the simulation setup.
200, with the CATE model being the sole model
IHDP. In the IHDP experiment, we worked with a fixed number of observations. The warm-start size was set to 20, and the acquisition step was set 10, meaning that data acquisition proceeded in increments of 10 data points. This setup was chosen to evaluate the model's performance with real-world dataset under controlled acquisition conditions.
F.2.2 HYPERPARAMETER SETTINGS
In the active learning framework, the overall process is divided into multiple iterative rounds. Each round consists of two sequential stages: Training and Acquisition ion. During the training stage, we update the model parameters using the currently labeled dataset. In the acquisition stage, the trained model, together with the proposed acquisition criterion, is used evaluate the utility of unlabeled points in the pool dataset.
GP. We perform GP regression without neural feature extractors in all experiments, including both synthetic simulations and the semi-synthett IHDP dataset. Instead, we directly apply bernel-based models implemented via GPyTorch (v1.12)². All models are trained using the Adam optimizer with fixed learning rate of 0.05 for 500 epochs. Full-batch training is used throughout. No early stopping is applied in the first 200 epochs, and validation performance is monitored every 10 epochs. Hyperparameters o² and any kernel parameters (e.g., bandwidths of the RBF kernel) are optimized
hhttps: :tttps///iiii com/ 1 uus- htt.g///pttrrrreeeelaaaee////gga
34

---

Published as a conference paper at ICLR 2026
by maximizing the exact marginal oog-likelihod, defined as:
log p Xr.) -- - ( + +02 yr log
+ +) 'T 1og(( (87)
2
where X e 1= are the inputs (including the covariates and treatment), yr E R are the observed outcomes, X,, is the kernel matrix computed via the RBF kernel (note that our kernel construction is detailed in Sec. 4), o² is the noise variance, and Ry is the number of training samples. Additionally, hyperparameters are updated following each acquisition of new labels. Unless otherwise specified, we use the standard radial basis function (RBF) kernel. We have also evaluated alternative kernel choices, including Matern and RQ kernels, in Fig. 16.
CME. We regularization parameter is a learnable parameter, initialized as À 0.01. Optimization is performed using the Adam optimizer with a learning rate of 0.05, and training is run for 100 epochs without early stopping or validation. To reduce computational cost in the active learning setting, we avoid retraining the CME model after each acquisition round by leveraging the closed-form conditional mean operator. In each round of CATE or ATT estimation, the kernel features of the adjustment variables s must be updated. Specifically, the CMEs /'s and /'saa' need be recomputed. This update can be done efficiently using the conditional embedding operator (details are in Sec. 4.2): instead of retraining, we only update the feature map & within the operator. This update requires only matrix multiplication and is thus computationally lightweight.
RBF kernel for the CME, with the bandwidth selected via the median heuristic. The
F.2.3 BASELINE METHODS
We use existing implementations for most of the baseline methods:
1. Random: In this case, during each round, we randomly choose the required my data points uniformly from the pool dataset.
2. Query-based Heterogeneous Treatment Effect estimation (QHTE): QHTE is core-set method that uses a distance-based acquisition function to strategically select data points, enhancing the estimation of heterogeneous treatment effects (Qin et al., 2021). We adopt this method as a baseline for the semi-synthetic dataset. To ensure a fair comparison, we also use a Gaussian Process (GP) as the regression function. Similar to QHTE, space. In the GP case, we define the distance between two points as the posterior covariance, where a smaller distance indicates stronger relationships. We greedily select ny data points from the pool dataset during each round to ensure they serve as an effective core-set.
select the core sets the feature
CausalBALD For CausalBALD (Jesson et al., 2021), we directly take the code of 4-BALD as the information gain method applied to the pool dataset. While CausalBALD is originally designed for the binary case, it can be extended to the continuous case by treating each subgroup independently. As reported in their paper, 4-BALD generally performs well, often achieving results comparable to the best methods.
4. Total Variance Reduction (TVR): For this method, we directly compute the posterior variance for all data points in the pool dataset (MacKay, 1992). We then calculate each data point's 'ss utility by measuring the total variance reduction, which is based on the sum of the posterior variances of all data points in the pool dataset.
For all baseline methods, ensure fair comparison, GP regression function as in our proposed method, including the learning rate, number of epochs, and validation dataset, as detailed in the implementation section above.
the same training configuration for the
F.3 HARDWARE
For the running time analysis and visualization experiments, we leverage a machine equipped with a 32-core Intel CPU and 128 GB of RAM. This setup provides sufficient resources for processing and handling the large-scale data required for our experiments. For running all the numerical results, we utilize Spartan, a high-performance computing platform. Each trial is typically run on an NVIDIA A100 (80GB) GPU, paired with an Intel CPU featuring 8 cores and 40 GB of memory.
'ttpps: / // thub. httpb//o///////////uuaaaaaa
35

---

Published as a conference paper at ICLR 2026
Algorithm 1 Active Learning Framework for CATE Estimation
Require: Initial labeled dataset Dr = {(a,,y)) =- Pool of unlabeled data Dp (xx))
N i-nz+1
Query budget n.g, Batch size n4, Utility function U()
Ensure: Final trained CATE model i
Train baseline model # Main active learning loop for each round Use the model from the previous round 4: Compute utility scores u, U feurrent for all Φ; Dp.
1: Train initial CATE model fo on Dr.
2: for r 1 to nn,n, do
3: Let Fourrent = Fy-1-
# Evaluate pool points # e.g., Top-k or
5: Select a batch of indices Σς of size ni from Dp based on scores {ur)-
greedy selection
6: Let the selected batch be X, {x; E
7:
Query outcomes Y, = {y/ | Ip} for the instances in X
# Oracle queries for labels # Add new data to
8:
Augment the labeled set: Dr- + Dr {(0;,,)) æ; E X6, yr Y,).
training set
9: Update the pool: Dp - Dp X
Remove queried points from pool Retrain CATE model f, on the updated labeled set Dy. Improve the model with new data
11: end for
12: return The final trained model Fng/uy-
F.4 METRICS
To evaluate the accuracy of our estimators, we use the Average Mean Squared Error (AMSE). This metric is computed over predefined set of "points of interest," which represent the specific subpopulations and treatment scenarios relevant to the research question. Since calculating the AMSE requires access to the ground-trut truth causal effect т, this evaluation approach is suitable for synthetic data settings where the true data generating process known. The general principle is to compute the squared error between the estimated effect 7 and the true effect across set of ny interest points. The structure of these points varies depending on the causal quantity being estimated.
AMSE for CATE. For the Conditional Average Treatment Effect, the estimator iCATE (a, 2) is function of both the treatment and the conditioning covariates The set of interest points is therefore a collection of crovaitet-covariiat pairs (a,, Zi) = {(a,, )- The AMSE defined as:
def 1 Σ (fcxm 2;) -CAATE (as, 2;))
AMSECATE
(88)
TI
AMSE for ATE and ATEDS. For the Average Treatment Effect (ATE) and ATE under Distribution Shift (ATEDS), the estimators fATE(a) and fATEDS(a) are functions of the treatment only. Con- sequently, the set of interest points is a collection of different treatment values a, {a The AMSE is defined as:
def 1 -
AMSEATE
(fxn (a)))
(89)
The formula for ATEDS analogous, using the respective estimator and true value.
AMSE for ATT. The Average Treatment Effect the Treated, ÊXTT (a, (a,a), a), a function of both the target treatment a and the conditioning treatment à. The set of interest points is therefore collection of pairs (az, àz)= {(a,, {(,,,,)) ". The AMSE is defined as:
del 1 Ai
AMSEATT
- Σ (fxIT(ai, ās) TXIT (xiT(a;; ?
(90)
nI 1m1
F.4.1 UNCERTAINTY REDUCTION AS AN EVALUATION CRITERION
While our acquisition strategy is designed reduce the uncertainty of the causal estimator, directly tracking this uncertainty reduction over acquisition rounds is an unsuitable primary metric for comparing different active learning methods. This is due to several fundamental challenges. First, the
36

---

Published as a conference paper at ICLR 2026
model is retrained after each batch acquisition. This process, which often includes hyperparameter re-optimization, recalibrates the model's 's posterior distribution. Consequently, the scale and meaning of uncertainty are not consistent from one round to the next, making direct numerical comparisons of uncertainty values across rounds unreliable. Second, different acquisition strategies are designed to reduce different forms of uncertainty. For instance, BALD- methods target the uncertainty in model parameters (0), whereas our method targets the posterior variance of the causal quantity (T) itself. Evaluating all methods using any single, internal uncertainty measure would unfairly favor the method that directly optimizes that specific measure. Therefore, following established practice in active learning, we evaluate performance using an external, task-based metric that is independent of any specific acquisition function. Since the ultimate goal of all compared methods is to produce the most accurate causal estimator possible, the AMSE the causal estimate is the most appropriate and equitable metric. It directly measures how effectively each strategy uses its budget to achieve the shared final objective.
F.5 COMPUTATIONAL COMPLEXITY
We analyze the computational complexity for single round of active learning. This process involves two main stages for each of the np points in the pool set Pp: (1) computing the posterior distribution of the causal quantity (e.g.. CATE), and (2) evaluating the acquisition function based this posterior. We denote the training set size as ay. The complexity of any neural network components is omitted from this analysis.
Stage 1: Posterior Computation Cost. The cost of computing the posterior is dominated by the number of "query points" required to estimate the causal quantity. Let be the number of such query points. For example, for a CATE 7(a, z), m would be the number of discrete values of a and z we are interested in. The initial, one-time cost for inverting the training kernel matrix is O(n)). ). After this, the cost for predicting the posterior for m query points is:
CDE with Monte Carlo Sampling: To estimate an integral (e.g.. over s), we sample n, points for each of the initial ng query points (e.g., ny different x values). This results in a total of MMC = ng n, effective query points. The cost to obtain their joint posterior covariance is dominated by:
(91) Case Method: The CME method analytically handles the integral, avoiding the need for sampling. The number of effective query points remains mcME ng. The CME estimator itself requires a one-time kernel matrix inversion, typically on the training data, with cost of (n²)), which is subsumed by the main GP inversion cost. The posterior computation cost is then dominated by:
O MMc mMC))
On TCME m.ME).
(92)
The key advantage of the CME method is that MCME < mwc. leading to significant reduction in complexity, especially when a large number of Monte Carlo samples nt, required for accuracy.
Stage 2: Acquisition Function Evaluation. After obtaining the m x m posterior covariance matrix E, for the causal quantity, we evaluate the acquisition function for each of the np pool points.
IG: This acquisition function is often related to the determinant of the posterior covariance matrix, representing the volume of the uncertainty ellipsoid. The cost of computing the determinant is Omm²). To find the best point, this must be done for all pool points, leading to a total complexity for one round of:
Onp (n + m² + m²)).
(93)
TVR: This acquisition function corresponds to the trace of the posterior covariance matrix, Tr(E,), E,). The cost of computing the trace is only (m)) The total complexity for one round is therefore much lower:
Onp (n² + m²))
(94)
37

---

Published as a conference paper at ICLR 2026
Batch Selection. When selecting a batch of n, points:
The simplest method is to compute the acquisition score for all np points and select the top rés- The cost is determined by the total complexity of the acquisition function evaluation, plus a negligible sorting cost of ((p log np).
Top-k Selection:
Greedy Selection; A greedy approach selects points sequentially, updating the selection criteria at each step. naive implementation would re-evaluate the acquisition function for all remaining pool points at each of the na steps. For a complex, set-based acquisition function like batch IG, this can be prohibitively expensive, with complexity scaling with ng np. More efficient greedy methods exist but beyond the scope of this analysis (Holzmüller al., 2023).
F.6 ADDITIONAL EXPERIMENTAL RESULTS
In this section, we present additional ablation results, including analysis of the robustness of our proposed methods with respect to different kernel choices, starting points, step sizes, and pool dataset sizes. All results are based the visualization datasets. Additionally, conduct various simulation setups and provide further results on both the simulations and the semi-synthetic IHDP data.
F.6.1 VISCALIZATION
Here, we present the full set of results, including the AMSE metrics and additional visualizations of the acquired data points, complementing the results presented in the main paper. Specifically, for the acquired data visualizations, we focus on the CATE case, where the condition variable is fixed at specific value, and we evaluate the performance across all possible treatments. The figure illustrates the joint distribution of the treatment and the adjustment variable S1. The blue points represent the scatter of the pool dataset, while the red points correspond to the scatter of the data points acquired by different methods. The red rectangle highlights the region of interest, which encompasses all possible treatments and the range of 51 conditioned on the fixed condition variable. From this figure, we can clearly observe the distribution shift between the target distribution, indicated by the CATE estimator, and the distribution of the pool dataset. In this scenario, our eight proposed methods consistently approximate the target distribution more accurately than all baseline methods, resulting in improved estimation performance.
F.6.2 ABLATIOON RESULTS
In this part, we present various ablation studies to evaluate the robustness of our data acquisition strategies. For active learning, the experimental setups control include: (1) the size of the pool dataset, in Fig. 15 (2) the number of warm-start data points, in Fig. 14, (3) different batch acquisition sizes, in Fig. 12. (4) using soft acquisition trick for batch acquisition, in Fig. 18 and (5) different noises 13. Regarding our model, since we use GPs as the regression function, we provide two more kernel choices, the Rational Quadratic (RQ) kernel and the Matérn kernel, to assess the model's reliability and sensitivity to different kernel selections, in Fig. 16.
Fig. 12 shows that smaller batch sizes improve performance by enabling frequent retraining, refining uncertainty estimates, and enhancing adaptive sampling. This supports our con- clusion that greedy acquisition approximates sequential single-point estimation. However, smaller batches also increase training itera- tions, balancing performance and efficiency.
CATE
CATE
Mumber aessind dbseeadioss
N.mberot artu doseevatinns
200
Robustness to Heteroscedastic Noise. To further evaluate the robustness of our frame- work against likelihood misspecification. specifically conditionally heteroscedastic er- rors, we conducted additional experiment using the visualization dataset setup. We modified the data generation process to introduce het- eroscedasticity, where the noise variance explicitly depends the covariate 81. Specifically, the
Figure 10: Full results of the VAMSE performance (with shaded standard error) for all methods on the Visualization dataset. From left to right: (I) In- distribution results; (2) Testing-distribution results.
38

---

Published as a conference paper at ICLR 2026
Teca TOCE
Teatmere
Tharmen
Treammn
Thearme
Traamma
F2L Copptaat
Figure 9: Acquired data points visualization.
outcome is generated as y= [(())+ +c, with <~ NN0, 0.5]s1)). The results are presented in Fig. 11. We observe that our proposed methods, particularly the CME-based strategies, consistently outperform the baselines. This demonstrates that our framework effectively handles scenarios where uncertainty varies across the input space. Even under this deviation from the fixed-variance assumption of the standard GP likelihood, the ActiveCQ strategies successfully identify and query informative regions, maintaining superior sample efficiency.
Across all results, our CME-based greedy data acquisition strategies, including the based and TVR-based methods, consistently demon- strate the best performance. In most cases, the naive method also performs comparably to the greedy approach, with the added benefit of sig- nificantly faster computation. In practice, when computational resources are limited, the naive method can be a practical alternative that deliv- ers strong performance.
CATE
CATE
4.3
4.4
JALL
D4.
THE DHL
Surber anpded deeavodioe
Marber of any digdint sodivn
Figure 11: Full results of the VAMSE perfor- mance (with shaded standard error) for all methods on the Visualization dataset. From left to right: (1) results for some methods for good visualization; (2) results for all methods.
The CDE MC-sampling methods generally achieve second-tier performance across various settings. This is because they aim to reduce the uncertainty of the target estimator, focusing on the constructed distribution rather than the observational data represented by the pool dataset. However, as previously discussed, the CDE method struggles in high-dimensional scenarios and is sensitive to noisy covariates that do not influence treatment assignment or outcome values.
In contrast, the CME-based method iteratively refines the learned representation by adjusting the bandwidth hyperparameter of the RBF kernel, enabling it to better capture the influence of relevant variables. Building on this insight, leveraging neural ketwork-based methods could further improve
39

---

Published as a conference paper at ICLR 2026
performance by filtering out irrelevant covariates more effectively (Xu et al., 2021). Alternativlly, one could apply independence tests between covariates, treatment, and outcomes to remove irrelevant variables (Zhang et al., 2018). However, this approach heavily depends on the reliability of the independence testing tool. it incorrectly excludes important variables, could adversely affect subsequent procedures. We leave these explorations for future work.
CATE
Surter diacieren andecatoo
Surte itse
Surber
Figure 12: One step acquisition Difffrent batch sizes. The v AMSE performance (with shaded standard error) for different batch sizes is presented. The first row illustrates the in-distribution performance, while the second row depicts the out-of out-t--distributiio performance. From left to right, the batch sizes are set to 5, 10, 15, 20, and 25.
CATE
3330
332
Humbere of acsine dhseveli
furbe d'ccuin Coneeatin
Sumber of sca Rservators
Hamber of eccureed diseration
CATE
CATE
CATE
CAFE
4.71
4.00
630
Sunder Connaatees
Figure 13: Different noises. The v AMSE performance (with shaded standard error) for different batch sizes is presented. The first row illustrates the in-distribution performance, while the second row depicts the citto--distribution performance. From left to right, the noises are normal, student, Laplace and uniform.
F.6.3 SIMULATION RESULTS
In this part, we present additional experimental results on simulation datasets, covering the CATE, ATE, ATT, and DS cases. Across all scenarios, we explore two primary treatment settings.
The first setting involves fixing the target treatment value, meaning we focus solely on estimating the causal effect for specific subgroups. This approach is particularly practical in real-world applications. For example, personalized medicine, we might estimate the effect of a fixed dosage of a drug on patients with specific genetic profile. Similarly, in educational interventions, impact of a fixed number of tutoring sessions on students with a certain performance baseline.
might evaluate the
The second setting, as discussed in the main paper, examines all possible causal effects across different treatment levels within specific subgroups. This broader perspective allows us to assess the variability of treatment effects, For instance, could evaluate how varying dosages of medication
40

---

Published as a conference paper at ICLR 2026
CATE
Kamter
Hamber
Figure 14: Different starting points. The v AMSE performance (with shaded standard error) for different warm starting points is presented. The first row illustrates the in-distribution performance. while the second row depicts the out- -itt----iittitttionn performance. From left to right, the starting poitns are set 10, 20, 30, 40, and 50.
Figure 15: Different sizes of pool datasets. The AMSS performance (with shaded standard error) for different Pool dataset sizes is presented. The first row illustrates the in-distribution performance, while the second row depicts the out-of-distribue sizes are set to 100, 150, 200, 300, 400, and 500.)
performance. From left to right, the Pool dataset
influence recovery rates for patients with a particular condition. Alternatiell, we might analyze how different levels of government subsidies affect the economic outcomes of households in given income bracket.
We would like clarify why reasons. The first reason pertains to the IG-based method, which requires calculating the determinant of the posterior covariance matrix. This determinant is influenced by the size of the pool dataset and the number of possible treatments, where the size is the product of these two factors. In our setup, we use 600 observations in the training dataset, and the starting point is set 100, meaning the pool dataset consists of 500 samples. If we were to use a continuous treatment, the worst-case scenario would involve 500 different treatment levels. This would result in covariance matrix of size 250, 000 x 250, 000, which is computationally prohibitive for calculating the determinant. While it is possible to approximate this by sampling covariates and treatments, which is reasonable approach in practice, it falls outside the scope of this paper. Therefore, we discretize the continuous treatment in most cases, excluding IG-based methods in the process.
discrete treatments in all cases, which is driven by main
The second reason relates to the simulation mechanisms in our dataset. When we focus on fixed treatment value, the binary case is straightforward, as we can easily segment the data into treated and untreated groups. However, in the continuous treatment case, the treatment determined by the covariates with added noise. To evaluate the performance of our estimator for a specific treatment, generating observations with fixed treatment value is challenging. To avoid this issue, discretize
41

---

Published as a conference paper at ICLR 2026
CATE
ON
4,000
DXTE
ONT
CAT
Mimber stsevators
Navber stservliiss
Figure 16: Different choices kernels. The VAMSE performance (with shaded standard error) for different kernel functions is presented. The first row illustrates the in-distribution performance, while the second row depicts the out-of-distribution performance. From left to right, the kernel functions are set to RBF, Matern and RQ
the treatment in the fixed-treatment case and also use discrete treatments for all possible treatment scenarios.
CATE
Figure 17: The / AMSE performance (with shaded standard error) for the CATE case on simulation data is presented. The first row shows the in-distribution performance, while the second row illustrates the -uttof-disttibution performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment, and all with continuous treatment.
CATE. In the CATE case, the results are shown in Fig. 17. We can see that our CME-based methods consistently show the best performance, and the CDE with MC sampling methods usually show the second-tier best performance. The benefit of our proposed methods over the baseline methods is substantial, since we fixed one value of the conditioning variable, which means that the gap between the interested distribution and the distribution of the pool dataset is quite large. Moreover, we would like to explain why we need to learn the conditional distribution Palz-
Even for the CATE case, to construct the interested distribution, the joint distribution Ps.x is decom- posed as Fs.z = Ps/z 12, meaning that we first sample the marginal distribution of F, to get the observations of the conditioning variable. This could be from delta distribution if we only care about a fixed value of the conditioning variable, which implies that we only care about the causal effect specific subgroup, from the observational distribution, which means that about all possible subgroups given different values the conditioning variable. Then, we can sample from the conditional distribution Paz obtain the observations of the adjustment variables.
42

---

Published as a conference paper at ICLR 2026
Notice that the conditional distribution remains unchanged in these cases, meaning that can directly take observations from the dataset as a surrogate for sampling from the conditional distribution. However, there are two main reasons why this is not sufficient. The first reason is that if we want to evaluate a subgroup, say = z", and this value does not appear the dataset, it is impossible to get any samples from the dataset. In this case, we need the flexible conditional distribution to sample from ==.=.. The second reason is that we can sample more observations from the learned distribution, which helps reduce the variance and improve accuracy
Also, note that if one cares about all possible subgroups, it may not be strictly necessary learn the conditional distribution. We believe that the primary benefit of our methods comes from the distribution shift caused by the treatment variable, from P,, to the interested distribution P.. This shift determined by whether one cares about the performance over fixed treatment or all possible treatments.
CATE
0.200
TVR, MONLS
TVR CNL.
1C MENLS
0.150
MENUG
AMSE
TVR_ONE_G
0.125
0.300
Here, we also incorporate the softmax-based batch acqui- sition strategy proposed by (Kirsch et al., 2023), which reweights the acquisition scores of candidate observations based on their sample importance within the pool dataset. This approach can accelerate the acquisition process and improve performance. We find that BALD-based meth- ods benefit from this strategy, achieving stable perfor- mance and demonstrating compatibility. However, our method consistently outperforms BALD, whether using Top-b selection or greedy acquisition, it directly tar- gets uncertainty reduction the target estimator. Notably, because our approach prioritizes predictive performance over reducing parameter uncertainty, the ceftmax-based reweighting appears to negatively affect its effectiveness.
0.075
0,000
0.025
100
Number of acquired observations
Figure 18: The VAMSE performance (with shaded standard error) for different kernel functions. Methods with "S" use the soft-max acquisition trick.
E
Figure 19: The V VAMSE performance (with shaded standard error) for the ATE case on simulation data presented. The first row shows the in-distribution performance, while the second row illustrates the out-of-distribution performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment, and all with continuous treatment.
ATE. For the ATE case, the results are shown in Fig. 19, We can observe that our methods consis- tently outperform the baseline methods across all different setups. As with the CATE case, we can conclude that when the interested distribution significantly differs from the observational distribution in the pool dataset, our methods perform much better than the baseline methods. Consequently, for all fixed treatment cases, methods yield superior results. However, for the all tressbeeettr in binary treatment scenarios, our methods show comparable performance to the baseline methods. This is because, to ensure that the overlap assumption holds, do not modify Pas significantly.
43

---

Published as a conference paper at ICLR 2026
Moreover, in the ATE case, the distribution shift arises solely from the shift in the treatment variables Fa. which is why all methods show similar results in this case.
AALL
AAGO
CEREE
Number eccurce deservatioe
Number of anquired deseredttt
Number AsQulie steereriee
Namber of (ow.eeeddoaeeeooo
ATT
ATT
ATT
ATTT
Number atarv
20 20 201
Figure 20: The VAMSE performance (with shaded standard error) for the ATT case on simulation data is presented. The first row shows the in-distribution performance, while the second row illustrates the out-of out-of-distriie of-distribution performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment.
ATT. For the ATT case, the results are shown in Fig. 20. We draw a similar conclusion before: our proposed methods consistently outperform the baseline methods, although the performance improvement is not substantial. We believe this is due the data simulation, where the covariates are high-dimensional, and ensure the overlap assumption, we avoid making significant distribution shifts in the treatment selection. This is further supported by the visualizations of the datasets shown in Fig. 8 and Fig. 7, which illustrate that Pajp does not change significantly when constructing the interested distribution.
Figure 21: The VAMSE performance (with shaded standard error) for the DS case on simulation data is presented. The first row shows the in-distribution performance, while the second row illustrates the out-of-distribution performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment, and all with continuous treatment.
For the ATEDS case (DS for short), the results are shown in Fig. 21. We believe this case provides the most compelling evidence to support the benefits of our proposed methods. Across all setups. methods consistently show significantly better results compared to the baseline methods. The performance improvement can be attributed not only to the changes in the treatment variable distribution P, but also to the distribution shifts in the covariates P... which we have intentionally
ATEDS.
44

---

Published as a conference paper at ICLR 2026
manipulated in this setting. Consequently, we observe a substantial performance improvement, which further highlights the importance of targeting the desired distribution when acquiring data points.
Also, note that for the ATE and DS cases, in the setup with fixed value treatment, the IG-based methods will perform identically to the TVR-based methods. This is because the target estimator collapses to the same formulation, where the target becomes point estimate and the posterior of this estimator is a normal distribution. In this case, minimizing entropy and minimizing variance become equivalent.
In conclusion, from the above results, we can derive the following key takeaway:
Takeaway Message: The benefits of CQ-specific data acquisition strategies are more pronounced when there significant distribution shift between the target distribution and the pool dataset distribution..
F.6.4 IHDP RESULTS
CATE. In this part, we present additional experimental results on the IHDP datasets, covering the CATE, ATE, ATT, and DS cases. Across scenarios, we explore two primary treatment settings shown before in the simulations. When studying the IHDP dataset, since all covariates in this dataset are fixed and the treatment assignment for the binary case is deterministic, it results in highly imbalanced data. However, for continuous treatment, assignment ourselves and must determine the outcome regression function.
design the specific form of the treatment
Figure 22: The v AMMSE performance (with shaded standard error) for the CATE case on IHDP data presented. The first row shows the in-distribution performance, while the second row illustrates the -itt---diitribttion performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment, and all with continuous treatment.
In the CATE case, the results are shown in Fig. 22. We observe that, in most cases, our CME-based method consistently outperforms the baseline methods. However, for the fixed-value treatment case, our method performs poorly due extreme data imbalance and the use of binary kernel, which inherently prevents the exploration of shared correlations between the two subgroups.
ATE and ATT. interesting observation is that, across all scenarios, none of the methods consistently outperform random acquisition. This outcome can be attributed to the minimal distribution shift in these cases. Additionally, the combination of high-dimensional features and limited number of observations hinders effective feature learning, which adversely impacts the precise representation of uncertainty in data points and, consequently, the overall performance.
the ATE and ATT cases, the results are presented in Fig. 23 and Fig. 24. An
For the ATE with distribution shift case, the results are shown in Fig. 25. Once again, our proposed methods demonstrate consistently better performance compared to all baseline methods. This improvement can be attributed to the large distribution shift, which amplifies the advantage of methods over the baseline approaches.
ATEDS,
45

---

Published as a conference paper at ICLR 2026
AT5
Remdon
Figure 23: The v AMSE performance (with shaded standard error) for the ATE case on IHDP data is presented. The first row shows the in-distribution performance, while the second row illustrates the -ii-tri-ittiittionn performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment, and all with continuous treatment.
ATT
ATT
ATT
ATT
4,00
5J813
10
Tel, 20
Norbor UCQUITE dd
Number AcQuITO dostWOtt
The AMSE performance (with shaded standard error) for the ATT case IHDP data is presented. The first row shows the in-distribution performance, while the second row illustrates out-of-distribution performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment.
Figure 24:
F.6.5 EXPERIMENTAL RESULTS ON LALONDE
Our first task was to actively estimate the CATE, Tab. presents the performance comparison be- tween our proposed methods and several strong baselines. The results demonstrate that our methods, particularly the CME-based variants (IG-CME and TVR-CME), consistently and significantly outper- form the baselines across all acquisition budgets. For instance, after acquiring only 60 labels, our TVR-CME method achieves an AMSE of 8.2 + 0.5, whereas the best best-performing baseline, Coresets, only reaches a similar error level after acquiring nearly 200 labels. This highlights our framework' superior sample efficiency in selecting informative individuals that are directly relevant to the CATE estimator.
The second task evaluated the framework's 's ability to handle a distribution shift when estimating the ATE, As shown in Tab. 2, our methods again demonstrate a clear advantage. They reduce the estimation error much more rapidly than the baselines in the initial stages of acquisition and maintain superior performance throughout the process. Interestingly, in this setting, both the MDN-based and CME-based variants of our approach perform comparably well after the first 60 samples, and both are significantly better than the baselines. This result validates our framework's 's ability to effectively target the relevant data distribution, even when it differs from the source population.
46

---

Published as a conference paper at ICLR 2026
Nambar
The v AMSE performance (with shaded standard error) for the DS case on IHDP data is presented. The first row shows the in-distribution performance, while the second row illustrates the ditto--distriuttion performance. From left to right, the settings include: fixed and binary treatment, fixed and discrete treatment, all with binary treatment, all with discrete treatment, and all with continuous treatment.
Figure 25:
Table 1: Performance comparison on the LaLonde dataset for active CATE estimation. Lower AMSE better.
Method
20
60
100
140
180
200
Random
21.7 0.7
1..1
13.6 0.5
12.1
11.0
10.8
Variance Reduction
21.7 0.7
16.0 1.0
12.7
1..1
10.7
Coresets
21.7 0.7
1.
0.6
1.81
8.6
0.2
H-BALD
2.7
1.20
9.0
8.0
7.0
6.6
IG-MDN (Ours)
0.7
1..9
1.6
8.9
8.3
8.1
TVR-MDN (Ours)
2.7
1..8
9.9
0.8
8.2
7.9
IG-CME (Ours)
2.7
0.5
6.2
0.5
5.5
5.9
TVR-CME (Ours) 0.7
8.2
6.1 0.5
0.5
0.5
5.8 0.4
Table 2: Performance comparison on the LaLonde dataset for active ATE with distribution shift estimation. Lower AMSE is better.
Method
20
60
100
140
180
200
Random
22.6 0.8
16.1 1.0
13.9 1.0
12.8 + 0.9
11.3 1.0
11.0 + 1.0
Variance Reduction
22.6 0.8
1.0
1.5
12.2 1.4
11.9 + 1.2
Coresets
0.8
1.22+1.2
1..8
10.6
9.2
8.9
4-BALD
22.6
1..0+1.6
8.2+
7.7+1.0
7.5
7.2
IG-MDN (Ours)
2..6
9.4
7.8
7.5+1.0
7.3
7.2
TVR-MDN (Ours)
2.8
9.1+1.0
7.3
7.1+100
7.1+100
7.1+1.0
IG-CME (Ours)
2.8
9.3+1.1
7.8
7.5+1.0
7.3
1.2
TVR-CME (Ours)
22.6
9.1
7.2
7.1
7.1
7.1
G DISCUSSIONS
G.1 SUMMARY OF CONTRIBUTIONS AND FUTURE DIRECTIONS
key challenge in causal inference is that target CQs are often defined by integrating over specific sub-populations that differ from the overall observed population. For instance, CATE is defined by the conditional distribution p(s]z), while ATT is defined by p(s(a). Consequently, the central principle of our work is that effective active learning strategy must target the uncertainty relevant to these specific sub- aab-poulatiion. Instead of selecting data to reduce the global uncertainty of the underlying outcome model f, our framework prioritizes samples that maximally reduce the posterior uncertainty
47

---

Published as a conference paper at ICLR 2026
of the final causal estimator & itself. While this principle of targeted uncertainty reduction establishes a strong foundation, our framework has several limitations that highlight promising avenues for future research.
Relaxing Identification Assumptions. Our current work operates under the standard assumption of no unmeasured confounders. A significant avenue for future work is to extend our active learning framework to more complex scenarios where this assumption is violated. This involves incorporating established identification strategies such as the front-door criterion (Dance et al., 2024; Chau et al., 2021b), instrumental variables (Xu et al., 2021), I), or proxy variables (Mastouri et al., 2021). Since many Bayesian and GP-based methods already exist for these settings, our framework can be naturally extended.
Model Architecture and model is a T-learner. An important extension would be to explore architectures that share statistical strength across treatment groups, such as S-learners or multi-task GPs (Alaa & Van Der Schaar, 2017). Furthermore, our current CME-based approach can be challenged by high-dimensional covariates. Integrating feature selection techniques (Witte et al., 2020) or neural deworkk-bssed feature extractors (Johansson et al., 2022) would be crucial step to enhance the framework's applicability to high-dimension data.
From a meta-learning perspective, our current
Scalability with Bayesian Neural Networks. While GPs offer theoretical guarantees and closed- form posteriors, their cubic complexity limits scalability. Replacing the GP with Bayesian Neural Network (BNN) (Jesson et al., 2021) is a compelling direction for future work. BNNs offer superior scalability for large datasets and greater flexibility in handling complex, high-dimensional inputs. They can learn feature representations directly from data, removing the need for kernel engineering, while still providing the necessary uncertainty quantification for active learning.
From Estimation to Policy Optimization. While our current focus is on minimizing the estimation error of the Causal Quantity (e.g.. îCATE) globally, a compelling direction for future work to target downstream decision-making directly. In many high-stakes applications, such as personalized medicine, the ultimate goal is to identify the optimal treatment policy (=2) that maximizes patient welfare. This decision-theoreee domain to accurate sign determination near the decision boundary. Our ActiveCQ framework provides unique advantage for this transition: unlike methods that rely on Enformation-thhonetie proxies (e.g., maximizing mutual information with model parameters (Jesson et al., 2021; Fawkes et al., 2025)), our derivation of the analytic posterior distribution for the causal estimator itself enables the direct calculation of the Expected Value of Information (EVOI). This capability allows for the design of acquisition functions that explicitly maximize the expected gain in downstream policy value, aligning with emerging frontiers in conffite adaptive trials (Klein et al., 2025).
objective shifts the priority from precise estimation across the entire
G.2 THE ADVANTAGES OF CME IN THE ACTIVE LEARNING SETUP
The choice of CMES over traditional CDEs offers several crucial advantages within our active learning framework. Foremost is its adaptability to learned features. In each round of the AL loop, as the GP retrains on new data and refines its -nduleddduced feature space, the CME can be re-estimated using this updated representation. This ensures the distribution embedding is always expressed in the most relevant, data-informed feature space, becoming progressively more aligned with the regression task.vSecond, CMES exhibit superior task alignment. Our ultimate goal is compute expectation (an integral), which is precisely what the CME is designed for, representing the distribution as a mean element in an RKHS. CDEs, in contrast, attempt to model the entire density function, a more general and often unnecessarily complex task, introducing extra overhead and potential for misspecification. The CME provides a more direct and efficient path to the quantity of interest.v particularly well-suited for high-dimensional settings. Unlike many density estimation techniques that suffer from the curse of dimensionality, CMES can gracefully handle higher-dimensional data. By imposing fewer restrictive assumptions on the distribution's 's form, they offer more flexible and robust alternative, making them a powerful and appropriate choice for the ActiveCQ task.
CMEs are
48


Reviews:
[{"review_id": "1KDwq4N3tT", "reviewer": "Reviewer_u6vV", "summary": "This paper proposes the first unified framework for active estimation of multiple causal estimands, including ATE, ATT, CATE and ATE under distribution shift. The framework is a Bayesian framework with uncertainty quantification, which is supported by solid theoretical grounding.   Empirical validation on several synthetic or semi-synthetic datasets shows the effectiveness of the framework.", "questions": "I would suggest using a alternative term for the word \"causal quantity\". In potential outcome framework, a more common definition is \"causal estimand\", indicating that we aim to estimate the value.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "1. This paper is novel in task formalization, unified modeling of causal estimands, and creative integration of Gaussian Processes and Conditional Mean Embeddings to balance uncertainty quantification with computational efficiency in active learning.\n2. The paper demonstrates exceptional quality through theoretical rigor, robust experimental design, and transparent reproducibility practices.\n3. The paper is exceptionally clear, making its technical content accessible to both experts in causal inference and researchers new to the field.", "weaknesses": "1. This paper acknowledges standard GPs’ cubic complexity and mentions efficient alternatives (sparse variational GPs, Random Fourier Features, Nyström). However, no empirical evaluation of these methods is provided.\n2. The CME’s regularization parameter $\\lambda=0.01$ is fixed with no sensitivity analysis.\n3. If I am right, the greedy strategy for batch selection still requires iteratively re-computing the posterior for all remaining pool points after each selection, which is not efficient enough.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "YbEtpxU0xq", "reviewer": "Reviewer_NNrP", "summary": "The paper introduces ActiveCQ, a pool-based active learning approach tailored towards estimating causal quantities of interest. The approach is observational, selecting “unlabeled” items from a pool to query their outcomes (labels), as opposed to carrying out interventions.  ActiveCQ follows a Bayesian active learning approach, where uncertainty about the causal quantity of interest is minimized. A key insight of the proposed approach is that if the outcome is modeled as a Gaussian process (GP), the causal quantity of interest is also a GP, with mean and covariance functions derived in equations (4), along with two sample-based estimators based on conditional density estimation and conditional kernel mean embeddings.  This insight allows utilizing acquisition functions from active learning with Gaussian processes, along with a (suitably adapted) theoretical analysis, to this novel setting.  The approach is evaluated in a quite extensive set of experiments.", "questions": "- The figures are tiny (especially Fig 3)... perhaps it would make sense to move some to the appendix?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The specific instance of active learning considered in this paper is interesting, practically relevant and has not been studied yet to the best of my knowledge\n- The approach is principled, and accompanied by a theoretical analysis\n- The experiments and baselines are quite comprehensive\n- The paper is quite well written (except some acronyms, such as SUTVA, should be defined to be more self-contained)", "weaknesses": "- The paper is primarily a natural composition of ideas that have individually been fairly well explored, thus perhaps not exceedingly surprising to researchers in the area.  Nevertheless, I find it makes a relevant contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "Q5GfD8kIgE", "reviewer": "Reviewer_TTvj", "summary": "The paper considers the problem of active learning when the goal is estimation of causal quantities like the average treatment effect (ATE) or conditional average treatment effect (CATE). More precisely, the authors consider settings in which there is available a (generally small) dataset containing outcomes $y$, treatments $a$, confounders $s$, and effect modifiers $z$, and also a larger dataset containing the same observables other than outcomes. One can obtain outcomes for selected observations in the second larger data set but at some cost. The problem is then to choose the data-point (or set of points) in the large sample, so that obtaining outcome information for that points leads to the greatest reduction in uncertainty in the target causal quantity of interest.\n\nThe approach of the authors is Bayesian, with Guassian process priors on the relevant features of the joint distribution of the observables and the outcomes assumed to be conditionally Gaussian with fixed conditional variance. The use of RKHS leads to tractable formulations for the posterior distribution of the causal quantities of interest. Points can be selected to achieve the greatest expected reduction in posterior uncertainty, as measured by either information gain or a total variance reduction criterion. A greedy algorithm is proposed for tractable selection of batches of points (selecting the optimal batch of $m$ points out of a large sample is computationally infeasible when $m$ is anything other than very small, for obvious combinatoric reasons). The performance of the method is assessed against benchmarks on semi-synthetic and simulated data and appears to perform very favourably.\n\nIn my reading, the key technical contribution here is as follows. Causal quantities depend on both the regression function ($E[y|a,z,s]$ in the authors' notation) and often also on the joint/conditional distribution of $z$, $s$ and (possibly) $a$. Two examples of causal quantities that depend on the latter are the ATE and CATE conditional only on $z$ (i.e., $E[y|1,z]-E[y|0,z]$). For standard active learning, it generally suffices to target the posterior of the regression function. The authors' approach is premised on the observation that many causal quantities of interest can be written as integrals of the regression function and the joint distribution or related objects. A key insight of the paper is then that, by expressing the joint distribution in terms of conditional kernel mean embeddings (CME) (an approach employed in some other recent papers on causal learning in RKHS), an analytic expression for the posterior of the causal quantities can still be obtained. The authors do consider an alternative based on kernel density estimation and Monte Carlo sampling, but the CME approach appears more tractable and elegant.", "questions": "I wondered which if any of the simulated/semi-synthetic datasets violate the assumption that the outcome is conditionally Gaussian with fixed variance? While the likelihood is misspecified if this is the case, the comparison on such potentially less favorable datasets might be quite informative as in the real world conditionally heteroskedastic errors may be very common. \n\nOne reason that the CATE may be of interest is that knowledge of the CATE can allow researchers to optimally allocate treatment in some new population on the basis of the covariates $z$. More precisely, one may hope to choose a treatment assignment rule $\\pi$, (where $\\pi(z)\\in \\lbrace 0,1 \\rbrace $) that achieves high expected utility compared to the baseline of no treatment in the new population. For example, equating the ITE with the utility gain from treatment, one would hope to achieve a large value of the objective  $\\ell[\\pi]=\\int E[y|a=\\pi(z),z]-E[y|a=0,z]d \\tilde{P}_z (z)$, where $\\tilde{P}_z$ is the distribution of $z$ in the new population which one may take as known. While this objective depends on the unknown CATE, given sample data, the posterior distribution of this object can be evaluated, and one could find the treatment allocation rule (perhaps in some restricted class) that maximizes the posterior expectation of this object (this is the approach in Bayesian statistical decision theory). This suggests an alternative (to information gain and the total variance) objective one could use to select the points whose outcomes one should sample, namely the supremum over $\\pi$ of the posterior expectation of $\\ell[\\pi]$. I wondered if the methods in the paper would still apply for this objective?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "I enjoyed reading the paper. I found it to be well-written. The discussion of the analysis and approach is thorough. The question is novel and the methods are innovative and practically relevant. The analysis is very elegant. I found myself convinced that the methods proposed by the authors are effective.", "weaknesses": "The could be a little more detail given regarding possible real-world applications. The figures are too small. The grammar in some parts of the appendices needs fixing.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "MIl4j9iEcN", "reviewer": "Reviewer_bTeT", "summary": "Strengths:This work makes both remarkable conceptual and methodological contributions: it delineates the new task of Active Estimation of Multiple Causal Quantities in a form that is immediately applicable to the real-world challenge of costly outcome measurement and begins closing the theoretical gap in CQ-aware active learning. On a methodological front, it would be difficult to find a more innovative way of marrying GP and CME than this paper, producing a coherent framework for uncertainty modeling in the RKHS, where analytic integration over target distributions replaces Monte Carlo sampling noise. Meanwhile,the paper sets a good basis in theory for consolidation and better portability to further acquisition strategies by summarizing the various causal quantities ATE, CATE, ATT, and DS into one integral formulation.Weaknesses:The paper remains thin on empirical validation; most experiments are run on simulated and, at best, semi-synthetic datasets (IHDP, LaLonde) without any observational reality, so it is challenging to talk about the paper with practical implications. Besides, the \"compatibility of the GP kernel space with the conditional embedding\" is claimed but not empirically tested, so this essential assumption is not empirically validated. Third, while the author blames the distributional mismatch for the poor performance of the methods reviewed, such a mismatch is not quantified, nor is any evidence provided that the performance boost observed for ActiveCQ results from addressing such an issue. Therefore, there is a logical gap in the causal reasoning", "questions": "na", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "na", "weaknesses": "na", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "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:

Decision: Accept (Poster)