EconBase
← Back to paper

Stochastic Treatment Choice with Empirical Welfare Updating

The exact contents of citations.db main_text.text for this paper — one flattened LaTeX string, title through conclusion, appendix excluded, unmodified except for removing email addresses. This is what our citation measures are computed over.

70,246 characters

Stochastic treatment choice with empirical welfare updating


\maketitle

\begin{abstract}
    This paper proposes a novel method to estimate individualised treatment assignment rules.
    The method is designed to find rules that are stochastic, reflecting uncertainty in estimation of an assignment rule and about its welfare performance.
    Our approach is to form a prior distribution over assignment rules, not over data generating processes, and to update this prior based upon an empirical welfare criterion, not likelihood.
    The social planner then assigns treatment by drawing a policy from the resulting posterior.
    We show analytically a welfare-optimal way of updating the prior using empirical welfare;
    this posterior is not feasible to compute, so we propose a variational Bayes approximation for the optimal posterior.
    We characterise the welfare regret convergence of the assignment rule based upon this variational Bayes approximation, showing that it converges to zero at a rate of $\ln\left(n\right)/\sqrt{n}$.
    We apply our methods to experimental data from the Job Training Partnership Act Study to illustrate the implementation of our methods.
    \vspace{\baselineskip}

    \textit{Keywords}: Empirical welfare maximisation, policy learning, PAC-Bayes learning, variational Bayes.
\end{abstract}


\section{Introduction}
\label{SEC:INTRODUCTION}

The principal goal of programme evaluation is to inform the social planner as to which individuals within a target population should receive which treatment.
When treatment effects are heterogeneous in individuals' observable characteristics, the social planner can improve social welfare by implementing an individualised treatment assignment rule based upon these characteristics.
The literature on statistical treatment choice initiated by \citet{Manski2004} studies how to estimate assignment rules based upon a finite sample and how to assess their welfare performance.
Given an experimental or observational sample, existing approaches -- including those proposed in \citet{athey2021policy}, \citet{Hirano2009}, \citet{Kitagawa2018a}, and \citet{Manski2004} -- yield deterministic assignment rules, which are functions mapping the individual's observable characteristics to a recommended treatment.
That is, individuals that share the same observable characteristics are all assigned the same treatment.
Such assignment rules are sharp and address the question of \textit{who should be treated?}
We adopt a broader view of the treatment choice problem by considering \textit{stochastic assignment rules} that map individual observable characteristics to a probability distribution over the different treatment arms, instead addressing the question of \textit{with what probability should an individual be treated?}

In static treatment choice problems with outcome distributions that exhibit the monotone likelihood ratio property, deterministic assignment rules form a class of admissible policies \citep{karlin1956theory,tetenov2012statistical}, such that restricting attention to this class is without loss (of welfare).
Once we allow the class of outcome distributions to be unconstrained though, there is little theoretical justification for focusing on deterministic rules.
In comparison to stochastic assignment rules, deterministic assignment rules have the following three potentially undesirable features.
First, deterministic assignment rules cannot incorporate  confidence or uncertainty about which treatment is best for each individual, with individuals typically assigned treatment if conventional point estimates suggest that treatment has a positive effect on average.
The strength of evidence in support of this conclusion, usually presented in the form of confidence intervals or p-values, is not generally acted upon;
what matters is whether empirical evidence supports a positive point estimate, and not
whether it is sufficient or insufficient to reject a non-positive effect.
Such a sharp dichotomy of assignment is naturally overconfident in its prescription, and there is no theoretical basis for the incorporation of confidence intervals or p-values into frequentist-based decision-making.
Stochastic assignment rules can represent such uncertainty, which arises due to the (finite sample) nature of experimental data or model misspecification, through their probability weighting of treatments.
Second, stochastic assignment rules facilitate future evaluation, since implementing a stochastic assignment rule can generate a new experimental sample in which treatment is randomised conditional on individual observable characteristics.
Third, unlike deterministic assignment rules for which the probability that a treatment is assigned changes discontinuously at some threshold, stochastic assignment rules feature assignment probabilities that smoothly change with respect to individual characteristics.
Such a feature is desirable if a fairness criterion requiring that individuals with similar characteristics have similar probabilities of treatment \citep{dwork2012fairness} is enforced.

This paper proposes novel and general methods for obtaining stochastic individualised assignment rules based on randomised control trial data.
Assuming that the social planner assigns individuals to a binary treatment, with her goal being to maximise additive (utilitarian) social welfare as in \citet{Manski2004}, we exploit an empirical analogue of the social welfare criterion to generate individualised assignment probabilities.
Specifically, we start with a prior distribution over a collection of deterministic assignment rules $\mathscr{G}$, each of which partitions the space of individual observable characteristics $\mathscr{X}$ into a group of characteristics $G$ and its complement $G^{c}$, thereby generating a deterministic assignment rule $g\left(x\right)=1\left( x \in G \right)$.
An individual is assigned treatment if their characteristics are such that $x\in G$, and is not assigned treatment if $x\in G^{c}$.
We then update the prior distribution based upon an empirical analogue of the social welfare criterion to obtain a posterior distribution over $\mathscr{G}$.
To generate a stochastic assignment, we draw a $g\in\mathscr{G}$ according to the posterior distribution over this collection, and implement the policy prescribed by $g$.
In this way, the stochastic assignment ruleapplied to an individual with characteristics $x$ treats her with the probability equal to the posterior probability that stochastic $G$ contains $x$.

One of the main contributions of this paper is that we derive an optimal updating procedure for obtaining the posterior distribution over $\mathscr{G}$.
This procedure minimises an upper bound on welfare regret and yields an exponential tilting of the prior over $\mathscr{G}$, where the tilting depends upon the empirical welfare criterion.
This novel updating formula resembles the quasi-posterior distribution that appears in the Laplace-type estimation studied by \citet{Chernozhukov03}, but differs in that the constant factor in the exponential tilting term is determined endogenously by the Lagrange multiplier of the optimisation.

Despite our analytical characterisation of the optimal posterior distribution, computation of this distribution or sampling of $g$ from it is not straightforward.
We therefore consider a variational approximation of the optimal posterior distribution by a parametric distribution.
In particular, as a specification of $\mathscr{G}$, we consider the class of Linear Eligibility Score rules that assign treatment if $x^{\top}\gamma$ (the linear score) exceeds some threshold $c$ (eligibility).
Building upon the LES class, we exploit the invariance of the welfare criterion to the ratio of $\gamma$ to $c$ (scale invariance) and approximate the optimal posterior distribution using a multivariate von Mises-Fisher distribution.

For the practice of reporting and communicating individualised allocations of treatment, our approach of obtaining a posterior distribution over policies is useful for generating some quantities that existing methods yielding deterministic assignment rules cannot produce.
First, the posterior probability that $g(x)=1$ offers a personalised probabilistic assessment that individuals with characteristics $x$ favour treatment over no treatment.
Reporting such a probability offers a novel alternative to the common practice of using the p-values of hypothesis testing to express confidence in a positive treatment effect, something which does not easily translate to a recommendation about what the social planner should do.
Second, viewing the posterior over $g$ as an inferential tool for the welfare-optimality of (deterministic) assignment policies within $\mathscr{G}$, we can obtain a credible region for the optimal policy by, for instance, selecting its highest posterior density region.
This approach to obtaining confidence sets for the optimal treatment assignment policy is an alternative to the frequentist approach that is studied in \citet{rai2019}.
Third, analogous to the practice of using a Bayesian posterior with a noninformative prior as a visual summary of the likelihood function, we can use our variationally approximated posterior over $g$ as a visual summary of the exponentiated empirical welfare criterion function.

To demonstrate how to implement our approach and what it delivers in practice, we apply our methods to the JTPA Study sample that is studied by \citet{bloom1997benefits}.
Given observations of prior earnings and years of education, we ask \textit{with what probability should an individual be treated?}
Restricting attention to linear assignment rules and a variational approximation of the optimal posterior distribution by a multivariate von Mises-Fisher distribution, we estimate a stochastic assignment rule that is more likely to allocate JTPA assistance to individuals with high prior earnings and fewer years of education.
\citet{Kitagawa2018a}, which similarly considers the JTPA Study sample, and estimates a deterministic assignment rule, serves as a useful benchmark for comparison.
Aside from the obvious difference that we estimate a stochastic rule (i.e., every individual has a non-trivial probability to be allocated JTPA assistance under our rule), our estimated rule allocates JTPA assistance to a smaller fraction of the population than the deterministic rule of \citet{Kitagawa2018a}, which targets individuals with low prior earnings and few years of education for treatment.
This difference reflects the shape of the empirical welfare criterion, which can be captured by our approach but is missed by deterministic policies that are obtained as the mode of the empirical welfare criterion.


