Good Allocations from Bad Estimates
Authors: Sílvia Casacuberta
Moritz Hardt
Venue: ICLR 2026 Poster (2026)
Keywords: Treatment Allocation
Treatment effects
Sample complexity
RCT
Primary area: alignment, fairness, safety, privacy, and societal considerations
TL;DR: Course treatment effect estimates suffice for near-optimal treatment allocations.

Abstract: Conditional average treatment effect (CATE) estimation is the de facto gold standard for targeting a treatment to a heterogeneous population. The method estimates treatment effects up to an error $\epsilon > 0$ in each of $M$ different strata of the population, targeting individuals in decreasing order of estimated treatment effect until the budget runs out. In general, this method requires $O(M/\epsilon^2)$ samples. This is best possible if the goal is to estimate all treatment effects up to an $\epsilon$ error. 
In this work, we show how to achieve the same total treatment effect as CATE with only $O(M/\epsilon)$ samples for natural distributions of treatment effects. The key insight is that coarse estimates suffice for near-optimal treatment allocations. In addition, we show that budget flexibility can further reduce the sample complexity of allocation.
Finally, we evaluate our algorithm on various real-world RCT datasets. In all cases, it finds nearly optimal treatment allocations with surprisingly few samples. Our work highlights the fundamental distinction between treatment effect estimation and treatment allocation: the latter requires far fewer samples.

Published as a conference paper at ICLR 2026
GOOD ALLOCATIONS FROM BAD ESTIMATES
Sílvia Casacuberta ++
Moritz H
1 Sttanford University
Max Planck Institute for Intelligent Systems, Tübingen Tübingen AI Center
ABSTRACT
Conditional average treatment effect (CATE) estimation is the de facto gold stan- dard for targeting a treatment to a heterogeneous population. The method esti- mates treatment effects up population, targeting individuals in decreasing order estimated treatment effect until the budget runs out. In general. this method requires O OMM/e² samples. This is best possible if the goal is to estimate all treatment effects up to an e error. In this work, we show how to achieve the same total treatment effect CATE with only ((M/c) samples for natural distributions of treatment effects. The key insight is that coarse estimates suffice for near-optimal treatment allocations. In addition, we show that budget flexibility can further reduce the sample complexity of allo- cation. Finally, we evaluate our algorithm on various real-world RCT datasets. all cases, it finds nearly optimal treatment allocations with surprisingly few sam- ples. Our work highlights the fundamental distinction between treatment effect estimation and treatment allocation: the latter requires far fewer samples.
error > 0 in each of M different strata of the
INTRODUCTION
Different groups in a population be it schools, counties, or age trackets-often respond differently to treatment (Gail Simon| 1985 Heckman & Robb. 1985) Imbens Angrist 1994 Banerjee & Dutlo 2009). This empirical fact has motivated significant body of work on estimating conditional average treatment effects (CATE), see, e.g., Imbens & Rubin (205)):[Ahhe) & Imbens (2016): Athey et al. (2019); Künzel et al. (2019).
key application of CATE estimation in relfare-maximizing treament allocation: assigning limited number of treatments those groups who benefit most from treatment. The optimal treatment allocation selects groups in descending order of treatment effect until the budget runs out. In practice, however, treatment effects first have to be estimated from data, such the responses from a randomized controlled trial (RCT) on the population. The standard method estimates the treatment effect in each group and allocates in descending order of estimated treatment effects. Good CATE estimates therefore seem to be the necessary first step of treatment allocation. Indeed, estimating single average treatment effect in one group up error O boils down to mean estimation and requires ((1/e²) samples. This sample size requirement holds robustly for almost all problem instances with few exceptions. By extension, given M -on-ovellap groups, CATE estimation requires O (M//2) samples-as as does finding (1 ooptmmaall treatment allocation. And so it appears that the two problems are essentially equivalent, at least in the worst-case.
In contrast, we show that for typical instances, we can find - 1-)--opptmall allocation with only OM//c) samples. Hence, for M O(1), treatment allocation typically requires quadratically fewer samples than treatment effect estimation. Whereas estimation has a robust quadratic lower bound, we show that the quadratic lower bound for allocation briule: It's easy to circumvent in theory with natural assumptions and it doesn't 't arise in any of the weal-world datasets examine. Our results follow from a simple but powerful win-win situation: If the treatment effect in a group is well above well below the optimal cut-off, we can figure that out with very few samples. Groups close to the threshold, on the other hand, have similar treatment effects. Therefore, we don't lose much if we mix them up. The only bad case arises when all units cluster around the threshold,
Work primarily done while interning at the Max Planck Institute for Intelligent Systems.

---

Published as a conference paper at ICLR 2026
not too close and not too far. We turn this intuitive observation into precise instance-dependent upper bound on the sample complexity of near-optimal treatment allocation. The formal argument is delicate, since we don't 't know the optimal threshold and we can only work from coarse estimates.
In a nutshell, our theory predicts that near-optimal allocation almost always has a linear- not -eterati---deeddeenee on 1/c. We thoroughly verify this prediction in five real-world RCT datasets. In all cases, we can find near-optimal allocations with even fewer samples than our up- per bound suggests. Fundamentally, our work highlights the stark difference between allocation and estimation: Nearly optimal allocations do not necessarily require highly accurate estimates.
1.1 OUR CONTRIBUTIONS
Consider a partition of the population into M groups and a budget K E (I....,M] that allows treating K out of M groups, such as schools, hospitals, or different age brackets. We refer to a group as a unit to indicate that for the purpose of treatment allocation we don't 't distinguish between individuals within a unit-we either treat all or none of the individuals within a unit. We assume that the cost of treating each unit is the same. Our goal is to identify the K units that would most benefit from receiving a specific treatment. Let r(u) E [0, 1] denote the average effect of treatment in unit E {I,.... M) J. A treatment allocation 2 S [M]] is any subset of K [M] units. The value achieved by an allocation i/ is the total sum Σnsu T(u) of treatment effects. An optimal allocation selects the K units with the highest r(u) values, breaking ties arbitrarily, achieving value V*.
In this paper, we study the sample complexity of finding a near-optimal allocation. We say that an allocation is (1 - -opptimal it achieves a value V that satisfies V/V" 2 - For bounded treatment effects, standard arguments show that we can always find (1-c)-optimal allocation from O(M/²²) samples. Moreover, there is a problem instance that requires (MM/²²) many samples. While these well-known bounds settle the worst-case sample complexity, our work shows that the typical sample complexity of treatment allocation is far lower. What matters for our theoretical results the shape of the distribution of treatment effects. For all reasonably smooth distributions- those that don't 't put excessive mass on small intervals we prove that O(M/c) many samples suffice to get a (1 opttmaal allocation.
Theorem (Informal). If the distribution of treatment effect values r(u) "smooth", we can obtain (1- -optimal allocation for any budget K S. M with (MMe)) many samples.
The key insight behind the theorem is that we only need to estimate the treatment effect values (()) up to accuracy م = @ Ve), rather than e, in order to obtain a near-optimal allocation. Intuitively, the problem of finding u, whether r(u) is above or below the cut-off TK. where TK is the smallest treatment effect of any unit in optimal allocation. For the units that have Tu)) value far from тк,
optimal allocation boils down to the problem of deciding, for each unit
determine so by estimating r(u) to low accuracy. For the units that are close to the threshold TK, mistakenly selecting one unit for another does not change the value of the allocation much.
We show that an accuracy level of p = @(vc) strikes the best balance between having the lowest possible level of accuracy while not losing too much value in the allocation. For this balance to hold, we need the number of units around the threshold TK to be reasonable.
We formalize this required notion of "smoothness" through our definition of p-regularity, which requires that any interval of length at least contain at most (ppM) units, that is, a multiple of what we would expect under the uniform distribution. Thus, p-regularity is a weak measure of closeness to the uniform distribution. Many natural distribution families, such as Gaussians or Beta distributions, are p-regular for a reasonable constant. The most regular distribution is the uniform distribution itself. Here, our analysis provides the biggest improvements. We prove that even in this case CATE estimation remains hard, requiring ((MMee)) samples. This formally proves a quadratie separation between treatment allocation and CATE estimation. In particular, CATE estimation remains hard even after imposing preguularity
Empirical evaluation. We consider data from five different randomized control trials (RCTs). For each dataset, we evaluate how many samples are needed to obtain near-optimal allocation. To do so, we treat the treatment effects estimated from the full RCT dataset as the ground truth. We then subsample the dataset to various sample sizes and compute the allocation value realized
2

---

