Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
65,550 characters · 19 sections · 75 citation commands
Optimal treatment assignment rules under capacity constraints
\address{Department of Economics, University of Rochester, Rochester, NY 14627, USA.} \email{[email removed]} \address{Department of Economics, University of Rochester, Rochester, NY 14627, USA.} \email{[email removed]}
When a social planner allocates goods such as vaccines or school vouchers, two key challenges often arise. First, the planner typically lacks knowledge of the true treatment effects. Second, even when treatment is clearly beneficial, capacity constraints---such as limited budgets or supplies---may prevent treating everyone. While many studies have addressed the first challenge, the second has received comparatively less attention. This paper studies a treatment assignment problem under such capacity constraints using Le Cam's limits of experiments framework.
We consider a planner who aims to maximize social welfare by designing treatment rules based on observable covariates subject to a capacity constraint. In unconstrained settings, the asymptotically optimal assignment rule can be determined pointwise: one simply treats individuals for whom the estimated (conditional) treatment effect is positive hirano2009asymptotics. However, under a capacity constraint, such pointwise rules may violate the overall treatment quota. A natural alternative is a plug-in rule that ranks individuals by their estimated effects and assigns treatment in descending order until the quota is met. But is this rule still optimal in the presence of constraints? A key challenge in answering this question stems from the need to coordinate treatment probabilities across the entire covariate distribution $F_{X}$, which makes it difficult to obtain an asymptotic representation of a treatment assignment rule.
To address this, we reformulate the planner's constrained maximization problem as an optimal transport problem. To illustrate, consider a binary treatment setting where a fixed fraction $p$ of the population can be treated. If the constraint binds exactly, then the distribution of treatment assignments must be a Bernoulli distribution with success probability $p$. The planner’s problem can then be framed as transporting the mass of $F_{X}$ into the Bernoulli distribution $F_{T}$, in a way that maximizes the social welfare. Since solutions to this optimal transport problem automatically satisfy the capacity constraints, this reformulation simplifies the analysis. This approach is also computationally attractive, as one can use existing software for optimal transport. In particular, when both $F_{X}$ and $F_{T}$ are discrete (or discretized), which is often the case in practice, the optimal transport problem reduces to a simple linear programming problem.
Based on this reformulation, we analyze the optimality of two canonical assignment rules: the plug-in rule and the Bayesian rule. The plug-in rule is known to be average optimal under point-identified models when the planner's utility function is smooth hirano2009asymptotics. The Bayesian rule is known to be average optimal under partially identified models when the planner's utility function is directionally differentiable christensen2023optimal.\footnote{Formally, christensen2023optimal adopt an optimality criterion that is weaker than average optimality. Under additional assumptions their results imply average optimality as well.} We show that both rules remain average optimal under capacity constraints, provided that the planner's utility function is smooth.
A meaningful distinction between the two rules arises when the planner's utility function is only directionally differentiable. This scenario is not merely a theoretical curiosity but arises in important practical settings. For instance, it occurs when the (conditional) potential outcome distributions may differ slightly between the target and training populations, and the planner adopts a utility function that is robust to such distributional shifts adjaho2023externallyvalidpolicychoice. In this setting, we show that the plug-in rule is generally not optimal, whereas the Bayesian rule remains optimal. This result echoes the findings of christensen2023optimal, where the lack of point identification similarly precludes full differentiability and renders the plug-in rule suboptimal.
Our simulation study supports these theoretical results. Specifically: (i) the Bayesian rule achieves a substantially lower risk than the plug-in rule in small samples for both the smooth and the directionally differentiable utility functions, (ii) the two rules behave similarly under the smooth utility function in larger samples, and (iii) the Bayesian rule continues to outperform the plug-in rule under the directionally differentiable utility function in larger samples.
We illustrate our methods using data from angrist2006long, who study the impact of receiving a randomly assigned voucher (which allows students to attend private secondary schools) on educational attainment seven years later. We hypothetically treat the marginal distribution of covariates (age and sex) in the observed sample as that of the target population, and compute both the plug-in rule and the Bayesian rule. The two rules produce identical allocations under the smooth utility function, but they produce different allocations under the directionally differentiable utility function.
This study builds on the literature on the statistical treatment assignment problems in econometrics, where the pioneering works include manski2004statistical and dehejia2005program. Within this expanding literature, our main contribution is to provide a decision-theoretic optimality result under capacity constraints on available treatments. Previous studies have established the decision-theoretic optimality of treatment rules in several settings including: (i) point-identified smooth (semi-)parametric models under local asymptotics hirano2009asymptotics,masten2023minimax, (ii) partially-identified smooth (semi-)parametric models under local asymptotics christensen2023optimal,kido2023locally,xu2024jmp, (iii) point-identified models with finite samples stoye2009minimax,stoye2012minimax,tetenov2012statistical,guggenberger2024minimax,kitagawa2024treatmentchoicenonlinearregret,chen2025note, (iv) partially-identified models with finite samples manski2007minimax,stoye2012minimax,yata2023optimal,ishihara2021evidence,fernandez2024robust,oleaDecisionTheoryTreatment2024. However, none of these studies consider the capacity constraints in the way we do. To the best of our knowledge, this is the first study to establish an optimality result under such constraints within the framework of hirano2009asymptotics, extended to cover non-binary treatments---both discrete and continuous.
Besides hirano2009asymptotics, the most closely related paper is christensen2023optimal. They extend hirano2009asymptotics to allow for partially identified parameters and (non-randomized) discrete actions, and show the asymptotic optimality of the Bayesian rule. Our setting differs in that the action space is the set of couplings of $F_{X}$ and $F_{T}$ (see notation below) equipped with the Wasserstein distance, which naturally allows for randomized assignment rules. This non-standard action space requires nontrivial extensions to analyze the asymptotic properties of the Bayesian rule, which we address by drawing on tools from optimal transport. In contrast to their framework, we focus on point-identified parameters.\footnote{xu2024jmp further extends the framework of christensen2023optimal to continuous decision problems using an expansion-based approach.}
There also exist studies that incorporate exogenously given constraints into the treatment assignment problem. bhattacharya2012inferring impose the capacity constraints in the way we do, but focus on the estimation and inference of a nonparametric plug-in rule. Other papers adopting the empirical welfare maximization approach allow for various types of constraints, including the capacity constraints kitagawa2018should,athey2021policy,mbakop2021model,sun2024empiricalwelfaremaximizationconstraints. kitagawa2018should show the optimality of their proposed rule in terms of the welfare convergence rate, which measures how quickly the average welfare achieved by the proposed rule converges to the maximum welfare under the true data generating process.
Some recent works utilize tools from optimal transport theory in the literature of treatment assignment problems. kidoDistributionallyRobustPolicy2022 and adjaho2023externallyvalidpolicychoice study the external validity of treatment choices by measuring the difference in potential outcome distributions between the training and target populations using the Wasserstein distance. hazard2025Who formulate a learning problem of optimal matching policies in a two-sided market as an empirical optimal transport problem, and derive a welfare regret bound for their estimated policy.
Our work also relates to the growing field of statistical methods for optimal transport problems chewi2024statisticaloptimaltransport, as we study the local asymptotic properties of transport maps of an optimal transport problem where the cost function is indexed by parameters that can be efficiently estimated.
The remainder of the paper is organized as follows. Section (ref) formulates the planner's problem and introduces the data generating process. Section (ref) introduces the decision theoretic framework and define the plug-in rule and the Bayesian rule. Then the optimality results are stated. Section (ref) provides a simulation study to evaluate the finite sample performance of rules. Section (ref) illustrates our methods using the data from angrist2006long. Finally, Section (ref) concludes. All of the proofs are relegated to Appendix.
A function $f:\Theta\subset\mathbb{R}^{k}\to\mathbb{R}$ is (Hadamard) directionally differentiable at $\theta_{0}$ if there is a continuous function $\dot{f}_{\theta_{0}}:\mathbb{R}^{k}\to\mathbb{R}$ such that \[ \lim_{n\to\infty}\left|\frac{f(\theta_{0}+t_{n}h_{n})-f(\theta_{0})}{t_{n}}-\dot{f}_{\theta_{0}}(h)\right|=0 \] for all sequences $\left\{ t_{n}\right\} \subset\mathbb{R}_{+}$ and $\left\{ h_{n}\right\} \subset\mathbb{R}^{k}$ such that $t_{n}\downarrow0$, $h_{n}\to h\in\mathbb{R}^{k}$ as $n\to\infty$ and $\theta_{0}+t_{n}h_{n}\in\Theta$ for all $n$. It is worth noting that this requires $\dot{f}_{\theta_{0}}$ need to be continuous, but not to be linear.
Let $P$ and $Q$ be Borel probability measures on $\mathcal{A}$ and $\mathcal{B}$, respectively. A joint distribution $\mu$ on $\mathcal{A}\times\mathcal{B}$ is called a coupling of $P$ and $Q$ if its marginals are $P$ and $Q$; that is, $\mu(A\times\mathcal{B})=P(A)$ and $\mu(\mathcal{A}\times B)=Q(B)$ for any measurable sets $A$ and $B$.
The setting of this paper closely follows that of hirano2009asymptotics. We consider a social planner who assigns a treatment $T$ to individuals based on their observable covariates $X$. Let $F_{X}$ denote the marginal distribution of $X$ in the target population, with support $\mathcal{X}$. We assume that $F_{X}$ is known to the planner. The planner can fractionally (probabilistically) assign treatment $T=t$ to an individual with covariate $X=x$.
Let $Y(t)$ denote potential outcomes under treatment $T=t$. In contrast to hirano2009asymptotics, we distinguish between the conditional potential outcome distribution in the target population and in the training population. We denote the conditional distribution of $Y(t)$ in the training population as $F_{t}(\cdot|x,\theta)$, where $F_{t}(\cdot|x,\theta)$ belongs to families of distributions indexed by a parameter $\theta\in\Theta\subset\mathbb{R}^{k}$. The planner must learn $\theta$ from the available data from experimental or observational studies.
The planner's utility for assigning treatment $T=t$ to an individual with covariate $X=x$ depends on the conditional distribution $F_{t}(\cdot|x,\theta)$ via a functional $w$. For the shorthand notation, we write \[ w(\theta,x,t):=w(F_{t}(\cdot|x,\theta)). \]
We consider two scenarios. First, $w(\theta,x,t)$ is fully differentiable in $\theta$. Second, $w(\theta,x,t)$ is only directionally differentiable in $\theta$. Two examples corresponding to each scenario are given as follows.
We consider a setting where the planner faces the capacity constraints on the available treatments. To illustrate, consider a simple binary treatment case, $T\in\mathcal{T}:=\left\{ 0,1\right\} $, where a fraction $p$ of the target population is to be treated. We assume that the capacity constraint binds exactly. Let $\mu(t|x)$ denote the conditional probability of assigning treatment $T=t$ to individuals with $X=x$. Then, under the capacity constraint, the planner's problem can be written as
We now show how to convert this constrained optimization problem into an unconstrained one. Observe that the distribution of treatment assignments must be a Bernoulli distribution with the success probability $p$. With this in mind, the planner’s problem can be seen as an optimal transport problem: the planner transports the mass of $F_{X}$ into $F_{T}$, the Bernoulli distribution, in a way that maximizes the social welfare.
Formally, let $\mathcal{M}_{a}$ be the set of all couplings of $F_{X}$ and $F_{T}$. Let $d$ be a distance function on $\mathcal{X}\times\mathcal{T}$, and define the Wasserstein distance of order 1 as
where $\Gamma(\mu,\nu)$ is the set of couplings whose marginals are $\mu$ and $\nu$. We focus on couplings that have a finite first moment: \[ \mathcal{M}:=\left\{ \mu\in\mathcal{M}_{a}:\int d((x_{0},t_{0}),(x,t))\mathrm{d}\mu<+\infty\right\} , \] for some arbitrary $(x_{0},t_{0})\in\mathcal{X}\times\mathcal{T}$. Then $(\mathcal{M},d_{W})$ becomes a metric space, and $d_{W}$ is finite on $\mathcal{M}$ villani2009optimal. Using this setup, the original constrained problem ((ref)) is equivalent to the following:
where
Hence, the action space of the planner is the space of couplings $(\mathcal{M},d_{W})$. We remark that $\mathrm{d}\mu(t|x)$ becomes a conditional probability measure when $F_{T}$ is continuous.
There are three important remarks regarding this optimal transport formulation. First, the capacity constraint is automatically satisfied by any coupling in $\mathcal{M}$, making the problem effectively unconstrained. Second, this reformulation is computationally attractive, as one can use an existing software for optimal transport. In particular, when both $F_{X}$ and $F_{T}$ are discrete---which is often the case in practice---the optimal transport problem reduces to a simple linear programming problem. Finally, this approach can easily accommodate non-binary treatment settings. We assume that $T$ follows a distribution $F_{T}$, determined by the capacity constraints, with support $\mathcal{T}$.
When the true (finite-dimensional) parameter $\theta_{0}$ is known to the planner, the optimal rules can be obtained by solving \[ \max_{\mu\in\mathcal{M}}W(\theta_{0},\mu). \] However, since the planner does not know $\theta_{0}$ in practice, she must select a rule in a data-driven manner. For this purpose, data $Z^{n}=(Z_{1},\dots,Z_{n})$, which are informative about $\theta$ (and hence about $F_{t}(\cdot|x,\theta)$) are available from a training population. We assume that the data $Z^{n}$ are i.i.d. with $Z_{i}\sim P_{\theta}$ on some space $\mathcal{Z}$ equipped with the Borel $\sigma$-algebra $\mathcal{B}(\mathcal{Z})$. We let $P_{\theta}^{n}$ denote the joint probability measure of $Z^{n}$. In Appendix (ref), we consider an extension in which the sampling distribution of data may depend on (possibly infinite-dimensional) nuisance parameters, as in a GMM model.
Following hirano2009asymptotics and among others, we use a local asymptotic framework where we perturb the data-generating process around the true one. Let $\Theta$ be an open subset of $\mathbb{R}^{k}$ and suppose that $\theta_{0}$ is the true parameter. Let $\theta_{nh}:=\theta_{0}+h/\sqrt{n}$. We assume that the sequence of experiments $\mathcal{E}_{n}=\left\{ P_{\theta}^{n}:\theta\in\Theta\right\} $ satisfies differentiability in quadratic mean (DQM) at $\theta_{0}$: there exists a function $s:\mathcal{Z}^{n}\to\mathbb{R}^{k}$ such that
where $s$ is the score function associated with $\mathcal{E}_{1}$. Let $I_{0}=\mathbb{E}_{P_{\theta_{0}}^{n}}[ss^{\prime}]$.
The planner's statistical treatment assignment rule (or just rule) $\mathcal{\mu}:\mathcal{Z}^{n}\to\mathcal{M}$ maps realizations of data into the coupling. Let \[ A_{0}:=\arg\max_{\mu\in\mathcal{M}}W(\theta_{0},\mu) \] be the set of couplings that maximize the welfare at $\theta_{0}$. We define the class of sequences of rules by
where $\stackrel{h}{\rightsquigarrow}$ denotes convergence in distribution along $P_{\theta_{nh}}^{n}$ with $Z^{n}\sim P_{\theta_{nh}}^{n}$ for each $n$, and $Q_{\theta_{0},h}$ is a (possibly degenerate) probability measure on $\mathcal{M}$. For technical reasons, we restrict our analysis to rules satisfying $\sqrt{n}P_{\theta_{nh}}^{n}\left(\mu_{n}(Z^{n})\notin A_{0}\right)\to0$, a condition also imposed by christensen2023optimal. This condition ensures that the treatment rule maximizes the welfare at the true parameter $\theta_{0}$ with high probability in a neighborhood of $\theta_{0}$, and that the probability of selecting a suboptimal coupling (i.e., $\mu_{n}\not\in A_{0}$) vanishes sufficiently fast.
Since $(\mathcal{M},d_{W})$ is compact (and thus complete and separable) by villani2009optimal, we obtain the following asymptotic representation theorem by van1991asymptotic.
This proposition states that any sequence $\left\{ \mu_{n}\right\} $ in $\mathcal{D}$ can be matched by some treatment rule $\mu_{\infty}$ in a limit experiment where we observe a single draw $\Delta$ from a shifted normal distribution and an independent uniform random variable $U$. This representation is useful for analyzing the asymptotic properties of rules, as the limit experiment is more analytically tractable than the original experiments $\mathcal{E}_{n}$.
We begin by introducing a decision theoretic framework to evaluate the performance of a sequence of rules $\left\{ \mu_{n}\right\} \in\mathcal{D}$. Let \[ W_{\mathcal{M}}^{*}(\theta):=\max_{\mu^{\prime}\in\mathcal{M}}W(\theta,\mu^{\prime}) \] denote the maximum attainable welfare at $\theta$. Following the literature, we define the welfare regret $W_{\mathcal{M}}^{*}(\theta)-W(\theta,\mu)$ as the loss incurred from choosing a coupling $\mu\in\mathcal{M}$. Accordingly, the risk associated with the map $Z^{n}\mapsto\mu(Z^{n})\in\mathcal{M}$ at $\theta$ is given by
where the expectation is taken with respect to the sampling distribution $P_{\theta}^{n}$ of $Z^{n}$. The planner's objective is to minimize the risk by constructing data-driven rules $\left\{ \mu_{n}\right\} \in\mathcal{D}$.
Let $\pi$ be any prior density function on $\Theta$ that is continuous and positive at $\theta_{0}$. A sequence of rules $\left\{ \mu_{n}^{*}\right\} \in\mathcal{D}$ is said to be average optimal if $\left\{ \mu_{n}^{*}\right\} $ attains the infimum of the asymptotic risk function:
Our main goal is to construct a sequence of rules that is average optimal. A natural candidate is the plug-in rule, which has been shown to be average optimal by hirano2009asymptotics in settings without capacity constraints. To formalize this rule, let $\hat{\theta}_{n}$ be a best regular estimator of $\theta_{0}$ such that \[ \sqrt{n}(\hat{\theta}_{n}-\theta_{nh})\stackrel{h}{\rightsquigarrow}N(0,I_{0}^{-1})\quad\text{as }n\to\infty. \] The maximum likelihood estimator or the Bayesian posterior mean estimator are typical examples of best regular estimators. The plug-in rule is the sequence $\{\mu_{n}^{P}(Z^{n})\}$, where for each $n$, \[ \mu_{n}^{P}(Z^{n})\in\arg\max_{\mu\in\mathcal{M}}W(\hat{\theta}_{n},\mu). \]
Another rule we examine is the Bayesian rule $\{\mu_{n}^{B}(Z^{n})\}$, defined for each $n$ by
where $\pi_{n}(\theta|Z^{n})$ is the posterior density obtained from a strictly positive, continuous prior density $\pi$ on $\Theta$. christensen2023optimal show the average optimality of such Bayesian rules in discrete choice problems when (i) the model is partially identified, and (ii) decision rules are not fractional.
We provide an optimality result for a Bayesian rule under the directional differentiability of $w$. We discuss asymptotic properties of the plug-in rule in Subsection (ref). We impose the following assumptions for results in this subsection. As in clarke2002information and christensen2023optimal, we say that a family $\mathcal{P}$ is locally quadratic if for any $\theta_{0}\in\Theta$, $D_{\mathrm{KL}}(p_{\theta}\mathrel{\Vert}p_{\theta^{\prime}})\le2(\theta-\theta^{\prime})^{\top}I_{0}(\theta-\theta^{\prime})$ holds for any $\theta$ and $\theta^{\prime}$ belonging to a neighborhood of $\theta_{0}$, where $D_{\mathrm{KL}}(p_{\theta}\mathrel{\Vert}p_{\theta^{\prime}})$ is the Kullback-Leibler divergence with respect to a common dominating measure $\nu$. Also, we say $\mathcal{P}$ is sound if weak convergence of $P_{\theta}$ to $P_{\theta^{\prime}}$ is equivalent to the convergence of $\theta$ to $\theta^{\prime}$ for probability measures $P_{\theta},P_{\theta^{\prime}}$ and parameters $\theta,\theta^{\prime}\in\Theta$.
The first three conditions are standard assumptions in local asymptotic frameworks. Conditions (iv) and (v) ensure that Schwartz’s theorem---which originally establishes posterior consistency in a space of density functions (see, for example, ghosh2003Bayesian and ghosal2017fundamentals)---can be applied in a parametrized setting, as in clarke2002information. We impose these assumptions to show that the Bayesian rule $\mu_{n}^{B}$ satisfies $\sqrt{n}P_{\theta_{nh}}^{n}\left(\mu_{n}^{B}(Z^{n})\notin A_{0}\right)\to0$. The same conditions are also imposed by christensen2023optimal.
For example, this condition is satisfied if $\mathcal{X}$ is a compact metric space, and $\mathcal{T}$ is a finite discrete space.
Note that discrete covariates are compatible with condition (ii), since the metric $d$ can be defined to incorporate the discrete metric.
Condition (i) imposes a uniform version of directional differentiability, rather than requiring it only pointwise in $(x,t)$. For condition (iii), a similar polynomial growth condition appears in the study of Bayes estimators van2000asymptotic. Choosing $p=1$ is sufficient for ((ref)) provided $\max_{(x,t)\in\mathcal{X}\times\mathcal{T}}\left\lVert \frac{\partial}{\partial\theta}\int y\mathrm{d}F_{t}(y|x,\theta_{0})\right\rVert <\infty$.
The order $p$ used in condition (ii) must align with the order in Assumption (ref) (iii).
Note that $W(\theta_{0},\mu)$ is constant over $A_{0}$. This condition requires that the value of $W(\theta_{0},\cdot)$ is uniformly separated between $A_{0}$ and $\mathcal{M}\setminus A_{0}$. The requirement arises because $\mathcal{M}$ is infinite; it is unnecessary when the action space is finite. To see an implication from this condition, note that the correspondence $A(\theta):=\arg\max_{\mu\in\mathcal{M}}W(\theta,\mu)$ is upper hemicontinuous at $\theta_{0}$ by the theorem of maximum of Berge. From this observation, one can show that Assumption (ref) implies that for sufficiently small $\varepsilon>0$ we have that $A(\theta)=A_{0}$ for all $\theta\in N_{\varepsilon}(\theta_{0})$, which means that $A(\theta)$ is invariant around the neighborhood of $\theta_{0}$.
We note that the sequence $\{\mu_{n}^{B}(Z^{n})\}$ of Bayesian rules may not be uniquely determined as our framework allows for multiple maximizers of the objective function. This non-uniqueness complicates the analysis since we cannot directly apply the argmax theorem, which is often used to study the asymptotic behavior of general argmax-functionals.
To deal with this, we utilize a penalized version of the Bayesian rule. Let $\nu\in\mathcal{M}$ be any fixed reference measure and $H:\mathcal{M}\to\mathbb{R}_{+}$ be a functional given by $\mu\mapsto\left(d_{W}\left(\mu,\nu\right)\right)^{2}$, which will serve as a penalty function of a maximization problem. For example, we can let $\nu$ be the product measure of $F_{X}$ and $F_{T}$. The functional $H$ has following properties:
Then we define the penalized Bayesian rule by \[ \mu_{n,\varepsilon}^{B}(z):=\arg\max_{\mu\in\mathcal{M}}\int\sqrt{n}W(\theta,\mu)\pi_{n}(\theta|z)\mathrm{d}\theta-\varepsilon H(\mu), \] for $\varepsilon>0$. Note that $\mu_{n,\varepsilon}^{B}$ becomes the unique maximizer of this penalized problem by the strict convexity of $H$. We then obtain a useful result on the penalized rules $\{\mu_{n,\varepsilon}^{B}\}$ by following the arguments of nutzIntroductionEntropicOptimal2022.\footnote{nutzIntroductionEntropicOptimal2022 provides corresponding results by choosing $H$ as the Kullback-Leibler (KL) information criterion between $\mu\in\mathcal{M}$ and any reference measure in $\mathcal{M}$. KL is nonnegative and strictly convex in $\mu$, but not continuous and bounded. Here we impose stronger requirements for the penalty function $H$, which is needed to handle weak convergence of functionals on $\mathcal{M}$ to study the asymptotic properties of rules. Accordingly, the mode of convergence of $\mu_{n,\varepsilon}^{B}(z)$ is modified to weak convergence from convergence in total variation, see nutzIntroductionEntropicOptimal2022.}
Thus, we can construct a unique $\{\mu_{n}^{B}(z)\}$ where $\mu_{n}^{B}(z)$ minimizes the penalty function $H$ over $\mathcal{M}_{opt}(z)$. The following result is stated in terms of $\{\mu_{n}^{B}(Z^{n})\}$ defined in this way.
Our proof strategy follows the approaches of hirano2009asymptotics, christensen2023optimal, and xu2024jmp, but requires suitable extensions since our action space $(\mathcal{M},d_{W})$ is more complicated than theirs. This gives rise to technical challenges specific to our framework, which we address by drawing on tools from optimal transport.
We will explore the asymptotic behavior of the plug-in rules to compare with the Bayesian rule.
As a part of the proof of Theorem (ref), we show that any average optimal rule $\{\mu_{n}(Z^{n})\}\in\mathcal{D}$ will be matched by the rule in the limit experiment $\mu_{\infty}$ where $\mu_{\infty}$ satisfies
with $\Delta\sim N(h,I_{0}^{-1})$ (see Lemma (ref)). The Bayesian rule satisfies this condition. We want to check whether the plug-in rule $\mu_{n}^{P}(Z^{n})$ satisfies this.
One can show that $\{\mu_{n}^{P}\}\in\mathcal{D}$, which implies $\mu_{n}^{P}(Z^{n})\in A_{0}$ with probability approaching to one along $P_{\theta_{nh}}^{n}$. Thus, for sufficiently large $n$, $\mu_{n}^{P}(Z^{n})$ equivalently solves
where the second equality follows because the value of $W(\theta_{0},\mu)$ is constant across $\mu\in A_{0}$. We will see that the maximization problem that plug-in rules solve weakly converges to a different maximization problem from the one that the matched rules of optimal rules must solve; i.e., ((ref)). We actually claim that
where $\Delta\sim N(h,I_{0}^{-1})$.
To see this, let $B_{n}(\mu):=\sqrt{n}\left[\int w(\hat{\theta}_{n},x,t)\mathrm{d}\mu-\int w(\theta_{0},x,t)\mathrm{d}\mu\right]$ and $B_{\infty}(\mu):=\int\dot{w}_{\theta_{0}}(x,t;\Delta)\mathrm{d}\mu$. We impose a high-level condition that the process $\left\{ B_{n}(\mu):\mu\in A_{0}\right\} $ is asymptotically tight. Also note that $\sqrt{n}(\hat{\theta}-\theta_{0})=I_{0}^{-1}S_{n}+o_{P_{\theta_{0}}^{n}}(1)$ with $S_{n}\stackrel{0}{\rightsquigarrow}N(0,I_{0})$ as $n\to\infty$ by the best regularity of $\hat{\theta}_{n}$. Combining Le Cam's third lemma and the delta method for the directionally differentiable functions fang2019inference yields \[ \sqrt{n}\left[\int w(\hat{\theta}_{n},x,t)\mathrm{d}\mu-\int w(\theta_{0},x,t)\mathrm{d}\mu\right]\stackrel{h}{\rightsquigarrow}\int\dot{w}_{\theta_{0}}(x,t;\Delta)\mathrm{d}\mu\quad\text{as }n\to\infty, \] where $\Delta\sim N(h,I_{0}^{-1})$. By the asymptotic tightness of $\left\{ B_{n}(\mu):\mu\in A_{0}\right\} $, we can extend this result to convergence in distribution of the process \[ B_{n}\stackrel{h}{\rightsquigarrow}B_{\infty}\quad\text{as }n\to\infty\text{ on }\ell^{\infty}(A_{0}), \] by vdvW1996. Then applying the continuous mapping theorem yields
which completes the argument.
It is evident that the solutions of RHS of ((ref)) need not to solve ((ref)), the maximization problem in the limit experiment that the matched rules of optimal rules must solve. Thus the plug-in rules might not be average optimal in general when $w$ is directionally differentiable.
If we strengthen the directional differentiability in Assumption (ref) to the full differentiability, then the plug-in rules become average optimal. To see this, notice that $\dot{w}_{\theta_{0}}(x,t;s)=\dot{w}_{\theta_{0}}(x,t)^{\top}s$ for some $\dot{w}_{\theta_{0}}(x,t)\in\mathbb{R}^{k}$ from the linearity of the directional derivative. Then the maximization problem ((ref)) is rewritten as \[ \max_{\mu\in A_{0}}\int\dot{w}_{\theta_{0}}(x,t)^{\top}\Delta\mathrm{d}\mu, \] which is the same as ((ref)), the maximization problem that the matched rules of the plug-in rules solve in the limit experiment.
This pattern is consistent with findings from the existing literature. As discussed earlier, hirano2009asymptotics show that the plug-in rule is average optimal in point-identified models when the utility function is fully differentiable. In contrast, christensen2023optimal demonstrate that the Bayesian rule is average optimal in partially identified models when the utility function is only directionally differentiable, and the plug-in rule fails to be optimal unless the full differentiability holds. In partially identified settings, directional differentiability is a natural and often unavoidable assumption, as full differentiability typically does not hold.
We conduct a simulation study to evaluate the performance of the Bayesian rule and the plug-in rule under the following conditions: (i) the welfare function is either smooth or only directionally differentiable, and (ii) the sample size is relatively small ($n=200$) and large ($n=500$).
We closely follows the data generating process described in Example (ref). For the training population, the latent variable is generated by \[ Y_{i}^{*}=X_{i}^{\top}\beta+\alpha T_{i}+u_{i}, \] where $X_{i}\in\mathbb{R}^{2}$ denotes the observable covariates and $T_{i}$ is the binary treatment that is randomly assigned. The first coordinate of $X_{i}$, interpreted as age, follows a truncated normal distribution with mean 4, standard deviation of 2, and is bounded on $[1,10]$. The second coordinate, interpreted as sex, is a binary variable assigned with equal probability. The observed outcome is \[ Y_{i}=\max\{0,Y_{i}^{*}\}. \] We set $\beta_{0}=(-2,-3)$, $\alpha_{0}=4$, and $u_{i}\sim N(0,\sigma_{0}^{2})$ with $\sigma_{0}=10$. The observed data is an i.i.d. sample $Z^{n}=\{(Y_{i},X_{i},T_{i})\}_{i=1}^{n}$. The parameters $\theta_{0}=(\beta_{0},\alpha_{0},\sigma_{0})$ can be estimated by the maximum likelihood using $Z^{n}$.
In this Tobit model, the conditional mean of the potential outcomes in the training population is given by \[ w(\theta,x,t)=(x^{\top}\beta+\alpha t)-(x^{\top}\beta+\alpha t)\Phi\left(\frac{-x^{\top}\beta-\alpha t}{\sigma}\right)+\sigma\phi\left(\frac{-x^{\top}\beta-\alpha t}{\sigma}\right), \] where $\Phi$ and $\phi$ denote the standard normal cdf and pdf, respectively. Following Examples (ref) and (ref), we specify the planner's utility function as \[ w_{R}(\theta,x,t,\varepsilon,\lambda)=\lambda w(\theta,x,t)+(1-\lambda)\max\left\{ w(\theta,x,t)-\varepsilon,0\right\} , \] where $\lambda\in[0,1]$ and $\varepsilon>0$. We note that $w_{R}$ is differentiable when $\lambda=1$, but only directionally differentiable otherwise. In what follows, we focus on the cases $\lambda=0,1$ and $\varepsilon=0.8$.
Figure (ref) plots the welfare contrasts $w_{R}(\theta,x,1,\varepsilon,\lambda)-w_{R}(\theta,x,0,\varepsilon,\lambda)$ for both males and females at $\theta_{0}$.
Under $\lambda=0$, kinks appear when $w(\theta,x,t)-\varepsilon<0$. For a given covariate $x$, the contrast is zero if both $w(\theta,x,1)<\varepsilon$ and $w(\theta,x,0)<\varepsilon$; that is, individuals with sufficiently low welfare are regarded as deriving no benefit from treatment.
The upper panels of Figure (ref) show the oracle (infeasible optimal) rule at $\theta_{0}$. The rule assigns treatment to females younger than approximately 6.5 and to males younger than approximately 5. Notably, this oracle rule remains unchanged across $\lambda=0$ and $\lambda=1$. It also remains the same at $\theta_{0}+h/\sqrt{n}$, for the range of local deviation parameters $h$ specified below.
We assume that the true distribution of covariates $X$ in the training population is known and that the target population shares the same distribution. Specifically, we define $F_{X}$ as the joint distribution of (a) the truncated normal distribution for the age variable and (b) the binary distribution for the sex variable. For computational purposes, we discretize $F_{X}$ into 99 bins, each corresponding to a distinct combination of age and sex. Each bin is assigned a probability mass according to $F_{X}$, representing the proportion of individuals falling into that bin.
Suppose that the planner has resources to allocate to 75% of the target population. Let \[ W(\theta,\mu)=\int w_{R}(\theta,x,t,\varepsilon,\lambda)\mathrm{d}\mu(x,t). \]
We evaluate performance under a sequence of perturbed DGPs: \[ \theta_{nh}=\theta_{0}+h/\sqrt{n},\quad\text{for }h\in H:=\{-2,-1.6,\dots,2\}, \] where $h/\sqrt{n}$ is added to $\theta_{0}$ element-wisely. Our goal is to compare the average risk \[ \int R(\mu_{n}^{Q},\theta_{nh})\mathrm{d}h=\int\mathbb{E}_{P_{\theta_{nh}}^{n}}\left[W_{\mathcal{M}}^{*}(\theta_{nh})-W(\theta_{nh},\mu_{n}(Z^{n}))\right]\mathrm{d}h \] for $Q=P,B$. The simulation proceeds as follows:
We use POT, an open-source Python library developed by JMLR:v22:20-451, to compute the plug-in rule and the Bayesian rule.
We study the cases of $n=200,500$, $J=2000$, and $L=2000$ for both $\lambda=0$ and $\lambda=1$. We first report the estimated risks, followed by comparisons of the resulting treatment allocations.
Figure (ref) shows the results for $n=200$. While our theory predicts the plug-in and Bayesian rules are asymptotically optimal under smooth welfare ($\lambda=1$), the simulation shows that the Bayesian rule performs better in small samples. We also observe that the Bayesian rule outperforms the plug-in rule under directionally differentiable welfare ($\lambda=0$).
Figure (ref) shows the results for $n=500$. The Bayesian rule still performs slightly better when $\lambda=1$, but the overall risk levels are substantially reduced, and the performance gap between the two rules narrows. This indicates that both rules are approaching optimality as the sample size increases from 200 to 500. When $\lambda=0$, the Bayesian rule continues to outperform the plug-in rule, which is consistent with our theoretical predictions: under the directionally differentiable welfare the Bayesian rule is optimal, but the plug-in rule may not be. Notably, the Bayesian rule performs particularly well when the values of $h$ are negative. In these cases, the welfare contrasts become smaller, making the assignment problem more challenging. This highlights the robustness of the Bayesian rule to local perturbations that make treatment decisions harder.
To gain further insight into the behavior of the two rules, we visualize the average allocations, $J^{-1}\sum_{j}\mu_{nh}^{Q,j}$, for $Q=P,E$, under $\theta_{0}$, $n=200$, and $\lambda=0$. Figure (ref) shows that the Bayesian rule deviates from the oracle rule only near the decision boundary, while the plug-in rule exhibits substantial deviations even away from it. This relative stability of the Bayesian rule contributes to a sizable risk reduction. A similar, albeit weaker, pattern is observed for $n=500$.
We illustrate our methods using data from angrist2006long continued from Example (ref).\footnote{For the replication dataset of the original article, see angristReplicationDataLongterm2019Am.Econ.Assoc..} As outcome variables, angrist2006long use the test scores in language and math. Since the estimation results are similar between these two, we use math scores as the outcome variable for illustration. We focus on the case where the observed test scores are censored at the tenth percentiles of the test score distribution among test-takers (denoted by $\tau$), in line with the original article, to address selection issues. In addition to the test scores, we observe treatment status, as well as age and sex as covariates. The sample includes 3,541 individuals overall, with 1,788 girls and 1,753 boys. Ages range from 10 to 17 with mean 12.7 and standard deviation 1.3. The maximum likelihood estimates are summarized in Table (ref).
We then hypothetically treat the marginal distribution of the covariates in the observed sample as that of the target population, and compute both the plug-in and the Bayesian rules as described in the previous section. In this example, the planner's utility function is given by:
where \[ w(\theta,x,t)=(x^{\top}\beta+\alpha t)+(\tau-x^{\top}\beta-\alpha t)\Phi\left(\frac{\tau-x^{\top}\beta-\alpha t}{\sigma}\right)+\sigma\phi\left(\frac{\tau-x^{\top}\beta-\alpha t}{\sigma}\right). \] Note that this is slightly different from the utility function in the previous section as the outcome variable is censored at $\tau\neq0$. In what follows, we focus on $\varepsilon=3.5$ and $\lambda=0,1$. We consider the case where we can assign vouchers for 50% of the target population.
Figure (ref) shows the allocations under smooth welfare ($\lambda=1$). As Table (ref) shows, age has a negative effect on outcomes. Accordingly, the plug-in rule allocates vouchers to younger individuals. Since the effect of sex is slightly negative, the plug-in rule prioritizes females over males, resulting in the allocation where vouchers are fully allocated to females aged 10-12, while not fully allocated to males at age 12 as the resource is exhausted due to the capacity constraints. In this setting, the Bayesian rule yields exactly the same allocation, which is natural since both rules are optimal under smooth welfare. This also aligns with the simulation result in the previous section: both rules perform similarly when the sample size is large enough.
Next, Figure (ref) shows the allocations under directionally differentiable welfare ($\lambda=0$). For the plug-in rule, the value of $w_{R}$ is censored by $\tau$ at age 13 for females and at 12 for males in this setting. As in the previous case, vouchers are fully allocated to females aged 10--12 and males aged 10--11. However, the remaining vouchers are randomly assigned, as the value of $w_{R}$ is just equal to $\tau$ for the rest. For the Bayesian rule, after integrating with respect to the posterior distribution, the value of $w_{R}$ is censored at $\tau$ at age 13 for both females and males. This leads to allocation to males at age 12 until the resource is exhausted, resulting in the same allocation as seen in Figure (ref). This illustrates that the plug-in and the Bayesian rules could generate different allocations under $\lambda=0$.
We studied the decision-theoretic optimality of treatment assignment rules under capacity constraints on available treatments. Since such constraints complicate the analysis of optimal rules, we transformed the planner's constrained maximization problem into the unconstrained one using tools from optimal transport theory. This reformulation allows us to search for optimal rules in terms of couplings that automatically satisfy the capacity constraints. We investigated two rules previously studied in the literature---the plug-in rule and the Bayesian rule. Both are average optimal when the planner's utility function is smooth; however, the plug-in rule may no longer be optimal when the planner's utility function is only directionally differentiable. A simulation study supports our theoretical predictions. We demonstrated our methods with a voucher assignment problems for private secondary school attendance using data from angrist2006long.
While we focused on average optimality as the optimality criterion, asymptotic minimax optimality is also a widely used benchmark in local asymptotics frameworks. kido2023locally provides an asymptotic minimax optimality result when the ATE is partially identified and there are no constraints on available treatments. In that setting, the plug-in rule becomes optimal only when the oracle rule is (Hadamard) differentiable. Given that hirano2009asymptotics show the minimax optimality of the plug-in rule under the full differentiability conditions, we conjecture that the modes of differentiability of the planner's utility function $w$ plays a key role in the minimax optimality of the plug-in rule in our setting.