Enhancing Treatment Effect Estimation via Active Learning: A Counterfactual Covering Perspective

Active Learning vs. Coreset Learning: The problem setting we described in lines 114-127 (left column) is an active learning loop that consists of 5 steps, where in Step 4 we identify a subset from the completely unlabeled pool set $\mathcal{D}=\{(\mathbf{x}_ {i},t_ {i})\}^{N}_ {i=1}$ for the oracle to obtain the labels (i.e., treatment outcome $y_i$). The newly labeled samples are then used to refine the estimation model in Step 5. As such, our setting differs from coreset learning that operates with all samples' labels $y_i$. As our labels are actively acquired, Algorithm 1 corresponds to Step 4 (unlabeled subset identification from $\mathcal{D}$) of our problem setting. In Eq.(8), though it is formulated using the labeled training sets $\mathcal{S}_ {0}$ and $\mathcal{S}_ {1}$, the calculation of Eq.(8) only requires covariate $\mathbf{x}_i$ and its corresponding treatment assignment $t_i$ (lines 202-203 right column), hence are applicable to all samples in the unlabeled pool set $\mathcal{D}$. In both Eq.(8) and Algorithm 1, we slightly abused the notation $\mathcal{S}$ to make the notations less cluttered, as their calculations do not actually rely on the access to treatment outcome labels $y_i$. To avoid potential confusion, we will update the description and notation in corresponding parts to highlight the active learning setting.

$\mathcal{R}$, and Class $\mathcal{H}$: Thanks for pointing out the typo "$\mathcal{R}$'', we will correct it to the feature space "$\mathcal{X}$''. The family of functions $\mathcal{H}$={$h \mid h: \mathcal{X} \rightarrow \mathbb{R}$} does contain the regression model $f$ as part of the function $h$. As per Theorem 3.3, we have $h_{f}(\mathbf{x},t):=\frac{1}{\kappa}l_{f}(\mathbf{x},t)\in \mathcal{H}$ ($t$ is a constant for any given $\mathbf{x}$), which contains the loss function $l_{f}$ to evaluate regression model $f$. The rationale for defining $\mathcal{H}$ is that the constant $\kappa_{\mathcal{H}}$ in our derived bound is dependent on $\mathcal{H}$. As per line 966, the term in Eq.(32c) is further upper-bounded by Integral Probability Metric (IPM) under the premise that $h_{f}(\mathbf{x},t)$ is in $\mathcal{H}$, which is a useful bounding technique also seen in (Shalit et al., 2017).

Corollary 3.4 and Noisy Response $y$: Thanks for improving the completeness of the Corollary, which should be claimed with the probability of least $1-\gamma$ since it is a further simplification of Theorem 3.3. Its main idea is to reveal what optimizable terms are closely related to risk convergence. For the gap $\Delta$, we follow the common practice in (Sener & Savarese, 2018) by assuming zero training loss and noiseless observation for the theoretical analysis purpose. Under the active learning setting, it is reasonable that the labels obtained from the oracle are sufficiently reliable, i.e., $y_i$ has minimal noise.

How We Query in Algorithm 1: As stated in lines 215-219 (right column), the intuition behind the query process is to greedily prioritize the reduction of the largest radius, which is a commonly used intuition also seen in (Tsang et al., 2005; Sener & Savarese, 2018) for the $k$-Center problem. As per Eq.(8), each of the four radii is defined in lines 202-204 (right column), thus, we just need to query the point (no label) that constitutes the radius $\delta$, and the query is explicitly defined in line 6 of the full Algorithm 1 (Appendix C) if we need to reduce $\delta_{(1,1)}$. The visualizations in Figure 1 are a graphical expression of Algorithm 1's intuition, i.e., the point at the other end of the radius is the next acquisition goal. Notably, Algorithm 1 is a naive instantiation of the covering radius-based active learning and comes with a strong (yet impractical) assumption on fully overlapping data distributions (lines 232-234 right column). Thus, we made its description relatively brief to pave the way for our proposed main algorithm FCCM (Algorithm 2), which relaxes the full coverage constraint to achieve stronger radius reduction under the same labeling budget. We will add this explanation when introducing Algorithm 1 to provide more intuitions.