study the strategies Published as a conference paper at ICLR 2026
Budget 30% (K=23)
Budget 30% (K=15)
Budget 30% (K=9)
Normallzes alloration value
aa-
Normalized allocation
ba
Normalzzdd allesation
0.8
9.7
6.7-
0.3
6.6-
0.5
Optimal
Optimal (1.0)
0.3
Optimal (1.0)
FulICATE
FulICATE
PulICATE
0.4
ALLOC
0.4-
ALLOC
Empirical data
Empirical data
0.3-
Empirical data
Sample 1536o
Sample 15les
15300
Sample
(a) STAR, schools.
(b) TUP, baseline poverty.
(c) NSW, baseline earnings.
Budget 30% (K=7)
Budget 30% (K=5)
Nomalized aalizzzd allocation value
ca
Normalized lllocattio
Optimal (1.0)
Optimal (1.0)
FullCATE
FulICATE
344
ALLOC
ALLOC
53
Empirical data
Empirical
5000
13630 Sample size
15000 20300
10000
19h0s
2010
Sample size
(d) Acupuncture, age.
(e) Post-op, BMI.
In all cases, realize close-to-optimal allocation with very few samples (blue), much less than the worst-case O (red) and even less than our theoretical bound (MM/e) (green).
Figure 1:
by our algorithm. As shown in Figure T in all cases we find near-optimal allocations with fewer samples than our theoretical upper bound predicts. This suggests that real-world treatment effect distributions meet the assumptions of our theory. As a robustness check, replicate the results across various different partitions of the populations and for varying budget sizes.
Discussion and limitations. There's immediate practical takeaway relevant future policy decisions about targeting welfare-promotee interventions. The standard sample size calculations for CATE estimation are excessively pessimistic for the purpose of treatment allocation. Indeed, we can find nearly optimal allocations from coarse treatment effect estimates. Perhaps counterintuitively, an RCT that is severely underpowered for CATE estimation can still yield excellent allocations. rule of thumb, M/c samples suffice for <) -optimal allocation.
While our smoothness requirement typically met practice, may not always hold. We address this limitation in two ways. First, expose exact instance-dependent optimality condition (Sec- tion 4.2) that can computed from coarse estimates as well. Thus, a policy maker can certify from few samples that solution is oear-optimal or-if it isnet-invest in additional samples. Alterna- tively, we consider strategies for the policy maker to mitigate cases where there are too many units around the optimal threshold тк without additional samples. Specifically, of underspending and overspending on the original budget K" in order to obtain a (1 - c)-optimal allocation with O(M/c) many samples. We show that, in practice, in the few cases where close- to-optimal allocation is not realized, we can find a very close threshold тк" at which we do realize with only (M//e) many samples. To summarize, our work strongly separates the sample com- plexity treatment effect estimation from that of treatment allocation. The latter requires far fewer samples than the former. This stark separation has the potential to inform future policy decisions.
1.2 RELATED WORK
Perdomo et al. (025) and Shirali et al. (2024) inspired our work by showing that effective resource
allocation need not require accurate predictions. These works bring into focus a critical examina-
tion of prediction in resource allocation (Barabas al. 2018: Perdomo 2024 Shirali et al.) 2025 Fischer- Abbigar et al. 2025 Mashiat et al., (025). Our results extend this line of work with fundamental observation: Extremely coarse treatment effect estimates still yield near-optimal
3

---

Published as a conference paper at ICLR 2026
allocations. Perhaps closest our work, Shirali et al. (2024) contrast unit-level allocations with individual-level targeting, showing that a measure of inter-unit inequality makes unit-level alloca- tions competitive with individual-level targeting. In our work, allocations are always unit-level; we study how accurate unit-level treatment effect estimates have to be for the purpose of allocation.
Another recent line of work studies different strategies for using RCT data effectively (Sharma & Wilder 2025 Cortes-Gomee et al., 2024 Wilder & Welle, 2025). In line with these results, our
work provides further evidence that investing in greater accuracy not always necessary in order to achieve a more efficient intervention. Unlike these works, our paper focuses on sample complexity upper bounds for treatment allocation, rather than on considering alternative strategies for estimating treatment effects. There is a vast causal inference literature on learning optimal targeting rules that address heterogeneity in the population (Manski, 2004, Qian Murphy 2011 Kitagawa Tetenov 2018 Kallus 2018, Athey & Wager, (2021). )). Much work tackles the case of large M using machine learning methods (Chernozhukov et al.] Nie & Wager 202T), which is not our focus. There's
also much work targeting rules subject to budget constraints, e.g., Bhattacharya & Dupas (2012); Luedike & Van Der Laan (2)16); Le et al. (2019). Our work also relates to the literature on bandits, specifically to the problem of identifying good arms (Audibert & Bubeck 2010) and the numerous
variants of this problem, such as identifying any subset of a set of good arms (Kalyanakrishnan ee al. 2012)] [Bubeck al., 2013 Chaudhuri & Kalyanakrishnan 2019 Rejwan & Mansour] (0220) In our setting, however, we are not interested in an adaptive procedure; we want our algorithm to be implementable in practice. The lower bound for CATE estimation runs through the bandit literature in various guises, see, c.g., Kaufmann et al. (2016); Katz-Samuut & Jamieson (2020).
NOTATION AND PRELIMINARIES
We have population X' divided into set U of M units. Each unit u E U is independent from the others, and we wish to select K out of the M units to carry out an intervention. We operate in the setting of a resource-constraines positive intervention: ideally, we would like to treat all units, but budget limits the number of units that we intervene on. In order to decide which units to select, a typical approach is for the policy-maker to estimate the average treatment effect T(u) in each unit u, and then choose the K units with the highest r(u) values. Throughout we assume bounded treatment effects ru) E [o. 1], normalized to the unit interval. Due to space constraints, defer all proofs of the main body to the appendix, while providing some proof sketches.
Estimating the average treatment effect within a unit-either via an RCT or an observational design-is the well-established problem of causal inference that is not the subject of this work. We therefore assume that we can get the average treatment effect r(u) up to an additive error é > 0 at the cost of O(1=²) samples. This assumption hides the challenges of causal inference, but exposes the sample complexity relevant for our argument.
Definition 1. Call f an (e, aaccuareee estimate of T if -|
< with probability 1- ò.
Definition 2 (Estimation oracle). Given a unit u E U and parameters e, > 0, an estimation oracle O returns an (c, -accurate estimate +(u) of r(u) at the cost of (ln(2/8) samples.
Throughout the paper, we always assume that the estimation oracle O is available to the algorithm. The sample complexity stated in Definition 2 follows from Hoeffding's ' inequality. The upper bound tight up to constants except in special cases. This follows from standard lower bounds for mean estimation in Bernoulli families (Le Cam 1973; Wasserman 2004] [sybakov| 2008). These lower bounds extend to estimating average treatment effects (Imbens & Rubin (015). Typically, in the setting of treatment effect estimation, we compute (e, 0) -accuraee estimate for every unit. We refer to this problem as the FullCATE problem.
Definition 3 (FulICATE problem). Given a population X of individuals divided into a set of M units U, solving the FullCATE(XX,4 c, problem consists of producing M estimates & such that {u)) is an (c. -acccuraee estimate of T (1) for each u E U.
We can compute the sample complexity of the FullCATE problem if we solve it by calling the estimation oracle with parameters 5c for each unit u U. making total of M independent calls. This upper bound is again tight, in general, up to constant factors when groups are non-overlapping, since each sample can only contribute estimating one of the treatment effects.
4

---

Published as a conference paper at ICLR 2026
Lemma 2. Having access to estimation oracle O, solve the FulICATE(X,U,C. prob- lem with a total of NFuICAATE = O(M ln(2/6)/e² samples from X.
2.1 ALLOCATION VERSUS ESTIMATION
Our goal is to find near-optimal allocation without incurring the cost of solving FullCATE. This means selecting the K < M units with the highest true treatment effect values r(u), where K is determined by the given budget. We think of the treatment
positive intervention, which is why we focus on identifying the units with the highes! values of (u)) For an integer K E ()....,M, we denote the K-th largest value of r(u) among the M possible treatment effect values by TK.
Definition 4 (Value). Given set U of M units and a budget K < M, an allocation function returns a set Ug l of K units. The value of the allocation function over an interval A A [0, 1], denoted V4,, (A), is defined as Vu, (A) = _nEu, T(u). The (total) value of the allocation function g for budget K is equal to Vu, ([0, 1J).
We abbreviate Vu, ((O, 1)) by Vu, if the interval A = [0. 1] can be inferred from the context. The optimal allocation for given budget is the one that realizes the highest value.
Definition 5. Given a budget to treat K units, the optimal allocation function is the one that selects the K values (u) that maximize the value Eveu T(u) over all subsets U of K units. We denote the set of units chosen by the optimal allocation function by U.
Throughout, assume that the true treatment effect values r(u) unique (otherwise, we can slightly perturb them), and thus the optimal allocation is unique.
Definition 6 (ALLOC problem). Given a budget K S M and parameters e, we say that a set Vu, (1,(().1))
Ug S u of K units is an (c, -optimml allocation ALLOC if Vur ([01I) 2 1 - e with probability least 5, where the probability is taken over the coins of g. Solving the L,, e,,,)) problem consists of selecting a set Ug of K units that is a (1 e, 6) -optimal allocation.
We usually drop the parameter 8, implying that the guarantee holds with high probability. Crucially, the value of an allocation is computed with the true values rather than with the estimated f values. The standard way of solving the ALLOC problem is through the FullCATE problem.
Lemma 3. Given any budget K- < M and parameters e, 8, we can solve the ALLOC(X, problem with N OM ln(2/8) (2) samples from X.
DIRECTLY TARGETING UNITS FOR ALLOCATION
Recall, the optimal allocation for a given budget K chooses the K units with the largest treatment effect value T (w)). It is clear that 2t {u r(u) 2 тK} (Claim B.I). Hence, solving the allocation problem reduces to solving threshold problem, where for given budget K, the task is to determine whether r(u) 2 TK or (u)) < TK for each u E U. In the former case, we decide to intervene on unit u; otherwise we do not. In order to solve this thresholding problem, do not need to estimate each of the r(u) values up to accuracy 6 we only need to determine whether Tu) lies above or below TK- The farther 7 (u) is from TK, the less accuracy need for our corresponding +(u) estimate. Conversely, we only need high accuracy in estimating fu) when r()) is close to the threshold TK.
3.1 LOW-ACCURACY ESTIMATION ALGORITHM
We introduce another parameter p that quantifies the eπrоr to which we estimate each treatment effect. Intuitively, we only try to determine whether r(u) is above + 2p or not. If we believe that r(u) belongs to the interval [TK, тк + 2p), then we "give up" and stop trying to determine the precise value of (u)). As we show, the value of that strikes the right balance between low sample complexity and a close-to-optimal allocation value is = O(ve). That is, instead of estimating all treatment effect values up to accuracy e, we only estimate them up accuracy of the order ve and then select the top K values based on these coarse estimates. This yields the following low-accuracy algorithm (Algorithm IT. which is non-adaptive RCT (further sample complexity reductions can be achieved by introducing adaptivity, for example with multiple stage RCT, but these are not usually implementable in practice).
that it can be implemented in typical one-round
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 Low-estimation non-adaptive allocation (LBA)
Input: M units, budget K, parameters E, 6, 7 > O, where 7 = (1).
For each unit u, obtain (p. >-acurrate estimate f(u) for TVE.
Output: A set ULEA consisting of the top K units sorted order f(u) values.
By definition, the algorithm selects all units such that i(u) 2 .x. First, by Hoeffding's bound and the choice of p. the algorithm requires a number of samples proportional to 1/e, rather than 1/e². Lemma 4. Given any X, X,U, K, Algorithm | requires ( (MMl(22M66/6 many samples from X.
By union bound, all units u satisfy (+(u) - r(u)] p with probabbility least 1- 8. Through careful analysis of the estimates, we also show the following:
Lemma 5. The output estimates ēK and the ULEA returned by Algorithm I satisfy the following properties: (I) Fix TKl S p. (2) All units such that r(u) > TK + 2p belong to ULEA (3) All units such that Tu) < TK- 2p do not belong to ULEA-
In other words, our algorithm selects the units optimally in the intervals +2p 1] and [0, TK 2p). We only need to worry about the interval surrounding TK. This motivates the following definitions:
Definition 7. We define the intervals A = (TK + 2p, 1] and D = [TK, -K- + 2p]- We also let ULEA- = {c E ULEA rT)) & A1), K1 = Hu r(u) E A1H1. and Ko = K K1
The guarantees of Lemma 5 imply that VeLea (A1) = Vu- (A,), and moreover allow us to lower bound VigEa ([O, 11) with (Tk - 2p) Ko, and upper bound (D) with (TK + 2 This allows us to get the following bound, where we use V(Ai) as a short-hand for Vinex (A1) = Vu- (A1).
Claim 3.1. VLLEA (10, 1])
4pKo
Vur ((0, 1]) 2 1 V(A,) (T + Ku
We let Pr. denote the discrete distribution of r(u) values. The quantity Pr,[TK, TK + 2p] and its relationship to 2pM is key to our analysis. To that end, we let Przx [TK: TK + 2pp = 0x and ((A) M. This leads to the following general accuracy bound for Algorithm
Claim 3.2. Given any X, U, K, Algorithm returns (1-m) 470k
vE)
-optimal approximation
+ TxOk
with O(M ln(2M/6)/e many samples.
In Section 4 we show how the expression Claim 3.2 simplifies to (1 - E) in the case of p-regular distributions, thus attaining a (1 - -P-OPT allocation.
3.2 QUANTILES SUFFICE
key insight in our analysis is the fact that for the problem of treatment allocation we only need to know quantiles of the distribution. In other words, it's enough to approximate the CDF of the distribution Pr, of treatment effect values.
Definition 8. Let F(() Pr (T() S t] be the CDF of the distribution of r(u) values, and let f denote the CDF corresponding to the distribution Pr, of the low-accuracy estimates f(u) produced by Algorithm ! That is, F(t) M [u: +(u)= S t)}.
We denote the PDF by ft). Given budget of K S. M, TK definitionally corresponds to the value in [0, 1] such that F(Tx) = 1 - K K/M; i.e., the K-th Mr- quantile of the distribution of r values. Similarly, 71 and Vu- be computed as permutation-invarian quantities derived from the CDF:
Claim 3.3. For any budget K < M, we can compute TK, T1 Vu- as: (1) Tk = F (1-K1- (2) = f (I), (3) Var ((0, 11) S"h dF))..
Moreover, the p-accurate estimates f(u) naturally provide approximation of f and F, which we can in turn use to check, from the low-accuracy estimates tx. the smoothness behavior of the treatment effect values around the optimal threshold:
Claim 3.4. For any E [0, 1). F(1 p) < F(I) À(t p).

---

Published as a conference paper at ICLR 2026
4 THE SAMPLE COMPLEXITY OF TREATMENT ALLOCATION
To illustrate the key ideas, we show that if the treatment effect values T(u) are uniformly distributed, our allocation algorithm achieves a (1 -ooptimll allocation with ((M//e) many samples. At the same time, any estimator that computes all M treatment effect values values r(u), each within e-aceuracy, must use least (MM/c² many samples. Therefore, the uniform distribution demon- strates a sample complexity gap between the problem of obtaining an optimal allocation and the problem of full CATE estimation. By uniformly distributed we mean that we place the 7 r()) values mniformly-space in [0, IJ and randomly permute them.
Theorem 6. Let the treatment effect values be uniformly distributed. Then: (I) We can obtain (1 c, 5) apprrxxiiatio of the optimal allocation with N = ( In(2M/6)/c) many samples from X for any budget K M. (2) FullCATE requires S?MM/2) samples.
Proof sketch. The sample complexity bound for Part (1) follows from Lemma The e-optimality of the realized allocation value follows a consequence of Claim 3.1 using the fact that 0g [T, TK + 2pp = 2p. given that the 7 ru) values are uniformly distributed by assumption. This allows us to obtain the term p² in the numerator of the expression 4p0x / (11 + (TK +2p))x)). ). which is the key to achieving a (1 - (1--1177 (v)))²) = (1 -e) --optimal allocation. The lower bound follows from minimax argument based LeCam's Theorem.
4.1 SMOOTH DISTRIBUTIONS
We can generalize our ((M/)) upper bound beyond the uniform distribution. All we require is for the D-interval to have "uniform-like" mass, that 0g A 2p and thus our proof of Theorem can still go through (i.c., to obtain the p² term in the numerator). Because TK could be anywhere, we require this condition on all intervals that have width at least 2p, a condition that we call p-regularity:
Definition 9. Given p O, we say that distribution Z on [0, is p-regular if for all intervals S C [0, 1] such that [S] 2 2p there exists a constant e @(1) satisfying Prz [S] S e U[S]. where 4 = Pru and U denotes the uniform distribution over [0, 1].
We can view this condition as a form of smoothness requirement. Note that this is a much looser notion than requiring. for example, O(p) total variation distance between Z and U, or with any other distance measure to the uniform that is not permutation invariant. Note that the value of e is always upper-bounded by the supremum of the density of the distribution Z; i.e., S sup, [(t).
Theorem 7. Let the treatment effect values be distributed according to a p-regular distribution for p ((ve). Then, can obtain (1 c, - approximmatio of the optimal allocation with N = O(M (22M/8)/c) many samples from X for any budget 0 < K < M.
Proof sketch. The p-regularity condition ensures that 0K (r,IITKTT TK + 2p] 2pc for e O(1). Following Claim 3.1, this allows us to again obtain the term p² in the numerator of the expression 4p0x/(71 + (Tx + 2p)0kk, thus ensuring that the ratio is upper bounded by
Note that for any distribution we can get (1 OPT approximation by setting ? = /n/ /Sc when calling Algorithm [ by letting c be the minimum value such that Pry [S] cc-4(S U(S for all intervals S such that (S] Z 2p. However, if is not @(1), then p might not be of the order ((ve). and so the sample complexity will no longer be of the order O(M/e). That is, for an arbitrary Pr,: Theorem 8. Let be the minimum value such that Pr-[S] < 2-2 all subsets 5 C [0, 1] such that |S| > 2p. For any budget K M, we call Algorithm with with p = VVe- /Sc. This gives (1 OPT approximation of the optimal allocation with ((M ( / p2pp) many samples.
We remark that Viz- is short-hand for Ve- ([0, 1)) and that is function of K. If the distribution is p-regular, then ( in which case require (MM/e) many samples, recovering Theorem 7 Note that solving the FullCATE problem to accuracy already required having a bound on тк (Claim 3p. if we want to ensure exactly that we get (1 -)-ottimal allocation. Similarly, here use an upper bound of 7 in order to ensure exact coppimallit.
Why p-regularity is not demanding property. As discussed, p-regularity is a permutation invariinn property that only depends on the shape of the CDF. It also ensures that the threshold Txk that deter- mines which units are selected for treatment does not yield highly arbitrary allocation. Empirically,
7

---

Published as a conference paper at ICLR 2026
as we extensively show in our experiments with real-world RCT data in Section 6 our algorithm highly succeeds in practice, further demonstrating that p-regularity is a natural condition. Moreover, typical families of distributions of treatment effect values are p-regular for reasonable values of y As we compute in Section Cection[..TTee using the CDF-derived quantities shown in Claim (3. we can explic- itly compute the values of Tk. c, Vu-, and y for typical distributions of treatment effect values from their CDF, which we do for the uniform, Beta, and Gaussian distribution families. In all of these cases, y ranges between 0.07 and 0.2.
4.2 GENERAL OPTIMALITY CONDITION
Theorem 7 is not an if and only if statement, and so p-regularity is a sufficient but not necessary condition for Algorithm [ to obtain a near optimal allocation. This means that, in some cases, we can obtain optimal allocations even if there is a lot of mass in the interval around the optimal threshold тк (e.g., if the units are extremely close to TK). We conclude this section by providing necessary and sufficient condition the CDF of the distribution of I values (for the worst-case high probability instance of Algorithm[|] D. thus providing a precise instane--dependent upper bound on the sample complexity of near-optimal treatment allocation. To do so, we require the following definition:
Definition 10. Given budger K S M, ak E [0. 1] corresponds to the smallest value such that ax]] Pr, [[TK, TK + +2pl]-
Pr, [[TK- 2p, TK-
Claim 4.1. For any budget K S M, Algorithm 7 called with parameters e, e, returns (1 - e, 0) OPT allocation if and only if the distribution of trealment effect values satisfies:
1T + tf(t) dt - -f
TK (t(t dt dt S & tf(t).
(1)
our estimation algorithm provides an approximation to the CDF of the distribution of treatment effect values. Specifically, we use it as certificate to check whether Equation D is satisfied (note that this is only a one-sided guarantee, given that we use various upper bounds). To do so, compute an upper bound of the expression (Vu- (D) VuRa \/Vr- ([O, 1]) using our i estimates. We do in Section C.2 by using Claim 3.4 In the same section, give an example of Claim in the case where the treatment effect values are uniformly distributed in [0, 1).
As we showed in Section 3.2,
5 FLEXIBLE BUDGET
When the original TK threshold has too many units around it for our algorithm to return an optimal allocation, we can slide the budget K up or down slightly into new budget K' so to find new threshold Tk¹ that lies in an interval of low probability mass (where p-regularity holds). In our experiments in Section 6 we find the closest K" to K that obtains (1 -)-ottmmal allocation with only OMM/c) samples, in the few cases where the original K does not give an optimal allocation.
Very dense distributions. Such budget K' does not always exist: e.g., if all T(u) values the same, then the interval [TK, TK + 2p] does not satisfy p-regularity for any K E [1, M]. However, in this case, Algorithm aaways returns an optimal allocation, since any selection of K units yields the optimal K- T(u) value (this is also shown by Claim 4.1 since (D= - Vuga = 0). In Section D we provide a "2 spikes" example where we show that no budget is able to achieve an optimal allocation. Nonetheless, we show that we almost achieve an optimal allocation, and in order to realize the required value, we can either promise a (1 OPT allocation for small s. or we can select more units for intervention. This type of resource augmentation is different from the "sliding budget" approach discussed at the beginning of this section. Here, we add more units to our allocation ULEA but we compete against the original budget K. For example, as we show in Section D for the 2 spikes case, increasing the budget to K' K(1 + 3e) ensures the e-optimality of a random allocation. In the appendix, we apply this overspending idea to distributions that are extra veK, units
not p-regular around the initial budget threshold Tk. As we show, adding suffices to obtain a (1 (1-a optomal allocation. In our empirical results in Section we obtain that in practice we only require adding one extra unit in order to realize the value of the optimal allocation.
Choosing an appropriate budget. In Section D we discuss various strategies for when the policy maker has some flexibility over the budget, where it is sensible to refine the number of units that
8

---

Published as a conference paper at ICLR 2026
School: Failure %
Poverty: Failure %
Earnings: Failure %
Age Groups: Failure %
BMI Groups: Failure %
Pellure %
Pallure %
Pallura %
Pallure %
Fallura %
10.4
10t4
10
10-
(a) STAR, schools.
(b) TUP, poverty.
(c) NSW, earnings.
(d) Acup., age.
(c) Post-op, BMI
School: Closest
Poverty: Closest K
Earnings: Closest
Age Groups: Closest
BMI Groups: Closest
Minimum
Miniuum
Minimum
Minimum
Minimum
0.5
0.5
10
10
10
10
10
10
10
10
10
10
10
10
10
10
10-7
Jra
Kra
(I) STAR, schools.
(g)
TUP, poverty.
(h) NSW, earnings.
(i) Acup., age.
i Post-op, BMI.
Figure 2: Top row: failure rate vs c. Bottom row: closest K' value vs c. Each row represents one dataset (STAR, TUP, NSW, Acup., Post-op) with one of the unit grouping methods that we study.
are selected for intervention after running Algorithm I These include: Detecting when a budget does not work by using the approximation F of the CDF (this is a one-sided test, as per the test of Claim D) sliding the budget from K to K' to find a p-regular interval, and adding few extra units to MLEA as to realize the optimal value (overspending), in the cases where the distribution is dense throughout a large range. We test the flexible budget strategies experimentally in Section
EXPERIMENTS WITH REAL-WORLD RCTS
We test allocation algorithm extensively on five different real-world RCT datasets across dif- ferent domains, these set of five datasets was recently studied in Sharma & Wilder (2025). These are: the (1) Tennessee's 's Student Teacher Achievement Ratio (STAR) project (Pate-Bain et al. 1997) (eduction), (2) Targeting the Ultra Poor (TUP) in India (Banerjee et al. 2021) (economic de-
velopment), (3) National Supported Work (NSW) demonstration (Dehejia Wahba 2002 1999. LaLonde, 1986) (labor economics), (4) Acupuncture dataset (Vickers et al. 2004)) (healthcare), and (5) Postoperative Puin dutasel (McHardy Chung 1999) (healthcare). We describe these detail,
as well our experimental set-up, in Section E and Section F
For each dataset, we group individuals in different ways using the covariates, ensuring that each group is somewhat balanced in terms of control vs treatment. Each different grouping method yields a partition L/ of M units. For each unit u, we define r(u) to be the average treatment effect within the unit across the entire dataset. Then, we run our semi-synthetic experiments follows: every time we sample from the population, we choose a unit u uniformly at random, and then receive a sample i.i.d. from Bern(7(u)). This gives us estimates +()) for each u. This semi-synthetic setting allows us to know the ground-truth of the value realized by the optimal allocation. We repeat the whole sampling procedure 50 times, to be able obtain confidence bounds. For the plots in Figure [ for various sample sizes, we compute the corresponding estimates f, and then the realized allocation value for various budgets K S M. For the plots in Figure [ we repeat the sampling simulation for different fixed values of c (and in turn repeating each simulation 50 times). For each, we call Algorithm and obtain a set of units ULEA for each budget K. For all budgets K E [1, M), we check whether not ULEA is an (1 e)-optimal allocation, and compute the percentage of budgets (out of a total of M possible budgets) that fail to be (1 - -)-optimal (top row in Figure I.. For the budgets that fail, we compute the closest budget K' that yields an (1 c)-optimal allocation (blue line, bottom row in Figure [. and the closest budget K" satisfying K' < K; i..., if we only consider underspending (orange line, bottom row in Figure [}. We see that, on average, the failure rate always says below 5%, and that 'K' - K is typically 1, even if we are only allowed to underspend (K' < K). We also try the resource augmentation approach of adding more units, and in all failed budget cases we always realize a (1 e) oottimal value by adding just one extra unit.
We provide all the details of the experiments and datasets along with the rest of figures (which include all of our grouping methods for each dataset, thus replicating Figures and extensively across all datasets and groupings) in the appendix in Sections and

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank Rediet Abebe, Florian Dorner, Kevin Jamieson, Juan Carlos Perdomo, Ali Shirali, and Bryan Wilder for helpful discussions, pointers, and feedback.
ETHICS STATEMENT
As we argue in this paper, treatment allocation is a problem of societal interest, given that policy makers use allocation algorithms and run randomized control trials in order to decide which groups of the population affect how these are performed in practice. In our paper, we only use publicly-ave publicly-available data from RCTs that have already been conducted, and so our study does not perform any real-world interven- tion. Still, we hope that our results can be beneficial in guiding treatment allocations in practice; specifically, our results suggest that good allocations can be attained with far fewer samples than the usual CATE estimation bound would suggest, and thus should be much cheaper and easier to implement in practice.
intervene on. Studying algorithms for such interventions can thus crucially
REPRODUCIBILITY STATEMENT
provide proofs for our mathematical statements in the appendix. We provide a detailed summary of the set-up of our experiments in Section and provide the code that we have used to generate all our plots and results in the supplementary material package for each of the five datasets. All the RCT datasets that explore in this paper publicly available the Internet.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Susan Athey and Guido W. Imbens. Recursive partitioning for heterogeneous causal effects. Pro- ceedings the National Academy of Sciences, 113(27)7353-7360, 2016. doi: 10.1073/pnas. 1510489113,
Susan Athey and Stefan Wager. Policy learning with observational data. Econometrica, 89(1): 133-161, 2021.
Susan Athey, Julie Tibshirani, and Stefan Wager. Generalized random forests. Annals of Statisties, 47(2):1148-1178 2019. doi:
Jean-Yvar Audibert and Sébastien Bubeck. Best arm identification in multi-armed bandits. In Conference on Learning Theory, pp. 13-p. PMLR, 2010.
Abhijit Banerjee, Esther Duflo, and Garima Sharma. Long-term effects of the targeting the ultra poor program. American Economic Review: Insights, 3(4))471-4866 2021.
Abhijit V Banerjee and Esther Duflo. The experimental approach to development economies. Anmual Review of Economics, 1(1):151-178, 2009, doi: 10. 1144////mue..ccoommss..05533 143233
Chelsea Barabas, Madars Virza, Karthik Dinakar, Joichi Ito, and Jonathan Zittrain. Interventions over predictions: Reframing the ethical debate for actuarial risk assessment. In ACM Conference on Fairness, Accountability and Transparency, pp. 62-76. PMLR, 2018.
Debopam Bhattacharya and Pascaline Dupas. Inferring welfare maximizing treatment assignment under budget constraints. Journal of Econometrics, 167(1):188-196, 2012.
Séebastian Bubeck, Tengyao Wang, and Nitin Viswanathan. Multiple identifications in multi-armed bandits. In International Conference on Machine Learning, pp. 258-265. PMLR, 2013.
Arghya Roy Chaudhuri and Shivaram Kalyanakrishnan. PAC identification of many good arms in stochastic multi-armed bandits. In International Conference on Machine Learning, pp. 991-1000. PMLR, 2019.
Victor Chernozhukov, Denis Chetverikov, Mert Demirer, Esther Duflo, Christian Hansen, Whitney Newey, and James Robins. Double/debiased machine learning for treatment and structural pa- rameters. The Econometrics Journal, pp. C1-C68, 2018.
Santiago
Naveen Raman, Aarti Singh, and Bryan Wilder. Data-driven design of ran- domized control trials with guaranteed treatment effects. In International Conference on Machine Learning. PMLR, 2024.
Rajeev H Dehejia and Sadek Wahba. Causal effects in nonexperimental studies: Reevaluating the evaluation of training programs. Journal of the American Statistical Association, 94(448): 1062, 1999.
Rajeev H Dehejia and Sadck Wahba. Propensity score-matching methods for nonexperimental causal studies. Review of Economics and Statistics, 84(1):151-161, 2002.
Unai Bischer-Abaigar Christoph Kern, and Juan Carlos Perdomo. The value of prediction in identi- fying the worst-off. In International Conference on Machine Learning. PMLR, 2025.
Mitchell Gail and Richard Simon. Testing for qualitative interactions between treatment effects and patient subsets. Biometrics, 41(2):361-372, 1985. doi: 10.2307/253086..
James J Heckman and Richard Robb. Alternative methods for evaluating the impact of interventions. In James Heckman and Burton Singer (eds.), Longitudinal Analysis of Labor Market Data, pp. 156-245. Cambridge University Press, New York, 1985, doi: 10. 10.1017/CB0978051 1661 180.006.
Guido W Imbens and Joshua D Angrist. Identification and estimation of local average treatment effects. Econometrica, 62(2))467-475 1994. doi: 10,2307/2951620,
Guido W. Imbens and Donald B. Rubin. Causal Inference Statistics, Social, and Biomedical Sciences. Cambridge University Press, New York, 2015.

---

Published as a conference paper at ICLR 2026
Nathan Kallus. Balanced policy evaluation and learning, Advances in Neural Information Process- ing Systems, 31, 2018.
Shivaram Kalyanakrishnan, Ambuj Tewari, Peter Auer, and Peter Stone. Pac subset selection in stochastic multi-armed bandits. ICML, volume 12, pp. 655-662, 2012.
Julian Katz-Samuels and Kevin Jamieson. The true sample complexity of identifying good arms. In International Conference on Artificial Intelligence and Statistics, pp. 1781-1791. PMLR, 2020.
Emilie Kaufmann, Olivier Cappé, and Aurélien Garivier. On the complexity of best-arm identifi- cation in multi-armed bandit models. The Journal of Machine Learning Research, 17(1):1-44 42, 2016.
Toru Kitagawa and Aleksey Tetenov. Who should be treated? empirical welfare maximization methods for treatment choice. Econometrica, 86(2):591-616, 2018.
Sören R. Künzel, Jasjeet S. Sekhon, Peter J. Bickel, and Bin Yu. Metaleamers for estimating het- erogeneous treatment effects using machine learning. Proceedings of the National Academy of Sciences, 1 116(1):4156-4165, 2019. doi: 10.1073/pnas. 111177//nas.10045971111 1804597116.
Robert LaLonde. Evaluating the econometric evaluations of training programs with experimental data. The American economic review, pp. 604-620, 1986.
Hoang Le, Cameron Voloshin, and Yisong Yue. Batch policy learning under constraints. In Inter- national Conference Machine Learning, pp. 3703-3712. PMLR, 2019.
Lucien Le Cam. Convergence of estimates under dimensionality restrictions. Annals of Statistics, (1):38-53, 1973. doi: 10. (0.1214/aos// 116340251177722511
Alexander R Luedtke and Mark Van Der Laan. Optimal individualized treatments in resource- limited settings. The International Journal of Biostatistics, 12(1):283-306 303, 2016.
Charles F Manski. Statistical treatment rules for heterogeneous populations. Econometrica, 72(4): 1221-1246, 2004.
Tasfia Mashiat, Patrick Fowler, and Sanmay Das. Who pays the rent? implications of spatial inequality for prediction-based allocation policies. In Proceedings of the AAAI/ACM Conference on AI, Ethics, and Society, volume 8, pp. 1686-1697, 2025.
FE McHardy and Chung. Postoperative sore throat: cause, prevention and treatment. Anaesthesia, 54(5):444-453, 1999.
Xinkun Nie and Stefan Wager. Quasi-oracle estimation of heterogeneous treatment effects. Biometrika, 108(2):299-319, 2021.
Helen Pain-Bain, Jayne foyd-Zaharias Van Cain, Elizabeth Word, and M Edward Binkley. STAR follow-up studies, 1996- 1997: The student/teacher achievement ratio (STAR) project. 1997.
Juan Carlos Perdomo. The relative value of prediction in algorithmic decision making. Proceed- ings of the 4/st International Conference on Machine Learning, pp. 40439 40460, 2024.
Juan Carlos Perdomo. Tolani Britton, Moritz Hardt, and Rediet Abebe. Difficult lessons on so- cial prediction from wisconsin public schools. In Proceedings of the 2025 ACM Conference on Fairness, Accountability, and Transparency, pp. (682-2704, 2025.
Min Qian and Susan A Murphy. Performance guarantees for individualized treatment rules. Annals of Statistics, 39(2): 1180, 2011.
Idan Rejwan and Yishay Mansour. Top-k combinatorial bandits with full-bandit feedback. In Algo- rithmic Learning Theory, pp. 752-776. PMLR, 2020.
Vibhhu Sharma and Bryan Wilder. Comparing targeting strategies for maximizing social welfare with limited resources. In International Conference Learning Representations. PMLR, 2025.
12

---

Published as a conference paper at ICLR 2026
Ali Shirali, Rediet Abebe, and Moritz Hardt. Allocation requires prediction only if inequality is low. In Proceedings of the 41st International Conference on Machine Learning, pp. 45114-45153, 2024.
Ali Shirali, Ariel Procaccia, and Rediet Abebe. The hidden cost of waiting for accurate predictions. In International Conference on Machine Learning. PMLR, 2025.
Alexandre B. Tsybakov. Introduction to Nonparametric Estimation. Springer Series in Statistics. Springer, New York, NY, 2008. doi: 10.1007/b13794.
Andrew J Vickers, Rebecca Rees, Catherine E Zollman, Rob McCarney, Claire M Smith, Nadia Ellis, Peter Fisher, and Robbert Van Haselen. Acupuncture for chronic headache in primary care: large, pragmatic, randomised trial. Bmj, 328(7442):744, 2004.
Larry Wasserman. All of Statistics: A Concise Course in Statistical Inference. Statistics. Springer, New York, NY, 2004. doi: 10.
Springer Texts in
Bryan Wilder and Pim Welle. Learning treatment effects while treating those need. Proceedings of the 26th ACM Conference on Economics and Computation, pp. 448-475 4888473, 2025.
13

---

Published as a conference paper at ICLR 2026
A DEFERRED PROOFS FROM SECTION 2
Lemma 9. We can compute an (e, (e.8)-accurate estimate (u) of u with (ll(2//))e²) samples.
Proof. By Hoeffding's inequality, given that r(u) E [0. 1] for all it follows that with probability at least 1 6,
Pr [lf(u) - (u) 2 <] S 2mm2
Hence, if we want F(u) be an (e, -accurae estimate of the true value r(u), need take m 2 ((n(2/00)2²) many samples from u.
Lemma 3. Given any budget K < M and parameters e, solve the cL, problem with = O(M ln(2/8)/e² samples from X.
Proof. By Lemma we can solve the FullCATE(XX 6) problem with parameters X,U, with (m (2 samples from X' We use parameters c' and 8/M, for é' be determined in the proof. This produces M estimates f such that f(u) is an (c', (eN)craaatuaat estimate of r(u) for each u E U. We then define the following allocation function ge select the units with the highest estimated values Fu)) This defines the set of units Ug.
By definition of тK. it follows that Vi- ((0, 1]) wrong, and by the e' accuracy guarantee we have that, in general, V2r ((,, 1J) Vu,, ((0, 1J) s 2K'', and equivalently
KTK- the worst case, all the units that we select
Vi,, (10, 11) Vur ((0, 1]]
2Kt
2K
2
2 1
1
Vur (1o 1))
KTK
TK
Thus, we set é' < TK when calling the FullCATE problem, we obtain (1 (1-c, 6)-OPP allocation function. Hence, we have solved the ALLOC problem, required.
B DEFERRED PROOFS FROM SECTION 3
Claim B.1. Given a budget K, u
{u r(u) 2 Tk].
Proof. By definition of TK. there are exactly K units that have 7(u) value higher than TK, all of which are included in 24. Hence swapping any unit in li" with one not in the L1" would decrease the value of the allocation, which would contradict the optimality of ALLOC".
Lemma 4, Given any X, X,U, K. Algorithm 7 requires O(M (n(2M/8) many samples from X.
Proof. By a Hoeffding bound using parameters and 8/M, letting (u) denote the number of samples that we require from unit u, we have that
Pr [[f(w) - T()) 2 p] s 2exp( (exp(-2N(1)) p²) = M N(u) (2
2p²
Therefore, adding the number of samples across the M units, we obtain that this low-accuracy estimation requires
M (222//8) Mm(2//8)
p
many samples from X.
Lemma 5. The output estimates fg and the set ULEA returned by Algorithm satisfy the following
properties: (1) fix Tkl S p. (2) All units u such that T(u) > Tk + 2p belong to ULE (3) All
units such that ruu) < TK 2p do not belong ULEN
14

---

Published as a conference paper at ICLR 2026
Proof. We prove each of the three parts separately.
(1) By definition TK. there are exactly K values T(u) such that r(u) 2 TK. By the p-accuracy estimates fu)) it follows that f(u) 2 7K p for all such u. Then, there K values f()) that are at least тк
guarantee
By definition of fx. then, must be that fK 2 TK p.
Symmetrically, ēK is definitionally the (N K 1)-th smallest number in the increasing sequence of values f(u), and similarly тк is the (N - K + 1)-th smallest number in the increasing sequence of values Tu)). Therefore, exactly N K values T(u) such that r(u) < Tx- By the p-aceuracy guarantee, +(u) < тк + for all such u. Thus there are - K values F(u) that are below тK + p- By definition of ēk. it must be that 7K 7K p.
Since ²K 2 TK and Êx- S TK + p. it follows that Ýx E [Tx - Ps 1, + p]. Hence |xx - Tkl S p- (2) By the p-accuracy guarantee, all units satisfying r(u) Tк 2p also satisfy f(u) > тк + p. By part (1) of this claim, it then follows that f(u) > fx. Hence, Algorithm treatment and they belong ULEA.
[
selects such u for
(3) Symmetrically, by the p-accuracy guarantee, all units u satisfying T(u) TK 2p also satisfy f(u) < тK - p. By part (1) of this claim, then follows that +(u) < fg. Hence, Algorithm does not select such u for treatment and so they do not belong to ULEA.
VLLEA (10, 1])
4pKo
Claim 3.1.
21
Vue (10, 1])
V(A1) (TK Ku
Before we prove Claim 3.1 we show some intermediate lemmas.
First, we use interval A to analyze the difference between ViLEa and Vu-.
Claim B.2. Vu- (10, 1]) = Vue ([TK, 1]).
Proof. As shown in Claim B.I the optimal allocation for budget K contains all units u such that T(u) 2 TK.
Claim B.3. ViLca (10, 1]) 2 VeLLa (A1) ((TK +2p,
Proof. Firstly, ViLea ([0, 1]) 2 VinEa (Ai) follows since A S (0, 1] and V-)) is an additive function. Secondly, ViLEa (A1) = V ((TK + 2p, 1)) follows by the definition of A
We can lower bound ViLsa (10, 1)) more precisely using the interval D |TK: TK 2p] and the set 7. Specifically, we write the values of the optimal allocation and our allocation
ULEA (see Definition using the D interval.
Claim B.4. (1) Vur (10, 1J) Ver (D) Vu ((TK 2p, 1))-
(2) VILeA (10, 1]) Eieur T(i) VILEA ((Tx + 2p, 11).
Proof. The first equality follows by Claim B.2 and the additivity of V(-). The second equality follows just from the additivity of V..).
We want to separate the budget between the number of units it treats over A₁ versus elsewhere, so we split K as follows:
Definition 11. Let K₃ = [ T(u)= E A1)), and Kg K K1.
We can now provide the following bounds:
Claim B.5. Vi (10, 2 (TK 2p)Kg,
Proof. By Lemma 5. all units LLLEN satisfy r(u) 2 TK 2p, and Ko $ K.
n
15

---

Published as a conference paper at ICLR 2026
Claim B.6. (D) & (TK Ko
Proof. By definition of A1, all units u such that r(u) E A1 satisfy satisfy 7(u) > TK + 2p. Hence, by definition of Ko, any unit u E Ko must satisfy r(u) TK + 2p.
VELEA (10, 11) Vu- (00, 1])
ViLea + Verge (10, 1J)
Vahna (A1) + (TK Ko
Claim B.7.
Vu- (A1) + Vu- (D)
Vue (A1) + (TK + 20
Proof. The first equality follows from Claim B.4 the inequality from Claims B.5 and B.6
To further simplify the expression, we use:
Claim B.8. VuLaa (A1) = Vu- (A1).
Proof. By definition of A,, all units u E A, satisfy 7(u)= > Tk 2p. By definition of TK. all such are part of Ut. As for our allocation, by 5)l such are also part of ULEA-
By the previous claim, we can use ((A)) as a short-hand for ViLLN (A1) (Ai( intermediate small claims, we can now show Claim 3.1
With these
Proof of Claim 3.1 This follows from re-arranging the expression in Claim B.7 Specifically,
VidNA (A1) (TK K.
V(Ar) (TK Ku
Vu- (A1) + (TK + Ko
V(Ai) + (TK + +2p)
V(Ai) (TK +2p) (TK + 2p) + (TK 20
V(A,) + (Tx + +p
(Tx + Ko (TK - Ko
V(A1) + (TK + 2p)
4pKo
1 V(A1) + (TK + Ko
For the proof of Claim 3.2 recall the formal definitions of 0g and 71: Definition 12. Let Pr, [TK, TK 2p] = 0g and let V(Ai) = 7 M.
Moreover, we can relate Ko the distribution Pr, of treatment effect values: Claim B.9. Ky Pr, [[T,, TK + 2p]] M.
Proof. All units selected with the Ko budget satisfy r(u) & A by definition, and hence they satisfy (u) S Tк + 2p. Since all units selected with the Kn budget are still part of U", it follows that they all satisfy r(u) 2 тK. Hence, [Kol = {u r()) E [TK, TK + 2p]]. Lastly, {[u 7(1)
[TK, TK + 2p]| = Pr [TK, TK + 2p)- N by definition of Pr,-
Claim 3.2. Given any X,U, K, Algorithm returns a (1-_) 470x
71 Tx vt) optimal approximation
with O(M (IMM((2M/886 many samples.
Proof. Using the Definition 12 on Claim 3.1 and dividing by N it follows that
Viana (10, 1))
4p0x
2
Vue (10, 11)
71 7 + (TK - 2p)0k.
16

---

Published as a conference paper at ICLR 2026
This is - e)-OPT approximation if
4p0k
S.
11 + (TK + 2p)0k
Hence, to get a c)-OPT approximation, we require
E71
0K 4p dTK + 2p)
Alternatively, by plugging in p TVE, we obtain the expression in the statement of Claim 3.2
Claim 3.3. For any budget K S M, can compute TK, n, Vur as: (1) TK F- K/M),
(1) = 1 M 22 (t), (3) Vu- (I0, 1]) = S+x ddF(t).
Proof. First, 7K definitionally the K-th largest ru) value. In terms of the CDF of treatment effect values Pr,, this means that тк corresponds to the value in [0, 1] satisfying F(TK) 1- K/M.
Second, " definitionally satisfies 7/1 = ((A))/N. By the definition of the value function and given that A (TK 2p, 1), it follows that
V(A)= = T(U).
Σ
u/T(w)EA
Therefore, this corresponds to taking the integral f2x+20 tf(t) dt, where accounts for the value of Tu)) and f(t) for the density of units for each value. Similarly,
Vu- 1]) .f tf(t) dt [ tf(t) dt,
given that by Claim B.2 Vir (10, 11) Vir (ITK, 11). Lastly, integrating (S(() di with the PDF equivalent integrating (t) with the CDF.
Claim 3.4. For any t E [0. 1), f(t p) s F(t) S F(t+ p).
Proof. If S t - p, by the p-accuracy guarantee follows that r(u) S f(u) + S t. Hence:
Ít p) (u: =(u)= <t -
p}] [u: T()) s <t} F(t).
M
M
Symmetrically, if ru) S 1, then by the p-accuracy guarantee it follows that f(u) S ((u)+/ S ltp.
Hence:
F(t) Hu: (u)
s stH
[[u f(u) St st+ +Hl +)
<
M
M
We can equivalently express Claim 1.4im3.4jj in terms of the number of units present in an interval. Specif- ically, for any interval in (0, 1] we have that:
Claim B.10. Given any interval (a, b] [0, 1].
Iu f(u) [a + p. pl Pr [la 6]]
M |u f(u) [a p,b p|/.
Proof. By the p-accuracy guarantee, we know that for any unit u, if f() E [a + b - pl, then the corresponding r(u) in [a, 6). Similarly, if 7(u) [a, b). then this implies that f(u) [a- p. p). Lastly, the number of (u) such that r(u) E (a, b) is precisely equal Pr, [[a, b]] M.
n1
17

---

Published as a conference paper at ICLR 2026
This allows us to low-accuracy estimates of the treatment effects to approximate the true probability mass in any interval by enlarging or shrinking the interval accordingly.
We use this approximation of the CDF to check the smoothness behavior of the treatment effect values from our low-accuracy estimates fx- Specifically, we are interested in the number of units in the D [TK, TK + 2p] interval. By Claim B.10 we have that
Pr [[TK TK + 2p] M S Iu: f(u) E [TK - TP -3pl/-
Moreover, by the p-accuracy guarantee, know that тк E [fx P. fx + pl. Hence, the number of units that have F(u) value in the interval [fx 2p, 5x + 4p] upper bounds the quantity
Pr [[TK.TK + 2p]] M,
and can be entirely computed from our low low-accuracy estimates f. This allows us to check whether the interval around the optimal threshold contains too many units, and thus fails to satisfy the smoothness condition required by p-regularity. As we discuss in Section 5 we can also use our low-accuracy estimation of the CDF to guide the choice of the budget, provided there is flexibility in the choice of the budget.
C DEFERRED PROOFS FROM SECTION 4
We begin by showing our upper-bound for p-regular distributions, and then show the upper-bound for uniform distributions corollary, given that the uniform distribution is p-regular.
Theorem 7. Let the treatment effect values distributed according to a p-regular distribution for P ((VE). Then,
obtain (1 apprxxxapiiniii of the optimal allocation with N = ((MM((M/8) many samples from X' for any budget O < K S.
Proof. For the first part, we use our Algorithm [ The sample complexity follows from Claim 4 As for the accuracy guarantee, by Claim 3.1 we know that we get (1 0PT approximation if
4p0k
S.
71 + (Tx + 2p)0k
Because Pr, is p-regular, we have that
0k Pr[TK,
2p] S 2pc
for some (1). As in Section bound by dropping (TK + 2p) from the denominator and requiring that
4 we bound the value of V ((A)). We obtain a more conservative
4p(2pc)
8cp2
S E.
71
71
Hence, we obtain a (1 OPT approximation by setting y = (21/(8c and p = YVG, which ensures that the previous inequality holds.
We can also keep the more precise ratio
4p0k * h++(x++2p00 +2p) Note that
?1 (Tk + 2p)0x 2 Vu- ([TK. 1J) Vur.
Note that Vu- is function of K as well (even if not indicated with subseript). Therefore, need to satisfy the inequality
4p0k
4p0k Vir
Sc.
71 + (TK + 2p)0k
which we do by setting 7 = / (8c) and p = nVe.
is clear from this proof why we need 0x < 2pc. If we only have the bound 0g < 1, then we get 4p/ Vu- I c, which in turn requires p to be of the order of e. rather than of fe. But 0x 2pc, then we get the term p² in the numerator, which is the crucial point that allows us to get p ((va) and thus an optimal allocation algorithm with sample complexity O(Mfe) rather than with (M/ee).
18

---

Published as a conference paper at ICLR 2026
Theorem 6. Let the treatment effect values be uniformly distributed. Then: (I) We can obtain
(1 e, c. apprrxxiiatin of the optimal allocation with N = O(M 1n(2M/6)/e) many samples from X for any budget K S M. (2) FullCATE requires ((M/ samples.
The sample complexity follows from Claim 4 As for the accuracy guarantee, by Claim 3.1 we know that we get (1 - -e, 0 approximation if
Proof. For the first part, we use our Algorithm I
4p0k
Sc.
71 + (TK + 2p)0k
For fixed M, recall that for uniform distribution of M treatment effect values, the r(u) equally-spaced in [0. 1] and randomly permuted among them. Thus, this discrete version of the uniform distribution is p-regular with p 1/M and c 2. This follows from the following argument: let S S [o, 1] be any interval of length É 2 2p. By the discrete equal spacing of the treatment effects, there are at most LM + 2 units in S, and so Pr[S] (EM + 2)/M f + 2/M
under this distribution. So for all intervals S with |S| 2 2p 2/M, it follows that
Pr[S] S é 2/M - 28 = =24[S]:
thus satisfying the definition of p-regularity with c = Given that is fixed, p = ((VE) and thus apply Theorem to this discrete uniform distribution, from which the upper bound follows.
For the lower bound, we provide a minimax argument based on LeCam's 's method. Given a parameter space θ for each 0 E θ we let Pe denote the distribution of the observed data under parameter 8. Then, LeCam's Theorem (also known as the two-point method) states that
inf sup max (Pr(w * O). Pr 0,0'
Vr +
KL(PolIPpr)
Poi
for any estimator & and parameter values 0,0' E θ, where Prp, denotes the probability when the data is distributed according to Po. and similarly for Po. Here, KL(Po]lPw) denotes the KL divergence between Po and Por.
Because we are working with the discrete uniform distribution, where the r(u) values are uniformly spaced in [0, 1], we consider the family of values = {1/M, 2/M Let be one of such values, bounded away from and from 1; without loss of generality, we can let p E [1/4, 3/ 4). For a unit u, we consider the two parameters
T(u) = p 0': 7()) = +
If (l) is an (e, 8) acccuraee estimate r(u). then definitionally must satisfy
Pr [[((u) (p - e) < <] 2 - 6, Pr (I+(w) (p + e)l S 2 1- 8. We consider the test V (w) 2 pl. For an (c, acccurate estimate, we have that
Pr( 1] o pr(w- = 0] 6.
That is, good estimator yields test with small error.
For each unit u, we observe N samples drawn i.dd. from Bern(T(u). By standard bounds on the KL divergence we have that
(Lennenpp c) Bermn + <)) < 13c..
Therefore, LeCam's Theorem tells us that for any estimator V
1
Pr[M
P
+]+ Pr = 0] 2
KL(Bem(p-e) (CBBmm(p-e))Bcm(ptt))
Hence we get that
28
2
1
13Ne
N
13 log lo/81/))).
2
Lastly, because the T(u) values are uniformly distributed over [0, 1] it follows that ( many units have 7(u) value bounded away from 0 and 1. (In the case where we consider E [1/4, 3/4], there are M/2 many such units.) Therefore, by applying the LeCam Theorem on each such unit and taking union bound over them it follows that the total number of samples required for producing (e. 8) -cccurate estimates for all units that have r(u) value bounded away from and I is N (MM for any estimator V and hence N S(M/e2).
wanted show. n
19

---

Published as a conference paper at ICLR 2026
While the upper bound of Theorem follows corollary of the p-regularity theorem (Theorem for the discrete uniformly-spaced placement of the treatment effect values, is illustrative to prove the upper bound for idealized version of the uniform distribution, where M + O.
In this case, we have that
0x Pr(TK, TK + 2p] 2p.
Recall that YI is such that ((A1) N1 We can compute V(A1) exactly (as per Claim 3.3 but for this proof is enough to bound ((A)). Given that A = (Tx + 2p, 1J. it follows that [ry[Ar] (Tx + 2p) TK 2p. Hence:
V(A1) N 2 (Tk + 2p) Pr[Ai] = (TK + 2p) - TK 2p),
given that all units in A1 have value least 7K + 2p. Using this bound the denominator of the LHS Equation 3 we get that
?1 + (Tx + 2p)0x 2 (TK 21 TK - 2p) + (Tx + 2p)(2p) (Tx + 2p) TK).
Hence by Equation we obtain - e, 0)-PT approximation if
4p0x
8p2 (TK + +pp)(1 -T) (TK + 2p) - Tx) S.
4p(2p)
71 (TK + 2p)0x
TK [K-
This is whenever p
Ve. In the case of the uniform distribution, we have that Tk - K/M, which allows us to compute the value of 7 exactly. We can also provide general bound for all budgets K: the function Tx -TK) for Ik E [0, 1] is maximized at /4, and so setting 4v2 guarantees that our Algorithm I achieves (1 - e, 6)PPT approximation.
8
= 7V for
Theorem 8. Let c be the minimum value such that Pr-[S] < c-((() for all subsets 5 C [0, such
that [S| 2p. For any budget K M, call Algorithm with with p = / /Sc. This gives
(1 -)PO approximation of the optimal allocation with (M (MIn(2M/6)) / p²) many samples.
Proof. While might not be of the order of @(1), we can still write 0x 2pc. Then, we get (1 0P approximation if
4p0
8cp²
71 (TK 2p)0x 71 (TK 2ppc) S.
We can similarly ensure the more conservative inequality 8cp² /m1 S by setting 7 / and
8c
p 7/c. Less conservatively, we can set = v VV/k/(BB))) given that 7 (TK +p)0k Vu-. D
C. EXAMPLES y TYPICAL DISTRIBUTIONS
For distribution Pr., recall that we denote the CDF by F, and so we denote the density function by f. Note that we can bound the value 71 using f:
L'
I (f(1)
1TH rk t2 if(1)
71
(f(1)
Fx +2p
TK
tf(t) dt -(TK 2p)0x
15(1) di (TK 2p)2cp.
TK
Recall that S", tf(() dt V.K. Indeed, in order compute the value of the optimal allocation, it suffices to have the CDF function, with any permutation of the units. Similarly, directly from the CDF. Following our expressions derived in Claim 3.3 we can compute the values of y, c, and Vu- for various typical distributions from their CDFs. In these examples, it is neater to integrate t (t) instead the CDF directly, but not that we only need knowledge of the CDF of the distribution (and the integration is equivalent).
compute тк
20

---

Published as a conference paper at ICLR 2026
Uniform distribution. For the uniform distribution, we have f(t) 1. Hence, 1, since it is the maximum possible value of f(t). Moreover, we have:
["' (f(() [ tdt
TK
Vur (10, 11) Moreover, TK K/M. Therefore, by Theorem 7 and Claim 8
JK
TK
it follows that
Vu-
(1 K/M)²
8c
For all K E [0, y is maximized at R 0.35.
Beta distribution. For a, 3. the PDF of the Beta distribution is given by
1
f(t: a, 9) B(a, B) " t)0-1,
where the Beta function B(a, 3) is the normalization constant, and the CDF is given by
F(t) B(t: a)) == L2(a,
B(a, 3)
where B(x; a, B) -o0 10- 1)8-1 di is known as the incomplete Beta function.
For a, B 1, the mode of the distribution is given by --1 (if 3, then the mode is equal to 1/2), and so
(a - 18 1)
c sup = f(
a + - = (a 8 -
As for Vi+, we have that
Vz- ((0, 1)) = f(/ 15(1) di = ["
10-1(1 1-1-) 1)8-1
JK
Bla, 3)
ta(1 t)8-1 dt
B(a, B) JTK
B(a L, 3) B(TKia 1,3) B(a, 3)
B(a+1,83
(1 ITC (a + 1,
Bla, 3)
a + (1 I (a 1,
Lastly, TK is the value such that F(TK) = K/M, or equivalently, such that S. S(1) dl 1 K/M, Given that F(t) Ie(a, 8), it follows that
(o,
For example, for Beta(2, 2) we have that f(t) 61(1 and sup, f(1) f(1/2) 1.5. We also have that (,(3 2) = 3 314, and so
2
Ir, (0+ 1,5) = ( - (4T -3 = 41- + kr
Vu- 4
4
Given that F(t) (2) 3t - 2t3, and that Tk is the value in [] such that F F(TK) =
1- K K/M, it follows that тк corresponds the unique root of the polynomial 2 37 + 1--KKM/ K/M in [0, 1]. For example, if K/M = 0.5, we get that V/y- x 0.34 and 7 < 0.17. For K/M = 0.25 we
get 7 S 0.12, and for K/M 0.75 we get 7 S 0.19.
For Beta(3. 3) we have that ft) 30t² (1 2)2 and so sup, f(t) f(1/2) = 1.875. We also have
that (,(33) = 10t4 15t 6 and so
Vue = 0 [1 I (a + 1, B)) 2
1
(10T
157 + +r+))
= (1 107 + 157 6T
1
2
21

---

Published as a conference paper at ICLR 2026
Given that F(t) = (3.33) = 10t - 15t + 6t and tk satisfies F(tx) 1 K/M, it follows that
TK corresponds to the root of the polynomial G - 155k + 10T - (1 K/M)= = 0 in [0. 1]. For
example, if /MM- 0.5, we get that Vu-- se 0.37 and y S 0.16. For K/M = 0.25 get y - 0.13,
and for K/M- = 0.75 we get 7 S 0.17.
For (eta 4) we have that f(t) 20r(1 t) and mode 2-1
244-2 4 and
S sup, f(t)
/(1/4) R 2.1. We also have that (3, 4) = 201 45/ + 361 10/ and so
Vu- - 1 (a+1, = - (1- 91--(205--555 45t 45:5--44---6--10000
2
1 (1- 11-200t +45t 366 + +106%)
6
3
Using F(I) 12(2, 4), it follows that тк corresponds to the root of the polynomial 10 207 K +
15T - 4T - -1+ K/M in j0, 1]. For example, if K K/M- 0.5, we get that Vu- R 0.24 and
7 0.12. For K/M- = 0.25 we get y s 0.09, and for K /MM 0.75 we get y 0.13.
Gaussian distribution. Given a normal distribution N(,, we truncate it to [0, 1] and re- normalize it. Let 6(t denote the PDF of ((,,0)) and (t) the CDF Let
1 Z x(b) (aa).
b
Then, the truncated PDF function is equal
0 1-
f(t)
oZ
where / E [0, 1). The CDF then equal to
+ (- ((a)) F(t)
7
for L E [0, 1]. Assuming that μ. E |0, 1]. we get that the maximum density occurs at (0),, where
x()
1
c= sup s(1) sup
oZ
oZv2T
Given budget K, the threshold Tk is the value satisfying F(TK) 1 K/M, and hence the value satisfying
① TK
(b) aZ.
Hence, тk- = = (- - oZ)). Lastly, we compute ViK Let yk (Tx - -1)/a.
Vu- = 1' tf((t)tt= =
['' t t 0
1x
Z TK
I" (p+ oy)d
yK
1 (n 1 Φly)dy + [ volu)dy)
Juk
1 (b(b) ( )) + (Ø(yK) - x(b))
where we used the change of variable y = بيا and t = p + σy and di ady.
For example, for = 0.5, a = 0.15, and budget K/M = 0.5, we get Tk zo 0.5, Vi- R 0.31, c 2.66, and 7 < 0.11. For μ 0.., o 0.2, and budget K/M 0.75, we get TK 0.19. Vu- Rs 0.23, c & 2.13, and < 0.13. For = 0.7, 0.10, and budget K/M = 0.25, we get TK 0.77, Vu- 0.19, e 3.98, and 7 S 0.07.
C.2 GENERAL CDF CONDITION
Claim 4.1. For any budget K < M, Algorithm 7 caalldd with parameters c, 5, p returns (1 - cc,)- OPT allocation if and only if the distribution of Treatment effect values satisfies:
STK g (f) di STx-29 (f(() dt dt. S f 15(1)
(1)
22

---

Published as a conference paper at ICLR 2026
Proof. By Claims B.7 and B.8, know that
ViLEA (10, 1))
VuLa (A1) + Ve (0.,1)
V(A1) + VuLes
Vur (I)
Vir- (A1) + Vu- (D)
V(A1) + (Vr-(1)
By re-arranging the expression in manner similar to Claim 3.1 (but without using lower and upper bounds), we get that:
ViA (10, 11)
V(A)) - (D) - Vu- (D) + Vig V(A,) + V22- (D)
Ve Viren
Vi2 (IO, 1T)
V(A)) + (D)
By definition of the value function and the PDF [(),, follows that
v(A)) [ tf(t)dt,
Vu- Vu- = f (2 tf(t)dt.
+2p
As for Vupaa by Claim 5 it follows that all units E ULEA satisfy T(u) 2 TK 2p. Therefore, in the worst case, Algorithm [ selects the Ko units with the lowest r(u) values in the interval D [TK, тK + 2p]. This is precisely what the definition of or captures (Definition 10 and so it follows that
TK aK
VuLen 2
(f(dd,
with the equality holding in the worst-case instance of Algorithm | (even with high probability). Hence, we get (1 (1-6, c, 8PP allocation whenever
f if(t)dt ST8-20 tf(t)dt
Se
Shc tf(t)dt
as we wanted to show.
We can use our low-accuracy estimates îK to check whether Equation I is satisfied (note that this is only one-sided guarantee; failing to satisfy it does not imply that the allocation achieved by Algorithm is not optimal, given that we use various upper bounds).
To do want to compute upper bound of the expression
Vu- Vu-(D) (D) VeLi
(4)
Vi- (0,
using our / estimates. To lower bound the denominator, the number of units that have ruu) E [Tx 1] lower-bounded by /u: f(u) E [TK 1]]. For each of the i(u) values that we find in this interval, by p-accuracy we know that (u) 2 f(u) - p. Hence,
Vu- (I0, 1J) N
Σ
7(u)
(x)
Because do not have access the true TK, but we know that |iK Txl S p by Claim 5 actually compute the lower bound as
Vee (10, 1))
Σ
f(u)
(5)
u=(())2fx22p
Next, upper bound V(D). The number of units that have ruu) E [TK, Tk + 2p) is upper-bounded by ju : (()) E [TK PK + 3p]]- For each the f(u) values that we find in this interval, by p- accuracy know that (u) < +(u) + p. Hence,
Vu- (D) S
Σ
+(u) + +
w: (u)E -ATTK--PTT*
Again because we do not have access to the true Tk but we do have the guarantee that | - S P. follows that
Vur (D) <
Σ
f(u)
(6)
24p
23

---

Published as a conference paper at ICLR 2026
Lastly, as for Vules we interested in lower bounding the number of units that belong to ULEA To do so, first lower bound the quantity Ko, which corresponds to upper bounding the quantity K1. We know that K2 S |uu 5 E |TK + p. 1]] S |u E [ 1]I Hence, Ko 2 K |u: E [ I|I- Recall that all units in LLEA satisfy 5 r(u) тK- 2p. To lower bound the value of Vifea we add up
the first |u: E [+x, 1]] units that have F(u) value above ēx- 3p, in increasing order.
Iu
Putting the three terms together, that is, our bounds for V4r- (J0. 1J) (Equation 5). Viz (D) (Equa- tion and Vigea (the sum of the first ---Iu |u + E [fx; 1] units that have (()) value above f 3p) computed with our low- awwaappooxiationnn i values, we have obtained an upper bound the i Claim
expression in Equation and thus to Equation 1
Hence, our upper bound is smaller than 6, we have a guarantee that Algorithm has returned an (1 c, 6)-PP approximation.
clear that, definitionally, we get (1 O)T allocation whenever we satisfy the inequality
Vilea [0.11)
2 6. The importance of Lemma 4.1 is that shows how we can check (one sided)
Vir- (10.1) whether we have obtained an optimal allocation using only our fx estimates.
We give an example of how to apply Claim 4.1 for specific distribution by repeating the accuracy and sample complexity caleulation for the case where the CATE values uniformly distributed, this time with a tight accuracy analysis rather than using lower and upper bounds for VA (f) and VA, (ALLOC*), ), respectively.
Tight accuracy calculation for the uniform distribution. Let the treatment effect values be dis- tributed as On, 1). Then, the PDF corresponds to the function ft) 1. We compute each of the three terms in the RHS of Equation | Given that the uniform distribution is symmetric around
TK for any value of TK, follows that ar- = 0 for all K.
Vir (10, 1J)
Σ
r(U)
t2
dt
2
1
(Tk)²
2TKP 2p²
2xTx +2p
Txx
2
[
Vur (D)
Σ
T(u)
t2p
ddl
{22
+2p
-2TKP +pp²
2 TI
(NT( +2p)
JK
-K-
-
ViLeA
E
T(u)
t
2TTKP- 22²
-2
2 -p
u:T()E 22.TK-
where the equality holds in the worst-case high probability output of Algorithm | otherwise it is lower bound. That is, the computation for Vue corresponds to the worst-case output of Algorithm as is the definition of ax. Then, V22- (Dz Vilia = 4p² and
ViLa (10,
Vu- (D)(DD) ViL Vie (10, 1]) V(A,) + Vi2 (D)
V22 (D) VIL Vyr (00,(011) 1])
4p²
(1- (Tx 1)/2
Hence, by setting p (1 (Tx)2) we obtain (1 e, 8P-OP approximation of ALLOC*
with O(M (MMIn(2M/0)/) many samples.
Remark C.1. These tight accuracy calculations assume a worst-case analysis (given the definition of the value of an allocation stated in Definition 6), rather than only achieving the looser notion of a -PP approximation on expectation.
D FLEXIBLE BUDGET
Very dense distributions. Consider the following "2 spikes" case, where M/2 units have 7(u) value equal to 1/2 - 2e, and the other M/2 units have r(u) value equal to 1/2 + 2e. This example is connected to the usual lower-bound that appears in the sample complexity of bandit algorithms [ubeeck et al. (2033); Kaufmann et al. (2016). We show that for this spikes case, no budget is able to obtain -opttmall allocation.
24

---

Published as a conference paper at ICLR 2026
For any budget K I M// we have that
K 2
1 2 + 22)
= K 4
Vu- (10, 11)
+ ck.
As for the low-accuracy estimation algorithm, given that our estimates are computed to p 7VE accuracy, the value of random allocation, which we denote by Vrand, upper-bounds the worst-case value realized by Algorithm [
On expectation, we have that
K
1
K
K 2
Vrand
-
2c +
+ 2e =
2
2
2
2
The worst-case value is naturally given by K(1/2- - 2c) = K/2 - 2cK. By definition, an (1 c, OPT allocation f must realize value
(1-6) K.K ( + 2c)
(I,, 1)) 2 -e]Vir [0,
K
+ 2cK K. ( + 26
2
K
( 26²
3E
+ K.
2
with probability at least 8. Hence, for any budget K, the expected value of random allocation never yields a (1 - e, o)-optimal allocation, and is always short of at least K ((3e/2 2 + 2 much value. Thus, the worst-case of Algorithm fails to satisfy the (1 -pptimatyyy requirement as well.
Nonetheless, we are quite close achieving the value required by an optimal allocation. This presents an interesting duality for the problem of allocation: either the distribution of treatment effect values is p-regular, very dense around the threshold, in which case [lgorithm[[] (and even random allocation) still does quite well. order to realize the required value, we can either promise (1 - kee - allocation, or we can increase the budget in order to include more units. This type of overspending/resouree augmentation is different from the modification of the budget discussed at the beginning of this section (which cannot work in this 2 spikes case, as we just showed). Here, we add more units to our allocation ULEA but we compete against the original budget K.
which case Algorithm (anhhmm[eeee returns an optimal allocation, or the distribution is
(and
For the former, we want to find the smallest A such that random allocation realizes (1-ke)-OPT allocation. This has value:
(1 k<) Ve- ((0, 1)) (1- (g+22) K
K
4eK
KEK
2k 2K
2
2
Given that the expected value of Vrand (10, 11) is K/2, we need satisfy
4eK ke K
2 0 = k 2
Indeed, 4e)-OPT allocation requires realizing value at least
K 2 2r 2c 822
K 2
(1 - 4c)4 (1 46) (0)
K
8
and so Vrand ((0, 11) > (1- Vey- (10, 1)), as desired. Performing a similar calculation with the worst-case scenario of Algorithm 1 we get K 8.
Overspending, Alternatively, we can take the dual approach and ask what is the smallest budget K' K that allows a random allocation to get at least (1 - c)-OPT allocation, where optimality is measured with respect to the original budget K. This is form of resource augmentation. For any budget K' > K, we showed that Vrand (0, 1]) = K2/2. We want K' that satisfies
K"
2 (1 - Vie (0,, 1)) = (1 - c) K (2 + +22).
2
25

---

Published as a conference paper at ICLR 2026
This means that K' must satisfy
K' 2 K 4eK eK 4e =
K Ke=
3eK -4
Hence, we need K' 2 K + 3eK ( + 3c).
We can apply this overspending/resource augmentation idea to a distribution that is not p-regular around the initial budget threshold тк- That is, suppose that the initial budget K' such that the interval [TK- 2p, TK 2p] contains a lot of probability mass. What is the minimum number S of extra units that we need spend on such that we achieve a (1 -)-OPP allocation, where optimality is measured with respect to the original budget K' The extra units that we add to ULEA are naturally added in descending order is, the final threshold is TK+S- Because we assume that the interval [TK - 2p, TK + 2p] is highly dense, we lower bound the value +()) of all of the extra units that we add by T 2p. Then:
value starting at fx. and we denote their added value by Vugi That
V (10, 11) Vil (10, 11)
V(Ai) + ViGen + Vigu V(A1) + Vic
Vuu (D) Viela Vug
V(A1) + V2i- (D)
(TX + 2p) - (TX - -2p) - VuS
V(4)) + (Tx + 20
(Tx + Kp (TK -2p) - (TK 2
V(Ai) + (TK + 2o
4pKo (TK 2p))
V(A1) + (TK + 2p Ku
By setting
4pKo - (TK - 2 V(A2) + (TK + Ko Se
we get that the minimum number of units that we need to add to our initial budget K is
4pp c(TK + 2p) Ko - cV
S
TK 2p
Asymptotically, this means that S is of the order of VeKo. As we will see in Section 6 in practice we only require adding one extra unit in order to realize the value of the optimal allocation.
Choosing an appropriate budget: guidelines for policy-makers. In the cases where the policy-
maker has some flexibility over the budget, it is sensible approach to refine the budget after running such that Tk lies in an interval (of width 2p) of uniform- like probability mass is desirable for several reasons. First, it ensures that optimal allocation is achieved, given that p-regularity is satisfied in the Diinterval. Second, it ensures that the allocation is less arbitrary, it minimizes the number of units that are very close the threshold that determines whether or not treatment is given. As shown in Section 4.2 the low-accuracy estimates f give an approximation of the CDF curve as well, and so we can upper bound the probability mass contained
our algorithm. Choosing a budget
in an interval [a, b] by counting the number of units u that have f(u) value in the interval [a (a-p. p, b+p). If the interval is too dense, the policy-maker can slide the budget up or down in order find threshold fx that lies in a www-mass interval.
While the spikes case demonstrates that we cannot always hope to find such an interval, this only occurs in the cases where the CDF satisfies a strong "anti-Lipschitz" condition, where changing K for much larger or much smaller K' induces very little difference between тк and Tgy. However, is rare to find such densely-packee packed distributions in real-world RCTss, and it would be highly arbitrary to decide on an allocation in the case of these distributions. As we will see Section 3 in our experiments for real-world data, we can always find a K' very close to K (even if we only willing to underspend) in the few cases where Algorithm T does not yield an optimal allocation. Moreover, as we have show in this section, high mass around Tx by overspending and adding some extra units to ULEA. We also test this method in our experiments in Section 6
obtain optimal allocations for distributions with
26

---

Published as a conference paper at ICLR 2026
Summarizing, if a budget does not work, the policy-maker can: (1) detect so correctly by using the approximation F of the CDF (this is one-sided test), (2) slide the budget from K to K' as to find smooth interval, or (3) if no interval is p-regular, then the distribution Pr, is very dense in the entire range, in which case Algorithm| 1 should obtain a (1 kcP allocation for small value of ni or, equivalently, we can add a few extra units to EL to realize the optimal value.
We test these strategies experimentally and report the results in Section F We find that, in practice, we only need to slide the budget by at most 2 units in order to find a K that obtains a (1 e)- allocation with ((M/E) many samples. As for the overspending/rrsouree augmentation strategy, we find that in all cases adding one extra unit (namely, the unit corresponding to fK+1) suffices to realize (1 - c)-optimal allocation.
E
EXPERIMENTAL DETAILS
E. RCT DATA
We test our algorithm on real-world RCT data in order to test its efficacy and study the p-regularity smoothness condition in practice. For robustness, we test our algorithm on data from five different real-world RCTs, which have been conducted across multiple domains. These datasets have recently been analyzed in the work of Sharma & Wilder (2025), and tested for different policies, which is why we choose these five datasets. We summarize each RCT dataset, following the descriptions given in Sharma & Wilder (2025) and in the original papers.
(I) Tennessee's 's Student Teacher Achievement Ratio (STAR) project (Pate-Bain al. 1997). Domain: education. This is a 4-year RCT conducted by the Tennessee State Department of Education to examine class size effects on student performance. The study involved 11,601 students across 79 schools. Students and teachers were randomly assigned to class types beginning in kindergarten and followed through grade 3. The original study considered three possible treatments: small class (xx13-1 13-17 students), regular class 2#222- 25 students), or regular class with a full-time aide. Following the analysis performed in Sharma & Wilder (2025). in order to keep the binary treatment/control study, treatment corresponds to having been assigned to small class, whereas control corresponds to having been assigned to a regular class. The outcome is a cumulative test-score measure in kindergarten (which should be positively affected by the treatment), which is composite outcome from four standardized test scores: reading, math, listening, and word skills.
We filter out out observations with missing test scores, treatment assignments, or school IDs, leaving a total of 3,712 students. Out of these, 1,989 are control and 1,723 are treated.
economic development. This RCT studies the long-term effects of providing large one-time capital grants to low-income households. The treatment given was one-time capital grant to ultra-poor households, and the control was to give no grant. The study measured how family income and overall consumption evolved over a period of 7 years. The outcome variable that we use in our experiments is the total household expenditure (which should be positively affected by the treatment). This is the difference between the endline consumption and the baseline consumption.
(2) Targeting the Ulra Poor (TUP) in India (Banerjee et al. 2021). Domain:
After filtering the dataset, we have 864 households, with 410 control and 454 treated.
(3) National Supported Work (NSW) demonstration (Dehejia & Wahba 2002. 1999. LaLonde 1986). Domain: labor economics. This study was carried out to analyze the impact of the National Supported Work Demonstration, which was a job training program, on the income of the partici- pants in the year 1978. This job-training program targeting disadvantaged workers. The dataset also collected the participants' income in the year 1975 as a bascline. The treatment was assignment to the NSW training program, and the control was assignment. The outcome variable that use is the annual carnings in 1978 (which should be positively affected by the treatment).
After filtering the dataset, have 722 individuals, out of which 425 are control and 297 are treat- ment.
(4) Acupuncture dataser (Vickers al. 2004). Domain: healthcare. This RCT evaluates the effect of acupuncture therapy on patients with chronic headache, with assessments at randomization, months, and 1 year. The treatment is accupuncture for chronic headaches, and control is the usual
27

---

Published as a conference paper at ICLR 2026
care. The outcome variable that use is the headache-sevee severity one year after randomization. Headache severity is measured using a discrete 0-5 Likert scale. Treatment is thus expected to lower the outcome, which is why we flip the sign in our figures for coherence with the other RCTs.
After filtering the dataset, we have 301 patients, with 140 control and 161 treatment.
(5) Postoperative Pain dataset (McHardy & Chung 1999). Domain: Healthcare. This RCT inves- tigates whether gargling a licorice solution prior to endotracheal intubation reduces postoperative sore throat, common side effect of thoracic surgery using double-lumen tubes. Hence, treatment corresponds to receiving licorice gargle prior to endotracheal intubation, and control corresponds placebo/standard care. The outcome variable measured on on a discrete 0-7 Likert scale. As in the case of the acupuncture dataset, treatment expected to lower the outcome, and so we also flip the sign in our figures for coherence with the other RCTs.
is throat-pain score 4 hours post-surgery,
After filtering the dataset, we have 233 patients, with 116 control and 117 treatment.
E.2 SEMI-SYNTHETIC EXPERIMENTS
Grouping methods. For each of the five RCT datasets, create set of units U by partitioning the individuals into groups from their covariates. For each dataset, we describe the various grouping methods that test in algorithm. For any grouping strategy, individuals and control individuals, and that the treatment rate is between 15% and 85%, to ensure balance within each unit. We only include the grouping methods that succeeded in creating a set of units satisfying these requirements, and discarded several other grouping methods based on the dataset variables.
require at least treated
(1) STAR dataset. We use the following grouping methods:
School groups. We cluster the students by school ID.
Performance groups. We group the students based baseline performance percentiles (test score averages).
Causal forest groups. We use random forest regressors to predict treatment effects (one for control and one for treatment), and then cluster the students based on predicted CATE values and covariates using K-means clustering. We cluster aiming for different number of groups (typically 30 and 50). We note that for short, but we are referring to CATE-based grouping obtained via random-fore and clustering; not through causal forest estimator.
the name "causal forest" in the figures
Propensity score groups. We perform a stratification based propensity scores using vrlss-validated logistic regression, dividing the propensity score distribution into equal- sized quantile- based strata (thus grouping students with similar propensity score).
(2) TUP dataset. We use the following grouping methods:
Baseline poverty groups. We group the households using baseline consumption as a poverty proxy, clustering them using K-means.
Demographic groups. We search for various demographic variables (such as "gender", "education", "age") and take their intersections.
Assets. We identify asxet-related variables (such "land" or "livestock") and group them using K -mmanss
Causal forest groups. We use random forest regressors to predict treatment effects, and then cluster the households based on the predicted CATE values and covariates using K-means.
Propensity score groups. We perform a stratification based on propensity scores using validaaal logistic regression.
(3) NSW dataset. We use the following grouping methods:
Baseline earnings groups. We divide the employed individuals into percentile- based carn- ings brackets.
28

---

Published as a conference paper at ICLR 2026
Age groups. We create percentile- based stratification of the age distribution.
Causal forest groups. We use random forest regressors to predict treatment effects, and then cluster the individuals based on the predicted CATE values and covariates using K-means.
Propensity score groups. We perform a stratification based on propensity scores using cross-validate logistic regression.
(4) Acupuncture dataset. We use the following grouping methods.
Age groups. We create percentile-base age brackets.
Age- Cee-Chronicity interaction groups. We compute a score for each patient combining their age and their chronicity variables, and then group them into equally-sized groups (in order of the scores).
Baseline headache. We group the patients by their initial headache score.
Chromicity groups. We group the patients by the length of chronic headache history.
Covariate forest groups. We perform a K K-manss clustering based on the covariate profiles.
Muliidimensional composite groups. We use the variables age, chronicity, baseline headache score, and combine them into an average score.
(5) Postoperative dataset. We use the following grouping methods.
BMI groups. We divide the patients into 30 equal-sized BMI brackets.
Age groups. We perform percentile-based stratification on the age variable.
Demographics. We combine different preoperative patient characteristics.
. Covariate forest. We perform K -mmans clustering on the patient characteristics.
Note that for the acupuncture and postoperative dataset we have less data than for the other datasets, which is why the group sizes that we obtain smaller and thus prone higher error.
Obtaining the CATE values. Each grouping method for each dataset yields partition of the dataset into a set 1i of M units. For each unit, we compute r(u) as follows: we compute the treated outcome mean (i.e., the mean of the outcome variable among the treated members of group u) and the control outcome mean (i.e., the mean of the outcome variable among the control members of group u). Subtracting the control outcome mean from the treatment outcome mean yields the treatment effect value for each unit. Lastly, we normalize all of the treatment effect values across the M units into [0, 1], yielding the final r(u) values.
Once we have the r(u) values, which allow us to compute the value of the optimal allocation, we simulate the sampling as follows. For every drawn sample of the population, we select one of the units u uniformly at random and draw a sample from the distribution Bern(r(u)). That is, each group represents an "arm" with Bernoulli reward equal its normalized CATE, In all experiments, we set the failure probability ö to 0.05.
E.3 FIRST APPROACH: GENERATING FIGURE
For this approach, we do not select specific value of c. Instead, we try different sample sizes, from N 100 N 20,000 across different budget constraints. Specifically, we try budgets K that are 10%, 20%, 30,, 50%, 70%, and 90% of M. For each sample size and each budget, we sample from the population for that many samples and compute the estimates + +(u) for each of the M units Then, select the K units with the highest estimates. Using the ground-truth values r(u), we compute the value of this allocation. We then plot the realized value of the allocation (y-axis) for each sample size (æ-axis), having plot for each fixed budget and grouping method. We repeat this sampling process 50 times, obtaining an average value the allocation for each sample size and confidence bounds.
In each figure, we also plot the value that we would expect from the worst-case bound O(M/ e²) and our bound of O(M/e). Specifically, from Lemma 3 we have that the expected value of the
29

---

Published as a conference paper at ICLR 2026
allocation using the worst-case bound is
(1- (1-V Min(2M/8)
Vu-..
From Theorem 7 we obtain that the expected value of the allocation using theoretic bound for p-regular distributions is
M In(2M/0)
Viu-
N
In Section F we display these plots for each of the grouping methods for each of the datasets. For each, we plot the relationship between the normalized allocation value and the sample size for four choices of a budget.
E.4 SECOND APPROACII: GENERATING FIGURE
For Figure 2 and its further displays Section F (the last two plots for each subsection), we carry out our analysis differently.
Here, we choose value of e, which ranges from 0.001 to 0.2. Based on our computation of the value of in Section C for various distributions, choose to call Algorithm | with y 0.5. This bound is best decided using an informed guess how the T values are distributed; if is too high then Algorithm will have higher error rate. For each value of number of samples in total that we use to construct our estimates f using our theoretical bound of N = M In (2M /c) Note that each grouping method yields a fixed number of units M. Sampling that many units, we obtain the coarse estimates i f(u). and for each possible budget K E (1,..., ......,/ M) compute the value realized by our allocation (where compute the value using the ground-truth estimates but they are selected using the low-accuracy estimates f). This gives the value VeLEa for each budget K.
compute the corresponding
Separately, for each possible budget K E {L...., }, we compute the value V/4- of the optimal allocation. Then, we check whether or not Vidiaa// VVe- 2 c, for each value of K and for the fixed value of c. We compute the rate of failed budgets (i.e., the proportion out of M) that do not achieve a - c-optmmal allocation. We repeat this process 50 times for each combination of K and c, giving us a robust average failure rate with confidence bounds. We do this for every grouping method and for every RCT dataset. The first row in Figure some of these failure rates. The second-to-last plot in each subsection of Section F shows these plots for all of the grouping strategies for each RCT dataset. In all cases, we see that the average failure rate is below 5%. This is higher for the acupuneture and postoperative datasets, since we work with fewer groups in all grouping strategies (because the dataset smaller). Moreover, we remark that in our experiments we test all budgets, and we do discard budgets that are very small (e.g.. K 1). As per the Vi2- term Theorem very small values of K increase the failure rate. The looser choice of > also increases the failure rate. For all these reasons, the failure rate that we obtain is higher than what it would be if we discarded small budgets and decreased r further, or if we had a higher number M of units.
Second, we test our flexible budget strategies discussed Section 5 In the cases where budget K fails to obtain (1 - -)-optimal allocation (recall that we are always using the fewer ((M/e /e) many samples to compute the allocation), we find the closest value of K' such that we obtain a (1 c)- optimal allocation, and then we report K K'". We also compute such a KK - K'| for K" < K; i.e., we are only willing to underspend. We repeat this process 50 times for each e, grouping method, and RCT. Some of the results are reported in the second row of Figure 2 the comprehensive set of plots is in Section F The last plot each subsection of Section F shows these plots for all of the grouping strategies for each RCT dataset. In all cases, including the subcases where only underspend, we see that, on average, K K'I < 2.
We also test overspending idea of adding extra units to allocation. We test it each of the failed budget settings. In all cases, without exception, we find that adding one extra unit (i.e., the unit that corresponds fx41). 1) realizes optimal allocation.
The code for this paper can be found at: httpp:: htt.://.....iiiiicaaac...
a olloc-o vS cate
30

---

Published as a conference paper at ICLR 2026
Use of LLMs. We have used LLMs to aid in exploring the five RCTs and their various features, in order to find sensible grouping strategies for each RCT. After implementing the code for the base simulation of our algorithm, used LLMs help adapt the algorithm to the specifies of each dataset. Lastly, we have used LLMs to double-check the integrals computed for the Beta and Gaussian distributions.
31

---

Published as a conference paper at ICLR 2026
F FURTHER EMPIRICAL RESULTS
F.I STAR DATASET
School Groups (M=78)
Normalized alocaton value
Budget = 20% (K=15)
Budget = 30% (K=23)
Budget = 50% (K=39)
Budget - 70% (K=54)
8.a
Normalizedd ellocation
Normal eizeelieedd eloocaaion
1 E
0.7-
0.3
(1.0)
Normalized
(1.0)
(8-0.5)
(8m0.5)
(8m0.5)
(8=0.5)
ALLOC (0=1.0)
ALLOC (0=10))
ALLOC 1gm1
ALLOC (fm1.0)
Empirical data
Empriiall data
Emprical date
Empirical dita
23000
10020
23850
10000
1,000
33000
Sample
Sample
Sample
Figure STAR dataset, school groups.
Groups (M=50)
Budget 20% (K=10)
Budget 30% (K=15)
Budget 50% (K=25)
Budget 70% (K=35)
Normaliced allocaton value
cormalized allocation value
0.9-
8.a
Normalized allocaion
allocation
(1.0)
Optimai (1.0)
Optimai (1.0)
Optimal (1.0)
RulICATE (p-0.5) 5)
FulICATE ($-0.5) 5)
FUlICATE ($-0.5)
FULICATE ($=0.5)
ALLOE (f-1.0)
ALLOC (f=10)
ALLOC ipmi
ALLOC (p=1.0)
Empirical data
data
data
Empiriaal dltta
23000
1992i Sample
ioloo
23600
1990 isien 2000
Tikan iidce 20000
Sample
Sample
Sample
Figure 4: STAR dataset, performance groups.
(30) (M=30)
Normaiized allocation value
30%
Budget
Budget
Normalized allocation
allocation
Mormallzed allocation
Optimal (1.0)
Optimal (1-0)
Optimal (1.0)
Optimal (1.0)
RulICATE ($-0.5)
RulICATE (B-0.5)
FUlICATE (B=0.5)
RULICATE (B-0.5)
ALLOC (f=1.0)
ALLOC
ALLOC (p=1.0)
ALLOC (p=1.0)
Empirical data
Erpiitail data
Empiicaal data
Empirical
1002
1080e
33000
5202
Labaa
33000
1002
10000
lidco 20030
100
1ib0o 19d6e 20000
Sample size
Sample size
Sample size
Sample size
Figure 5: STAR dataset, causal forest 30.
Causal Forest (50) (M=47)
Budget (K=9)
Budget
(K=14)
Budget (K=23)
Budget (K=32)
Norma illocaation vea
Normaalzed bllocatio
Norma Illocation
Illocation
5 C
5. i C
Ca
Optimal (1.0)
Optimal (1.0)
Optimal (1.0)
Optimal (1.0)
FUlICATE
FUlICATE ($-0.5)
FUlICATE ($=0.5)
FULICATE ($=0.5)
ALLOC (f=1.0)
ALLOC (5=1.)
ALLOC 1p=1.00
ALLOC (p=1.))
Empirical data
Emprical data
Emprical data
Empiriaal ditta
1032
Jidee Fibo
30000
Laiaa
33800
Lanan
iideo
Sample size
Sample size
Sample size
Sample sice
Figure 6: STAR dataset, causal forest 50.
32

---

Published as a conference paper at ICLR 2026
Propensity Score (M=50)
Norralized ailocation value
Budget = 20% (K=10)
Budget - 30% (K=15)
Budget = 50% (K=25)
Budget - 70% (K=35)
Normalized allocation
Normallzed allocation
Normallzedd ellocation
H E
6.2-
0.3
(1.0)
(1.0)
(1.0)
(8-0.5)
(8m0.5)
(3m0.5)
(8m0.5)
ALLOC (f-1.0)
ALLOC (f-
ALLOC (f=1.0)
ALLOC (p=1.0)
Empirical data
Empriicll dita
Emprical diaa
Empirical diat
Fillo
Sample
Sample
Sample
Sample
Figure 7: STAR dataset, propensity score.
School: Fallure rate vs.
Performance: Failure rate
Causal Forest 30: Fallure rate
Failure rate
1.5
Failure rate
Failure rate
1.0
0.5
0.90
10
10-1
9. 3
7
10
90
10
10-1
(a) STAR, schools.
(b) STAR, performance.
(c) STAR, causal forest 30.
Causal Forest 50: Failure rate
Propensity Score: Fallure rate vs.
Failure rate
Failure rate
1.5
1.0
2
0.5
90
10
10
0.0
10
10-4
10
(d) STAR, causal forest 50.
(e) STAR, propensity score.
School: Closest K'
Performance: Closest K' g
Causal Forest 30: Closest K' E
1.54
Minimum K
1.0
K"
Minimum
Minimum
0.5
0.08
10
10
-0-3
10
10
0.0
10
10
(a) STAR, schools.
(b) STAR, performance.
(c) STAR, causal forest 30.
Causal Forest 50: Closest K' vs
Propensity Score: Closest K' vs.
Any K
1.5
Any
K'
Minimum
Minimum
1.0
1
0.5
-
-0
10
0.0
20
10
10-1
(d) STAR, causal forest 50.
(c) STAR, propensity score.
33

---

Published as a conference paper at ICLR 2026
F.2
TUP DATASET
Baseline Poverty (30) (M-30)
Budget 20% (K-6)
Budget 30% (K-9)
Budget 50% (K=15)
Budget 70% (K=21)
Normalizod allocatinn
Normllzzd aloocation
Normal allocation
Normal alfocation
Optimal 11.0)
Optimal 11.0)
Optimal 11.0)
Optimal 11.0)
ALLOC
ALLOC
C 2
ALLOC
ALLOC
Empinical data
Empinical data
Empinccal data
Empincal data
sbt 1280e 1o00o
23600
sib
1sd0o
50b1
1sdco
1ob00 11deo
Sample size
Sample size
Sample size
Sample size
Figure 10: TUP dataset, baseline poverty 30.
Poverty (50) (M=50)
Budget
Budget
Budget
Budget
Normaizzed allocation
Normalized allocation
Normalll Normalized allocationn
to
Optimal (1.0)
Normalized
Optimal 11.0)
Optimal 11.0)
Optimal 11.0)
PUICATE
RUICATE
RJICATE
RUICCTTE
Empirical data
Empirical data
Empirical data
Empiriaal
3902
1300e
13000
23000
5902
12000
13000
23000
5002
12000
33000
20000
5000
10000
33060
20000
Sample size
Sample size
Sample size
Sample size
Figure 11: TUP dataset, baseline poverty 50.
Asset (M=23)
Budget 20% (K-4)
Budget 30% (K-6)
Budget 50% (K-11)
Budgat 70% (K-16)
Boomllieed bilocation value
Normalicee alicaeeeed ellocation
ellocation
elfocation
ALLOC
ALLOC
ALLOC
ALLOC
Empirical
Empirical
Empirical data
Empirical data
62
15600
23600
12000
15600
23800
12000
20000
Sample
Sample
Sample
Sample
Figure 12: TUP dataset, asset groups
Asset Groups (50) (M=32)
Budget (K=6)
Budget 30% (K=9)
Budget (K=16)
Budget (K=22)
2.0-
Nortraized bllocation
Bllocation
allocation
5. C
0.
Optimal (10)
Optimal (1.0)
5
Optimal (1.0)
Optimal (1.0)
RJICCATE
2
Norma
2
Empirical data
Empirical data
Empirical data
Empirical data
ribon 30000
loloo 30800
lideo
1400
Jaban lidce
Sample size
Sample size
Sample size
Sample sice
Figure 13: TUP dataset, asset groups 50.
34

---

Published as a conference paper at ICLR 2026
Causal Forest (30) (M=30)
Budget 20% (K=6)
Budget 30% (K=9)
Budget 50% (K=15)
Budget - 70% (K=21)
Norralized allocation
Normalized allocation
Normallzes ellocation B
Normallzzdd allocation
(1.0)
(1.0)
(8m0.5)
(8m0.5)
(8m0.5)
(8m0.5)
ALLOC (f-1.0)
ALLOC (f-10) 1.0
ALLOC Ig=
ALLOC (p=1.0)
Empirical data
Empriiall dida
Empricall diat
Empirical diat
riloo
Sample
Sample
Sample
Sample
Figure 14: TUP dataset, causal forest 30.
Propensity Score (30) (M-29)
Budget 20% (K=S)
Budget 30% (K=8)
Budget 50% (K-14)
Budget - 70% (K=20)
Normalzzo ailocation
Normall allocation
Normell allecation C
Normall allocation
(1.0)
Optimal
Optimal
Optimal
ALLOC
ALLOC
ALLOC
ALLOC
Empirical data
Empinical data
Empiniaal data
Empincal data
12l0o
inloo
inkeo
Tako Sample
Sample size
Sample
Sample
Figure 15: TUP dataset, propensity score.
Baseline Poverty: Failure rate vs
Assets: Failure vs.
Demographics: Failure rate vs.
0.20
Failure rate
Failure rate
Failure rate
4
0.15
0.10
0.05
-0-1
10-2
10
40-5
10-2
0.09
10-2
10
(a) TUP, haseline poverty.
(b) TUP, assets.
(c) TUP, demographics.
Causal Forest 30: Failure rate
Propensity Score: Failure rate vs.
4
Failure rate
Failure rate
-0
-0-
10-1
9
10-2
10
(d) TUP, causal forest 30.
(c) TUP, propensity score.
35

---

Published as a conference paper at ICLR 2026
Baseline Poverty: Closest K' vs.
Assets: Closest K' vs. E
Demographics: Closest K' vs. E
1.5
0.20
Minimum K
Any K
Minimum K
Any
Any
K Ks
Minimum
1.0
1
0.10
0.5
0.05
9
10
10
0.0
10
10
0.00
10
10
10
E
(a) TUP, baseline poverty.
(b) TUP, assets.
(c) TUP, demographics.
Causal Forest 30: Closest
Propensity Score: Closest K' vs.
Minimum K
K
Any
1.5
Minimum
1.0
0.5
0.90
10
10
0.0
10
E
(d) TUP, causal forest 30.
(c) TUP, propensity score.
F.3 NSW DATASET
Baseline Earnings (M=30)
Budget= 20% (K=6)
Budget 30% (K=9)
Budget 50% (K=15)
Budget = 70% (K=21)
Norralized alocssion value
C.a
Normallized alcccation
8.5-
12
Normalized
ALLOC
ALLOC:
ALLOC
ALLOC
0.3
Empiriaa
Empincal
Empincaal data
C
Empirical data
23000
13030
23000
1,0000
1,000
Sample
Figure 18: NSW dataset, baseline earnings.
Age Groups (M=18)
Budget (K=3)
Budget 30% (K=5)
Budget (K=9)
Budget = (K=12)
Norma allocation vea
Normalized bllocation
Norma Illocatton
allocaaion
2 5. i 1
0.a
Optimal (10)
Optimal (1.0)
5
Optimal (1.0)
Optimal (1.0)
PJICATE
RJICCTE
RJICCATE
2
Empirical data
Empirical data
Empirical data
Empirical data
ribon 30000
loloo 30800
lideo
1400 Jaban lidee
Sample size
Sample size
Sample size
Sample size
Figure 19: NSW dataset, age groups
36

---

Published as a conference paper at ICLR 2026
Causal Forest (M-22)
Norralized ailocation value
Budget 20% (K=4)
Budget
30% (K=6)
Budget = 50% (K=11)
Budget - 70% (K=15)
Normalized allocation
Normallia allocation
Normall ellocation
6-2-
0.3
(1.0)
Optimal
ALLOC
ALLOC
ALLOC
ALLOC
Empincal data
Empincal data
Empincall data
Empiriall data
Tiloo
Sample
Sample
Sample
Sample
Figure 20: NSW dataset, causal forest.
Propansity Score (M=43)
=udget= 20% (K=B)
Budget 30% (K=12)
Budget 50% (K=21)
Budget = 70% (K=30)
Rorallced Dilocation value
Noraalized allocation
Normalized ellocation
Normalizedd ellocation
PUICATE
PUICATE
PUICATE
PUICATE
ALLOC
ALLOC
ALLOC
ALLOC
Empirical
Empirical
Empirical data
Empirical data
33600
13000
15000
33636
18000
15600
20000
19000
35600 2000
Sample
Sample
Sample
Sample
Figure 21: NSW dataset, propensity score.
F.4 ACUPUNCTURE DATASET
Age Groups (M=24)
Budget (K=4)
Budget
value
Budget
Budget
Normalized alocation
12 t
Optimal (1.0)
Jormalized
Optimal 11.0)
Optimal 11.0)
Optimal (1.0)
RUICATE
RJICATE
RJICATE
RJICATE
Normal
+
Empirical data
+ Empirical data
Empirical data
Empirical data
2362e
15600
23600
5301
12000
25000 23800
5001
12000
isdco
20000
5000
18600
iséco 20000
Sample size
Sample size
Sample size
Sample size
Figure 23: Acupuncture dataset, age groups.
Age-Chronicity Interaction (M-26)
Budget (K=5)
Budget 30% (K=7)
Budget (K=13)
Budget (K=18)
Norma allocation ve
Norma bllocation
Norma Illoccatio
22 5. 1
Optimal (1.)
Optimal (1.0)
5
Optimal (1.0)
Optimal (1.0)
RUICCATE
Empirical data
Empirical datt
Empirical data
Empirical data
30000
Foloo 30800
lideo
1400
lidee
Sample size
Sample size
Sample size
Sample sice
Figure 24: Acupuneture dataset, age-chronicity interaction.
37

---

Published as a conference paper at ICLR 2026
Baseline Eamings: Failure rate
Age Groups: Fallure rate vs,
Causal Forest 30: Failure rate vs.
10
Fallure rate
Failure rate
Fallure rate
6
4
90
10
10-1
---
10
10
10-3
10
10
(a) NSW. baseline earnings.
(b) NSW, age groups.
(c) NSW, causal forest 30.
Causal Forest 50: Fallure rate vs.
Propensity Score: Fallure rate
Fallure rate
Failure rate
-
10
10
90
10"
10
(d) NSW, causal forest 50.
(e) NSW, propensity score.
Baseline Earnings: Closest K"
Age Groups: Closest K'
Causal Forest 30: Closest K'
1.5
Minimum K"
1.0
Any K
WK
Minimum K"
2
Any
Minimum K
0.5
0.5
0.0
10
%
10
'10
0.0
10
(f) NSW, baseline carnings.
(g) NSW, age groups.
(h) NSW, causal forest 30.
Causal Forest 50: Closest vs.
Propensity Score: Closest
& 1.5
Any
1.5
Aey K
Minimum
Minimum
1.0
0.5
0.5
0.0
10
10
0.0
10
10
(i) NSW, causal forest 50.
(i) NSW, propensity score.
38

---

Published as a conference paper at ICLR 2026
Baseline Headache (M=25)
Budget 20% (K=5)
Budget 30% (K=7)
Budget 50% (K=12)
Budget - 70% (K=17)
Norralized bilocation value
Normall allocation value
6.4-
Normalized allocation
Normallzed ellocation
E
(1.0)
Optimll
ALLOC
ALLOC
ALLOC
ALLOC
Empincal data
Empincal data
Empincall data
Empiriall data
Tiloo
Sample
Sample
Sample
Sample
Figure Acupuncture dataset, baseline headache.
Chronicity Groups (M=19)
Budget (K-3)
Budget (K=5)
Budget (K=9)
Budget
2.0.
Hormaized allocation
Hormalized allocation
Normalized allocatioon
6.--
Optimal (1.0)
Optimal 11.0)
Normalized lized
Optimal 11.0)
Optimal (1.0)
FuICATE
FUIGATE
RUTCATE
PUTICATE
Empirical data
Empinical dat
Empinical data
Empirical data
1032
1i66e rolon
23000
13ian
23602
19han
vidco
1idce
2000
Sample size
Sample size
Sample size
Sample size
Figure 26: Acupuncture dataset, chronicity groups.
(M=15)
Budget 20%
30%
50%
(K=10)
looaliieed allocatio value
Normalized aloocatio
8.a.
Optimal (1.0)
Optimal 11.00
Optimal (1.0)
cormaliedd
Optimal (1.0)
FUICATE
RJICATE
RJICCTE
RJICCATE
ALLOC
ALLOC
ALLOC
ALLOC
Empiniaa data
Empiniaa data
C
Empinical
Empirical data
5008
23000
13600
23000
5002 10000
33000
23000
3002
10000
33000
20000
3000
18000
33000 20000
Samplo size
Sample
Sample
Sample
Figure 27: Acupuncture dataset, covariate forest.
Multidimensional Composite (M=29)
Budget = 20% (K=5)
Budget 30% (K=8)
Budget 50% (K=14)
Budget = 70% (K=20)
Normaliced bilocation value
ellocation 12
Normallized ellocation
0.5:
12
Normalized
Optimal
FUICATE
FUICATE
ALLOC
ALLOC
ALLOC
ALLOC
Empincal date
Empincal date
Empincall data
Empirical data
Sample
Sample
Sample
Figure 28: Acupuncture dataset, multidimensional.
39

---

Published as a conference paper at ICLR 2026
Age-Chronicity Interaction: Failure rate vs. E
Age Groups: Failure rate
Baseline Headache: Failure rate
Failure rate
Fallure rate
7.5
4
Fallure
5.0
2.5
10
10
90
10
10
E
0.0
10
(b) Acup., age chronicity interac-
Acup., age groups.
tion.
(c) Acup., baseline hcadache.
Chronicity Groups: Fallure rate vs. Covariate Forest 50: Fallure rate
Fallure rate
6
Failure rate
4-
2-
98
10
46
10
10
(d) Acup., chronicity groups.
(c) Acup., covariate forest 50.
Age-Chronicity Interaction: Closest vs.
Age Groups: Closest vs. e
2
Any
Baseline Headache: Closest vs.
Minimum K"
1.5
Any K
Minimum
Any
Minimum&
1
%
10
10
0.0
10
10
10
10
10
(b) Acup., age chronicity interac-
(a) Acup., age groups
tion.
(c) Acup., baseline headache.
Multidimensional Composite: Closest vs. E
1.5
Chronicity Groups: Closest K' vs.
Covariate Forest 50: Closest K' vs.
1.0
Any
Minimum K"
1.0
Ank
Minimum
sk
KE
1.0
Minimum
0.5
0.5
0.00-
0.0
10
10
0.0
10
10
0.0
10-4
10
8
(f)
Acup.,
multidimensional
(d) Acup., chronicity groups
(e) Acup., covariate forest 50.
composite.
40

---

Published as a conference paper at ICLR 2026
F.5
POSTOPERATIVE DATASET
BMI Groups (M=19)
Budget (K=3)
Budget (K=5)
Budget (K=9)
Budget (K-13)
2.3
2.0
Normalzzd allocation
Normallzed alloccation
0.3
FuIGATE
Optimal (1.0)
Normaltzed
PuICATE
Optimal 11.0)
Normallzed
Optimal 11.0)
PuTCATE
PUTCATE
Optimal 11.0)
0.4
Empirical data
Empiniaal data
Empincal data
Empirical data
3952
1360e 13000
23000
3202
12000
33000
23000
3002
12000
33d0o
20000
5000
10000
33d00
20000
Sample size
Sample size
Sample $120
Sample size
Figure 31: Postoperative dataset, BMI.
Groups
Budget 20% (K=2)
Budget 30% (K=4)
Budget 50% (K=7)
Budget (K%)KK))
lormaiized ailocation
0.a
2.
12 12
PUICAAT
(1.0)
Normalized
RUICCTT
11.0)
Normalized
C 1
RUICCATE
11.0)
Normallcen
RLIICATE
11.0)
ALLOC
ALLOC
ALLOC
ALLOC
Empincal data
Empincal
Empincal
C
Empirical
5002
33000
13000
23000
5002
13900
33000
23000
3002
1,900
33000
20000
3000
18000
33000
20000
Salmplo
Sample
Sample
Sample
Figure 32: Postoperative dataset, age groups.
Demographics (M-13)
Budget= 20% (K=2)
Budget 30% (K=3)
Budget 50% (K=6)
Budget 70% (K=9)
NomT allocation
0.2-
2.-
allocation 2 B
allocation
6.2-
2.3
2.
ALLOC Empiriaal data
ALLOC
ALLOC
ALLOC
Empincal data
Empiniall data
Empincall data
Sa Sample
Sample
Sample
Sample
Figure 33: Postoperative dataset, demographics.
Covariate Forest (M=17)
Budget (K=3)
Budget (K=5)
Budget (K-8)
Budget
2.0.
Horralized allocation
leeeellllddddeeeeeeee ellocation
0.5-
FUICATE
Optimal (1.0)
Normalized
RUIGATE
Optimal 11.0)
RUICATE
Optimal 11.0)
Optimal (1.0)
PuICCATE
Empirical data
Empirical data
Empinical data
Empirical data
5052
2362e 12d0o
23000
5202
12000
2sd0o 23800
5002
12000
asdeo
20000
5000
12600
asdeo 20000
Sample size
Sample size
Sample size
Sample size
Figure 34: Postoperative dataset, covariate forest.
41

---

Published as a conference paper at ICLR 2026
BMI Groups: Failure rate
Age Groups: Failure rate
Fallure rate
Fallure rate
2
2
90-
10
10
10-3
10-2
10-1
(a) Post-op, BMI.
(b) Post-op, age groups.
Demographics: Failure rate
Covariate Forest 30: Failure rate
6
Fallure rate
Failure rate
90-
10
10
%
10
10
(c) Post-op, demographics.
(d) Post-op, covariate forest 30,
BMI Groups: Closest K'
Age Groups: Closest
3
Minimum K'
Any
1.0
2-
Minimum
0.5
90-1
10-1
0.00-
10
10
(e) Post-op, BMI
Post-op, age groups
Demographics: Closest K'
Covarlate Forest 30: Closest K' vs.
Minimum K"
1.0
Any
Minimum K
1.5
Any
1.0
0.5
0.5
0.0
10
10
0.0
0
10
10
(g) Post-op, demographics.
(h) Post-op, covariate forest. 30.
42


Reviews:
[{"review_id": "BxmygB9P3t", "reviewer": "Reviewer_92Sv", "summary": "This paper fills the gap between estimating CATE and making decisions on the allocations. The authors show that while estimating all CATEs within $\\epsilon$ accuracy requires $O(M/\\epsilon^2)$ samples, achieving a near-optimal $(1-\\epsilon)$ treatment allocation typically needs only $O(M/\\epsilon)$ samples under mild distributional assumptions. In general, I personally see the results quite interesting and insightful.", "questions": "See above.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "1. The paper makes a clear and elegant theoretical distinction between estimation and allocation. The reduction of the sample complexity from $M/\\epsilon^2$ to $M/\\epsilon$ is insightful and exciting.\n2. Practical relevance: Direct implications for RCT and policy design: significant reduction in sample cost.\n3. The proofs are clean and well-structured and the theoretical results are rigor.", "weaknesses": "In general, I enjoy reading the paper a lot. I do not have major concerns.\n\n1. Comparison to bandit best arm identification could be expanded. The link is conceptually strong. Particularly, recently, there are some works on good arm identification. Some ideas are very similar, although they are not is a causal inference setting.\n\n2. Policy implication is strong (“RCTs underpowered for CATE estimation can still yield good allocations”), but guidance on how to detect $\\rho$-regularity or compute sample sizes in practice is missing.", "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": "an9Pw9Kltr", "reviewer": "Reviewer_mQ5Z", "summary": "The authors provide new theory showing that optimal treatment allocation can be achieved with fewer samples than classic predict-then-optimize approaches (FullCATE in the paper) would require. Their theory is built upon the insight that accurate estimates of CATE are only needed around the treatment allocation threshold.", "questions": "All the items listed in the Weaknesses section may be interpreted as questions by the authors.\n\nTypo: line 353 \"notion that requiring\"", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The authors provide interesting insights about sample size requirements for optimal treatment allocation.\n- The authors substantiate their claims with extensive theoretical analysis", "weaknesses": "Some other related work exists that uses similar insights about the problem of optimally allocating treatment, though often without an extensive theoretical analysis. For example, some work has argued that when trying to find the optimal treatment allocation, accurate CATE estimation is not always the most effective [1,2]. \n\nWhile the authors provide a very extensive theoretical analysis, they only briefly explain the potential impact of their contributions. For example, I find it hard to understand what practitioners should do with this new information. It would be helpful if the authors could discuss this.\n\nI wonder why the authors decided to put the individuals into different groups and do the analysis based on these groups. As the number of groups decreases, the number of samples needed also decreases, but the quality of the overall allocation will also go down (because as you make the groups more granular, you will find more heterogeneity between groups, allowing for better decision-making). How should you choose the number of groups in practice if you have very little a priori information about the CATE distribution?\n\nRelated to the previous point, I do not understand how the different groups are created in the experiments. To me, this seems like the most crucial part when evaluating treatment allocation quality.\n\nI wonder why the authors did not use datasets that are often used in Uplift Modeling [1] and treatment effect estimation [3]. \n\n[1] Devriendt, F., Van Belle, J., Guns, T., & Verbeke, W. (2020). Learning to rank for uplift modeling. IEEE Transactions on Knowledge and Data Engineering, 34(10), 4888-4904.\n\n[2] Fernández-Loría, C., & Provost, F. (2022). Causal decision making and causal effect estimation are not the same… and why it matters. INFORMS Journal on Data Science, 1(1), 4-16.\n\n[3] Curth, A., Svensson, D., Weatherall, J., & Van Der Schaar, M. (2021, August). Really doing great at estimating cate? a critical look at ml benchmarking practices in treatment effect estimation. In Thirty-fifth conference on neural information processing systems datasets and benchmarks track (round 2).", "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": "p2bVk8xiFq", "reviewer": "Reviewer_zW6P", "summary": "The authors investigate sample complexity bounds for estimating CATE across M different groups when trying to perform allocation to a K subset of them. Standard CATE analysis demonstrates a $\\frac{M}{\\epsilon^2}$ bound, as each group needs to learn a separate CATE, then the K-highest CATE values are selected. However, the authors argue for a $\\frac{M}{\\epsilon}$ bound by performing coarser estimation for treatments near the boundary. Essentially, the authors demonstrate that if CATE estimates can be learned within $\\sqrt{\\epsilon}$, then any estimation mistakes will not be very costly. Such learning is possible under certain assumptions on $\\tau$; for example, that $\\tau$ is smooth or is near-uniform. The authors conclude with experiments on real-world RCTs to demonstrate the efficacy of their experiments.", "questions": "1. In practice, on the datasets listed in the experiments section, how does the sample complexity of the CATE-style selector ($\\frac{1}{\\epsilon^2}$ compare with the selector proposed", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. **Work tackles an important problem** - The authors tackle the important problem of $K$ selection from $M$ groups in a causal setting. Such a problem can be seen across a variety of real-word situations, and is especially prevalent in the world of policy. This can help more efficiently allocate resources and avoid unnecessary experiments. \n2. **Analysis is intuitive and clean** - The authors present a clean and intuitive reason why their proposed selector should outperform baselines. By avoiding the need to recompute CATE for each of the $M$ groups, the authors are able to achieve a better sample complexity, due to their ability to adaptively explore in some sense. \n3. **Characterizes when their new bound is possible** - The authors describe several scenarios where there newly proposed estimators reach the desired $\\frac{1}{\\epsilon}$ bound, and describe why such assumptions are not onerous. For example, the authors describe a class of smooth distributions for $\\tau$ that allow for such bounds, and they express an if and only if condition based on the CDF of $\\tau$. \n4. **Extension to Flexible Budget** - In Section 5, the authors include a discussion of flexible budgets, where an alternative budget $K'$ is used near $K$ that achieves better sample complexity performance. The authors sketch when such a method does and does not work, dependent on the distribution of $\\tau$.", "weaknesses": "1. **Experiments are not Extensive** - The experiments in Section 6 are condensed to a half a page (with some extra material in the Appendix). As a result, it's hard to understand some of the results. For example, why is the failure percentage not monotonic in $\\epsilon$; presumably with increasing $\\epsilon$, it is less a stringent failure threshold, so it is surprising that this pattern is exhibited across datasets. Additionally, there is little comparison with the $\\frac{1}{\\epsilon^2}$ method. Finally, what does the actual distribution of $\\tau$ look like in practice; are the assumptions validated?", "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)