\subsection{Literature review}
\label{SEC:LIT}

This paper contributes to the growing literature on statistical treatment choice initiated by \citet{Manski2004}.
Exact minimax regret assignment rules are studied in \citet{ishihara2021}, \citet{schlag2006eleven}, \citet{Stoye2009}, \citet{Stoye2012}, \citet{tetenov2012statistical}, and \citet{yata2021}.
\citet{Hirano2009} analyses asymptotically-optimal assignment rules in limit experiments, and \citet{bhattacharya2012} considers capacity constrained policies.
\citet{Kitagawa2018a} proposes \textit{Empirical Welfare Maximisation} (EWM) methods for individualised assignment, which maximise a sample analogue of the social welfare function over a constrained class of policies.
Similar approaches have been studied in the literature on machine learning and personalised medicine, as in \citet{Beygelzimer2009}, \citet{Swaminathan2015}, \citet{Zadrozny2003}, \citet{Zhang2012}, and \citet{Zhao2012}.
Recent advances in learning individualised assignment policies include \citet{adjaho2022externally}, \citet{athey2021policy}, \citet{dadamo21}, \citet{han2022}, \citet{kido22}, \citet{Kitagawa2017}, \citet{KST21}, \citet{liu22}, \citet{Mbakop2018}, \citet{nie2019learning}, \citet{Sakaguchi2019},  \citet{sasaki20}, \citet{sun21}, and \citet{viviano21}, to list but a few works.
The assignment rules estimated in these works are all deterministic.

There are some earlier works that investigate the decision-theoretic justification for stochastic (fractional) assignment and the welfare performance of these rules, with \citet{manski20092009} providing a detailed review of settings where stochastic rules are preferable.
When the welfare criterion is only partially identified, minimax regret-optimal rules are stochastic given knowledge of the identified set \citep{Manski2000,manski2005social,manski2007identification, Manski2007}, which remains true even after taking into account uncertainty of estimates of the bounds \citep{Stoye2012, yata2021, manski2022identification}.
As shown by \citet{manski2007admissible} and \cite {manski20092009}, stochastic assignment rules can also be justified by a nonlinear welfare criterion in a point-identified setting.
\citet{kitagawa2022} shows that, for a wide class of nonlinear welfare regret criteria, admissible assignment rules are stochastic (fractional).
In particular, \citet{kitagawa2022} shows that the minimax squared-regret rule is stochastic, with the probability of assignment equal to the posterior probability of a positive treatment effect under the least-favorable null.
\citet{kitagawa2022} proposes using this probability as a measure of the strength of evidence for a positive treatment effect, replacing the commonly used p-value of a hypothesis test.
In contrast, this paper obtains the probability of assignment from a posterior probability distribution over assignment rules, rather than over the treatment effect parameters, with a quasi-likelihood built upon the empirical welfare criterion.
\citet{kock2022treatment} obtains a stochastic assignment rule in a setting where the oracle optimal rule is fractional due to nonlinearity in the social planner's chosen welfare criterion.
In contrast to the static treatment assignment problem, dynamic treatment assignment problems analysed in the multi-arm bandit literature often consider stochastic assignments that balance the exploitation versus exploration trade-off, such as the posterior probability matching algorithm of \citet{thompson1933likelihood} does. Thompson sampling algorithms build upon the standard Bayesian posterior distribution for treatment effects such that the allocation algorithm crucially relies on a parametric specification of the data generating process, which our approach does not require.

\citet{Chamberlain2011} and \citet{Dehejia2005} approach the treatment choice problem from a Bayesian perspective.
In their framework, the potential outcome distributions are parametric, and it is over the parameters of these distributions that a prior is imposed.
For the standard mean welfare criterion, the Bayes optimal allocation rule is deterministic.
Our approach differs from these works in that we do not assume a prior distribution over the data generating process.
We instead impose few restrictions on the data generating process, and form prior and posterior distributions over the parameters that index assignment rules.
Our approach can be advantageous when compared to \citet{Chamberlain2011} if the social planner is concerned about potential misspecification of the likelihood.
If the likelihood is misspecified, the resulting Bayes-optimal assignment rule can be suboptimal even for large samples.
In contrast, our approach is guaranteed to yield a distribution over policies that is guaranteed to concentrate on welfare-optimal policies without requiring a specification for the data generating process.

Our approach is also related to that of \citet{Bissiri2016} and \citet{csaba2020learning}, where loss function-driven (quasi-Bayes) updating rules are proposed.
Rather than follow their approach by adopting exponentiated loss as a quasi-likelihood and solving the quasi-Bayesian decision problem, we obtain an optimal learning rule by minimising a high probability upper bound on welfare regret.
This way of establishing optimality is similar to the structural risk minimisation approach of \citet{Vapnik1998} and the \emph{Probably Approximately Correct} (henceforth, PAC) analysis proposed by \citet{Valiant1984}, which was extended to the study of randomised predictors in \citet{McAllester1999}, and \citet{Shawe-Taylor1997}, constituting the development of PAC-Bayes theory.
For classification and regression problems, various PAC-Bayes bounds on prediction generalisation errors are obtained in \citet{Begin2014,Begin2016}, \citet{catoni2007pac}, \citet{Derbeko2004}, \citet{Germain2009}, \citet{McAllester2003}, \citet{Pentina2015}, and \citet{Seeger2002}, and can accommodate quasi-Bayesian procedures similar to ours.
See \citet{Guedj2019} for a recent review of this literature.
To our knowledge, the PAC-Bayes bounds that we derive for treatment choice are new to the literature and offer a contribution of independent interest. We also note that \citet{pellatt2022pac} makes use of PAC-Bayes theory to analyse the treatment allocation problem with stochastic assignment rules under a budget (or resource) constraint.

Although the treatment choice problem is distinct from prediction problems -- as is discussed in \citet{Kitagawa2018a} -- the EWM approach for treatment choice is closely related to the cost-sensitive binary \emph{classification} problem, as first pointed out by \citet{Zadrozny2003}.
The PAC-Bayes classification analysis with variational posterior approximation that is proposed by \citet{Alquier2016} is, therefore, closely related to our analysis.
There are, however, important differences with \citet{Alquier2016}.
First, we make use of the approach proposed by \citet{Begin2016}, which allows for the construction of a variety of different bounds via a general convex function.
This introduces the complication that classification is not standard (i.e., cost is homogeneous) and is instead cost-sensitive.
Introducing heterogeneity in the cost of misclassification leads to a non-trivial challenge in deriving the PAC bounds.
To address these complications we leverage results in \citet{Maurer2004} for continuous loss functions over the unit interval.
Second, \citet{Alquier2016} considers approximating the optimal posterior distribution by a Gaussian distribution.
We exploit the scale invariance property of the welfare criterion and approximate the optimal posterior distribution by a multivariate von Mises-Fisher distribution over the hypersphere.


\section{Model and Setup}
\label{SEC:MODEL}


\subsection{Notation and setting}
\label{SEC:NOTATION}

We let $D \in \{1,0\}$ be a binary treatment;
we let $Y_{1}\in\mathbb{R}$ and $Y_{0}\in\mathbb{R}$ be the potential outcomes associated with the two treatment states;
and we let $X\in\mathscr{X}$ be a vector of observed characteristics, which we refer to as covariates, where $\mathscr{X}\subset\mathbb{R}^{m-1}$.
We suppose that experimental data comprising $n$ independent and identically distributed observations concatenated as $\left(Y_{i},D_{i},X_{i}\right)$ are available, where $Y_{i} = D_{i}\cdot Y_{1i} + \left(1-D_i\right)\cdot Y_{0i}$ is the post-treatment observed outcome of individual $i$.
We denote the joint distribution of the experimental sample, which we reiterate is a probability distribution over $\left(Y_{i},D_{i},X_{i}\right)$, by $P^{n}$.
The joint distribution of $\left(Y_{1},Y_{0},D,X\right)$ that induces the independent and identical distribution of the observations $\left(Y_{i},D_{i},X_{i}\right)$, $i=1, \dots, n$ is referred to as the data generating process and denoted by $P$.
We assume that the covariates consist only of those characteristics that the planner can use to discriminate between individuals in the target population, with budgetary, ethical or legal considerations precluding the use of other characteristics.

