Budgeted Multi-Source Counterfactual Annotation for Off-Policy Evaluation

AI summary

Overview

Research area: Off-policy evaluation (OPE) in contextual bandits, specifically the statistical design of data acquisition — budgeted collection of counterfactual annotations from heterogeneous, imperfect sources (domain experts, LLMs) — combined with discrete optimization theory.

Technical level: Advanced. The paper combines unbiasedness and MSE decomposition proofs, threshold and regime analysis of a local objective, an integer program with a majorization-minimization (MM) solver, dynamic programming, and a sensitivity bound for estimated source profiles. Readers need comfort with estimator variance decomposition and discrete optimization.

Scope (one sentence): Given a fixed budget, a set of annotation sources with different costs and error profiles, and logged bandit data, the paper asks which context-action pairs to annotate and from which sources so as to minimize the part of the OPE estimator's mean squared error that the annotation plan controls.

What This Paper Is About

Off-policy evaluation estimates how well a new policy would perform using only data logged by a different (behavior) policy. When the target policy wants to take actions the logged data barely covers, standard reweighting becomes high-variance and reward models must extrapolate. Counterfactual annotations — a clinician saying what would likely have happened under an alternative treatment, or an LLM predicting a student's response to a different activity — can fill those gaps, but experts are expensive and models are noisy, biased, or cheap in ways that vary by context-action pair. The paper's goal is to decide, under a hard budget, exactly which annotations to buy from which source.

Key Contributions

1. Annotation-value regimes. The authors derive a first-annotation threshold — the exact condition under which buying a single annotation does not increase the local objective — and a four-regime characterization of when annotations help immediately, help only up to a finite amount, help only after a sufficiently large batch, or do not help locally.

2. Budgeted allocation algorithm. They show the allocation problem is coupled across actions within a context, nonconvex, and neither submodular nor supermodular, then develop a majorization-minimization (MM) algorithm with dynamic-programming subroutines that monotonically improves the original objective.

3. Sensitivity to estimated source profiles. A second theorem translates errors in estimated annotation bias and excess variance, plus the solver's optimization error, into a bound on excess OPE mean squared error, separating the gain from better profile estimates from the gain from a better solver.

4. Empirical validation. Comparisons against greedy allocation and certified solver benchmarks, profile-estimation sensitivity, and validation of the final OPE estimator on finite pilots and independent evaluation samples, across synthetic clinical and LLM-annotated education bandits.

Main Findings

• Unbiasedness holds, but annotation bias still costs. The annotation-augmented estimator (called DM⁺-IS in the paper) is unbiased under the stated sample design; annotations enter only the reward estimate, not the importance-sampling correction. Annotation bias therefore does not shift the estimate, but it enters the variance through the term J₁ − J₂.

• MSE decomposes cleanly. The total mean squared error equals C₀ + J(n)/N, where C₀ does not depend on the allocation and J(n) = J₁(n) − J₂(n) + J₃(n) does. Because C₀ is allocation-invariant, minimizing J exactly minimizes fixed-profile MSE. J₁ captures accumulated reward-model bias, J₂ is the cross-action coupling term within a context, and J₃ captures reward-model variance.

• A first annotation is not always worth buying. The first annotation, going from n = 0 to n = 1, does not increase the local objective if and only if the squared bias is at most η times ((n_f + 1)σ² / n_f − Δ), where η = 1 − π_b(a|s). The same source can be worth buying at one context-action pair and not another, because its bias is judged against the variance reduction available at that pair. If Δ ≥ (n_f + 1)σ²/n_f, the first annotation strictly hurts for every pair with nonzero bias.

• Four regimes of annotation value. The sign of the local objective's derivative is governed by an affine term, producing: (I) annotations locally harmful (monotonically increasing); (II) decrease then increase, with a unique continuous minimizer; (III) monotonically decreasing, so keep adding; (IV) increase then decrease, with a unique continuous maximizer — where a strict improvement over no annotation is possible only if ε² < ησ²/n_f. Regime II means spending the whole available budget on one pair can be counterproductive.