Random Variable $T$: Thanks for pointing out the typo, we will correct $t$ to $T$ where needed.

Q1: For instance, when performing A/B tests to compare two software versions, the service provider needs to assign two distinct user groups to different versions (i.e., treatments). To understand which software version offers a better experience for a specific user profile (i.e., features), it is necessary to collect users' feedback via customer surveys. Though the features and treatments associated with each user are known to the service provider, collecting all treatment outcomes -- users' satisfaction levels is time-consuming and economically unrealistic (e.g., paid surveys). In this case, active learning can be used to select a small subset of the most informative users, from whom the treatment outcomes can be obtained at a much lower cost.

Q2+W1: We provide the figures under the 2% increment evaluation scheme, which is the most fine-grained setting based on the 50 query steps in total in the anonymous link https://anonymous.4open.science/r/To-Reviewer-FYYH (click into folder ``To-Reviewer-FYYH''). Note that on the IBM dataset, all methods yield an impractically large prediction error at the initial acquisition steps (<20% budget), thus we show the results from 20% for better clarity.

Q3+W3: Our main rationale is to avoid making the search space prohibitively large -- up to $\mathcal{O}(n^4)$ if each radius has $n$ candidate settings. The key insight to reduce the search complexity is that, because $P(\mathcal{A})=\frac{1}{4}(P(\mathcal{A}^{t=1}_ {F})+P(\mathcal{A}^{t=0}_ {F})+P(\mathcal{A}^{t=1}_ {CF})+P(\mathcal{A}^{t=0}_ {CF}))$, each radius is independent and each of the sub-terms increases monotonically with the radius, then so does the mean coverage $P(\mathcal{A})$. Thus, by the independence and monotonicity, the search space greatly reduces to $\mathcal{O}(n)$ while the search of four radii stays in the same direction, thus we let the four radii share the same setting. When will different radii help? Though the shared value among four radii allows for effective and efficient hyperparameter tuning, it could potentially fail if the distribution discrepancy between the two treatment groups is very large. This is because the counterfactual covering radius can be $\delta_{(t,1-t)}\gg\delta_{(t,t)}$ for the counterfactual coverage $P(\mathcal{A}^{t}_{CF})$ to get close to full, lowering the utility of the uniform radii value. As such, if the distribution discrepancy between treatment groups is large, we can then identify a different value for each covering radius to maintain the high coverage under the linear complexity due to the independence and monotonicity. We will add this discussion to an updated version of our paper.

Q4+W4: The reason for omitting 0% in IBM is that the estimation risk at the first few steps is too large to be practical for all methods. Thus, the reported results start from 20% for clarity. We will append a more explicit explanation on the figure caption to eliminate potential confusion.

Q5+W2: Thanks for the suggestion, we will update the table to show the PEHE results for both FCCM and FCCM$^-$ in our paper. Please check the anonymous link https://anonymous.4open.science/r/To-Reviewer-FYYH (click into folder ``To-Reviewer-FYYH'') for the detailed results.

Probability Threshold $\lambda$ ($\gamma$ in Our Paper): We will add a discussion of the magnitude of the probability threshold $\gamma$. It is noted that higher confidence gives a larger upper bound due to the increase in the third term of the bound, but it remains a constant during our optimization.

SUTVA: We will update it with the full spell, namely Stable Unit Treatment Value Assumption (SUTVA).

Typos: Thanks for pointing them out, and we will correct them correspondingly. For ``by change of variable reverting y to x'', we intended to express that the integral over the domain $\mathcal{Y}$ is changed to the domain $\mathcal{X}$ by the change of variable. We will modify the description to clear the confusion.

Potential Failure Modes: FCCM is designed to better handle the partially overlapped data for a quicker bound reduction while maintaining high coverage (Figure 3). As such, in scenarios where the two treatment groups have no overlapping regions (e.g., biased treatment assignment), the data acquisition of FCCM will be challenged as there are no overlapping, counterfactual pairs to be identified. One possible remedy is to operate FCCM in the latent space where the inter-group distributions are aligned by methods like (Zhang et al., 2020; Wang et al., 2024), but it also offsets the model-independent advantage of FCCM. We will add this discussion in an updated version.