Throughout our analysis, we maintain several assumptions.
We follow \citet{Manski2004} and the subsequent literature in supposing that the social welfare criterion is that of a utilitarian social planner who aims to maximise the average level of individual outcomes.
We note that other criteria could also be implemented, such as inequality-averse social welfare and Gini social welfare.\footnote{
For example, \citet{Kasy2016} and \citet{Kitagawa2017} study a setting where the social welfare function is a weighted average of the outcomes with rank-dependent weights, which includes the Gini social welfare function as a special case.}
\begin{assumption}[External validity]
\label{ASS:EV}
The population to which policy is to be applied -- the target population -- has the same distribution over $\left(Y_{1},Y_{0},X\right)$ as the marginal distribution of $\left(Y_{1},Y_{0},X\right)$ that is obtained from the data generating process.
\end{assumption}
\begin{assumption}[Unconfoundedness]
\label{ASS:U}
The data generating process satisfies $\left(Y_{1},Y_{0}\right)\perp\!\!\!\perp D|X$.
\end{assumption}
\cref{ASS:EV,ASS:U} are satisfied, for instance, if the experimental data is extracted directly from the target population and the treatment is, conditional on the covariates, randomly assigned,\footnote{
\citet{Kitagawa2018a} considers a setting where the marginal distribution of $X$ differs between the population of interest and the data generating process.
\citet{adjaho2022externally} and \citet{kido22} study settings that differ also in terms of the distribution of potential outcomes.}
independently of the potential outcomes \citep{Rosenbaum1983}.
\begin{assumption}[Bounded outcomes]
\label{ASS:BO}
There exists a constant $0<M<\infty$ such that the support of Y is contained in $[0,M]$.
\end{assumption}
\begin{assumption}[Strict overlap]
\label{ASS:SO}
There exists a $\psi\in\left(0,1/2\right)$ such that the propensity score satisfies, for all $x\in\mathscr{X}$, $e\left(x\right)\in[\psi,1-\psi]$, where $e\left(x\right)=\mathrm{E}_{P}\left(D|X\right)$.
\end{assumption}
We normalise the outcome variable to the interval $[0,M]$.
As is discussed in \citet{Kitagawa2018a} and \citet{Swaminathan2015}, policies that maximise an empirical welfare criterion are not invariant to positive affine transformations of outcomes, which is the case for the empirical welfare criterion that we consider in this paper.
Given \cref{ASS:BO,ASS:SO}, we can define
\begin{equation}
\label{EQ:WEIGHTSFORIMPLEMENTATION}
    H
    \equiv
    \frac{\psi\cdot Y/M}{e\left(X\right)\cdot D + \left(1-e\left(X\right)\right)\cdot\left(1-D\right)},
\end{equation}
which is confined to the unit interval, and which we interpret as weights and that are motivated by an unbiased estimator of the (scaled) expected potential outcomes.
We solve the planner's problem using these transformed outcomes. This transformation of outcomes does not affect the welfare ranking over assignment policies, both in the population and according to in-sample welfare criteria, yet facilitates the proof of our PAC-Bayes bounds shown in Theorem \ref{THM:PAC} below.

Let $G\subset\mathscr{X}$ specify a set of individuals -- identified by their individual characteristics -- to be treated deterministically.
We refer to $G$ as a policy.
Adopting an additive utilitarian perspective, the average level of social welfare attained by $G$ is proportional to
\begin{equation}
\label{EQ:ORIGINAL}
    W\left(G\right)
    \equiv
    \mathrm{E}_{P}\left(Y_{1} \cdot 1\left(X \in G\right)+Y_{0} \cdot 1\left(X \notin G\right)\right).
\end{equation}
Given \cref{ASS:EV,ASS:U,ASS:SO} and that $Y=Y_{0}+D\cdot\left(Y_{1}-Y_{0}\right)$, we can re-write \cref{EQ:ORIGINAL} as
\begin{align}
    W\left(G\right)
    &=
    \mathrm{E}_{P}\left(\frac{Y\cdot D}{e\left(X\right)} \cdot 1\left(X \in G\right)+\frac{Y\cdot \left(1-D\right)}{1-e\left(X\right)} \cdot 1\left(X \notin G\right)\right),\\
    &=
    \mathrm{E}_{P}\left(\frac{Y\cdot\left(1-D\right)}{1-e\left(X\right)}\right)+\mathrm{E}_{P}\left(\left( \frac{Y\cdot D}{e\left(X\right)}-\frac{Y\cdot\left(1-D\right)}{1-e\left(X\right)}\right)\cdot 1\left(X\in G\right)\right).
\end{align}
Accordingly, the sample analogue of \cref{EQ:ORIGINAL} can be written as
\begin{align}
    W_{n}\left(G\right)
    &\equiv
    \frac{1}{n}\sum^{n}_{i=1}\left(\frac{y_{i}\cdot d_{i}}{e\left(x_{i}\right)} \cdot 1\left(x_{i} \in G\right)+\frac{y_{i}\cdot\left(1-d_{i}\right)}{1-e\left(x_{i}\right)} \cdot 1\left(x_{i} \notin G\right)\right),\\
    &=
    \frac{1}{n}\sum^{n}_{i=1}\frac{y_{i}\cdot\left(1-d_{i}\right)}{1-e\left(x_{i}\right)}+\frac{1}{n}\sum^{n}_{i=1}\left(\frac{y_{i}\cdot d_{i}}{e\left(x_{i}\right)}-\frac{y_{i}\cdot\left(1-d_{i}\right)}{1-e\left(x_{i}\right)}\right)\cdot 1\left(x_{i}\in G\right),
    \label{EQ:TOREWRITE}
\end{align}
where $W_{n}\left(G\right)$ is an unbiased estimator for the true level of welfare that arises from the implementation of a particular $G$.
Given the additive social welfare criterion, the maximal welfare level can be attained by a deterministic policy.
Hence, as far as the population welfare maximisation problem is concerned, the social planner wants to select the $G$ that maximises $W\left(G\right)$.

Each $G$ can be associated with a binary function $g$ that indicates membership in $G$.
We refer to $g$ as a deterministic asssignment rule, or simply as an assignment rule, with $\mathscr{G}$ constituting the class of assignment rules.
With this notation to hand, we can write \cref{EQ:TOREWRITE}, with some abuse of notation, as
\begin{align}
W_n\left(g\right)
&=
\frac{1}{n}\sum^{n}_{i=1}\frac{y_{i}\cdot\left(1-d_{i}\right)}{1-e\left(x_{i}\right)}+\frac{1}{n}\sum^{n}_{i=1}\left(\frac{y_{i}\cdot d_{i}}{e\left(x_{i}\right)}-\frac{y_{i}\cdot\left(1-d_{i}\right)}{1-e\left(x_{i}\right)}\right)\cdot g\left(x_{i}\right),\\
&=
\frac{1}{n}\sum^{n}_{i=1}h_{i}\cdot\frac{M}{\psi}\cdot\left(e\left(x_{i}\right)\cdot d_{i}+\left(1-e\left(x_{i}\right)\right)\cdot\left(1-d_{i}\right)\right)\cdot\left[\frac{d_{i}}{e\left(x_{i}\right)}\cdot g\left(x_{i}\right)+\frac{1-d_{i}}{1-e\left(x_{i}\right)}\cdot\left(1-g\left(x_{i}\right)\right)\right],\\
&=
\frac{1}{n}\sum^{n}_{i=1}h_{i}\cdot\frac{M}{\psi}\cdot\left[d_{i}\cdot g\left(x_{i}\right)+\left(1-d_{i}\right)\cdot\left(1-g\left(x_{i}\right)\right)\right],\\
&=
\frac{1}{n}\sum^{n}_{i=1}h_{i}\cdot\frac{M}{\psi}\cdot 1\left(g\left(x_{i}\right)=d_{i}\right),\\
&=
\frac{1}{n}\sum^{n}_{i=1}h_{i}\cdot\frac{M}{\psi}-\frac{1}{n}\sum^{n}_{i=1}h_{i}\cdot\frac{M}{\psi}\cdot 1\left(g\left(x_{i}\right)\neq d_{i}\right),
\label{EQ:WEIGHT}
\end{align}
where $h_i$ is the realisation of $H$, as defined in \cref{EQ:WEIGHTSFORIMPLEMENTATION}, for observation $\left(y_{i},d_{i},x_{i}\right)$.
We observe from \cref{EQ:WEIGHT} that $W_{n}\left(g\right)$ depends upon $g$ only through its second term, such that
\begin{equation}
\label{EQ:OBJECTIVE}
    \underset{g}{\operatorname{argmax}}\,W_n\left(g\right) = \underset{g}{\operatorname{argmin}}\,\frac{1}{n}\sum_{i=1}^n h_{i}\cdot 1\left( g\left(x_{i}\right) \neq d_{i} \right).