• Multiple sources can beat any single source. The multi-source optimum is never worse than the best single-source optimum, and mixing two sources with opposite-signed biases at a fixed total annotation count can shrink the reward-model bias below what either source achieves alone. Gains over the best single source can be modest when one source dominates and mixing offers little cancellation.

• Cross-action coupling changes the threshold. When the aggregate reward-model bias in a context is zero and α₃ > 0, including J₂ relaxes the first-annotation threshold to ε² ≤ (n_f + 1)σ²/n_f − Δ, removing the η factor and permitting larger admissible bias.

• Headline empirical reductions. In a synthetic clinical bandit with 100 contexts, 5 actions, M = 1,000 fitting records and budget κ = 10,000, the four-source MM allocation reduced scaled MSE from 𝒱(0) = 84.23 to 66.90, a 20.58% reduction relative to no annotation. The abstract additionally reports a 10.77% reduction for the LLM-annotated education bandit.

• Regimes matched in simulation. Across 100 independently sampled environments, Source 1 decreased monotonically (Regime III), Source 2 first increased then decreased (Regime IV), Source 3 decreased then increased (Regime II), and Source 4 increased monotonically (Regime I). The decomposition showed variance reduction (J₃) competing with bias accumulation (J₁) and correction (J₂).

• Source-subset sweep. Evaluating all 2^K − 1 = 15 nonempty subsets of the four sources, the four-source setting achieved the lowest scaled MSE among the evaluated subsets.

• Runtime. One MM update costs O(T_Bis |S| |A| (K² N_max T_swap + N_max²)), with the K² N_max T_swap term from local source search and N_max² from dynamic-programming aggregation; the bound grows linearly in the number of context-action pairs.

Methodology in Plain English

The setup is a contextual bandit with finite contexts and actions. Two separate datasets are used: one for fitting the reward model, one for evaluation. Logged data only reveal rewards for actions the behavior policy actually took, so the authors add annotations from K sources that estimate what the reward would have been for other actions. Each source has a known cost per annotation, a bias, and an excess variance relative to the true reward noise.

The annotations are folded into a single averaged reward estimate per context-action pair, combining factual observations and annotations. Crucially, this annotated reward estimate is used only inside the reward model, not in the importance-weighting step, which is what keeps the final estimator unbiased even when the annotations are biased. The authors then work out how the estimator's mean squared error splits into a part that does not depend on the annotation plan and a part that does.

That allocation-dependent part is the objective. The authors analyze it in stages. First, one pair with one source, to derive the threshold for buying a first annotation and the four behavioral regimes — essentially asking whether the bias the annotation adds outweighs the variance it removes. Then they scale up, showing the full problem is coupled across actions in the same context, nonconvex, and not submodular, so local per-pair rules cannot be applied independently.

To solve it, they use majorization-minimization: at each iteration they replace the one coupling term (a squared aggregate bias) with a tight affine upper bound that is exact at the current allocation, which makes the surrogate separable by context-action pair. Given the surrogate, they search a price λ on the budget by bisection; for each λ, they solve a local source-mixing problem per pair using pairwise-exchange local search, then use dynamic programming to decide how many annotations each pair gets within each context, respecting per-context and per-pair capacities. Keeping the current allocation as a fallback guarantees the objective never increases.

Experiments are run on a synthetic clinical bandit and a semi-synthetic education bandit built from real student interaction data, comparing against greedy allocation and certified exact solvers, and testing how sensitive the final result is to errors in the estimated source profiles.

Why This Matters

Impact on research. Most OPE work assumes the data (or annotations) are already in hand and focuses on the estimator. This paper moves upstream and formalizes the acquisition stage, giving a statistical account of when a counterfactual annotation is worth its price and a discrete-optimization method for spending a fixed budget across heterogeneous sources. It extends prior single-source annotation work to multiple sources and provides an explicit bound for the realistic case where source quality is estimated from a small pilot rather than known.