Practicality of Strong Ignorablity (SI): It is noted that the focus of this paper is active label acquisition, it can then integrate with the selection bias-aware method to make predictions. The validity of the SI can be approximated by carefully selecting sufficient relevant covariates and constructing a more balanced dataset.

Practicality of Lipschitz Continuity: Theorem 3.3 assumes the Lipschitzness of the $p^{t}(y|\mathbf{x})$, this can be a practical assumption if the regression model $f$, e.g., a well-regularized neural network (NN), learns a smooth mapping from $\mathbf{x}$ to $y$, which implies the Lipschitzness for $p^{t}(y|\mathbf{x})$. Also, the NN $f$ is differentiable w.r.t. $\mathbf{x}$, thus the squared loss $l_{f}$ is also bounded and sufficiently differentiable w.r.t. $\mathbf{x}$, further implying the Lipschitzness of $l_{f}$.

Q1: Elaboration on The Sentence: Despite that $\mu\rho$-BALD ($\mu\rho$) highlights searching for samples within the non-overlapping region, its acquisition criterion naturally prefers samples that lead to higher estimated uncertainty by its criterion. However, the criterion design of $\mu\rho$ is not density-aware. For example, in early acquisition steps, two unseen data points $a$ and $b$ can be of the same uncertainty, while $a$ is an outlier yet data $b$ belongs to a dense cluster (not seen). Then, prioritizing the acquisition of $b$ over $a$ can help generalize the estimator to more samples, which means training on $b$ can eliminate more estimation risk when data is limited. However, $\mu\rho$ does not differentiate the importance of $a$ and $b$, and thus "account for less risk'' in such cases. How FCCM addresses this challenge: Pertaining to the previous example, the factual covering from FCCM can prioritize the acquisition of $b$ over $a$, as $b$ has more neighbors (i.e., higher out-degree) than $a$ in the graph constructed in Algorithm 2. To enhance the distribution alignment, FCCM makes use of the weight $\alpha$ imposed on the counterfactual edges (line 308 left column). For example, assume $c$ is in the same treatment group as $b$, while $b$ and $c$ have similar numbers of neighbors in the graph. Then, if $c$ has more neighbors from the counterfactual group than $b$ does, then the priority of $c$ is higher than $b$ owing to the additional bonus (with $\alpha>1$) on counterfactual edges when calculating the out-degree. For cases with significant non-overlap (almost no overlap), please refer to our response to potential failure modes. Why $\mu\rho$ queries non-overlapping samples in Figure 6(c): As stated in line 85, $\mu\rho$ relies heavily on the accuracy of estimated uncertainty (model-dependent), which can be erroneously high for certain samples outside the overlapping region, boosting the acquisition metric score. Thus, a few such samples are selected as per Figure 6(c). In contrast, FCCM is model-independent, and the acquisition score is robustly computed via the out-degree of nodes (Algorithm 2).

Q2: Please see the anonymous link: https://anonymous.4open.science/r/To-Reviewer-hAfw for sensitivity analysis results on $\delta$ and $\alpha$ (click into folder "To-Reviewer-hAfw''). Note that the acquisition on treatment sample $t=1$ is insensitive on $\delta_{(0,0)}$ and $\delta_{(0,1)}$ in our setting, as all control samples ($t=0$) are seen. For $\alpha$, our setting of $\alpha=2.5$ has an overall lower error across different acquisition budgets. We will include all results in an updated version.

Q3: Theorem 3.3 (Lipschitz continuity): If the Lipschitzness does not hold, the multiplicative constant will be unbounded, and thus the reduction of the radii may not help control the risk upper bound. Theorem 4.1 (SI): SI provides an ideal scenario for acquiring the counterfactual samples, where a quick bound reduction by Algorithm 1 is seen in Figure 2(a). If it cannot be guaranteed in real-world data, e.g., CMNIST, the reduction of the bound will be significantly slower and less effective for Algorithm 1, as shown in Figure 2(b). Thus, it motivates us to propose Algorithm 2 which can handle compromised data distributions more effectively via a slight trade-off in coverage.