\end{equation}
Accordingly, we define
\begin{equation}
    R\left(g\right)
    \equiv
    \mathrm{E}_{P} \left(H\cdot 1\left(g\left(X\right) \neq D\right)\right),
    \label{EQ:LESRISK}
\end{equation}
which we term the \textit{welfare risk} of $g$, and its empirical analogue
\begin{equation}
    R_{S}\left(g\right)
    \equiv
    \frac{1}{n} \sum_{i=1}^{n}h_{i} \cdot 1\left(g\left(x_{i}\right) \neq d_{i}\right),
    \label{EQ:EMPIRICALLESRISK}
\end{equation}
which we term the \textit{empirical welfare risk} of $g$.
In view of \cref{EQ:OBJECTIVE}, the social planner's objective is to minimise \cref{EQ:LESRISK} via \cref{EQ:EMPIRICALLESRISK} in $g$, following the empirical risk minimisation principle of \cite{Vapnik1998}.

One special set of policies that we draw particular attention to is the \textit{Linear Eligibility Score} (LES) class that is defined in \citet{Kitagawa2018a}, and that we denote by $\mathscr{F}$.
Assignment rules in this class are indexed by a finite-dimensional parameter vector $\gamma$ and a threshold $c$, and are associated with a binary function $f_{\beta}$ that satisfies, for all $x\in\mathscr{X}$,
\begin{equation}
\label{EQ:UNIQUENESSOFINDEX}
    f_{\beta}\left(x\right)
    \equiv
    1\left(x^{\top}\gamma\geq c\right),
\end{equation}
where we take $\beta$ to include both $\gamma$ and $c$ (i.e., $\beta$ is an $m$-dimensional vector).
Each LES rule induces a partitioning of the covariate space into two half-spaces, such that individuals in the upper contour set receive treatment and individuals in the lower contour set do not.
By restricting $\beta$ to the unit hypersphere (i.e., the Euclidean length of $\beta$ is one), we guarantee that each policy is associated with a unique $\beta$.
In what follows, we exploit the interchangeability of $\beta$ and the LES rule that it indexes, adopting $\beta$ as the argument of the loss functions that we consider.
For instance, and again with some abuse of notation, whenever we focus on the LES class of decision rules we write
\begin{equation}
    R\left(\beta\right)
    =
    \mathrm{E}_{P}\left(H\cdot 1\left(f_{\beta}\left(X\right)\neq D\right)\right),
\end{equation}
and
\begin{equation}
    R_{S}\left(\beta\right)
    =
    \frac{1}{n} \sum_{i=1}^{n}h_{i} \cdot 1\left\{f_{\beta}\left(x_{i}\right) \neq d_{i}\right\},
\end{equation}
respectively, in place of \cref{EQ:LESRISK,EQ:EMPIRICALLESRISK}.


\subsection{Posterior over policies as a stochastic assignment rule}
\label{SEC:POSTERIOR}

We now adapt \cref{EQ:LESRISK,EQ:EMPIRICALLESRISK}  to handle stochastic assignment rules.
We let $\Pi$ denote a probability distribution over $\mathscr{G}$ that we interpret as a posterior distribution, assuming that $\mathscr{G}$ can be embedded in a measurable space.\footnote{
For $\mathscr{G}$ to be embedded in a measurable space, $\mathscr{G}$ cannot be too rich. We defer to \citet{molchanov2005} and \citet{gunsilius2019path} for further discussion of this point.}
We let $\mathscr{M}$ denote the collection of all posterior distributions.
\begin{definition}[Posterior assignment rule]
Let $\Pi$ be a probability distribution over $\mathscr{G}$ that is constructed upon observing the sample.
The posterior assignment rule under $\Pi$ is a stochastic assignment rule that assigns individuals with $x\in\mathscr{X}$ to treatment with probability $Q^{\Pi}\left(x\right)\equiv\int_{\mathscr{G}}g\left(x\right)\cdot\mathrm{d}\Pi$.
\end{definition}
To implement posterior assignment rules in practice, we randomly draw a $g$ from $\mathscr{G}$ according to $\Pi$ for each individual in the target population.
In this way, similar individuals, who can have similar assignment probabilities, can be assigned to different treatment arms.
Moreover, this approach does not require computation of the probability of treatment.
\begin{definition}[Expected welfare risk under $\Pi$]
We define the expected welfare risk under $\Pi$ as
\begin{equation}
\begin{aligned}
R^{\Pi}
&\equiv
\int_\mathscr{G}R\left(g\right)\cdot\mathrm{d}\Pi\left(g\right),\\
&=
\mathrm{E}_{P} \left(H \cdot\left(D\cdot\left(1-Q^{\Pi}\left(X\right)\right)+ \left(1-D\right)\cdot Q^{\Pi}\left(X\right)\right) \right),
\end{aligned}
\end{equation}
with its empirical analogue taking the form
\begin{equation}
\begin{aligned}
    R_{S}^\Pi
    &\equiv
    \int_\mathscr{G}R_{S}\left(g\right)\cdot\mathrm{d}\Pi\left(g\right),\\
    &=
    \frac{1}{n}\sum_{i=1}^{n} h_{i} \cdot\left(d_{i}\cdot\left(1-Q^{\Pi}\left(x_{i}\right)\right)+ \left(1-d_{i}\right)\cdot Q^{\Pi}\left(x_{i}\right)\right).
    \end{aligned}
\end{equation}
\end{definition}
The interpretation of $R^{\Pi}$ is the average welfare loss that the social planner expects from stochastic implementation of $g$ in $\mathscr{G}$ in the target population when $g$ is distributed according to $\Pi$.

We reiterate that stochastic assignment is achieved by randomly drawing $g$ according to $\Pi$.
This way of selecting assignment rules is reminiscent of the \textit{Gibbs classifier} in statistical learning theory \citep{Germain2009} and might offer a computational advantage over other methods if drawing $g$ according to $\Pi$ is easier than maximising empirical welfare or finding the mode of $\Pi$, say.
An advantage of stochastic assignment is the possibility for sequential treatment evaluation:
the induced assignment of treatment and non-treatment to individuals in the target population by $\Pi$ is random conditional on $X$, which allows for estimation of the causal effect of treatment in future studies.
In this sense, stochastic assignment is well suited to balancing existing evidence about what constitutes the optimal assignment for each individual against the benefit of further exploration of the treatment effect \citep{manski2000using}.
We do not, however, study this channel and focus on a purely static problem in this paper.

Our framework allows the social planner to hold some prior as to what constitutes the best policy.
We differentiate the subjective beliefs that the social planner holds, which we encode using the prior distribution $\pi_{0}$, and their updated beliefs following their observation of sample data, which we encode using $\Pi$.
In the analysis that follows, we provide finite sample regret guarantees in the form of PAC-Bayes bounds for stochastic assignment.
The approach discussed here differs significantly from other Bayesian treatment choice settings, such as those discussed in \citet{Chamberlain2011}, because we do not impose any kind of prior belief on $P$.
We instead choose to model the beliefs that the social planner has regarding the optimal policy.
Using a decision procedure that is free from specification of the likelihood comes with a desirable robustness property, as we discuss in due course.

Aside from its more conventional role as a means of expressing existing information about what constitutes the best policy, $\pi_{0}$ can also play several other roles within our framework.
For instance, $\pi_{0}$ can also embed any constraints that are imposed upon the set of policies through truncation of its support.
Such a zero density condition can be imposed in lieu of an explicit restriction on $\mathscr{G}$ and is easy to implement in practice via rejection sampling, with only those policies that satisfy any budgetary, ethical or legal constraints being retained under the sampling procedure.
Moreover, $\pi_{0}$ can also be used to describe the status quo, with restrictions on the shape of $\pi_{0}$ governing how much policy can deviate.
These interpretations of $\pi_{0}$ naturally extend to $\Pi$.


\section{Optimal stochastic assignment and convergence of welfare}
\label{SEC:OPTIMALITY}


\subsection{Bounding expected welfare risk}
\label{SEC:DEVIATIONS}

Seeing as experimental data provides only an insight into the welfare performance of any policy in the target population, a very natural question to ask is how much we can expect $R^{\Pi}_{S}$ to differ from $R^{\Pi}$ for any given $\Pi$.
We provide an answer to this question here.
\begin{theorem}
\label{THM:PAC}
Suppose that \cref{ASS:EV,ASS:U,ASS:BO,ASS:SO} are satisfied, and $n\geq8$.
Then, for any $0<\epsilon<1$ and $\Pi\in\mathscr{M}$ that is absolutely continuous with respect to $\pi_0$, the following inequality holds with probability at least $1- \epsilon$ in terms of $P^{n}$:
\begin{equation}
\label{EQ:PAC-BOUND}
    R^{\Pi} \leq R_{S}^{\Pi}+\sqrt{\frac{1}{2 n}\cdot\left(\mathrm{KL}\left(\Pi \| \pi_0\right)+\ln \left(\frac{2 \sqrt{n}}{\epsilon}\right)\right)}.