Real-world applications:

• Healthcare. A clinician assesses what would likely have happened under an alternative treatment, letting a new treatment policy be evaluated without deploying it — the paper's synthetic clinical bandit is modeled on this setting.
• Education technology. An automated tutor estimates how a student would respond to a different instructional activity; the paper builds its education bandit from 6.1 million ASSISTments student–problem interactions, with 100 context clusters and 10 problem-type actions, using LLM-predicted first-response times converted to reward.
• Annotation pipeline budgeting. Teams that buy labels from a mix of expensive human experts and cheap LLM annotators can use the regime analysis to decide when a second opinion from an opposite-biased source is worth more than more samples from the same one.
• Any high-stakes logged-data decision system — recommendation, pricing, or operations — where exploring an untaken action is expensive or risky, and where the target policy assigns weight to poorly covered actions.

Industry relevance. Counterfactual annotation via LLMs is now cheap enough to buy at scale, which makes the budget-allocation question commercially real. The paper's approach is designed for practical constraints: per-sample annotation limits, heterogeneous per-source costs, and source profiles estimated from a pilot rather than known exactly — and it explicitly separates the benefit of better profile estimates from the benefit of a better solver, which is exactly the tradeoff an engineering team faces.

Future Directions

• Repeated annotations per sample. The formulation enforces a one-annotation-per-factual-sample protocol via the capacity constraints, which the authors note can be relaxed if repeated annotations are allowed. The allocation algorithm and the regime analysis would need to be revisited under that relaxation.
• Beyond the DM⁺-IS estimator. The objective is derived specifically for the annotation-augmented doubly robust estimator studied here. Whether analogous thresholds and allocation algorithms carry over to other OPE estimators that use annotations differently is open.
• Tighter solvers and scalability. The local source-mixing step uses pairwise-exchange local search with a proven approximation bound rather than exact enumeration, because enumeration grows combinatorially in K. Better or exact methods for the per-pair source mixture, and for many sources, are natural next steps.
• Robust allocation under profile uncertainty. Theorem 2 bounds MSE loss from estimated bias and excess variance. Designing allocations that are robust to adversarial or heavily misspecified profiles — rather than just quantifying the damage — is a further step the paper's sensitivity analysis invites.

Target Audience

Researchers and graduate students working on off-policy evaluation, contextual bandits, and causal inference from logged data; optimization researchers interested in nonconvex, non-submodular integer allocation with practical structure; and applied scientists or ML engineers in healthcare, education technology, or recommendation systems who need to decide how to spend a finite annotation budget across a mix of human and LLM labelers. Readers unfamiliar with importance sampling, doubly robust estimation, or majorization-minimization will need background reading, since the paper moves quickly into variance decompositions and threshold conditions.

Authors’ abstract

Off-policy evaluation (OPE) estimates the value of a target policy from logged data, but limited behavior-policy coverage can force high-variance reweighting or reward-model extrapolation. Counterfactual annotations can add evidence about unobserved actions, yet practical sources, including domain experts and large language models (LLMs), may be costly, biased, or noisy. We study budgeted acquisition of such annotations for contextual-bandit OPE. Given source-specific costs and error profiles, we formulate an integer allocation problem over context-action pairs and annotation sources to minimize the component of estimator variance that depends on the annotation plan. We characterize when annotations are valuable through a first-annotation threshold and local annotation-value regimes. For the coupled multi-source problem, we develop a majorization-minimization algorithm with dynamic-programming subroutines that monotonically improves the objective. Experiments in synthetic clinical and LLM-annotated education bandits show that our allocation method reduces fixed-profile mean squared error (MSE) by 20.58% and 10.77%, respectively, relative to no annotation.