We will adopt all the constructive suggestions to update the paper.

Q1: $\Gamma$ is a pseudo-operator defined in the "Input'' section of the Algorithm 1, i.e., $\Gamma=\arg\max\min d(\cdot,\cdot)$. We use $\Gamma$ as a shorthand for the distance-based radius defined underneath Eq. (8) (line 203 right column) to keep Algorithm 1 concise. Take line 6 of Algorithm 1 as an example, to reduce radius $\delta_{(1,1)}$, by applying $\Gamma$ as the selection criteria, the selected sample from $\mathcal{D}_ {1} \setminus\mathcal{S}_ {1}$ is $a=\arg\max_{a'\in\mathcal{D}_ {1}\setminus\mathcal{S}_ {1}}\min_{j\in \mathcal{S}_ {1}} d(\mathbf{x}^{t=1}_ {a'},\mathbf{x}^{t=1}_ {j})$, where $D_ {1}$ and $S_ {1}$ are respectively the unlabeled and labeled sets for $t=1$.

Q2: As per line 273 (left column), the informative bound in Theorem 3.3 consisting of four covering radii is derived under the full coverage, i.e., all actively acquired samples should together cover 100% of the full data space (we will mathematically explain the full coverage in Q3). Thus, the straightforward solution in Algorithm 1 strictly maintains 100% coverage requirement using currently the smallest radii (in Definition 3.1 and 3.2) at every acquisition step, while progressively minimizing the four radii via newly acquired samples. However, due to this full coverage requirement, Algorithm 1 cannot work effectively on real-world datasets where inter-group distributions do not align well (Figure 2 and lines 255-269 left column). Thus, in Section 4.2, we further build Algorithm 2 (FCCM) upon Algorithm 1 as a solution. Clarifying the statement: Instead of computing the minimal radii in Algorithm 1, Algorithm 2 fixes the radii by a small constant when performing data acquisition. Though the 100% coverage requirement cannot be fully satisfied in early acquisition steps, as the acquisition expands, we can effectively increase the coverage of the data space. So as shown in Figure 3(b), FCCM can achieve up to 25% bound reduction compared to Algorithm 1 with only a negative $1\%$ coverage gap (99% coverage). Thus, in line 238 (right column) we state ``given the relatively smaller fixed covering radius'', we can '' approximate (approaching) the full coverage''. We will add more explanations to those parts to better motivate Section 4.2 in an updated version.

Q3: For each treatment group $t$,$P(\mathcal{A}^{t}_ {F})$ is the proportion of the data points from the unlabeled pool set $\mathcal{D}_ {t}$ that has been covered by the training set $S_{t}$ in the covering ball $\mathcal{A}^{t}_ {F}$, i.e., $P(\mathcal{A}^{t}_ {F})=\frac{|\mathcal{A}^{t}_ {F}|}{|\mathcal{D}_ {t}|}$. Analogously, $P(\mathcal{A}^{t}_ {CF})=\frac{|\mathcal{A}^{t}_ {CF}|}{|\mathcal{D}_ {1-t}|}$, then the mean coverage rate $P(\mathcal{A})=\frac{1}{4}(P(\mathcal{A}^{t=1}_ {F})+P(\mathcal{A}^{t=0}_ {F})+P(\mathcal{A}^{t=1}_ {CF})+P(\mathcal{A}^{t=0}_ {CF}))$. We will add this formal definition to complement our textual descriptions in lines 266-271 (right column) to make it easier to follow. A further note is that, for Algorithm 1, we have the full coverage, i.e., $P(\mathcal{A})=1$, strictly satisfied at every acquisition step, while Algorithm 2 -- FCCM is to maximize the mean coverage rate $P(\mathcal{A})$ to approximate the full coverage ($P(\mathcal{A})=1$) given the four fixed small covering radii to constitute a lower bound value.