\end{equation}
\end{theorem}
We present a proof of this theorem in \cref{PROOF:PAC}.
The proof builds upon \citet[\S Lemma 3]{Begin2016}, which offers a flexible approach that allows for the recovery of many different PAC-Bayes bounds.
The approach centres around a convex function of $R^{\Pi}$ and $R^{\Pi}_{S}$, which we specify so as to recover the form that is presented in \citet[see \S (8) in \citealp{Begin2016}]{McAllester2003}.
We leverage results in \citet[\S Lemma 3 and \S Theorem 1]{Maurer2004}, exploiting the properties of Bernoulli random variables and convex functions, to adapt \citet{Begin2016} and the bound that is presented therein to allow for heterogeneous cost (in lieu of the standard binary loss function that is prevalent in the classification literature and that is studied in \citealp{Begin2016}).

We note that \cref{THM:PAC} holds for any rules that update $\pi_{0}$ and deliver $\Pi$, and that we have not committed to any particular updating rule to obtain \cref{EQ:PAC-BOUND}.
A feature of \cref{EQ:PAC-BOUND} is that the regularisation term -- the square root term containing the Kullback-Leibler divergence of $\Pi$ from $\pi_{0}$ -- enters additively, which we find is a convenient feature for establishing convergence of our variational approximation (regularisation can otherwise be effected multiplicatively).\footnote{
See \citet{Begin2016} for the implications of different convex functions.}
The regularisation term, by design, prevents overfitting.
To illustrate this point, suppose that $\pi_{0}$ is uniform over $\mathscr{G}$:
the best response of the social planner absent regularisation is to concentrate probability mass on the optimal (in-sample) $g$ as suggested by data, such that $\Pi$ is degenerate.
In the presence of regularisation, however, this is no longer a best response, since the Kullback-Leibler divergence infinitely penalises degeneracy \textit{vis-\`{a}-vis} uniformity.
Rather, the best response of the social planner is to allocate probability mass on all $g$ in $\mathscr{G}$, albeit concentrating more mass on those $g$ that are associated with low empirical welfare risk.
Put differently, the regularisation term controls how far away from $\pi_{0}$ a stochastic assignment rule can be, with this difference governed by the number of observations in the sample.

The Vapnik-Chervonenkis (VC) dimension is the standard measure of complexity in the statistical learning literature.
We instead associate complexity with the Kullback-Leibler divergence.
If one is willing to impose distributional constraints on a posterior assignment rule, an advantage of the PAC-Bayes approach and of using the Kullback-Leibler divergence is that complexity is then purely in terms of the selected (stochastic) assignment rule $\Pi$ rather than in terms of the class of possible stochastic assignment rules $\mathscr{M}$ or the class of underlying deterministic rules $\mathscr{G}$.
As such, \cref{THM:OPTIMAL} does not explictly require any assumption about the VC dimension of $\mathscr{G}$.
The influence of VC dimension for $\mathscr{G}$ is implicit in our setting:
the upper bound on the difference between $R_{S}^{\Pi}$ and $R^{\Pi}$ implied by \cref{EQ:PAC-BOUND} is governed by the Kullback-Leibler divergence, which is increasing in the dimension of the support of $\Pi$ and $\pi_{0}$, and so is non-decreasing in the complexity of $\mathscr{G}$.


\subsection{Optimal updating rule}
\label{SEC:UPDATING}

In a standard Bayesian setting, unknown parameters index the distribution of data and inference on parameters is conducted with respect to the posterior distribution.
Typically, the posterior distribution is constructed from a well-defined likelihood function via Bayes' theorem.
We leverage \cref{THM:PAC} to construct $\Pi$ from $\pi_{0}$ and $R^{\Pi}_{S}$.
This approach is valid since \cref{THM:PAC} holds for all $\Pi\in\mathscr{M}$.
We emphasise that the posterior distribution that we construct is over assignment rules rather than over the data generating processes, which is distinct from \citet{Bissiri2016} and \citet{csaba2020learning}.

Following \citet{McAllester2003} and \citet{Germain2009}, we define an optimal posterior distribution, which we denote by $\Pi^{*}$, as a distribution over $\mathscr{G}$ that minimises the right hand side of \cref{EQ:PAC-BOUND}. That is, $\Pi^{*}$ minimises
\begin{equation}
     R_{S}^{\Pi} +\sqrt{\frac{1}{2 n}\cdot\left( \mathrm{KL}\left(\Pi \| \pi_0\right) \text+\ln\left(\frac{2 \sqrt{n}}{\epsilon}\right)\right)}.
\end{equation}
\begin{theorem}
\label{THM:OPTIMAL}
The optimal posterior $\Pi^{\ast}$ over $\mathscr{G}$ satisfies
\begin{equation}
    \mathrm{d}\Pi^{*}\left(g\right)
    \equiv
    \frac{\exp \left(-\chi\cdot R_{S}\left(g\right)\right)}{\int_{\mathscr{G}} \exp \left(-\chi\cdot R_{S}\left(g\right)\right)\cdot \mathrm{d}\pi_{0}\left(g\right)} \cdot \mathrm{d}\pi_{0}\left(g\right),
\end{equation}
where
\begin{equation}
\label{EQ:OPTIMAL-CHI}
\chi
\equiv
4 n\cdot\sqrt{\frac{1}{2 n}\cdot\left(\mathrm{KL}\left(\Pi^* \| \pi_0\right)+\ln\left( \frac{2 \sqrt{n}}{\epsilon}\right)\right)}.
\end{equation}
\end{theorem}
We present a proof of this theorem in \cref{PROOF:OPTIMAL}.

The posterior distribution that we derive is analogous to the optimal posterior in \citet{McAllester2003} with the difference that our observations are mapped to the unit interval rather than to $\{-1,1\}$, which is the standard support in classification.
This particular distribution is common in the statistical mechanics literature and is a Boltzmann (or Gibbs) distribution, and has the form of exponential tilting of the prior, where the exponential tilting term involves the negative empirical welfare risk.
The degree of tilting depends upon the magnitude of $\chi>0$, which is the inverse of the Lagrange multiplier of the associated minimisation problem and that corresponds to the root of \cref{EQ:OPTIMAL-CHI}.
The Lagrange multiplier controls the extent to which $\pi_{0}$ is updated by empirical welfare risk in minimising the upper bound for $R^{\Pi}$.

Although \cref{THM:OPTIMAL} offers an analytical characterisation of the optimal posterior, we are unable to obtain a closed-form expression for the optimal posterior density.
Whilst \cref{EQ:OPTIMAL-CHI} does suggest a means to compute this density, it is likely that, in practice, this computation is difficult to perform with any precision.
Given that we have, however, established that this density exists, we can consider approximating it using the \emph{variational approximation} method, as is considered in \citet{Alquier2016}.


\subsection{Variational approximation of the optimal stochastic assignment rule}
\label{SEC:APPROXIMATION}

We develop a variational approximation of the optimal posterior density.
Variational approximation is useful in situations where Gibbs distributions are difficult to sample from directly,
such as is the case for graphical models where Markov Chain Monte Carlo (MCMC) sampling is costly \citep{Wainwright2008}.
In variational approximation, we choose to approximate the optimal posterior distribution via a family of distributions of our choice, $\mathscr{V}\subset\mathscr{M}$.
This allows us to develop an analytically tractable upper bound for the welfare regret attained by the resulting stochastic assignment rule.

Aside from guaranteeing tractability in estimation, variational approximation can also be motivated as a convenient way to impose constraints on the set of policies.
For instance, a fairness criterion requiring that individuals with similar characteristics have similar probabilities of treatment \citep{dwork2012fairness} can be enforced by specifying that $\mathscr{V}$ is a continuous family, and by limiting the concentration of the density.
A similar approach can be used if $\mathscr{V}$ is a parametric family to limit how much policy can deviate from the status quo, by fixing the parameters of the posterior distribution or restricting them to some set, say.

We approximate $\Pi^{\ast}$ (implicitly defined in \cref{THM:OPTIMAL}) by minimising the right-hand side of \cref{EQ:PAC-BOUND} with respect to posterior distributions in $\mathscr{V}$, defining
\begin{equation}
    \tilde{\Pi}
    \equiv
    \underset{\Pi \in \mathscr{V}}{\operatorname{argmin}} \,\left\lbrace R_{S}^{\Pi} + \sqrt{\frac{1}{2 n}\cdot\left(\mathrm{KL}\left(\Pi \| \pi_{0}\right)+\ln \left(\frac{2 \sqrt{n}}{\epsilon}\right)\right)}\right\rbrace.
\end{equation}
We then use this optimal variational posterior distribution to define a new bound for welfare regret from which we can characterise its convergence rate.
\begin{lemma}
\label{THM:VARIATIONAL}
With probability $1-\epsilon$ in terms of $P^n$, the expected welfare risk under the optimal variational posterior satisfies
\begin{equation}
    R^{\tilde{\Pi}} \leq \inf _{\Pi \in \mathscr{V}} \left\{ R^{\Pi} +\frac{a\left(\lambda, n\right)}{\lambda}+\frac{\mathrm{KL}\left(\Pi \| \pi_0\right)}{\lambda}+\frac{\ln \left(\frac{2}{\epsilon}\right)}{\lambda}+\sqrt{\frac{1}{2 n}\cdot\left(\mathrm{KL}\left(\Pi \| \pi_{0}\right)+\ln \left(\frac{4 \sqrt{n}}{\epsilon}\right)\right)} \right\},
\end{equation}
where $\epsilon\in\left(0,1\right)$ and $\lambda > 0$ are arbitrary constants and $a$ is a function that depends upon a positive constant $\lambda$ and the sample size $n$.
\end{lemma}
In what follows, we let $\mathscr{V}$ be a variational family of distributions that assigns positive density to assignment rules in $\mathscr{F}$ only.
Equivalently, we let $\mathscr{V}$ be a directional family that assigns positive density to unit vectors on the hypersphere, recalling that every policy in the LES class can be uniquely associated with a unit vector (see \cref{EQ:UNIQUENESSOFINDEX} and surrounding discussion).
A particularly tractable directional family, and one that we use, is the von Mises-Fisher family of distributions, which is characterised by a probability density function satisfying, for all $\kappa>0$ and $m$-dimensional unit vectors $\mu$,
\begin{equation}
    \mathrm{d}\Pi\left(\beta;\kappa,\mu\right)
    \equiv
    \frac{\kappa^{m/2-1}\cdot \exp\left({\kappa\cdot\mu^{\top}\beta}\right)}{\left(2\pi\right)^{m/2}\cdot I_{m/2-1}\left(\kappa\right)}\cdot\mathrm{d}\beta,
\end{equation}
where $I_{\nu}\left(z\right)$ is a modified Bessel function of the first kind with order $\nu$ and argument $z>0$.
A von Mises-Fisher distribution is the analogue of a multivariate Gaussian distribution on the unit hypersphere.
We refer to $\kappa$ as the concentration parameter and to $\mu$ as the mean direction (the location parameter), noting that the distribution becomes degenerate as $\kappa\rightarrow\infty$ and uniform as $\kappa\rightarrow 0$.

We now introduce some further notation that facilitates our analysis.
First, we let $\overline{R} = \inf_{\Pi \in \mathscr{M}} R^{\Pi}$ denote the minimum expected welfare risk amongst assignment rules in $\mathscr{F}$ and $\overline{\beta}$ denote the vector that parametrises the assignment rule that induces $\overline{R}$.
We also add the following assumption that restricts the marginal distribution of $X$.
\begin{assumption}[Margin assumption]
\label{ASS:CO}
There exists a constant $c>0$ such that, for any $m$-dimensional unit vectors $\beta^{\dagger}$ and $\beta^{\ddagger}$, $P\left(\langle X, \beta^{\dagger}\rangle\cdot\left\langle X, \beta^{\ddagger}\right\rangle \leq 0\right)
\leq
c\cdot\left\|\beta^{\dagger}-\beta^{\ddagger}\right\|_2$.
\end{assumption}
This assumption is satisfied whenever $X$ has bounded density on the unit hypersphere and is also present in the analysis of \citet{Alquier2016}.
An interpretation of this assumption is that the proportion of individuals whose treatment status switches is continuous with respect to the linear eligibility score coefficients, with an implication being that $R\left(\beta\right)-\overline{R}\leq 2c\cdot\|\beta - \overline{\beta}\|_2$.

We now present a high-probability uniform upper bound for the welfare regret of the stochastic assignment rule obtained by variational approximation via the von Mises-Fisher family of distributions.
\begin{theorem}
\label{THM:RATE}
Suppose that \cref{ASS:EV,ASS:U,ASS:BO,ASS:SO,ASS:CO} are satisfied, that $\pi_0$ is a uniform distribution over the unit hypersphere, and that $\Pi$ is a von Mises-Fisher distribution.
Then for $n\geq8$,
with probability at least $1-\epsilon$ in terms of $P^n$,
\begin{equation}
\label{EQ:DECAY}
    R^{\tilde{\Pi}} - \overline{R}
    \leq
    M\cdot\frac{\ln\left(n\right)}{\sqrt{n}},
\end{equation}
where $M$ is a universal constant.
\end{theorem}
We present proof of this theorem and provide an analytical expression for the universal constant in \cref{PROOF:RATE}.

The uniform upper bound on welfare regret that is defined by \cref{EQ:DECAY} decays at a rate of  $\ln\left(n\right)/\sqrt{n}$.
This rate is slightly slower than the welfare regret convergence rate of the EWM (deterministic) assignment rule studied in \citet{Kitagawa2018a}.
A simple comparison of these rates, however, is not quite meaningful for the following reason:
we do not know if the convergence rate of $\ln\left(n\right)/\sqrt{n}$ that is obtained in \cref{THM:RATE} is sharp or not.
\cref{THM:RATE} requires \cref{ASS:CO}, whilst \citet{Kitagawa2018a} does not impose this assumption in showing that $1/\sqrt{n}$ is the minimax optimal rate of welfare regret convergence.\footnote{
\citet[\S Theorem 2.3 and \S Theorem 2.4]{Kitagawa2018a} establishes that the minimax-optimal rate under a stronger condition than our \cref{ASS:CO} is $1/n^{2/3}$.
This stronger condition embeds a margin assumption that implies our \cref{ASS:CO} but also embeds the requirement that the first-best treatment rule is contained in the set of admissible decision rules -- that $\mathscr{F}$ contains the deterministic assignment rule that minimises expected welfare risk amongst all assignment rules.
We do not make any assumption about whether $\mathscr{F}$ (or indeed $\mathscr{G}$ in earlier parts of our analysis) contains the first-best assignment rule, or whether this rule is deterministic or stochastic.}
We do not know what the minimax optimal rate of welfare regret convergence is when \cref{ASS:CO} is additionally imposed and, hence, cannot rule out the possibility that the convergence rate of \cref{THM:RATE} can be improved upon and made faster than $1/\sqrt{n}$.
We leave further investigation of this matter for future research.

It is worth emphasising that the regret convergence result of \cref{THM:RATE} imposes weak restrictions on the distribution of data (Assumptions \ref{ASS:BO} - \ref{ASS:CO}) and does not require the specification of a likelihood function or of regression equations.
This contrasts with other approaches such as a Bayesian approach, where misspecification of likelihood can lead to non-convergence of the welfare regret even when Bayes optimal policies are constrained to deterministic ones in $\mathscr{F}$.

Our approach is similar to \citet{Alquier2016} in which the families of distributions for variational approximation are multivariate Gaussian distributions on Euclidean space with flexible covariance matrices.
In our approach, we stipulate a class of von Mises-Fisher distributions, which are the hyperspherical analogues of multivariate Gaussian distribution with diagonal covariance matrices featuring a constant variance element.
Since the empirical welfare criterion for LES rules is invariant to the scale of $\beta$, it is natural to consider von Mises-Fisher distributions rather than Gaussian ones. The scale invariance of von Mises-Fisher distributions can simplify optimisation of the variational approximation by reducing the set of optima to a singleton.


\section{Implementation}
\label{SEC:IMPLEMENTATION}

To implement our procedure, we restrict attention to $\mathscr{F}$ and let $\mathscr{V}$ be the von Mises-Fisher family of distributions.
Our goal is then to minimise the objective function,
\begin{equation}
\label{IMPLEMENTATION:OBJECTIVE}
    R^{\Pi}_S+\sqrt{\frac{1}{2n}\cdot\left(\text{KL}\left(\Pi|\pi_0\right)+\ln\left(\frac{2\sqrt{n}}{\epsilon}\right)\right)},
\end{equation}
with respect to $\kappa$ and $\mu$, which are the parameters of our chosen variational family (the concentration parameter and the mean direction, respectively).
Here, we reiterate that $\epsilon$ relates to the probability with which our high probability bounds hold.
We assume that $\pi_0$ is the uniform distribution over the sphere and set $\epsilon$ equal to the 5\% level throughout.

We propose numerically minimising the objective function, approximating $R^{\Pi}_{S}$ using Monte Carlo draws of $\beta$ from the von Mises-Fisher distribution for a given realisation of data and for fixed values of the parameters of the von Mises-Fisher distribution.
We let $\{\beta^j : j=1,...,J\}$ be the pseudo-random draws that we obtain, and compute
\begin{equation}
\label{EQ:RISK-IMPLEMENTATION}
    \hat{R}^{\Pi}_{S}
    \equiv
    \frac{1}{n}\sum_{i=1}^{n}h_{i}\cdot\left(D_{i}\cdot\left(1-\hat{\Pi}_{i}\right)+\left(1-D_i\right)\cdot\hat{\Pi}_{i}\right),
\end{equation}
where, for all $i=1,...,n$,
\begin{equation}
    \hat{\Pi}_{i}\equiv
    \frac{1}{J}\sum_{j=1}^{J}1\left(x_{i}^{\top}\beta^j\geq 0\right).
\end{equation}
Fast pseudo-random sampling of von Mises-Fisher random vectors is possible using the rejection sampling method of \citet{wood1994simulation} or the inversion method of \citet{kurz2015stochastic}.
The analogue of \cref{IMPLEMENTATION:OBJECTIVE} that we then minimise is
\begin{equation}
\label{IMPLEMENTATION:OBJECTIVE2}
\hat{R}^{\Pi}_S+\sqrt{\frac{1}{2n}\cdot\left(
\nu\cdot\ln\left(\frac{\kappa}{2}\right)-\ln\left(I_{\nu}\left(\kappa\right)\right)-\ln\left(\Gamma\left(\nu+1\right)\right)+\frac{I_{\nu+1}\left(\kappa\right)}{I_{\nu}\left(\kappa\right)}\cdot\kappa
+\ln\left(\frac{2\sqrt{n}}{\epsilon}\right)\right)},
\end{equation}
which follows from our assumption of a uniform prior and the form that the Kullback-Leibler divergence takes under that assumption \citep{KitagawaRowley}.

As is shown in \citet{KitagawaRowley}, the Kullback-Leibler divergence of the von Mises-Fisher distribution from the uniform distribution over the hypersphere does not depend upon $\mu$ and increases at the logarithmic rate in $\kappa$.
In contrast, although $R^{\Pi}_{S}$ is a function of $\mu$ and $\kappa$, the nature of this relationship is not clear.
That several combinations of $\mu$ and $\kappa$ induce local minima of $R^{\Pi}_{S}$ cannot be ruled out for instance.
Moreover, since we construct $\hat{R}^{\Pi}_{S}$ based upon random draws of $\beta$ from $\Pi$, our objective function is not a smooth function of $\mu$ and $\kappa$.
Given this, we suggest that a grid-based search for the minimum of the objective function is appropriate, with this search limited to values of $\kappa$ between zero and some specified upper limit.
Further insight about the behaviour of the objective function is provided in an online appendix.


\section{Empirical illustration}
\label{SEC:ILLUSTRATION}

We illustrate our procedure using data from the National Job Training Partnership Act (JTPA) Study.
Applicants to the Study were randomly allocated to one of two groups.
Applicants allocated to the treatment group were extended training, job search assistance and other services provided by the JTPA over a period of 18 months.
Applicants allocated to the control group were excluded from JTPA assistance.
Along with information collected prior to the commencement of the intervention, the Study also collected administrative and survey data relating to applicants' earnings in the 30 months following its start.
Further details about the data and the Study can be found elsewhere \citep[see, for instance,][]{bloom1997benefits}.
We restrict attention to a sample of 9,223 observations for which data on years of education and pre-programme earnings amongst the sample of adults (aged 22 years and older) used in the original evaluation of the programme and in subsequent studies (\citealp{bloom1997benefits,heckman1997matching,abadie2002instrumental}) is available.
Applicants in this sample were assigned to the treatment group with a probability of two thirds.
Like \citet{Kitagawa2018a}, we define $D$ to be the initial assignment of treatment, rather than the actual take-up due to the presence of non-compliance in the experiment.
We study stochastic assignment rules.

We follow \citet{Kitagawa2018a} in considering total individual earnings in the 30 months after programme assignment as our principal welfare measure.
Moreover, we focus exclusively on the class of linear rules,
\begin{equation}
\label{EMPIRICS:FCLASS}
\begin{aligned}
    \mathscr{F}
    &=
    \left\{x : \beta_{0}+\beta_{1}\cdot\text{prior earnings}+\beta_{2}\cdot\text{education}\geq 0\left\vert\,\|\left(\beta_{0},\beta_{1},\beta_{2}\right)\|_{2}=1\right.\right\},\\
    &=
    \left\{x : \beta_{0}+\beta_{1}\cdot x_{1}+\beta_{2}\cdot x_{2}\geq 0\left\vert\,\|\left(\beta_{0},\beta_{1},\beta_{2}\right)\|_{2}=1\right.\right\},\\
\end{aligned}
\end{equation}
that are studied in that paper.

\begin{figure}
    \centering
    \caption{Variation in treatment propensity across individuals in the JTPA Study sample}
    \label{FIG:PS}
    \includegraphics[width=\textwidth,height=3.5in]{BD-JTPA.pdf}
    \caption*{\footnotesize\textit{Notes:}
    JTPA Study sample.
    This figure illustrates the treatment propensity of individuals under the posterior assignment rule that is induced by $\left(\kappa^{*},\mu^{*}\right)$.
    Each point represents the individual characteristics of an individual or several individuals in the (crosses denote individuals with zero in-sample weight).
    For comparison, individuals to the left of the solid diagonal line are assigned treatment under the optimal deterministic assignment rule of \citet{Kitagawa2018a}.}
\end{figure}

To implement our procedure, we map prior earnings and education to the unit interval,\footnote{
We map each variable to the unit interval by dividing through by its maximum in the sample.
\citet{Kitagawa2018a} also does this.
Such a change of units is useful when the domain of one variable is much larger than the domain of another and the respective coefficients on the two variables reflect this.
For instance, in our sample, every individual has between seven and 18 years of education, and no individual earned more than \$63,000 prior to the start of the intervention.}
and calculate $H$ as outlined in \cref{EQ:WEIGHTSFORIMPLEMENTATION}.
We perform this calculation without adjusting post-programme earnings by the average cost of JTPA assistance (\$774 per individual) for treated individuals.\footnote{
We adjust post-programme earnings by the average cost of JTPA assistance in an online appendix.}
We then utilise a grid search approach over the parameters of the von Mises-Fisher distribution, specifying a reasonably fine grid over the unit sphere and over a finite subset of the reals.\footnote{
We design our grid so as to place an upper limit on the great-circle distance between any point on the sphere and its closest point on the grid.
Our grid comprises a total of 10,116 directional vectors combined with a sequence of evenly-spaced concentrations on the $[0,5]$ interval.
For reference, the surface area of the sphere is $4\pi$, which means that our grid has a density of approximately $0.001$.}

For each point on our grid, we draw 1,000 values of $\beta$ from the corresponding von Mises-Fisher distribution and approximate empirical welfare risk as per \cref{EQ:RISK-IMPLEMENTATION}.
We then substitute these values into \cref{IMPLEMENTATION:OBJECTIVE2} to provide an estimate of the objective function.\footnote{
In \cref{IMPLEMENTATION:OBJECTIVE2}, given \cref{EMPIRICS:FCLASS} and its restriction of $\beta$ to the unit sphere, $\nu=1/2$ .
More generally, maintaining our convention of defining $\mathscr{X}\in\mathbb{R}^{m-1}$ and adding to this vector a constant (i.e., an intercept), $\nu=m/2-1$.}

We find that the objective function is minimised (amongst the class of von Mises-Fisher distributed linear assignment rules) by the stochastic assignment rule with $\kappa=1.550$ and $\mu=\left(0.883,0.442,0.158\right)$,\footnote{
This directional vector can be represented by an azimuth of $27^\circ$ and an inclination of $81^\circ$ using spherical coordinates.}
which we label as $\kappa^{*}$ and $\mu^{*}$, respectively.

\begin{figure}
    \centering
    \caption{Behaviour of the objective function at $\mu^{*}$ given variation in $\kappa$}
    \label{FIG:KAPPA}
    \includegraphics[width=\textwidth,height=3.5in]{OR-JTPA.pdf}
     \caption*{\footnotesize\textit{Notes:}
    JTPA Study sample.
    This figure illustrates the shape of the objective function and its risk component at $\mu^{*}$ as $\kappa$ is varied;
    high concentration is associated with low risk but incurs a large penalty for divergence from the uniform prior. }
\end{figure}

The optimal stochastic assignment rule, on average, assigns treatment to individuals in the JTPA Study sample around 83\% of the time.
This probability is not, however, uniform, and there is some variation in the probability with which distinct individuals are assigned treatment.
This variation in assignment propensity can be seen in \cref{FIG:PS}, which plots the individual characteristics of all individuals in the sample.
The propensity with which individuals with distinct characteristics are assigned treatment is represented by the color of each point, and the weight given to individuals in the sample with particular characteristics is represented by the size of each point.
The weight attached to a given point is proportional to the sum of post-programme earnings over all individuals with those characteristics.\footnote{
To simplify \cref{FIG:PS}, we scale the weights such that they sum to one.}
\cref{FIG:PS} shows that the optimal stochastic assignment rule is more likely to assign treatment to an individual with few years of education and high prior earnings than an individual with more years of education and lower prior earnings, with the assignment probability ranging from 78\% to 84\%.
The deterministic assignment rule of \citet{Kitagawa2018a}, in contrast, assigns only individuals with few years of education and low prior earnings to treatment, with around 93\% of individuals assigned treatment.
We plot this deterministic rule as a useful benchmark for comparison in \cref{FIG:PS}.

It is important to emphasise that it is the regularisation term and, in particular, the Kullback-Leibler divergence that limits the value of $\kappa$ at the optimum, and leads to an interior probability of assignment for all individuals.
This can be seen in \cref{FIG:KAPPA}, which plots the USD equivalent of the objective function (left-hand axis, solid line) and of empirical welfare risk (right-hand axis, dashed line) for a range of $\kappa$, holding fixed $\mu$ at $\mu^{*}$.
We observe that empirical welfare risk decreases as $\kappa$ increases, remaining low and constant once its value is sufficiently large.\footnote{
That empirical welfare risk decreases for small to moderate values of $\kappa$ is specific to the data and chosen $\mu$, and is arguably also attributable to the lack of consideration given to the cost of treatment.}
The intuition here is that large values of $\kappa$ lead to stochastic assignment rules that mimic deterministic ones;
we expect $\mu$ to eventually coincide with the deterministic assignment rule of \citet{Kitagawa2018a} as the value of the concentration parameter approaches infinity,since there does not exist a (linear) deterministic rule that can improve upon this.
Tempering this preference towards large values of $\kappa$ is the Kullback-Leibler divergence of the von Mises-Fisher distribution from the uniform distribution, which is increasing at the logarithmic rate in $\kappa$ and generates the difference between the objective function and empirical welfare risk in \cref{FIG:KAPPA}.
As $\kappa$ increases, the regularisation term begins to dominate.

\begin{figure}
    \centering
    \caption{Deterministic assignment rules and empirical welfare risk}
    \label{FIG:DETERMINISTICRULES}
    \includegraphics[width=\textwidth,height=3.5in]{DHM-JTPA.pdf}
    \caption*{\footnotesize\textit{Notes:}
    JTPA Study sample.
    This figure illustrates the risk that is associated with (deterministic) assignment rules in $\mathscr{F}$.
    A spherical coordinate mapping is implemented.
    The intersection of the two white lines is located at $\mu^{*}$.
    The intersection of the two black lines is located at the optimal deterministic assignment rule of \citet{Kitagawa2018a}, which attains the minimal regret amongst all deterministic linear rules.}
\end{figure}

\begin{figure}
    \centering
    \caption{Behaviour of the objective function at $\kappa^{*}$ given variation in the mean direction $\mu$}
    \label{FIG:MU}
    \includegraphics[width=\textwidth,height=3.5in]{HM-JTPA.pdf}
    \caption*{\footnotesize\textit{Notes:}
    JTPA Study sample.
    This figure illustrates the risk that is associated with (stochastic) assignment rules in $\mathscr{V}$;
    the concentration parameter is fixed at $\kappa^{*}$ whilst $\mu$ is varied.
    The intersection of the two white lines is located at $\mu^{*}$.
    The intersection of the two black lines is located at the optimal deterministic assignment rule of \citet{Kitagawa2018a}, which attains the minimal regret amongst all deterministic linear rules.}
\end{figure}

To better understand why $\mu^*$ does not coincide with the deterministic assignment rule of \citet{Kitagawa2018a} holding fixed $\kappa$ at $\kappa^*$, we refer to \cref{FIG:DETERMINISTICRULES}, which plots empirical welfare risk for all vectors on the unit sphere -- i.e., the empirical welfare risk associated for each deterministic assignment rule in $\mathscr{F}$.
\cref{FIG:DETERMINISTICRULES} utilises the spherical coordinate system
\begin{equation}
    \left(\beta_0,\beta_1,\beta_2\right)=\left(\cos\left(\theta\right)\sin\left(\phi\right),\sin\left(\theta\right)\sin\left(\phi\right),\cos\left(\phi\right)\right),
\end{equation}
where $\theta\in\left[-180^\circ,180^\circ\right)$ is the azimuth and $\phi\in\left[0,180^\circ\right]$ is the inclination.
It is perhaps convenient to think of the azimuth as related to longitude and the inclination as related to latitude.
For non-trivial $\kappa$, the von Mises-Fisher distribution allocates probability mass to the sphere in such a way that its density contours are concentric about $\kappa$, with points closer to $\mu$ more likely to occur.
The deterministic assignment rule of \citet{Kitagawa2018a} can be seen from \cref{FIG:DETERMINISTICRULES} to be located on the boundary between a high risk region (no-one treated) and a moderate risk region (everyone treated).
As such, a stochastic assignment rule with non-trivial $\kappa$ and $\mu$ located at this point would approximately allocate probability mass to each of these regions in equal amounts.
By shifting $\mu$ towards the centre of the moderate risk region, we allocate relatively more mass to rules that induce moderate risk and less mass to rules that induce high risk, which reduces empirical welfare risk overall.

This pattern underlies what we observe in \cref{FIG:MU}, which plots empirical welfare risk for all $\mu$ on the sphere holding fixed $\kappa$ at $\kappa^{*}$.
Despite the apparent discontinuity of risk over deterministic assignment rules, empirical welfare risk (and the objective function) appear to vary smoothly.


\section{Conclusion}
\label{SEC:CONCLUSION}

Central to our analysis of the treatment choice problem is the question of \textit{how should the social planner allocate individuals to treatment with a given probability rather than with certainty?}
To answer this question, we focus on stochastic assignment rules that we formulate as posterior distributions obtained from well-defined prior distributions via a PAC-Bayes approach.
These distributions are able to accommodate any initial belief that the social planner holds about what constitutes the best treatment, as well as any budgetary, ethical or legal constraints that are imposed and that can arise due to concerns about fairness or about maintaining the status quo.
We establish that it is possible to obtain the minimum expected welfare risk under a variational approximation of the optimal posterior distribution, which we also characterise.
In keeping with the notion that variational approximation replaces a general class of assignment rules with a simpler set of policies, we focus on stochastic assignment rules that can be expressed as density functions over the LES class of (deterministic) assignment rules.
We exploit the scale invariance of the LES class to restrict attention to distributional families on the unit hypersphere, selecting the von Mises-Fisher family of distributions as our chosen variational family.
We demonstrate how our methods can be used in an empirical setting by estimating which individuals should be entered onto a job training programme, using data from the well-known JTPA Study.

Our research suggests several further questions that remain unanswered.
\textit{How does the choice of variational family affect the rate of convergence?}
The von Mises-Fisher family of distributions is not the only directional family -- other families such as the Matrix Langevin and Kent families allow for richer covariance structures.
Whilst more general distributions admit more complex stochastic assignment rules, does this additional generality come at an additional cost?
\textit{Do other parametrisations of the class of assignment rules exist that can achieve faster decay rates?}
We establish that the optimal stochastic assignment rule in our chosen variational family yields a convergence rate for welfare regret of $\ln\left(n\right)/\sqrt{n}$.
Ignoring the logarithm in the numerator, this rate coincides with the minimax optimal rate for deterministic assignment rules in the absence of the margin assumption, as shown in \citet{Kitagawa2018a}.
Our convergence rate result, however, relies upon Assumption \ref{ASS:CO} that constrains the marginal distribution of $X$.
We do not know whether this assumption leads to the minimax optimal rate being faster than $1/\sqrt{n}$ and whether the proposed method can attain a faster convergence rate than $\ln\left(n\right)/\sqrt{n}$.