arXiv:2605.06608v1 [stat.ML] 7 May 2026
DARTS: Targeting Prognostic Covariates in Budget-Constrained Sequential Experiments
Kateryna Husar Department of Statistical Science Duke University Durham, NC 27705
Alexander Volfovsky Department of Statistical Science Duke University Durham, NC 27705
Abstract Randomized controlled trials typically assume that prognostic covariates are known and available at no cost. In practice, obtaining high-dimensional pretreatment data is costly, forcing a trade-off between covariate-adaptive precision and a measurement budget. We introduce Dynamic Adaptive Rerandomization via Thompson Sampling (DARTS), which treats covariate acquisition as a sequential optimization problem embedded within a design-based causal inference task. A budgeted combinatorial Thompson sampler learns which covariates are most prognostic across successive batches; selected covariates then drive rerandomization and regression adjustment to reduce batch-level average treatment effect variance. Our primary theoretical contribution is a decoupling result: adaptive covariate selection based on past batches preserves batch-level randomization validity, and the cumulative inverse-variance weighted estimator achieves at least nominal asymptotic coverage. We further derive a Bayes risk bound for the acquisition layer that matches the minimax lower bound up to logarithmic factors. Empirically, DARTS systematically concentrates the budget on informative features, significantly closing the efficiency gap to oracle designs while maintaining strict inferential validity.
1
Introduction
A common challenge in experimental design is the “Individual Trial Fallacy”: while complete randomization ensures unbiasedness over the design space, any single realized allocation may exhibit significant conditional bias due to covariate imbalance, particularly in small samples or high-dimensional settings. To mitigate this, researchers utilize pretreatment data to enforce balance during the design phase (e.g., via rerandomization or stratification) or to adjust for imbalances during analysis. But the assumption that high-quality pretreatment data are readily available is often unrealistic. In practice, acquiring such data is rarely cost-free, constrained by medical diagnostic expenses, survey fatigue, or strict computational overhead [24, 13, 8]. Furthermore, the efficacy of covariate-adaptive methods diminishes as the number of prognostic factors grows. When the covariate space is high-dimensional relative to the sample size, achieving simultaneous balance across all dimensions becomes mathematically improbable and computationally expensive. If the most prognostic factors are known a priori, existing literature provides robust frameworks for their inclusion [18, 11, 29, 2, 27, 26, 28, 12, 17]. Yet, in many settings, these factors are unknown. In such cases, we must learn which variables to prioritize while adhering to a strict measurement budget. Instead of examining all of the data at once, it is natural to consider sequential batches and incorporate learning steps. Moreover, we often already encounter settings where data arrive in groups: participants in studies might continuously enroll over months or years, or interventions are implemented district-by-district. By treating the experiment as a series of batches, we can iteratively update our understanding of the covariate-outcome relationship. Preprint.
A substantial literature addresses related problems across three communities. We propose the Dynamic Adaptive Rerandomization via Thompson Sampling (DARTS) framework, which treats covariate acquisition as a sequential optimization problem embedded within a valid randomized trial design, and organize the relevant work to clarify where DARTS sits relative to each. Covariate-adaptive randomization and design-based inference. The foundational approach of Morgan and Rubin [18] and subsequent work (Bugni et al. [2], Yang et al. [26, 27], Zhang and Yin [28], Zhou et al. [29]) address covariate balance under a known prognostic set, either by filtering allocations via Mahalanobis distance or by tilting assignment probabilities as units arrive; Li and Zhao [12] and Molitor and Gold [17] additionally incorporate anytime-valid inference. DARTS is complementary: it answers the prior question of which covariates to balance on when they are unknown and costly to measure. Bandits with knapsacks and combinatorial semi-bandits. The budgeted combinatorial bandit framework of Badanidiyuru et al. [1], extended by Das et al. [5] to the semi-bandit setting with LP relaxations, provides the algorithmic backbone for DARTS’s acquisition layer. Liu and Ročková [15] introduce Thompson sampling for variable selection using Beta posteriors; DARTS extends this to a budgeted combinatorial setting with an explicit causal inference objective. The key question is not how to run a budgeted semi-bandit efficiently, but whether doing so inside a sequential RCT preserves causal validity. Adaptive feature acquisition and active experimental design. A growing literature treats feature acquisition as a sequential decision problem (Janisch et al. [8], Li and Oliva [13], Shim et al. [23], Covert et al. [4]), learning policies that minimize prediction error under feature costs. These methods target predictive accuracy for individual units and carry no randomization validity requirement: the acquisition policy may freely use outcome data. DARTS maintains a strict separation between acquisition policy and outcomes. Contributions. DARTS’s contributions are ordered as follows: First and foremost, we introduce a new problem formulation: adaptive covariate acquisition as part of sequential RCT design under a measurement budget. Second, we establish a design-based validity guarantee: adaptive covariate selection based on past batches does not compromise batch-level randomization validity (Theorem 1), and the cumulative inverse-variance weighted estimator achieves at least nominal asymptotic coverage (Theorem 2). Third, we provide supporting bandit theory for the acquisition layer: a Bayes risk bound and a minimax lower bound for the binary-reward case, whose rates match up to a logarithmic factor, and an extended bound for the fractional reward scheme used in practice (Theorems 3 and 4). Fourth, we validate empirically that adaptive covariate acquisition meaningfully closes the gap to oracle rerandomization across a range of DGPs and budget regimes, and that the Thompson sampler learns to concentrate on the true prognostic covariates.
2
Sequential Randomized Trial
Randomized control trial. We consider an experiment over T batches of n units each. Let Zi (t) ∈ {0, 1} denote the treatment assignment for unit i in batch t, with N1 (t) = N0 (t) = n/2 being the number of treated and control units respectively. For each unit, we define a pair of potential outcomes: Yi (1, t), the outcome that would be observed under treatment, and Yi (0, t), the outcome that would be observed under control. We impose Stable Unit Treatment Value Assumption (SUTVA): potential outcomes of unit i depend only on its own treatment, and treatment is well-defined and uniform across units. The observed outcome for unit i defined by Yi (t) = Zi (t)Yi (1, t) + (1 − Zi (t))Yi (0, t), where the only source of randomness comes from treatment assignment. Our target estimand is the Average Treatment Effect (ATE), τ = E[Yi (1, t) − Yi (0, t)], common across all batches. Since both potential outcomes are never jointly observed, a natural within-batch estimator is the difference-in-means: n n X X DiM 1 1 τ̂t = N1 (t) Zi (t)Yi (t) − N0 (t) (1 − Zi (t))Yi (t). i=1
i=1
This is trivially unbiased, but can be imprecise when an unlucky allocation produces covariate imbalance between arms. 2
Rerandomization, covariate adjustment, and precision gains. To improve precision within each batch, we can leverage pre-treatment covariates Xi (t) ∈ Rp . In the design phase, we can rerandomize using a generalization of [18]: for batch t, we draw K candidate allocations and select z ∗ = arg min dM (z, X(t)) = (X̄1 (t) − X̄0 (t))⊤ Σ̂(t)−1 (X̄1 (t) − X̄0 (t)), z
where X̄1 (t), X̄0 (t) are the treatment and control covariate means and Σ̂(t) is the pooled sample covariance within batch t. We accept randomization Z(t) for batch t by tossing a fair coin between z ∗ and 1 − z ∗ to preserve exact marginal balance. In the analysis phase, Lin-adjusted OLS on Y (t) ∼ Z(t) × X(t) yields τ̂t with HC2 variance estimate v̂t . Combining rerandomization and Lin adjustment never hurts precision under a mirror-symmetry condition [11] which our minimization-based design satisfies. Combining batch estimates. Across batches, we combine estimates using precision-weighted meta-analysis. After observing batch t, the cumulative ATE estimate is updated as: −1 −2 µ̂0 = 0, (1) σ̂t2 = σ̂t−1 + v̂t−1 , µ̂t = σ̂t2 µ̂σ̂t−1 + v̂τ̂tt , 2 t−1
which is the standard inverse-variance weighted estimator. This sequential update is equivalent to pooled OLS on all observed data when the batch-level variance estimates are well-calibrated. Each batch contributes to the cumulative estimate in proportion to its precision, so batches with better-balanced covariates and hence smaller v̂t naturally receive greater weight. Theorem 1 (Unbiasedness of sequential ATE estimator). Let Ft−1 denote the sigma-field generated by all covariate measurements, treatment assignments, and outcomes through batch t − 1. Suppose the selected covariate set St is chosen by an Ft−1 -measurable rule. Assume SUTVA and that the units in batch t are drawn i.i.d. from the target population independently of Ft−1 . Let Pt (Z(t) = z | X(t), Ft−1 ) = Pt (Z(t) = 1 − z | X(t), Ft−1 ) for every feasible assignment vector z. Then the unadjusted difference-in-means estimator satisfies E τ̂tDIM | Ft−1 = τ . Consequently, adaptive covariate selection based only on previous batches does not bias the batch-level design-based estimator. Proof of Theorem 1 is provided in Appendix A.1; the argument is batch-local and does not require the covariate set to be fixed across batches. Moreover, the combined across-batch estimator is asymptotically valid. Theorem 2 (Validity and conservativeness of the sequential ATE estimator). Define µ̂T and σ̂T as above. Under sufficient high-level regularity conditions (stated in Appendix A.3) ensuring that estimated weights are equivalent to oracle inverse-variance weights (A1–A2); no single batch dominates the asymptotic variance (A3); rerandomization does not inflate variance relative to simple random assignment(A4–A5); and any conditional bias is asymptotically negligible (A6), µ̂T −τ d − → N (0, c−2 ) σ̂T
for some c2 ≥ 1. Consequently, P |µ̂T − τ | ≤ zα/2 σ̂T → P (|Z| ≤ zα/2 c) ≥ 1 − α, so the interval µ̂T ± zα/2 σ̂T has at least nominal asymptotic coverage. Full proof can be found in Appendix A.4. While exactly unbiased under rerandomization, the difference-in-means estimator can be made more precise using regression adjustment [11]. For the Lin-adjusted estimator, the usual agnostic regression adjustment has a finite-sample bias of order O(n−1 ) (see Appendix A.2). In the equal-batch-size case with, σt2 ≍ n−1 , and w̄t ≍ n, the accumulated bias is negligible relative to the standard error when T = o(n). Thus ordinary Lin adjustment is covered in regimes where the batch size grows faster than the number of batches. More generally, the theorem requires only the stated negligible-bias condition. The two inferential statements play different roles. The first is an exact design-based result: adaptive covariate selection does not bias the batch-level difference-in-means estimator because the selected set is fixed before the current assignment is randomized. The second is an asymptotic meta-analysis result: the cumulative estimator uses same-batch estimated precision weights, so exact unbiasedness is not available, but conservative coverage follows when those weights are consistent for a design-agnostic reference variance and the accumulated centering bias is negligible. 3
Unknown and costly covariates. In practice, the most predictive covariates are rarely known in advance and costly to measure, forcing the experimenter to identify which variables are worth buying within a hard budget B across all T batches. The sequential structure offers a natural resolution: we learn which covariates are most prognostic during the experiment, updating beliefs after each batch. To measure importance, we consider efficiency gain from rerandomization and Lin’s adjustment strategies, which scales with R2 , the proportion of variance in the potential outcomes explained by St [11, 18]. More predictive covariates yield larger reductions in the standard error of the batch-level ATE estimator, making predictiveness the natural criterion for covariate selection. This motivates treating the prognostic importance of each covariate as the quantity to be learned: if covariates are independent, a more predictive covariate contributes more to R2 and hence delivers greater efficiency gains. We evaluate the impact of this proxy in the empirical results.
3
Variable selection via combinatorial semi-bandit with knapsacks
We frame identification of prognostic factors problem as a sequential variable selection problem, following the approach of Liu and Ročková [15], extending it to accommodate budget constraints and an explicit causal inference objective. 3.1
Problem formulation
We assume a known pool of p pre-treatment candidate covariates, where each variable j ∈ [p] = {1, . . . , p} has an associated measurement cost cj ∈ [0, 1]. These variables correspond to the set of arms available to the agent, and each arm j has a stochastic reward with fixed but unknown distribution with mean θj∗ ∈ [0, 1]. We are given a total budget B, and each round can select at most K ≤ p arms. At each time period t ∈ [T ], the agent selects a subset of arms St ⊂ [p], observing a binary reward ∗ rj (t) ∈ {0, 1} from each arm j ∈ St with E[r P j (t)] = θj (semi-bandit feedback). The global P reward of a super-arm St is defined as R(St ) = j∈St rj (t), with expected value rθ∗ (St ) = j∈St θj∗ . Following Liu and Ročková [15], we treat θj∗ as a proxy for the prognostic importance of covariate j for the control potential outcome Yi (0): an arm with high mean reward is one whose covariate is strongly predictive of baseline outcomes. The goal is to maximize the total expected reward subject to the overall budget constraint B. 3.2
From unconstrained to budget-constrained selection
The oracle solution under no budget is to measure every covariate for every unit in every batch. In practice, of course, budget constraints make this infeasible, and the selection problem becomes non-trivial. A natural attempt is to allocate a fixed per-batch budget B/T , spending it greedily within each round. Das et al. [5] show in Lemma 1, however, that per-batch budget schemes can perform arbitrarily badly: if the oracle set of covariates — those with the highest prognostic value — has total cost exceeding the per-batch budget, the algorithm will never be able to select them, and their rewards will never be observed regardless of how much total budget remains. A purely greedy strategy — always pulling the highest-ranked arms first — fails because early reward estimates are noisy: budget may be exhausted on arms that appeared promising before reliable rankings emerge, with no resources remaining to correct mistakes. While Explore-Then-Commit (ETC) is a natural alternative, its uniform exploration rapidly exhausts strict budgets on high-dimensional noise covariates. Because the optimal commit point relies on unknown parameters (gaps, noise, costs), ETC cannot dynamically truncate its search. Instead, it remains trapped in unguided exploration until budget depletion, yielding performance similar to a random allocation. The resolution, proposed by Badanidiyuru et al. [1] and extended by Das et al. [5], is to pose the selection problem as the integer program (3), relax it to a linear program, and work with the dual of that relaxation. Full problem formulation and extension following Das et al. [5] is presented in Appendix A.5. 4
3.3
Thompson sampling and regret bound
Under binary rewards, arm j yields rj (t) ∈ {0, 1} with E[rj (t)] = θj∗ , placing us squarely in the Beta-Bernoulli conjugate setting. We maintain independent Beta posteriors θj ∼ Beta(αj , βj ), initialized at αj (0) = βj (0) = 1 (unless prior knowledge is available) and update after each round via the standard conjugate update: αj (t) ← αj (t − 1) + rj (t),
βj (t) ← βj (t − 1) + (1 − rj (t)).
(2)
At each round t, the idealized CBwK-LP-TS policy draws posterior samples θ̃j (t) ∼ Beta(αj (t − 1), βj (t − 1)) for all j ∈ [p] and uses these in place of the UCB estimates of Das et al. [5], with all other aspects of the CBwK-LP selection procedure unchanged. We call this algorithm CBwK-LP-TS, adapting the UCB-based CBwK-LP-UCB of Das et al. [5] to Thompson sampling. The idealized policy assumes the LP dual solutions are known at each round, so that the only source of suboptimality is uncertainty about the true arm means θ⋆ . The Bayes learning regret of Theorem 3 below therefore measures exclusively the cost of substituting the posterior draw θ̃t for θ⋆ in the selection map. Theorem 3 (Bayes Risk of idealized CBwK-LP-TS). Let π T S denote the idealized CBwK-LP-TS algorithm with independent Beta(αj (0), βj (0)) priors on θj∗ . Assume semi-bandit feedback with binary rewards rj (t) ∈ {0, 1} satisfying E[rj (t) | θj∗ , Ft−1 ] = θj∗ , arms with costs cj ≥ cmin > 0, PT budget B, and horizon T . Define the regret Regret(T, π) = t=1 [rθ∗ (St∗ ) − rθ∗ (St )]. Then ! r BayesRisk(T, π T S ) ≤ O
B p log(pT ) · min KT, cmin
,
where K ≤ min(p, ⌊B/cmin ⌋) is the maximum feasible super-arm size. Detailed proof can be found in the Appendix A.6. Our proof leverages the Russo-Van Roy decomposition [22], which bounds the Bayes risk using an artificial sequence of upper confidence bounds (UCBs). If the action map is an LP/primal-dual map with a known deterministic planning gap ALP as defined by Das et al. [5], the LP benchmark bound adds E[ALP ] to the displayed learning term. Detailed discussion on planning gap is included in Appendix A.7. Lower bound analysis We establish a lower bound by considering a uniform-cost instance where arm identification must come entirely from reward observations, and applying the gap-free minimax lower bound of Kveton et al. [9] for combinatorial semi-bandits with linear rewards. Theorem 4 (Minimax Lower Bound). For any policy π, consider the regime with K < p, B ≥ p·cmin , K ≤ ⌊B/cmin ⌋, and T ≥ ⌈p/K⌉. Then there exists a Combinatorial Semi-Bandit with Knapsacks instance with p arms, costs cj ∈ (0, 1], budget B, and horizon T such that ! r B p · min KT, cmin
E[Regret(T, π)] ≥ Ω
.
The conditions of the theorem ensure that the effective pull budget is at least p, placing the problem in the nontrivial regime where all arms could in principle be sampled and regret is driven by statistical identification rather than by an insufficient horizon or budget. Full proof is provided in Appendix A.9. Theorem 3 bounds the Bayes risk of CBwK-LP-TS via the framework of Russo and Van Roy [22], while Theorem 4 establishes a frequentist minimax lower bound. These operate under different statistical frameworks: the Bayes risk averages over a prior on θ∗ , whereas the minimax bound holds for a specific hard instance. p Nevertheless, in the non-trivial regime of the theorem, the polynomial terms match up to a log(pT ) factor. Establishing a formal frequentist upper bound for CBwK-LP-TS remains an avenue for future work. 3.4
Practical considerations
Predictiveness proxy and fractional rewards Efficiency gains from rerandomization and regression adjustment are typically measured using R2 . To avoid treatment effect contamination, we estimate predictiveness using control units only via LASSO regression of Y (0, t). 5
Under standard sparsity conditions, LASSO provides consistent prediction in high dimensions [16, 25], and its coefficients serve as a proxy for prognostic importance [10]. We define fractional rewards rj (t) =
|λ̂j | maxk∈St |λ̂k |
,
with rj (t) = 0 if λ̂j = 0. Treating rj (t) as pseudo-Bernoulli observations in a Beta update yields a pseudo-posterior that aggregates predictive signals over time, down-weighting variables that receive consistently low or unstable rewards. We formally characterize the regret bound of this fractional scheme in Appendix A.8, establishing that the algorithm retains sublinear Bayes risk when the pseudo-posterior approximation error is sufficiently small. Such theory–implementation separation is common in posterior-sampling methods: regret analyses are often derived for tractable posterior updates, while practical systems use approximate or modified posterior-sampling rules whose behavior must be checked separately [3, 21, 20, 19]. Empirically, as shown in Figure 1, the fractional scheme achieves strictly lower regret than a binary reward update. It reduces regret relative to the oracle policy faster, motivating its use in the full DARTS algorithm.
Figure 1: Cumulative regret relative to Oracle DARTS across 1000 replications. Solid lines show medians; shaded bands show 95% intervals from run percentiles. DARTS with fractional and binary rewards exhibit sublinear growth, with fractional Note on independence. The bandit framework rewards converging faster. assumes independent arm rewards. LASSO partially mitigates correlation violations by typically selecting a sparse set of representatives from collinear groups. Under a budget constraint, this is highly desirable: it prevents wasting budget on correlated substitutes while using fewer degrees of freedom without sacrificing predictive performance. Back to Experimental Design. At each round t, once the super-arm St is selected, the algorithm transitions from the bandit problem back to the causal inference problem. The selected covariates {Xij (t) : j ∈ St } are measured for each unit i in batch t, yielding the covariate matrix X(t) ∈ Rn×|St | , which is then used for treatment assignment and ATE estimation as described in Section 2. Outcome observation and ATE estimation. Once treatment is assigned, outcomes are observed and the batch-level ATE is estimated via Lin-adjusted OLS on the same |St | variables used for rerandomization, with HC2 heteroskedasticity-robust variance estimate v̂t , as described in Section 2. No additional measurement cost is incurred. The cumulative estimate µ̂t is then updated via the precision-weighted meta-analysis in (1). Lin’s regression adjustment fits separate within-arm regressions, consuming 2|St | additional parameters relative to the unadjusted estimator. When |St | is large relative to n, OLS loses degrees of freedom and the HC2 variance estimate can become unstable. We therefore implicitly assume |St | ≪ n throughout. In practice this can be enforced by choosing larger batch sample sizes or additional budget constraint, which limits the number of covariates measured per batch. When the budget permits selecting many covariates, the analyst may wish to add an explicit degree-of-freedom constraint of the form |St | ≤ n/k for some k > 2 as an additional knapsack constraint in Problem 3. In our implementation we use the Moore-Penrose pseudoinverse to handle near-rank-deficient designs gracefully when estimating Σ̂(t)−1 , but this is a numerical safeguard rather than a substitute for adequate sample size.
Reward update. Finally, the observed batch is used to update the Thompson sampler. We fit a LASSO regression on the control units of batch t and compute fractional rewards rj (t) = |λ̂j |/ maxk |λ̂k | for each j ∈ St , then apply the pseudo-posterior Beta update (2). The updated posteriors {Beta(αj (t), βj (t))}pj=1 inform arm selection in round t + 1, closing the adaptive loop between variable selection and experimental design. 6
The complete procedure is summarized in Algorithm 1. Algorithm 1 Dynamic Adaptive Rerandomization via Thompson Sampling (DARTS) Require: Covariate matrix X, budget B, costs {cj }pj=1 , batch size n, number of batches T Ensure: Sequential ATE estimates {τ̂t }Tt=1 , posterior inclusion probabilities {πj }pj=1 1: Initialise: αj ← 1, βj ← 1, νj ← 1 for all j ∈ [p]; νbudget ← 1 p 2: Set B0 ← min(B, T ), c̃j ← cj · B0 /B, ε ← log(p + 1)/B0 3: for t = 1, . . . , T do 4: // Arm Selection (CBwK-LP-TS) 5: Sample θ̃j ∼ Beta(αj , βj ) for all j ∈ [p] 6: Compute effective costs c̃eff j ← c̃j · νbudget + νj for all j ∈ [p] 7: Select St ⊆ [p] greedily by θ̃j /c̃eff j , subject to remaining budget Bt−1 8: Update shadowP prices: νj ← νj (1 + ε) and νbudget ← νbudget (1 + ε)c̃j for all j ∈ St 9: Bt ← Bt−1 − j∈St c̃j 10: // Randomisation 11: Observe X(t) ∈ Rn×|St | 12: Draw K candidate assignments 13: Set Z(t) by fair coin flip between z ∗ and 1 − z ∗ , where z ∗ = arg minz dM (z, X(t)) 14: Observe Yi (t) = Zi (t)Yi (1) + (1 − Zi (t))Yi (0) for all i 15: // Estimation 16: Fit τ̂t ← β̂Z from regressing (Y (t) on Z(t), X(t) and their interaction; compute HC2 variance v̂t 17: if t = 1 then 18: µ̂1 ← τ̂1 , σ̂12 ← v̂1 19: else −1 −2 2 + v̂t−1 , µ̂t ← σ̂t2 µ̂t−1 /σ̂t−1 + τ̂t /v̂t 20: σ̂t2 ← σ̂t−1 21: end if 22: // Thompson Sampler Update 23: Observe reward rj for all j ∈ St 24: αj ← αj + rj , βj ← βj + (1 − rj ) for all j ∈ St 25: end for 26: Return: µ̂T , σ̂T , πj ← αj /(αj + βj ) for all j ∈ [p]
4
Simulation Studies
We evaluate DARTS across synthetic settings testing the learning mechanism, budget constraint, and inferential validity. Full reproducibility details are in Appendix B. 4.1
Setup
Covariates and outcome models. We consider two DGPs with q = 20 truly predictive covariates out of p ∈ {100, 1000} candidates. The linear DGP draws independent Gaussian covariates with a sparse linear outcome surface. The Liang DGP [14] introduces correlated covariates via a shared random effect and a nonlinear outcome, partially violating the arm-independence assumption. In both DGPs, τ = 4. Budget and batch parameters. The Monte Carlo grid uses T = 100 batches of n ∈ {100, 1000} units, equal costs, and budgets B ∈ {1000, 2000, 20000}, averaged over 1000 replications. We compare DiM (complete randomization, no adjustment); Random (rerandomization and Lin adjustment on a randomly drawn budget-feasible covariate subset, no learning); DARTS (Algorithm 1); and Oracle (rerandomization and Lin adjustment on the true q = 20 covariates, budget-free). The comparison to ARMM and MADCovar uses 1000 replications of the Liang DGP with T = 200 batches, n = 1000, B = 2000, and independent variable costs cj ∼Uniform(0, 2) — a more difficult set up due to correlation and linearity violations. We report medians and 95% bands from run percentiles. Two robustness checks with heterogeneous treatment effects and oracle-costly cost structures are in Appendix B.6. 7
4.2
Results
Rerandomization with adjustment. Table 1 in the appendix confirms, across 1000 replications, that combining rerandomization with Lin adjustment consistently dominates rerandomization alone. The advantage narrows when p is large relative to n, as OLS adjustment consumes degrees of freedom, but never substantially harms precision. All subsequent comparisons use the combined estimator. Does DARTS learn the right covariates? Figure 2 shows the distribution of final posterior inclusion probabilities πj = αj /(αj + βj ) across 1000 replications, separately for true signal (j ≤ 20) and noise (j > 20) covariates. Under the Liang DGP with B = 2000, p = 100, n = 1000, the signal and noise distributions are completely separated (signal mean 0.34, noise mean 0.11), confirming that the Thompson sampler systematically concentrates budget on prognostic covariates. Separation is weaker under B = 1000 and p = 1000, where the sampler must identify 20 signal arms from 1000 candidates within a limited horizon, an honest boundary case where DARTS degrades gracefully toward the random baseline. When budget is 20, 000 while p = 100, essentially removing any monetary constraint, all methods perform similarly, since all get access to every oracle covariate. Across the 1000-replication method comparison, the share of batch budget allocated to oracle variables increases monotonically over batches (Figure 6a, Appendix B.7), confirming that the Thompson sampler progressively concentrates resources on the truly prognostic covariates as posterior beliefs sharpen with accumulated evidence.
(a) Liang, B = 2000, p = 100, n = 1000: signal and noise distributions completely separated.
(b) Liang, B = 2000, p = 1000, n = 100: separation collapses as the learning problem becomes intractable.
Figure 2: Distribution of final posterior inclusion probabilities πj = αj /(αj + βj ) across 1000 replications for true signal (j ≤ 20, solid blue) and noise (j > 20, dashed yellow) covariates. The Thompson sampler concentrates budget on prognostic covariates when the problem is tractable, and spreads it randomly when the candidate pool is too large relative to the horizon. Does learning translate to precision gains? Table 2 shows that SE(DiM) ≥ SE(Random) ≥ SE(DARTS) > SE(Oracle) holds for most DGPs and budgets. The main exceptions occur in high-p, low-n settings, where the algorithm has insufficient data to learn effectively, and at high budgets, where all methods gain access to all oracle variables. DARTS achieves up to RE = 3.56 over DiM versus RE = 1.26 for Random, isolating the Thompson sampler’s contribution. After a 30-batch burn-in, batch-level rewards correlate positively (r = 0.977) with SE reduction relative to DiM (Figure 6b,Appendix B.7). Moreover, later batches correspond to higher SE reduction. This confirms that the LASSO fractional rewards proxy translates learning into precision gains. Coverage is at the nominal level throughout, consistent with Theorem 2. Comparison to ARMM and MADCovar. We compare against ARMM [26] and MADCovar [17] over 1000 replications as described in Section 4.1. Both methods are given a randomly preselected subset of covariates per simulation — drawn until the budget per batch is met, averaging around 12 — the realistic scenario where an analyst pre-commits without knowledge of which are prognostic, while DARTS sees the full pool and must learn under the same budget. Table 3 summarizes performance, where Rel. RMSE = RMSEDiM /RMSEmethod with values above 1 indicating improvement over DiM. DARTS achieves Rel. RMSE = 1.415, compared to 1.018 and 1.036 for ARMM and MADCovar respectively, closing a substantial share of the gap to Oracle DARTS (Rel. RMSE = 1.812). Figure 3 reinforces this: panel (a) shows that DARTS accumulates strictly less regret than ARMM and MADCovar throughout the horizon, and panel (b) shows that DARTS achieves narrower CI widths over the last 50 batches, with both differences remaining consistent across replications as seen by the non-overlapping 95% bands. The advantage of DARTS 8
reflects the value of adaptive acquisition when the prognostic set is unknown — when given oracle covariates, ARMM and MADCovar perform comparably to Oracle DARTS.
(a) Cumulative regret relative to Oracle DARTS. Solid lines show medians; shaded bands show 95% intervals from run percentiles.
(b) 95% CI width over the last 50 batches. Solid lines show medians across 1000 replications; shaded bands show 95% intervals from run percentiles.
Figure 3: Comparison of DARTS against ARMM [26] and MADCovar [17] across 1000 replications of the Liang DGP (n = 1000, T = 200, p = 100, B = 2000, variable costs). ARMM and MADCovar are given 12 randomly pre-selected covariates; DARTS learns from the full candidate pool. Solid lines show medians across 1000 replications; shaded bands show 95% intervals from run percentiles. All CIs are 95% valid at fixed T ; HC2 standard errors throughout.
5
Limitations and discussion
5.1
Limitations
While DARTS consistently improves precision in unknown-covariate settings, we note several limitations regarding its reward proxy, theoretical guarantees, and implementation constraints. First, the theoretical regret guarantees rely on an idealized setting. Theorem 3 assumes exact BetaBernoulli conjugacy. While Corollary 2 bounds the pseudo-posterior approximation error of fractional rewards, this bound depends on an unobservable total variation distance. The guarantee is robust in sparse settings but should be treated as indicative in dense-signal regimes. Second, the framework currently supports fixed-horizon inference. Theorem 2 establishes asymptotic coverage only for the fixed-horizon estimator µ̂T . The confidence intervals are valid only at the pre-specified end of the trial; sequential peeking or early stopping invalidates coverage. Finally, the empirical comparisons to ARMM and MADCovar are informationally asymmetric by design — DARTS learns from the full candidate pool, while baselines use randomly pre-selected covariate sets to accurately reflect standard practice. This asymmetry is inherent to the problem setting: ARMM and MADCovar are not designed to perform covariate acquisition and take a fixed covariate set as given, so the comparison is best understood as a demonstration of the value of adaptive acquisition over naive pre-selection rather than a head-to-head competition between methods solving the same problem. 5.2
Discussion and future work
DARTS demonstrates that it is possible to close the efficiency gap to oracle rerandomization without sacrificing inferential validity, gracefully degrading to a completely randomized baseline when the learning problem is intractable. To build on this framework, future work should directly address the limitations of the independent arm assumption. In settings where covariate clusters are known a priori, integrating alternative reward functions like group LASSO, or extending the framework to combinatorial bandits with correlation structure [6], would allow the Thompson sampler to share information across correlated arms. Furthermore, extending the inverse-variance weighted estimator to utilize confidence sequence constructions would enable anytime-valid inference, a natural evolution for sequential and adaptive trial designs. 9
References [1] Ashwinkumar Badanidiyuru, Robert Kleinberg, and Aleksandrs Slivkins. Bandits with knapsacks. Journal of the ACM (JACM), 65(3):1–55, 2018. [2] Federico A. Bugni, Ivan A. Canay, and Azeem M. Shaikh. Inference under covariate-adaptive randomization. Journal of the American Statistical Association, 113(524):1784–1796, 2018. An empirical evaluation of thompson sam[3] Olivier Chapelle and Lihong Li. pling. In Advances in Neural Information Processing Systems, volume 24, pages 2249–2257, 2011. URL https://proceedings.neurips.cc/paper/2011/hash/ e53a0a2978c28872a4505bdb51db06dc-Abstract.html. [4] Ian Connick Covert, Wei Qiu, Mingyu Lu, Na Yoon Kim, Nathan J White, and Su-In Lee. Learning to maximize mutual information for dynamic feature selection. In International Conference on Machine Learning, pages 6424–6447. PMLR, 2023. [5] Debojit Das, Shweta Jain, and Sujit Gujar. Budgeted combinatorial multi-armed bandits. arXiv preprint arXiv:2202.03704, 2022. [6] Samarth Gupta, Jinhang Zuo, Carlee Joe-Wong, Gauri Joshi, and Osman Yağan. Correlated combinatorial bandits for online resource allocation. In Proceedings of the Twenty-Third International Symposium on Theory, Algorithmic Foundations, and Protocol Design for Mobile Networks and Mobile Computing, pages 91–100, 2022. [7] Peter Hall and Christopher C Heyde. Martingale limit theory and its application. Academic press, 2014. [8] Jaromír Janisch, Tomáš Pevnỳ, and Viliam Lisỳ. Classification with costly features using deep reinforcement learning. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pages 3959–3966, 2019. [9] Branislav Kveton, Zheng Wen, Azin Ashkan, and Csaba Szepesvari. Tight regret bounds for stochastic combinatorial semi-bandits. In Artificial Intelligence and Statistics, pages 535–543. PMLR, 2015. [10] Quinn Lanners, Harsh Parikh, Alexander Volfovsky, Cynthia Rudin, and David Page. Variable importance matching for causal inference. In Uncertainty in Artificial Intelligence, pages 1174–1184. PMLR, 2023. [11] Xinran Li and Peng Ding. Rerandomization and regression adjustment. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 82(1):241–268, 2020. [12] Xinran Li and Anqi Zhao. Design-based theory for causal inference from adaptive experiments. arXiv preprint arXiv:2602.21998, 2026. [13] Yang Li and Junier Oliva. Active feature acquisition with generative surrogate models. In International conference on machine learning, pages 6450–6459. PMLR, 2021. [14] Faming Liang, Qizhai Li, and Lei Zhou. Bayesian neural networks for selection of drug sensitive genes. Journal of the American Statistical Association, 113(523):955–972, 2018. [15] Yi Liu and Veronika Ročková. Variable selection via thompson sampling. Journal of the American Statistical Association, 118(541):287–304, 2023. [16] Nicolai Meinshausen and Bin Yu. Lasso-type recovery of sparse representations for highdimensional data. The Annals of Statistics, 37(1):246–270, 2009. [17] Daniel Molitor and Samantha Gold. Anytime-valid inference in adaptive experiments: Covariate adjustment and balanced power. arXiv preprint arXiv:2506.20523, 2025. [18] Kari Lock Morgan and Donald B Rubin. Rerandomization to improve covariate balance in experiments. The Annals of Statistics, 40(2):1263–1282, April 2012. doi: 10.1214/12-AOS1008. 10
[19] My Phan, Yasin Abbasi-Yadkori, and Justin Domke. Thompson sampling and approximate inference. In Advances in Neural Information Processing Systems, volume 32, pages 8801–8811, 2019. URL https://proceedings.neurips.cc/paper/2019/hash/ f3507289cfdc8c9ae93f4098111a13f9-Abstract.html. [20] Chao Qin, Zheng Wen, Xiuyuan Lu, and Benjamin Van Roy. An analysis of ensemble sampling. In Advances in Neural Information Processing Systems, volume 35, 2022. URL https://proceedings.neurips.cc/paper_files/paper/2022/hash/ 874f5e53d7ce44f65fbf27a7b9406983-Abstract-Conference.html. [21] Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. In International Conference on Learning Representations, 2018. URL https://openreview.net/forum? id=SyYe6k-CW. [22] Daniel Russo and Benjamin Van Roy. Learning to optimize via posterior sampling. Mathematics of Operations Research, 39(4):1221–1243, 2014. [23] Hajin Shim, Sung Ju Hwang, and Eunho Yang. Joint active feature acquisition and classification with variable-size set encoding. Advances in neural information processing systems, 31, 2018. [24] Peter D Turney. Types of cost in inductive concept learning. arXiv preprint cs/0212034, 2002. [25] Larry Wasserman and Kathryn Roeder. High dimensional variable selection. Annals of statistics, 37(5A):2178, 2009. [26] Haoyu Yang, Yichen Qin, Fan Wang, Yang Li, and Feifang Hu. Balancing covariates in multiarm trials via adaptive randomization. Computational Statistics & Data Analysis, 179:107642, 2023. [27] Haoyu Yang, Yichen Qin, Yang Li, and Feifang Hu. Sequential covariate-adjusted randomization via hierarchically minimizing mahalanobis distance and marginal imbalance. Biometrics, 80(2): ujae047, 2024. [28] Hengtao Zhang and Guosheng Yin. Response-adaptive rerandomization. Journal of the Royal Statistical Society: Series C (Applied Statistics), 70(5):1281–1298, 2021. [29] Quan Zhou, Philip A. Ernst, Kari Lock Morgan, Donald B. Rubin, and Anru Zhang. Sequential rerandomization. Biometrika, 105(3):745–752, 2018.
A
Technical appendices and supplementary material
A.1
Proof of theorem 1
Fix batch t and condition on Ft−1 . Since St is chosen using only previous batches, the selected covariate set is fixed under this conditioning. Let Gt = Ft−1 ∨ σ(X(t)) denote the sigma-field that additionally contains the selected covariates for batch t. Write Yi (1, t) = Yi (0, t) + τi (t), where τi (t) = Yi (1, t) − Yi (0, t). The unadjusted difference-in-means estimator satisfies X X Yi (0, t) − N01(t) Yi (0, t) τ̂tDIM − τ = N11(t) i:Zi (t)=1
i:Zi (t)=0
|
{z
}
(I)
+ N11(t)
X
τi (t) − τ .
i:Zi (t)=1
|
{z
}
(II)
11
Conditional on Gt , the selected covariates and potential outcomes in batch t are fixed. By the mirror-symmetric rerandomization rule, Pt (Z(t) = z | Gt ) = Pt (Z(t) = 1 − z | Gt ). Thus the treated and control labels are exchangeable under the conditional randomization distribution. Therefore E[(I) | Gt ] = 0. The same symmetry implies that the treated units are, in expectation under the conditional randomization distribution, representative of the finite batch with respect to the individual treatment effects: X E N11(t) τi (t) Gt = τ̄t , i:Zi (t)=1
Pn where τ̄t = n−1 t i=1 τi (t) is the finite-batch average treatment effect. Hence E[(II) | Gt ] = τ̄t − τ. Taking expectations over the i.i.d. sampling of the batch, which is independent of Ft−1 , gives E[τ̄t | Ft−1 ] = τ. By iterated expectation, E[τ̂tDIM − τ | Ft−1 ] = 0, which proves the claim. A.2
Asymptotic unbiasedness under covariate adjustment
Under the rerandomization procedure and regularity conditions (i)–(iii) below, the Lin-adjusted estimator satisfies |E[et | Ft−1 ]| = O(n−1 ), where n is the batch size. Regularity conditions: (i) Covariates bounded: ∥Xi ∥ ≤ C < ∞. (ii) The finite Pare uniformly ⊤ population covariance matrices n−1 X X converge to positive definite limits Σz for i z i i:Zi =z z ∈ {0, 1}. (iii) Potential outcomes satisfy E[|Yi (z)|2 ] < ∞. To simplify notation, we assume without loss of generality that the covariates have already been centered within each batch such that X̄t = 0. The Lin-adjusted estimator is defined as: ⊤ ⊤ τ̂t,Lin = (Ȳt,1 − β̂t,1 X̄t,1 ) − (Ȳt,0 − β̂t,0 X̄t,0 ),
where β̂t,z are arm-specific OLS coefficients. We decompose the error et = τ̂t,Lin − τ as follows: et = (Ȳt,1 (0) − β1∗⊤ X̄t,1 ) − (Ȳt,0 (0) − β0∗⊤ X̄t,0 ) | {z } (I)
+ τ̄t,1 − τ | {z } (II)
+ (β1∗ − β̂t,1 )⊤ X̄t,1 − (β0∗ − β̂t,0 )⊤ X̄t,0 , |
{z
(III)
}
where βz∗ denote the population limits of the regression coefficients that exists under conditions (ii) and (iii). The expectations of Terms (I) and (II) vanish exactly following the symmetry and sampling arguments established in the rerandomization proof (A.1). By the marginal 1/2 symmetry of the mirrorflip procedure, P (Zi (t) = 1 | X(t), Ft−1 ) = 1/2. Given that potential outcomes and βz∗ are fixed conditional on X(t), the expected adjusted baseline outcomes are equal across arms, so E[(I) | X(t), Ft−1 ] = 0. By the marginal 1/2 symmetry, E[τ̄t,1 | X(t), Ft−1 ] = τ̄t , the batch average. Taking the outer expectation over the i.i.d. sampling of units gives E[τ̄t | Ft−1 ] = τ , thus E[(II) | Ft−1 ] = 0. 12
Term (III) captures the finite-sample bias of the OLS adjustment. Under condition (ii), the coefficient estimation error satisfies E[∥β̂t,z − βz∗ ∥2 | Ft−1 ] = O(n−1 ). Under rerandomization and withinbatch centering, the group-mean imbalance satisfies E[∥X̄t,z ∥2 | Ft−1 ] = O(n−1 ). Applying the Cauchy-Schwarz inequality in L2 : h i E |(βz∗ − β̂t,z )⊤ X̄t,z | Ft−1 ≤ E[∥β̂t,z − βz∗ ∥2 | Ft−1 ]1/2 · E[∥X̄t,z ∥2 | Ft−1 ]1/2 = O(n−1 ). Therefore, |E[(III) | Ft−1 ]| = O(n−1 ). Combining these results, the sequence {et } satisfies |E[et | Ft−1 ]| = O(n−1 ). A.3
Regularity conditions for Theorem 2
As before, for each batch t, let Gt = Ft−1 ∨ σ(X(t)) be the sigma-field containing previous batches and the selected covariates for the current batch. Let τ̂t be the batch-specific ATE estimator and let et = τ̂t − τ . Define !−1 T PT X wt τ̂t 2 t=1 µ̂T = PT w , σ̂T = wt , wt = v̂t−1 . t=1
t
t=1
Let
σt2 = Var(τ̂t | Gt ) denote the conditional variance under the actual rerandomization design, and let 2 σt,◦ = Var◦ (τ̂t | Gt ) −2 denote the corresponding design-agnostic variance under simple random assignment. Set w̄t = σt,◦ . Let PT w̄t2 σt2 VT = Pt=1 2. T ( t=1 w̄t ) We impose the following conditions.
A1 (Relative weight consistency) maxt≤T A2 (Estimated-weight remainder)
wt w̄t − 1
p
− → 0.
PT
t=1 (wt − w̄t )et = op
q
PT
2 2 t=1 w̄t σt
.
w̄2 σ 2
A3 (Uniform asymptotic negligibility) maxt≤T PT t w̄t2 σ2 → 0. s s s=1 P 2 2 σ w̄ p t t t A4 (Variance ratio convergence) P − → ℓ ∈ (0, 1] t w̄t 2 A5 (Conservative reference variance) σt2 ≤ σt,◦
for all t. q PT PT 2 σ2 . A6 (Negligible centering bias) t=1 w̄t E[et | Gt ] = op w̄ t=1 t t A7 (Bounded moments) There exist δ > 0 and M < ∞ such that supt E[|et |2+δ | Gt ] ≤ M. Condition (A3) is a uniform asymptotic negligibility (UAN) condition for the martingale difference triangular array {w̄t (et − E[et | Gt ])}. Combined with the bounded moments Condition (A7), these jointly imply the conditional Lindeberg condition required by the Martingale CLT. Conditions (A1)–(A2) permit replacement of wt with w̄t . Condition (A6) rules out asymptotically non-negligible centering bias. Conditions (A4)–(A5) ensure conservative variance estimation under rerandomization. (A5) is satisfied whenever the within-batch assignment mechanism is a rerandomization scheme in the sense of Morgan and Rubin [18]. In the finite-population design-conditional setting, Li and Ding [11] show that such schemes reduce or preserve variance relative to simple random assignment. (A5) imports this conclusion into the superpopulation framework as a primitive assumption on the joint conditional distribution of τ̂ given Gt . 13
A.4
Proof of theorem 2
By relative weight consistency, PT p t=1 wt PT − → 1. w̄ t t=1
By the estimated-weight remainder condition, T X
v u T uX (w − w̄ )e = o t w̄2 σ 2 . t
t
t
p
t
t=1
t
t=1
Therefore, µ̂T − τ =
PT p t=1 w̄t et P + op ( VT ). T w̄ t t=1
By the negligible centering-bias condition, T X
v u T uX w̄t E[et | Gt ] = op t w̄t2 σt2 .
t=1
t=1
Thus the centered triangular array w̄t {et − E[et | Gt ]} is a martingale difference array with respect to {Gt } with sum of conditional variances equal to PT 2 2 t=1 w̄t σt by construction. To verify the Lindeberg condition, note that for any ε > 0, by Markov’s inequality and Condition (A7): " # ! T X |w̄t {et − E[et | Gt ]}| w̄t2 {et − E[et | Gt ]}2 P 2 2 pP E Gt ·1 >ε 2 2 s w̄s σs s w̄s σs t=1 PT 2+δ 21+δ M t=1 w̄t ≤ · → 0, P 1+δ/2 εδ ( s w̄s2 σs2 ) where the last step follows from conditions (A3) and (A7). The Lindeberg condition is therefore satisfied. By (A6) and Slutsky’s theorem, the centering bias is negligible, and the Martingale CLT [7] gives: PT d w̄t et √Pt=1 − → N (0, 1). T 2 2 t=1 w̄t σt
Equivalently, µ̂√ T −τ d − → N (0, 1). VT
Let 2 σT,◦ =
T X
!−1 w̄t
.
t=1 −2 Because w̄t = σt,◦ and using (A5), VT = 2 σT ,◦
PT PT 2 2 2 2 t=1 w̄t σt,◦ t=1 w̄t σt P P ≤ = 1. T T w̄ t t=1 t=1 w̄t
p
2 By Condition (A4), VT /σT,◦ − → c−2 , where c2 = ℓ−1 ≥ 1. Relative weight consistency also yields P −1 2 p ( t wt ) σ̂T = → 1. P 2 −1 − σT ,◦ ( w̄t ) t
By Slutsky’s theorem, µ̂T −τ d − → N (0, c−2 ). σ̂T
Since c−2 ≤ 1, the Wald interval based on σ̂T has asymptotic coverage at least 1 − α. 14
A.5
LP dual problem and shadow prices
The integer programming program can be formulated as follows: p T X X
max a(1),...,a(T )
subject to
aj (t)θj∗
t=1 j=1 p T X X
aj (t)cj ≤ B,
t=1 j=1
T X
(3) aj (t) ≤ T,
aj (t) ∈ {0, 1}
∀j ∈ [p], t ∈ [T ],
t=1
where aj (t) = 1 indicates covariate j is measured in batch t. The mean rewards θj∗ ’s are unknown and must be learned sequentially, compounding the difficulty of Problem (3) beyond its NP-hard nature. The LP relaxation replaces the binary constraint aj (t) ∈ {0, 1} with a continuous allocation variable ζj ≥ 0: p p X X max ζj θj∗ subject to ζj Mij ≤ B ′ ∀ i ∈ {1, . . . , p + 1}, (4) ζ≥0
j=1
j=1
where B ′ = min(B, T ), c̃j = cj ·B ′ /B are rescaled costs, and M encodes both the global budget and per-arm frequency constraints on a common scale. By Lemma 3 of Das et al. [5], OPTLP ≥ OPTIP , so benchmarking against OPTLP provides a valid upper bound on the integer optimum. Shadow prices and the dual problem Rather than solving (4) at each round, we work with its dual. The dual of (4) introduces a shadow price νi ≥ 0 for each resource constraint i. Intuitively, a shadow price measures how much additional reward one extra unit of budget would yield: strictly positive when the constraint is binding, zero when slack. We maintain a vector of dual shadow prices ν ∈ Rp+1 , one per arm and one for the global budget constraint, initialized to 1 and updated multiplicatively after each round: νj ← νj (1 + ε),
νp+1 ← νp+1 (1 + ε)c̃j
∀j ∈ St ,
p
where ε = log(p + 1)/B0 is the multiplicative weight learning rate. Each time an arm is pulled, its shadow price increases, making it less attractive in subsequent rounds. The update has a natural interpretation: having spent heavily on arm j in recent rounds, the algorithm treats the remaining budget as more valuable and demands a higher return per unit cost before pulling j again. On average, this prevents early over-spending and ensures the global budget is not breached. The effective cost of arm j at round t is then: c̃eff j (t) = c̃j · νp+1 (t) + νj (t), which inflates the nominal cost of frequently-pulled arms. Looking ahead over the remaining H = T − t + 1 rounds, the super-arm is selected greedily by the bang-per-buck ratio: n o θ̃j (t) St = Greedy , B t−1 , c̃eff (t) j
j∈[p]
where Greedy sorts arms in decreasing order of θ̃j /c̃eff j and selects greedily subject to the remaining budget Bt−1 . A.6
Proof of theorem 3
The Bayes risk over T periods is defined as the expected cumulative regret against the optimal subset St∗ : T X BayesRisk(T, π T S ) = E[rθ∗ (St∗ ) − rθ∗ (St )], t=1
where the expectation is over the prior on θ∗ and the algorithm’s randomness. Let Ft−1 be the filtration generated by the implemented Beta policy, π T S , up to round t. To analyze the Thompson Sampling policy without directly tracking the posterior quantiles, we apply the fundamental identity of Russo and Van Roy [22] (Proposition 1). Because Thompson Sampling 15
ensures that the distribution of the selected arm St matches the posterior distribution of the optimal arm St∗ , we can introduce any valid Ft−1 -measurable upper confidence bound sequence {Ut | t ∈ N} as an analytical tool. We can then decompose the Bayes Risk into a confidence width term (Term I) and a UCB failure term (Term II): BayesRisk(T, π T S ) = E
T X
[Ut (St ) − rθ∗ (St )] + E
t=1
T X [rθ∗ (St∗ ) − Ut (St∗ )]. t=1
= Term I + Term II Since we have semi-bandit feedback, we can construct independent confidence bounds. Let θ̂j (t) = P 1 s<t,j∈Ss rj (s) (with θ̂j (t) = 0 if nj (t) = 0), where nj (t) denotes the number of times arm nj (t) j was pulled up to time t. For each j ∈ [p] and round t, we define: ( ( ) ) r r Uj,t = min 1, θ̂j (t) +
log(pT 2 ) 2nj (t)
log(pT 2 ) 2nj (t)
Lj,t = max 0, θ̂j (t) −
,
,
with Uj,t = 1 and Lj,t = 0 when nj (t) = 0. Since rj (t) ∈ [0, 1] are bounded, by Hoeffding’s inequality (assuming E[rj (t) | θj∗ , Ft−1 ] = θj∗ ) for each arm j at each round t: ! r P (θj∗ > Uj,t | Ft−1 ) = P
θj∗ > θ̂j,t +
=P
θj∗ − θ̂j,t >
log(pT 2 ) 2nj (t) Ft−1
r
log(pT 2 ) 2nj (t) Ft−1
!
n o 2n (t) log(pT 2 ) ≤ exp − j 2nj (t) = pT1 2 , and, symmetrically P (θk∗ < Lj,t | Ft−1 ) ≤ 1/(pT 2 ). We define the super-arm bounds additively by linearity of the reward: X X Ut (S) = Uj,t , Lt (S) = Lj,t . j∈S
j∈S
Term II bound: The event rθ∗ (St∗ ) > Ut (St∗ ) requires at least one arm j ∈ St∗ to satisfy θj∗ > Uj,t . By the union bound over at most K arms in St∗ : P (rθ∗ (St∗ ) > Ut (St∗ ) | Ft−1 ) ≤
X
P (θj∗ > Uj,t | Ft−1 ) ≤ K pT1 2 .
j∈St∗
Since each arm contributes at most 1 to the reward, rθ∗ (St∗ ) − Ut (St∗ ) ≤ |St∗ | ≤ K. Thus, T T X X ∗ ∗ K K2 ∗ E [rθ (St ) − Ut (St )] ≤ K pT 2 = pT . t=1
t=1
This term vanishes as T → ∞. Term I bound: Let Et = {Lj,t ≤ θj⋆ ≤ Uj,t ∀j ∈ [p]} denote the event that all confidence bounds hold at round t. By union bound and Hoeffding’s inequality: P (EtC | Ft−1 ) = P (∃j : θj∗ < Lj,t ∨ θj∗ > Uj,t | Ft−1 ) ≤
p X
P (θj∗ < Lj,t ∨ θj∗ > Uj,t | Ft−1 )
j=1
≤ p pT2 2 = T22 . This complementary event contributes at most K · 2/T 2 per round in expectation. Thus, the total PT 2K 2K expected penalty is at most t=1 2K T 2 = T 2 · T = T = O(K/T ). 16
Therefore, on high probability event Et : Ut (St ) − rθ∗ (St ) ≤ Ut (St ) − Lt (St ) =
X
(Uj,t − Lj,t ) =
X
r 2
log(pT 2 ) 2nj (t) .
j∈St
j∈St
Because semi-bandit feedback increments nj (t) independently for each arm j, we can exchange the order of summations over time and arms. T T X r X X 2 ) E [Ut (St ) − rθ∗ (St )] ≤ 2 log(pT 2nj (t) t=1
t=1 j∈St
=
≤
p
p
2 log(pT 2 )
2 log(pT 2 )
p X T X
√1
j=1 t=1 p q X
nj (t)
nj (T )
2
harmonic sum bound
j=1
v 2 u u X p q p u = 2 2 log(pT 2 )t nj (T ) j=1
v u X p u p ≤ 2 2 log(pT 2 )tp nj (T )
Cauchy-Schwarz inequality
j=1
p p = 2 2 log(pT 2 ) pQ(T ), Pp where Q(T ) = j=1 nj (T ) is the total number (quantity) of arms pulled across all rounds and arms. The algorithm enforces two physical constraints on the number of arms that can be pulled: 1. at most K arms per found for T rounds, so Q(T ) ≤ KT ; 2. each pull of arm j costs at least cmin , and the total budget is B, so cmin · Q(T ) ≤ B, giving Q(T ) ≤ B/cmin . B Since both must hold, Q(T ) ≤ min KT, cmin . We get E
T X p [Ut (St ) − rθ∗ (St )] ≤ 2 2 log(pT 2 ) ·
r
B p · min KT, cmin .
t=1
Combining this with the vanishing terms for event EtC and term II, we get the bound in the theorem. A.7
Planning gap
Let B ′ = min(B, T ) and let L(θ⋆ ) = OPTLP (θ⋆ ) denote the LP optimal value under true means θ⋆ . Following Das et al. [5], the deterministic planning gap of the CBwK-LP primal-dual policy under known means is ! r log(p + 1) p ⋆ ⋆ ADas (T, θ ) = CD L(θ ) + ′ , B′ B where CD > 0 is a universal constant. This bounds how far the primal-dual policy π LP (θ⋆ ) falls short of OPTLP (θ⋆ ) when arm means are known, capturing the resource-allocation loss from online dual updates rather than statistical uncertainty about θ⋆ . Taking expectations over the prior gives ! r log(p + 1) p GapLP (T ) = E[ADas (T, θ⋆ )] = CD E[L(θ⋆ )] + ′ . B′ B p When B ′ ≫ p, this term is of order O(L(θ⋆ ) log p/B ′ ) and is negligible relative to the learning regret of Theorem 3. 17
To obtain a bound against the LP benchmark, let OPTLP (θ⋆ ) and OPTIP (θ⋆ ) denote the values of the LP relaxation and integer program under true means θ⋆ respectively (Appendix A.5); by Lemma 3 ⋆ ⋆ of Das et al. [5], hPOPTLP (θ )i≥ OPTIP (θ ), so benchmarking against OPTLP is conservative. Let T ⋆ ⋆ Rθ⋆ (π) = Eπ t=1 rθ (St ) denote the expected total reward of policy π under true means θ , and let π LP (θ⋆ ) denote the CBwK-LP primal-dual policy with θ⋆ supplied in place of posterior samples. Decomposing the total gap, OPTLP (θ⋆ ) − Rθ⋆ (π T S ) = OPTLP (θ⋆ ) − Rθ⋆ (π LP (θ⋆ )) + Rθ⋆ (π LP (θ⋆ )) − Rθ⋆ (π T S ) , | {z } | {z } planning gap
learning gap
the first term is controlled by the known-means analysis of Das et al. [5]; taking expectations over the prior on θ⋆ gives GapLP (T ) = E[ADas (T, θ⋆ )], where ADas is defined in Appendix A.5. The second term is bounded by Theorem 3, since BayesRisk(T, π T S ) = E[OPTLP (θ⋆ ) − Rθ⋆ (π T S )] controls the learning gap after taking expectations over the prior. Corollary 1 (LP-benchmark Bayes risk of CBwK-LP-TS). Under the conditions of Theorem 3, BayesRiskLP (T ) := E OPTLP (θ⋆ ) − Rθ⋆ (π T S ) ! r B ≤ GapLP (T ) + O . p log(pT ) · min KT, cmin Since OPTIP (θ⋆ ) ≤ OPTLP (θ⋆ ), the same bound holds against the integer optimum. A.8
Fractional Rewards
Corollary 2 (Bayes Risk under Fractional Rewards). Under fractional rewards and the pseudoposterior Beta update (2), ! r T X TS B BayesRisk(T, π ) ≤ O p log(pT ) · min KT, cmin +K E[εt ], t=1
where εt = dTV (Π̃t , Πt | Dt−1 ) is the total variation distance between the pseudo-posterior and the true posterior at round t. The additional term is small when rewards concentrate near 0 or 1, in which case the fractional algorithm achieves the same asymptotic rate as the binary case. PT The additional term K t=1 E[εt ] captures the price of using a pseudo-posterior. This term is small when the fractional rewards concentrate near 0 or 1 — that is, when covariates are either clearly predictive or clearly not — since in this regime the fractional and binary updates agree approximately and the pseudo-posterior is close to the true posterior. In high-dimensional settings where the signal is sparse, this condition is empirically plausible. A.9
Proof of theorem 4
Consider an instance where cj = cmin for all j ∈ [p]. With uniform costs, the bang-per-buck ratio θj /cj = θj /cmin is strictly monotone in θj . Identification of the optimal super-arm must come entirely from reward observations. Assign each arm a Bernoulli reward distribution with mean θj∗ ≥ 1/2 > 0. Because every arm has strictly positive expected reward, pulling an additional arm always increases expected reward, so any algorithm pulling fewer than K arms in any round incurs strictly additional regret relative to the optimal policy. Without loss of generality we therefore restrict attention to algorithms that pull exactly K arms per round on this instance. The condition K < p ensures there exist suboptimal arms, making the identification problem nontrivial. Because each arm pull costs exactly cmin regardless of which arm is selected, the total budget B induces a deterministic stopping time: the experiment terminates after exactly ⌊B/cmin ⌋ total arm pulls independently of the algorithm’s choices. Combined with the horizon constraint of KT total pulls, the effective total pull budget is: j k B Q = min KT, cmin . 18
Since algorithms pull exactly K arms per round on this instance, Q total pulls correspond to t∗ = ⌊Q/K⌋ full rounds. A lower bound on regret over t∗ rounds is a valid lower bound over the full horizon, since rounds beyond t∗ are either infeasible due to budget exhaustion or contribute non-negatively to cumulative regret. We apply the lower bound of Kveton et al. [9], whose Proposition 2 is established for the case K | p to construct a disjoint partition of arms into non-overlapping super-arms, ensuring independence between pulls. For general p, we restrict to a sub-instance of p′ = ⌊p/K⌋ · K arms, which satisfies K | p′ and recovers the same asymptotic bound since p′ ≥ p/2. The problem now reduces to a stochastic combinatorial semi-bandit with p′ base arms, super-arm P size K, linear reward rθ∗ (S) = j∈S θj∗ , rewards bounded in [0, 1], semi-bandit feedback over t∗ rounds. Proposition 2 of Kveton et al. [9], which holds for any horizon t > 0, proves a gap-free lower √ bound of Ω( Kp′ t∗ ), under the only non-trivial setting of Q ≥ p where B ≥ p · cmin by assumption, and T ≥ p/K so that we can visit every arm. We get q p p Q E[Regret(T, π)] ≥ Ω K · p′ · K = Ω( p′ Q) = Ω( pQ), √ √ where the last equality holds since p′ = ⌊p/K⌋ · K ≥ p/2, so p′ = Ω( p). Substituting Q = min(KT, ⌊B/cmin ⌋) yields the final bound: ! r B . E[Regret(T, π)] ≥ Ω p · min KT, cmin This lower bound isolates the statistical learning complexity under uniform costs. Instances with severe cost heterogeneity may be strictly harder, potentially yielding a tighter lower bound than established here. Characterizing whether cost heterogeneity introduces an additional dominant complexity term remains an open question.
B
Simulations
B.1
Data-Generating Processes
Linear DGP. Covariates are drawn independently as X ∼ N (0, Ip ). Outcomes are Yi (0) = Xi⊤ β + 2Xi,1 Xi,5 + εi , εi ∼ N (0, 1), with βj ∼ N (2, 0.1) for j ≤ 10, βj ∼ N (−2, 0.05) for j ∈ {11, . . . , 20}, and βj ∼ N (0, 0.01) for j > 20. The weak but nonzero coefficients for j > 20 reflect a realistic setting where noise covariates carry negligible predictive signal rather than being exactly inert. This is intentional: DARTS need not recover the exact support of β, but rather concentrate budget on the covariates that most reduce Var(τ̂t ) — a strictly easier and more practically relevant objective. iid
Liang DGP. Each covariate is generated as Xij = (ei + zij )/2 where ei , zij ∼ N (0, 1), then scaled to [0, 1] and mean-centered. Outcomes follow 10X
Yi (0) = 1+Xi,2 + 5 sin(Xi,3 − Xi,4 ) + 2Xi,5 + Xi⊤ β + εi , 2 i,1
with coefficients as in the linear DGP. Treatment and ATE. In both DGPs, Yi (1) = Yi (0) + τ with τ = 4. B.2
Simulation Regimes
Monte Carlo grid (Section 4.2). We use T = 100 batches of n =∈ {100, 1000} units, p ∈ {100, 1000}, budgets B ∈ {1000, 2000, 20, 000}, and equal costs cj = 1 for all j. Both DGPs are evaluated. Results are averaged over 1000 replications with seeds set.seed(s) for s = 1, . . . , 1000. Method comparison (Section 4.2). We use T = 200 batches of n = 1000 units, p = 100, and budget iid B = 2000. Variable costs cj ∼ Uniform(0, 2) and the covariate set for ARMM and MADCovar are redrawn in each replication using the replication seed set.seed(s) for s = 1, . . . , 1000, which also governs the data and assignment randomness; on average, approximately 3 of the 12 selected covariates overlap with the true oracle set. 19
B.3
Hyperparameters and Implementation
Beta p priors are initialized at αj (0) = βj (0) = 1. The shadow-price learning rate is ε = log(p + 1)/B0 where B0 = min(B, T ). Rerandomization uses K = 5,000 candidate assignments per batch; the accepted assignment is chosen by a fair coin flip between the minimizer and its complement. The LASSO penalty is selected by 5-fold cross-validation with the lambda.min rule on control units. If LASSO selects no covariate, all rewards for that batch are set to zero and the prior is unchanged. Covariate matrices with near-singular sample covariance are handled via the Moore–Penrose pseudoinverse (ginv in R). ARMM uses ρ = 0.75 and ridge stabilization λ = 10−4 . MADCovar uses plain OLS nuisance models with pmad = 0.5. B.4
Compute Resources
Single-run analyses and diagnostic plots were conducted locally on an Apple M1 processor (8 cores) running R version 4.5.0. Monte Carlo replications were parallelized across independent jobs on an institutional high-performance computing cluster using bash job submission scripts, with each replication assigned an independent random seed. For the primary method comparison (n = 1000, T = 200, p = 100, K = 5,000 rerandomization candidates), cluster jobs were allocated 1 CPU core each. The algorithm reached an observed peak memory consumption of approximately 1.4GB. Each replication required approximately 35 CPU minutes. To optimize cluster throughput for the extensive Monte Carlo grid, the 1000-replication studies were distributed as an array of 100 jobs, with each job sequentially executing 10 replications on a single CPU core. The heaviest grid settings (n = 1000) peaked at approximately 1.1GB of RAM per job. The smaller grid settings (n = 100, T = 100) required substantially less compute time and had a peak memory footprint of only 250MB. All code was implemented in R using the tidyverse, glmnet, car, and MASS packages. B.5
Inference
Batch-level ATE estimates use Lin’s regression adjustment with HC2 heteroskedasticity-robust standard errors. The cumulative estimate µ̂T is the inverse-variance weighted combination of batchlevel estimates as in equation (1). All reported confidence intervals are fixed-T Wald intervals at the 95% level, µ̂T ± 1.96 σ̂T . B.6 B.6.1
Robustness Checks Heterogeneous Treatment Effects
We modify the method comparison DGP to introduce individual-level treatment effect heterogeneity. The outcome surface for Yi (0) follows the Liang structure on covariates 1–20, while the individual treatment effect is τi = 2 + 3 sin(πXi,6 Xi,7 ) + 4(Xi,8 − 0.5)2 + 2Xi,9 + Xi,10 , a Friedman-style function of covariates 6–10. The true ATE is τ = E[τi ] ≈ 5.6451, with substantial individual-level variation (between 2.64 and 8.28). All other design parameters match the method comparison (T = 200, n = 1000, p = 100, B = 2000, variable costs,set.seed(7)). This setup tests whether DARTS, which uses a reward signal based on Yi (0, t) predictiveness among control units, still concentrates budget on useful covariates when some oracle variables are effect modifiers rather than purely prognostic. Results are presented as a single illustrative run. B.6.2
Oracle-Costly Covariates
We use the Liang DGP with τ = 4 but assign differential costs: oracle covariates (j ≤ 20) have cost cj = 1.1 and noise covariates (j > 20) have cost cj = 0.8, with B = 2000. This tests whether DARTS still concentrates budget on prognostic covariates when doing so is relatively more expensive. All other design parameters match the method comparison (set.seed(7)). In both robustness checks (Figures 4,5), DARTS closes a substantial share of the efficiency gap to Oracle DARTS, consistent with the main method comparison results. 20
(b) Final ATE estimates with 95% CIs at T = 200.
(a) Cumulative regret relative to Oracle DARTS. Single run, set.seed(7).
Figure 4: Heterogeneous treatment effects robustness check. Liang outcome surface on covariates 1–20 with Friedman-style CATE on covariates 6–10; true ATE ≈ 5.645. Single run (n = 1000, T = 200, p = 100, B = 2000, variable costs). ARMM and MADCovar receive randomly preselected covariates; DARTS learns from the full pool. Oracle variants use the true 20 signal covariates. All CIs are 95% valid at fixed T ; HC2 standard errors throughout.
(b) Final ATE estimates with 95% CIs at T = 200.
(a) Cumulative regret relative to Oracle DARTS. Single run, set.seed(7).
Figure 5: Oracle-costly covariates robustness check. Liang DGP with τ = 4; oracle covariates (j ≤ 20) assigned cost cj = 1.1 and noise covariates cost cj = 0.8, rescaled by their maximum with B = 2000 rescaled accordingly. Single run (n = 1000, T = 200, p = 100). ARMM and MADCovar receive randomly pre-selected covariates; DARTS learns from the full pool. Oracle variants use the true 20 signal covariates. All CIs are 95% valid at fixed T ; HC2 standard errors throughout. B.7
Diagnostic Plots
Oracle budget share. Figure 6a shows the median share of batch budget allocated to oracle covariates (covariates 1–20) across 1000 replications of the method comparison, with 95% bands from run percentiles. The share increases monotonically from near-random in early batches toward a stable plateau, confirming progressive concentration of resources on prognostic covariates. Reward–SE correlation. Figure 6b plots batch-level LASSO fractional rewards against the SE ratio σ̂DiM,t /σ̂DARTS,t for each batch after a 30-batch burn-in, pooled across 1000 replications. The positive correlation confirms that batches in which DARTS selects more prognostic covariates — as measured by the reward signal — yield larger reductions in batch-level standard error relative to DiM.
21
(a) Median share of batch budget allocated to oracle covariates (covariates 1–20) across 1000 replications. Solid line shows median; shaded band shows 95% interval from run percentiles.
(b) Batch-level LASSO fractional reward against SE ratio σ̂DiM,t /σ̂DARTS,t after a 30-batch burn-in, pooled across 1000 replications. Each point is one batch from one replication. Correlation coefficient is 0.977.
Figure 6: Diagnostics from the 1000-replication method comparison (Liang DGP, n = 1000, T = 200, p = 100, B = 2000, variable costs). Left: the oracle budget share increases monotonically, confirming progressive concentration of resources on prognostic covariates. Right: the positive reward–SE correlation confirms that the LASSO fractional reward proxy translates learning into precision gains. Table 1: Rerandomisation vs. rerandomisation with adjustment in DARTS designs. DiM and Oracle are benchmark rows; Oracle is separated by a partial rule within each setup block. Bold marks the best non-benchmark value within each setup. RE vs. DiM is the ratio of DiM MSE to the method’s MSE, estimated across 1,000 simulation replications; % Red. is the percentage reduction in MSE relative to DiM. B = total budget, p = covariates, n = units/batch, T = batches. Setup
Method
MSE
RE vs DiM
% Red.
LIANG | B=1000, p=100, n=100, T =100 DiM 0.0302 3×10−4 Rerand (DARTS) 0.0253 −9×10−4 Rerand+Adj (DARTS) 0.0245 −1×10−3
9×10−4 6×10−4 6×10−4
1.0000 1.4216 1.5251
0.00 29.65 34.43
1×10−3
2×10−4
3.8874
74.28
LIANG | B=2000, p=100, n=100, T =100 0.0308 −4×10−4 DiM Rerand (DARTS) 0.0249 −1×10−3 Rerand+Adj (DARTS) 0.0218 −6×10−4
9×10−4 6×10−4 5×10−4
1.0000 1.5355 1.9918
0.00 34.88 49.79
1×10−3
2×10−4
4.0421
75.26
LIANG | B=2000, p=100, n=1000, T =100 DiM 0.0086 −5×10−4 Rerand (DARTS) 0.0064 −1×10−4 Rerand+Adj (DARTS) 0.0057 0
1×10−4 0 0
1.0000 1.7856 2.3177
0.00 43.99 56.85
0
3.6377
72.51
9×10−4 9×10−4 9×10−4
1.0000 1.0144 0.9573
0.00 1.42 −4.46
Oracle
Oracle
Oracle
Emp SD
0.0153
0.0153
0.0045
Bias
1×10−4
LIANG | B=2000, p=1000, n=100, T =100 DiM 0.0296 −0.0003 Rerand (DARTS) 0.0294 0.0004 Rerand+Adj (DARTS) 0.0302 0.0002
Continued on next page.
22
Table 1 continued from previous page. Setup
Method
Emp SD
Bias
MSE
RE vs DiM
% Red.
Oracle
0.0153
−0.0012
2×10−4
3.7321
73.21
LIANG | B=2000, p=1000, n=1000, T =100 DiM 0.0083 4×10−4 Rerand (DARTS) 0.0078 −1×10−4 Rerand+Adj (DARTS) 0.0075 0
1×10−4 1×10−4 1×10−4
1.0000 1.1470 1.2309
0.00 12.81 18.76
0
3.4878
71.32
1×10−4 0 0
1.0000 1.5726 3.2967
0.00 36.41 69.67
0
3.9257
74.53
0.0337 0.0288 0.0290
1.0000 1.1681 1.1606
0.00 14.39 13.84
0.0012
0.0022
15.5885
93.59
LINEAR | B=2000, p=100, n=100, T =100 DiM 0.1871 −0.0010 Rerand (DARTS) 0.1514 0.0087 Rerand+Adj (DARTS) 0.1206 0.0053
0.0350 0.0230 0.0146
1.0000 1.5223 2.3994
0.00 34.31 58.32
−0.0007
0.0020
17.3388
94.23
LINEAR | B=2000, p=100, n=1000, T =100 DiM 0.0581 −0.0014 Rerand (DARTS) 0.0518 0.0011 Rerand+Adj (DARTS) 0.0308 0.0007
0.0034 0.0027 0.0009
1.0000 1.2605 3.5633
0.00 20.55 71.93
0.0006
0.0002
17.5896
94.31
LINEAR | B=2000, p=1000, n=100, T =100 DiM 0.1869 0.0046 Rerand (DARTS) 0.1837 −0.0006 Rerand+Adj (DARTS) 0.1898 0.0018
0.0349 0.0337 0.0360
1.0000 1.0358 0.9701
0.00 3.46 −3.08
−0.0015
0.0023
15.3760
93.50
LINEAR | B=2000, p=1000, n=1000, T =100 DiM 0.0588 −0.0059 Rerand (DARTS) 0.0600 0.0025 Rerand+Adj (DARTS) 0.0529 0.0007
0.0035 0.0036 0.0028
1.0000 0.9677 1.2480
0.00 −3.37 19.84
0.0003
0.0002
16.9663
94.10
Oracle
0
0.0045
LIANG | B=20000, p=100, n=1000, T =100 DiM 0.0089 −2×10−4 Rerand (DARTS) 0.0071 −2×10−4 Rerand+Adj (DARTS) 0.0049 0 Oracle
0.0045
1×10−4
LINEAR | B=1000, p=100, n=100, T =100 0.1835 0.0043 DiM Rerand (DARTS) 0.1698 0.0024 Rerand+Adj (DARTS) 0.1704 0.0026 Oracle
Oracle
Oracle
Oracle
Oracle
0.0465
0.0449
0.0139
0.0477
0.0144
23
Table 2: Performance of rerandomisation with adjustment (DARTS) versus rerandomisation with randomly selected covariates (Random), with DiM and Oracle benchmarks. Oracle is separated by a partial rule within each setup block. Bold marks the best non-benchmark RE vs. DiM within each setup. Median SE is the median estimated standard error across replications; Coverage is the empirical 95% CI coverage rate; RE vs. DiM is the ratio of DiM MSE to the method’s MSE, estimated across 1,000 simulation replications. Setup
Method
Median SE
Coverage
RE vs DiM
LIANG | B=1000, p=100, n=100, T =100 DiM 3×10−4 0.0302 9×10−4 −4 Random 3×10 0.0280 8×10−4 −3 DARTS −1×10 0.0245 6×10−4
0.0294 0.0277 0.0231
0.949 0.946 0.943
1.0000 1.1623 1.5251
1×10−3
2×10−4
0.0146
0.939
3.8874
LIANG | B=2000, p=100, n=100, T =100 DiM −4×10−4 0.0308 9×10−4 −4 Random 2×10 0.0290 8×10−4 −4 DARTS −6×10 0.0218 5×10−4
0.0294 0.0270 0.0208
0.933 0.933 0.939
1.0000 1.1323 1.9918
1×10−3
2×10−4
0.0146
0.937
4.0421
LIANG | B=2000, p=100, n=1000, T =100 DiM −5×10−4 0.0086 1×10−4 −4 Random 1×10 0.0075 1×10−4 DARTS 0 0.0057 0
0.0085 0.0077 0.0057
0.946 0.960 0.956
1.0000 1.3261 2.3177
1×10−4
0.0045
0.948
3.6377
LIANG | B=2000, p=1000, n=100, T =100 DiM −0.0003 0.0296 9×10−4 Random 0.0023 0.0315 1×10−3 DARTS 0.0002 0.0302 9×10−4
0.0294 0.0293 0.0285
0.945 0.939 0.937
1.0000 0.8810 0.9573
2×10−4
0.0146
0.942
3.7321
LIANG | B=2000, p=1000, n=1000, T =100 DiM 4×10−4 0.0083 1×10−4 Random 1×10−4 0.0085 1×10−4 DARTS 0 0.0075 1×10−4
0.0085 0.0083 0.0076
0.959 0.944 0.954
1.0000 0.9525 1.2309
0
0.0045
0.953
3.4878
LIANG | B=20000, p=100, n=1000, T =100 DiM −2×10−4 0.0089 1×10−4 −4 Random −3×10 0.0045 0 DARTS 0 0.0049 0
0.0085 0.0047 0.0047
0.948 0.953 0.940
1.0000 3.8763 3.2967
1×10−4
0.0045
0.952
3.9257
0.1824 0.1736 0.1610
0.952 0.957 0.944
1.0000 1.1413 1.1606
0.0022
0.0442
0.934
15.5885
LINEAR | B=2000, p=100, n=100, T =100 DiM −0.0010 0.1871 0.0350 Random −0.0095 0.1726 0.0298
0.1824 0.1677
0.944 0.945
1.0000 1.1714
Oracle
Oracle
Oracle
Oracle
Oracle
Oracle
Bias
−0.0012
Emp SD
0.0153
0.0153
0.0045
0.0153
0.0045
0.0045
MSE
0
0
0
LINEAR | B=1000, p=100, n=100, T =100 DiM 0.0043 0.1835 0.0337 Random −0.0008 0.1718 0.0295 DARTS 0.0026 0.1704 0.0290 Oracle
0.0012
0.0465
Continued on next page.
24
Table 2 continued from previous page. Setup
Method
Bias
Emp SD
MSE
Median SE
Coverage
RE vs DiM
DARTS
0.0053
0.1206
0.0146
0.1142
0.937
2.3994
Oracle
−0.0007
0.0449
0.0020
0.0442
0.953
17.3388
LINEAR | B=2000, p=100, n=1000, T =100 DiM −0.0014 0.0581 0.0034 Random 0.0011 0.0518 0.0027 DARTS 0.0007 0.0308 0.0009
0.0583 0.0583 0.0583
0.953 0.977 0.998
1.0000 1.2605 3.5633
0.0006
0.0002
0.0142
0.956
17.5896
LINEAR | B=2000, p=1000, n=100, T =100 0.0046 0.1869 0.0349 DiM Random 0.0019 0.1895 0.0359 DARTS 0.0018 0.1898 0.0360
0.1825 0.1852 0.1846
0.944 0.953 0.943
1.0000 0.9736 0.9701
−0.0015
0.0023
0.0446
0.931
15.3760
LINEAR | B=2000, p=1000, n=1000, T =100 DiM −0.0059 0.0588 0.0035 Random 0.0025 0.0600 0.0036 DARTS 0.0007 0.0529 0.0028
0.0583 0.0583 0.0583
0.946 0.939 0.964
1.0000 0.9677 1.2480
0.0003
0.0143
0.953
16.9663
Oracle
Oracle
Oracle
0.0139
0.0477
0.0144
0.0002
Table 3: Performance at T =200 batches of n=1000 units, p=100 covariates, B=2000, variable costs iid cj ∼ Uniform(0, 2). True ATE = 4. Summary across 1000 MC replications. DiM is the baseline benchmark; oracle methods use the true 20 signal covariates and are separated by a partial rule. Bold marks the best non-benchmark, non-oracle value. Rel. RMSE = RMSEDiM /RMSEmethod ; values >1 indicate improvement over DiM. CI Width = 2 × 1.96 × Med. SE. All CIs are 95% (valid at fixed T ). Method
Mean ATE
Bias
RMSE
Med. SE
CI Width
Rel. RMSE
−4
DiM ARMM MADCovar DARTS
4.0003 4.0000 4.0000 3.9998
+3×10 0 0 −2×10−4
0.0058 0.0057 0.0056 0.0041
0.0059 0.0056 0.0056 0.0041
0.0230 0.0221 0.0219 0.0162
1.000 1.018 1.036 1.415
Oracle DARTS Oracle MADCovar Oracle ARMM
3.9999 4.0001 3.9999
−1×10−4 +1×10−4 −1×10−4
0.0032 0.0032 0.0032
0.0032 0.0031 0.0032
0.0125 0.0122 0.0125
1.812 1.812 1.812
25
Table 4: Final ATE estimates after T =200 batches of n=1000 units, p=100 covariates, B=2000, iid variable costs cj ∼ Uniform(0, 2). True ATE = 4. Single run, set.seed(7). DiM is the baseline benchmark; oracle methods use the true 20 signal covariates and are separated by a partial rule. Bold marks the best non-benchmark, non-oracle value. SE is the HC2 posterior standard error; CI Width is 2 × 1.96 × SE; Rel. SE = σ̂DiM /σ̂method (values >1 indicate lower estimated variance than DiM). All CIs are 95% (valid at fixed T ). Method
ATE
Bias
SE
CI Width
Rel. SE
−2
DiM ARMM MADCovar DARTS
4.01222 4.00945 3.99432 4.00322
+1×10 +9×10−3 −6×10−3 +3×10−3
0.00579 0.00557 0.00550 0.00393
0.02270 0.02181 0.02154 0.01539
1.0000 1.0404 1.0535 1.4743
Oracle DARTS Oracle ARMM Oracle MADCovar
4.00319 4.00458 4.00236
−2×10−3 +5×10−3 +2×10−3
0.00319 0.00318 0.00311
0.01249 0.01246 0.01220
1.8175 1.8221 1.8598
26