Survey calibration for causal inference: a simple method to balance covariate distributions

Survey calibration for causal inference: a simple method to balance covariate distributions

\onehalfspacing

Keywords: calibration estimators, quantile estimation, balancing weights, observational data

JEL: C18, C21, C83.

\doublespacing

Introduction

Recent literature on causal inference for observational studies includes several approaches to balancing whole distributions of covariates rather than moments (e.g. means, variances). In particular, hazlett_kernel_2020 proposes kernel entropy balancing (KEB), which consists in making the multivariate density of covariates approximately equal for the treated and control groups when the same choice of kernel is used to estimate these densities. zhao_covariate_2019 proposes a covariate balancing rule that modifies the balancing propensity score model by reproducing the kernel Hilbert space and is implemented within the framework of the tailored loss function approach. Finally, sant2022covariate proposes an integrated propensity score model, which aims to minimise imbalances in the joint distribution of covariates.

In this paper we propose a simple method, which consists in balancing control and treatment distributions at specific quantiles along with moments, if needed. Drawing on the theory of calibration estimators in survey sampling deville1992calibration and, in particular, on quantile calibration estimators harms2006calibration, we develop methods for observational studies. The proposed method involves adding new variables to be balanced based on pre-treatment covariates $\boldsymbol{X}$ using a modified Heaviside function (or its approximation), thus reducing bias and the root mean square error. The procedure does not require a prior knowledge of the covariate distributions, does not involve integration or tuning parameters such as bandwidths, and is computationally simple as it applies local linear approximation (for example through step-wise regression).

In this paper, we extend two methods for balancing covariate moments between treatment and control groups, namely entropy balancing proposed by hainmueller2012entropy, and the covariate balancing propensity score proposed by imai2014covariate. As our method does not add new restrictions on assumptions, nor is it restricted to specific treatment designs, it can be further applied to other methods. For example, our approach has already been extended to multi-category designs and other methods by Noah Greifer (Harvard University; Institute for Quantitative Social Science), who has already implemented our approach in his R package WeightIt (ver. 0.14.2.9003, weightit) for binary and multi-category (for selected also continous) treatments using the non-parametric covariate balancing propensity score fong2018covariate, inverse probability tilting graham2012inverse, optimisation-based weighting wang2020minimal and energy balancing weighting huling2023independence, huling2020energy. Our approach can be further generalised to other designs (e.g. continuous, longitudinal), assumptions (e.g. treatment is endogenous) or methods (e.g. synthetic control method as an alternative to gunsilius_distributional_2023).

The paper is structured as follows. Section (ref) presents the theory underlying calibration estimators for means/totals and quantiles. Section (ref) describes the proposed approach for the entropy balancing hainmueller2012entropy and the covariate balancing propensity score imai2014covariate, which we refer to as distributional entropy balancing (DEB) and distributional propensity score (DPS) methods, respectively. Section (ref) presents results of two simulation studies aimed at validating the two methods: DEB is used to estimate the average treatment effect on the treated (ATT) and the quantile treatment effect on the treated (QTT), while DPS is used to estimate the average treatment effect (ATE) and the quantile treatment effect (QTE). Section (ref) presents an empirical example aimed at estimating the ATE and QTE of 401(k) retirement plans on asset accumulation benjamin2003does and compares our results with those obtained by sant2022covariate. The paper ends with a conclusion and additional results are presented in the Appendix.

Theoretical basis from survey sampling

Calibration estimator for a total

Let $\boldsymbol{X}$ be a random auxiliary variable and $Y$ be the target random variable of interest. In most applications, the goal is to estimate a finite population total $~{\tau_{Y}=\sum_{k\in U}Y_{k}}$ or the mean $\bar{\tau}_{Y}=\tau_{Y}/N$ of the variable of interest $Y$, where $U$ is the population of size $N$. The Horvitz-Thompson estimator is a well-known estimator of a finite population total, which is expressed as $~{\hat{\tau}_{Y\pi}=\sum_{k=1}^{n}d_{k}Y_{k}=\sum_{k\in s}{d_{k}Y_{k}}}$, where $s$ denotes a probability sample of size $n$, $d_{k}=1/\pi_{k}$ is a design weight, and $\pi_{k}$ is the first-order inclusion probability of the $k$-th element of the population $U$. This estimator is unbiased for $\tau_{Y}$ i.e. $E\left(\hat{\tau}_{Y\pi}\right)=\tau_{Y}$.

Let $\boldsymbol{X}_{k}^{\circ}$ be a $J_{1}$-dimensional vector of auxiliary variables (benchmark variables) for which $~{\tau_{\boldsymbol{X}}=\sum_{k\in U}\boldsymbol{X}_{k}^{\circ}=\left(\sum_{k\in U}X_{k1},\ldots,\sum_{k\in U}X_{kJ_{1}}\right)^T}$ is assumed to be known. In most cases, in practice the $d_{k}$ weights do not reproduce known population totals for benchmark variables $\boldsymbol{X}_{k}^{\circ}$. It means that the resulting estimate $\hat{\tau}_{\boldsymbol{X}\pi}=\sum_{k\in s}{d_{k}\boldsymbol{X}_{k}^{\circ}}$ is not equal to $\tau_{\boldsymbol{X}}$. The main idea of calibration is to look for new calibration weights $w_{k}$ that are as close as possible to original design weights $d_{k}$ and reproduce known population totals $\tau_{\boldsymbol{X}}$ exactly. In other words, in order to find new calibration weights $w_{k}$ we have to minimise a distance function $D\left(\boldsymbol{d},\boldsymbol{v}\right)=\sum _{k\in s}d_{k}\hspace{2pt} G\hspace{0pt}\left(\frac{v_{k}}{d_{k}}\right) \to \textrm{min}$ to fulfil calibration equations $\sum_{k\in s}v_{k}\boldsymbol{X}_{k}^{\circ} = \sum_{k\in U}\boldsymbol{X}_{k}^{\circ}$, where $\boldsymbol{d}=\left(d_{1},\ldots,d_{n}\right)^T$, $\boldsymbol{v}=\left(v_{1},\ldots,v_{n}\right)^T$ and $G\left(\cdot\right)$ is a function which must satisfy some regularity conditions: $G\left(\cdot\right)$ is strictly convex and twice continuously differentiable, $G\left(\cdot\right)\geq 0$, $G\left(1\right)=0$, $G'\left(1\right)=0$ and $G''\left(1\right)=1$. Examples of $G\left(\cdot\right)$ functions are given by deville1992calibration. For instance, if $G\left(x\right)=\frac{\left(x-1\right)^{2}}{2}$, then using the method of Lagrange multipliers the final calibration weights $w_{k}$ can be expressed as $w_{k}=d_{k}+d_{k}\left(\tau_{\boldsymbol{X}}-\hat{\tau}_{\boldsymbol{X}\pi}\right)^T\left(\sum_{j\in s}d_{j}\boldsymbol{X}_{j}^{\circ}\boldsymbol{X}_{j}^{\circ T}\right)^{-1}\boldsymbol{X}_{k}^{\circ}$. It is worth adding that in order to avoid negative or large $w_{k}$ weights in the process of minimising the $D\left(\cdot\right)$ function, one can consider some boundary constraints $L\leq \frac{w_{k}}{d_{k}}\leq U$, where $\ 0\leq L\leq 1 \leq U,\ k=1,\ldots,n$. The final calibration estimator of a population total $\tau_{Y}$ can be expressed as $\hat{\tau}_{Y\boldsymbol{X}}=\sum_{k\in s}w_{k}y_{k}$, where $w_{k}$ are calibration weights obtained after selecting a given $G\left(\cdot\right)$ function.

Note that the Kullback-Leibler (KL) distance function, given by $~D(d,v) = \sum_k q_k \{v_k \log(\frac{v_k}{d_k}) -v_k + d_k\}$, reduces to $D(d,v) = \sum_k v_k \log(\frac{v_k}{d_k})$ if $q_k =1$ and $\sum_k v_k = N$. This distance function is known as entropy divergence in hainmueller2012entropy, minimum entropy distance in deville1992calibration, exponential tilting (EL) in kim2010calibration or generalized exponential tilting in wu2016calibration. After applying the KL distance function one obtains $G(x)=x\log(x) - x + 1$, which is the well known raking ratio method for categorical data and is available in statistical software for sample surveys (e.g. in the calibrate function with calfun="raking" argument in the R \texttt{survey} package; cf. r-survey).

deville1992calibration and kim2010calibration proved that the use of the KL function is asymptotically equivalent to a regression estimator (i.e. using $G(x) = \frac{1}{2}(x-1)^2$ function) for sample surveys (i.e. a calibration estimator can be efficient if $Y_k$ is linearly related to $\boldsymbol{X}_k$). This was proved once again by zhao2016entropy in the context of observational studies.

Calibration estimator for a quantile

harms2006calibration considered a way of estimating quantiles using the calibration approach, which is very similar to that proposed by deville1992calibration for a finite population total $\tau_{Y}$. By analogy, in their approach it is not necessary to know values for all auxiliary variables for all units in the population. It is enough to know the corresponding quantiles for the benchmark variables. Let us briefly discuss the problem of finding calibration weights in this setup.

We want to estimate a quantile $Q_{Y,\alpha}$ of order $\alpha \in \left(0,1\right)$ of the variable of interest $Y$, which can be expressed as $Q_{Y,\alpha}=\mathrm{inf}\left\{t\left|F_{Y}\left(t\right)\geq \alpha \right.\right\}$, where $F_{Y}\left(t\right)=N^{-1}\sum_{k\in U}H\left(t-Y_{k}\right)$ and the Heaviside function is given by

We assume that $\boldsymbol{Q}_{\boldsymbol{X}^*,\alpha}=\left(Q_{X^*_{1},\alpha},\ldots,Q_{X^*_{J_{2}},\alpha}\right)^{T}$ is a vector of known population quantiles of order $\alpha$ for a vector of auxiliary variables $\boldsymbol{X}_{k}^{*}$, where $\alpha \in \left(0,1\right)$ and $\boldsymbol{X}_{k}^{*}$ is a $J_{2}$-dimensional vector of auxiliary variables. Note that we distinguish $\boldsymbol{X}^{\circ}_k$ and $\boldsymbol{X}^{*}_k$ auxiliary variables for which we wish to reproduce totals and quantiles respectively. This is because, in general, the numbers $J_{1}$ for $\boldsymbol{X}^{\circ}_k$ and $J_{2}$ for $\boldsymbol{X}^{*}_k$ may be are different. It may happen that for a specific auxiliary variable its population total and the corresponding quantile of order $\alpha$ will be known. However, in most cases, quantiles will be known for continuous auxiliary variables, unlike totals, which will generally be known for categorical variables. In most observational studies unit level data is available for all $\boldsymbol{X}$ variables thus possibility of calibration to specific quantiles is easier than for sample surveys.

In order to find new calibration weights $w_{k}$ that reproduce known population quantiles in a vector $\boldsymbol{Q}_{\boldsymbol{X}^*,\alpha}$, an interpolated distribution function estimator of $F_{Y}\left(t\right)$ is defined as $ \hat{F}_{Y,cal}(t)=\frac{\sum_{k \in s} w_{k} H_{Y, s}\left(t, Y_{k}\right)}{\sum_{k \in s} w_{k}} $, where the Heaviside function in formula ((ref)) is replaced by the modified function $H_{Y, s}\left(t, Y_{k}\right)$ given by

where $L_{Y, s}\left(t\right)=\max \left\{\left\{Y_{k} \mid Y_{k} \leqslant t, k \in s\right\} \cup\{-\infty\}\right\}$, $U_{Y, s}\left(t\right)=\min \left\{\left\{Y_{k} \mid Y_{k}>t, k \in s\right\} \cup\{\infty\}\right\}$ and $\beta_{Y, s}\left(t\right)=\frac{t-L_{Y, s}\left(t\right)}{U_{y, s}\left(t\right)-L_{Y, s}\left(t\right)}$ for $k=1,\ldots,n$, $t \in \mathbb{R}$. A calibration estimator of quantile $Q_{Y,\alpha}$ of order $\alpha$ for variable $Y$ is defined as $\hat{Q}_{Y,cal,\alpha}=\hat{F}_{Y,cal}^{-1}(\alpha)$, where a vector $\boldsymbol{w}=\left(w_{1},\ldots,w_{n}\right)^{T}$ is a solution of optimization problem $D\left(\boldsymbol{d},\boldsymbol{v}\right)=\sum _{k\in s}d_{k}\hspace{2pt} G\hspace{0pt}\left(\frac{v_{k}}{d_{k}}\right) \to \textrm{min}$ subject to the calibration constraints $\sum_{k\in s}v_{k}=N$ and $\hat{\boldsymbol{Q}}_{\boldsymbol{X}^*,cal,\alpha}=\left(\hat{Q}_{X^*_{1},cal,\alpha},\ldots,\hat{Q}_{X^*_{J_{2}},cal,\alpha}\right)^{T}=\boldsymbol{Q}_{\boldsymbol{X}^*,\alpha}$ or equivalently $\hat{F}_{X_{j}^*,cal}\left(Q_{X^*_{j},\alpha}\right)=\alpha$, where $j=1,\ldots,J_{2}$.

As in the previous case, if $G\left(x\right)=\frac{\left(x-1\right)^{2}}{2}$ then using the method of Lagrange multipliers the final calibration weights $w_{k}$ can be expressed as $w_{k}=d_{k}+d_{k}\left(\mathbf{T_{a}}-\sum_{k\in s}{d_{k}\boldsymbol{a}_{k}}\right)^{T}\left(\sum_{j\in s}{d_{j}}\boldsymbol{a}_{j}\boldsymbol{a}_{j}^{T}\right)^{-1}\boldsymbol{a}_{k}$, where $\mathbf{T_{a}}=\left(N,\alpha,\ldots,\alpha\right)^{T}$ and the elements of $\boldsymbol{a}_{k}=\left(1,a_{k1},\ldots,a_{kJ_{2}}\right)^{T}$ are given by

with $j=1,\ldots,J_{2}$. Alternatively, one can consider the logistic function instead of (ref)

where $X^*_{kj}$ is the $k$th row of the auxiliary variable $X^*_j$ ($j=1,...,J_2$), $N$ is the population size, $Q_{X^*_j, \alpha}$ is the known population $\alpha$-th quantile, and $l$ is a constant set to a large value (e.g. 1,000).

In the next sections we describe how this method can be applied to make causal inferences in observational studies. We focus on four causal parameters: ATT, QTT, ATE and QTE.

Proposed approaches

Setup

Let us assume that $\mathcal{D}_k = \{0, 1\}$ is a treatment indicator variable, sample $s_{0}$ denotes the control group of size $n_0$, $s_{1}$ denotes the treatment group of size $n_1$ and $Y_k$ denotes the variable of interest, where $Y_k(1)$ and $Y_k(0)$ are the potential outcomes for the treatment group and the control group, respectively. The realised outcome is $~{Y_k=\mathcal{D}_k{Y_k(1)}+(1-\mathcal{D}_k)Y_k(0)}$ and $\boldsymbol{X}_k^{\circ}$ is vector of pre-treatment covariates whose support is denoted as $\mathcal{X}$. Let $p(\boldsymbol{x})=\mathbb{P}(\mathcal{D}_k=1|\boldsymbol{X}_k^{\circ}=\boldsymbol{x})$ be the propensity score where $x \in \mathcal{X}$ and for $\delta = \{0, 1\}$ the distribution and the quantile of the potential outcome $Y_k(\delta)$ is given by $F_{Y_k(\delta)}(y) = \mathbb{P}(Y_k(\delta) \leq y)$ and $q_{Y_k(\delta)}(\alpha) = \mathrm{inf}\left\{t\left|F_{Y_k(\delta)}\left(t\right)\geq \alpha \right.\right\}$.

Let us further assume that the researcher is interested in estimating the average treatment effect on the treated $ATT=\mathbb{E}[Y(1) \mid \mathcal{D}=1]-\mathbb{E}[Y(0) \mid \mathcal{D}=1]$, the quantile treatment effect on the treated $QTT(\alpha) = q_{Y(1) \mid D=1}(\alpha)-q_{Y(0) \mid D=1}(\alpha)$, an overall average treatment effect $ATE=\mathbb{E}(Y(1)-Y(0))$ or the quantile treatment effect $QTE(\alpha) = q_{Y(1)}-q_{Y(0)}$.

In this paper we follow a commonly used identification strategy in policy evaluation rosenbaum1983central, firpo_efficient_2007: 1) given $\boldsymbol{X}_k^{\circ}$, $(Y_k(1),Y_k(0))$ are jointly independent from $\mathcal{D}_k$ (conditional ignorability), 2) $\forall_{\boldsymbol{x} \in \mathcal{X}}\; p(\boldsymbol{x})$ is bounded away from zero and one, 3) uniqueness of quantiles. The identification strategy is the same as that used in the literature mentioned above, since we adopt the same assumptions.

For the ATT and QTT the counterfactual mean and $\alpha$-quantile can be estimated as

$$

$$

where $\omega_k$ is a weight chosen for each control unit.

For ATE one can use the approach suggested by rosenbaum1987model, i.e.

and for QTE with $\delta \in \{0, 1\}$, $F_{Y(\delta)}(y)$ is identified by

where $1\{.\}$ is the indicator function, which means that $QTE(\alpha)$ can also be written as functionals of the observed data sant2022covariate.

Distributional entropy balancing

hainmueller2012entropy proposed entropy balancing (EB) to reweight the control group to the known characteristics of the treatment group to estimate ATT and QTT. This method can be summarised as follows when only the first moments are constrained

where $v_k$ is defined as previously, $d_k >0$ is the base weight for unit $k$ set to e.g. $d_k=1/n_0$ and $m_j$ is the mean of the $X^{\circ}_j$-th covariate in the treatment group. As in the case of calibration, $\omega_k$ are solutions to (ref).

As we discussed in section (ref) the EB method is a variant of calibration approach where the KL distance function is used. Therefore, we can be simply extended (ref) to achieve not only the mean balance but also the distributional balance via quantiles. Instead of using known or estimated population totals $\boldsymbol{Q}_{\boldsymbol{X}^*,\alpha}$, we can use treatment group quantiles denoted by $\boldsymbol{q}_{\boldsymbol{X}^*,\alpha} = \left(q_{X^*_{1}, \alpha},\ldots,q_{X^*_{J_{2}},\alpha}\right)^{T}$, where the same $\alpha$ is applied for all $\boldsymbol{X}^{*}$ variables and the definition of the vector $\boldsymbol{a}_{k}=\left(1,a_{k1},\ldots,a_{kJ_{2}}\right)^{T}$ changes to

with $j=1,\ldots,J_{2}$ where $n_1$ is the size of the treatment group. Alternatively, one can use a modified (ref) given by

$$ a_{kj} = \frac{1}{1+ \exp\left(-2l\left(X^*_{kj}-q_{X^*_j, \alpha}\right)\right)}\frac{1}{n_1}. $$

Our proposal, which leads to distributional entropy balancing (hereinafter DEB), consists of extending the original idea by adding additional constraint(s) on the weights on $\boldsymbol{a}_k$, as presented below where the same $\alpha$ is applied for all $\boldsymbol{X}^*_j\; j=1,...,J_2$ variables

$$

$$

This approach can be easily extended for vector $\boldsymbol{\alpha}$, say quartiles $(0.25, 0.5, 0.75)$ or deciles $(0.1, \ldots, 0.9)$. Our approach is similar to that proposed by hazlett_kernel_2020, who extended (ref) by replacing the first condition by $\sum_{k \in s_0 } v_{k}\phi\left(X_{k}\right) =\frac{1}{n_{1}} \sum_{k \in s_1} \phi\left(X_{k}\right)$, where $\phi\left(X_{k}\right)$ are the basis functions for the kernel function (in particular the Gaussian kernel). Our approach is simpler as we locally approximate this relationship with a step-wise (constant) regression because the EB method assumes linear model of $Y_k$ and balancing variables.

Remark 1: zhao2016entropy showed that EB method is doubly robust with respect to linear outcome regression and logistic propensity score regression. Including quantiles in the constraints is simply adding new variables to outcome and propensity score regression (as higher order of moments). This leads to step-wise linear/logistic regression, which can approximate non-linear relationships in both models. Effectiveness of this approach will be showed in the simulation study.

Remark 2: Instead of modelling the whole distribution, one can start by adjusting the medians or quartiles. The number of quantiles for $\boldsymbol{X}^*_k$ can vary. In the simulation study we show that even a small number of quantiles significantly improves the estimates, especially in the presence of non-linear relationships.

Remark 3: This approach assumes that the distributions of $\boldsymbol{X}_k^*$ between the control and treatment groups have the same support, i.e. it is possible to generate a vector $\sum_k a_{kj} > 0$.

Our approach can be further applied to hierarchically regularised entropy balancing, as proposed by xu_hierarchically_2023, or the empirical likelihood method, as recently discussed by zhang_calibration_2022. Development version the WeightIt weightit package in R rcran currently supports binary and multi-category treatments.

Distributional propensity score method

imai2014covariate proposed the covariate balancing propensity score (CBPS) to estimate the (ref), where unknown parameters of the propensity score model $\boldsymbol{\gamma}$ are estimated using the generalized method of moments as

where $p(\dot)$ is the propensity score with unknown vector of parameters $\boldsymbol{\gamma} \in \mathbb{R}^{J_1}$ (including the intercept). Equation (ref) balances means (if $f(\boldsymbol{X}^{\circ}=\boldsymbol{X}^{\circ})$; or other moments, if specified) of the $\boldsymbol{X}^{\circ}$ variables, which may not be sufficient if the variables are highly skewed or we are interested in estimating DTE or QTE.

We propose a simple approach based on the specification of moments and $\alpha$-quantiles to be balanced. Instead of using the matrix $\boldsymbol{X}^{\circ}$ (including constant), we propose either using $\boldsymbol{X}^*$ solely (i.e. balancing only quantiles) as given below or $\boldsymbol{X}^{\circ}$ and $\boldsymbol{X}^*$ jointly (i.e. balancing both quantiles and moments). Both cases we denote distributional propensity score (DPS) method as in both cases the goal is to balance distribution of $\boldsymbol{X}^*$ variables via $\alpha$-quantiles.

To describe the main idea let's focus on the first case where only $\alpha$-quantiles $\boldsymbol{X}^*$. In such case the equation (ref) is replaced by

where rows of matrix $\boldsymbol{A}$ are given by $\boldsymbol{a}_{k}=\left(1,a_{k1},\ldots,a_{kJ_{2}}\right)^{T}$ and elements $\boldsymbol{a}_{k}$ for treatment units are given by

where $n_1$ is the size of the treatment group and $q_{X^*_{j,\alpha}}$ is $\alpha$-quantile for $X_j^*$ for treatment group and for control units

Alternatively, the logistic function (ref) can be used. Note that the elements of $\boldsymbol{A}$ for treatment group will sum up to the selected $\alpha$ orders of the quantiles as $n_1 \to \infty$ (though for small sample sizes, they may not sum up to the specified $\alpha$). As a result, the propensity score weights balance the $\alpha$ orders of the treatment and control groups and, as shown in harms2006calibration, the $\alpha$ quantiles of the selected variables.

The second case of the DPS method combines both $\boldsymbol{X}^{\circ}$ and $\boldsymbol{X}^*$ where $\boldsymbol{A}$ equation (ref) is then replaced by $\boldsymbol{X}$ as in

where $\boldsymbol{X} = [\boldsymbol{A} \; \boldsymbol{X}^{\circ}]$ and the intercept in the $\boldsymbol{X}^{\circ}$ is removed as $\boldsymbol{A}$ already contains it. DPS presented in (ref) will preserve not only means but also $\alpha$-quantiles of selected variables if $f(\boldsymbol{X}) = \boldsymbol{X}$.

Remark 4. The ATE or QTE estimator based on the CBPS or DPS method is doubly robust if either the outcome or the propensity score is correct. If we use the DPS method presented in (ref) we assume that $\mathbb{E}(Y|\boldsymbol{X})$ is approximated by step-wise linear regression with steps created by $\alpha$-quantiles through $\boldsymbol{A}$. The same applies to $p(\boldsymbol{X}, \boldsymbol{\gamma})$, where $\boldsymbol{A}$ is used to add the step-wise linear part to e.g. logistic regression. It means that the inclusion of $\alpha$-quantiles makes it possible to approximate non-linear relationships of $Y$ or $\mathcal{D}$ and $\boldsymbol{X}^\circ$ with step-wise linear models.

As our method simply involves adding new covariates (like adding higher moments), it is possible to use standard procedures to estimate $\boldsymbol{\gamma}$ parameters with a just-identified or over-identified set of equations (e.g. generalised method of moments, empirical likelihood). Since we do not change the estimation method itself, any method proposed in the literature can be applied. Furthermore, this matrix can be plugged into a high-dimensional setting with variable selection as in ning2020robust, non-parametric CBPS as in fong2018covariate or the improved CBPS proposed by fan2016improving. Our approach is not limited to binary cases but can also be used for multi-category or continuous treatments or longitudinal settings. In this way, one can achieve distribution balancing without developing more sophisticated methods than those proposed by imai2014covariate.

The development version of the WeightIt R package currently supports only non-parametric CBPS for binary and multi-category treatments but our approach can be used either by creating matrix $\boldsymbol{A}$ independently or by applying our R package jointCalib containing the joint_calib_cbps function, which relies on the CBPS package cbps-pkg. In the next section we verify the DEB and DPS methods in simulation studies.

Simulation results

Simulation for the DEB method

To show the effectiveness of our approach we follow the simulation procedure described by hainmueller2012entropy. We generate 6 variables: three ($X_1, X_2$ and $X_3$) from a multivariate normal distribution $\operatorname{MVN}(\boldsymbol{0}, \boldsymbol{\Sigma})$, where

$$ \boldsymbol{\Sigma} =

, $$

$X_4 \sim \operatorname{Uniform}[-3,3]$, $X_5 \sim \chi^2(6)$ and $X_6 \sim \operatorname{Bernoulli}(0.5)$. The treatment and control groups are formed using

$$ \mathcal{D} = \boldsymbol{1}[X_1 + 2X_2 - 2X_3 - X_4 -0.5X_5 + X_6 + \epsilon > 0]. $$

We consider three designs: Design 1 (D1): $\epsilon \sim N(0,30)$, Design 2 (D2): $\epsilon \sim N(0,100)$ and Design 3 (D3): $\epsilon \sim \chi^2(5)$ scaled to mean 0.5 and variance 67.6; and three outcome designs:

$$

$$

where $\eta \sim N(0,1)$. In the simulation study we consider equal sample sizes $n_0=n_1=1000$. As the definition of $\mathcal{D}$ can lead to unequal sample sizes, we use simple random sampling with replacement from the simulated treatment and control groups to meet the requirement of $n_0=n_1=1000$.

In the simulation study, we use three methods: entropy balancing (EB), kernel entropy balancing (KEB), distributional entropy balancing (DEB) with balancing means and quartiles (denoted as DEB MQ) and DEB with balancing means and deciles (denoted as DEB MD) of $X_1$ to $X_5$.

Table (ref) contains results for ATT and QTT($\alpha$) for $Y_3$ for all designs where $\alpha \in \{0.10, 0.25,0.5,\allowbreak0.75\allowbreak,\allowbreak0.90\}$. Note that we use partially overlapping $\alpha$ for balancing and QTT. Results for $Y_1$ and $Y_2$ are presented in the Appendix (ref). In all studies we report Monte Carlo $\text{Bias}=\bar{\hat{\theta}} - \theta$, $\text{Variance}=\frac{1}{R-1}\sum_{r=1}^R\left(\hat{\theta}_r - \bar{\hat{\theta}}\right)^2$ and root mean square error $\text{RMSE}=\sqrt{\text{Bias}^2 + \text{Variance}}$, where $\bar{\hat{\theta}}=\frac{1}{R}\sum_{r=1}^R \hat{\theta}_r$, and $\theta$ is the known effect (i.e. ATT, QTT, ATE or QTE) and $R$ is the number of simulations set to 500.

For all designs, as expected, KEB yields better results in terms of RMSE (mainly thanks to small variance) for $Y_3$, while results for $Y_1$ and $Y_2$ vary. The proposed estimators are better than KEB for $Y_2$ under all three designs, and for $Y_1$ results obtained using DEB are comparable or slightly better than those for KEB.

Compared with EB, DEB performs better in terms of the RMSE, which is smaller for ATT and QTT(0.25) to QTT(0.90), with a small increase / and slightly higher / for QTT(0.10) in D1 and D3. The proposed approach improves the estimates of ATT by almost halving the variance of EB for D1 and D2 and significantly reducing the bias for the non-linear case (D3). For D3, an increase in variance is observed for DEB with mean and deciles. DEB MQ and MD are more efficient compared to EB as $\alpha$ increases.

For the non-linear case, DEB MD yields an almost unbiased estimate of the cost of increasing the variance, since the bias for $\alpha=0.90$ in D3 is around 0.83, while for EB it is over 3.6 and the variance is 20.7 and 12.7, respectively. For D2, both the bias and variance decrease compared to their corresponding values for EB, while for D1 only the variance decreases leading to a decrease in RMSE.

The results suggest that the proposed approach offers more efficient ATT and QTT estimators, especially for non-linear cases in comparison to EB. It should be noted that the DEB approach significantly improves estimates of the upper part of the distribution, which can be beneficial in economic studies. Furthermore, estimation techniques for KEB take a significant amount of time even for small sample sizes (several minutes), while the proposed approach takes less then a few seconds.

Simulation for the DPS method

In the next simulation, we follow imai2014covariate and sant2022covariate. We generate four variables $\boldsymbol{X} \sim \operatorname{MVN}(\boldsymbol{0}, \boldsymbol{\Sigma})$ where $\boldsymbol{\Sigma}$ is an $4 \times 4$ identity matrix (for correctly specified models). Next, we generate $\boldsymbol{W} = (W_1, W_2, W_3, W_4)^T$ with $W_1 =\exp(X_1/2)$, $W_2=X_2/(1+\exp(W_1))$, $W_3=(X_1X_2/25 + 0.6)^3$ and $W_4=(X_2 + X_4 + 20)^T$ (for mis-specified models where instead of $\boldsymbol{X}$ we observe $\boldsymbol{W}$). The true propensity score of the treatment status $\mathcal{D}$ is given by

$$ p(\boldsymbol{X})=\frac{\exp \left(-X_{1}+0.5 X_{2}-0.25 X_{3}-0.1 X_{4}\right)}{1+\exp \left(-X_{1}+0.5 X_{2}-0.25 X_{3}-0.1 X_{4}\right)}, $$

and the treatment status $\mathcal{D}$ is generated $\mathcal{D}=\boldsymbol{1}\{p(\mathbf{X})>U\}$, where $U\sim \text{Uniform}[0,1]$. The potential outcomes $Y(1)$ and $Y(0)$ are given by $ Y(1)=210+m(\boldsymbol{X})+\varepsilon(1)$ and $\quad Y(0)=200-m(\boldsymbol{X})+\varepsilon(0)$, where $m(\boldsymbol{X})$ is defined as $m(\boldsymbol{X})=27.4 X_{1}+13.7 X_{2}+13.7 X_{3}+13.7 X_{4}$ and where $\varepsilon(1)$ and $\varepsilon(0)$ are independent $N(0,1)$ random variables. We focus on ATE and QTE($\alpha$) where $\alpha$ is defined as in DEB. The true effect equals 10 for ATE and all QTE. Standard errors were estimated using the same method as in sant2022covariate to make the results comparable between the methods. We compare the following approaches:

To compare balance of distributions we use the same metrics as sant2022covariate, i.e.

where $ \operatorname{DistImb}(\boldsymbol{X}, \boldsymbol{\gamma})=\mathbb{E}\left[\left(\omega_{1}\left(\mathcal{D}, \tilde{\boldsymbol{X}} ; \boldsymbol{\gamma}\right)-\omega_{0}\left(\mathcal{D}, \tilde{\boldsymbol{X}} ; \boldsymbol{\gamma}\right)\right) 1\left\{\tilde{\boldsymbol{X}} \leqslant \boldsymbol{x}\right\}\right], $ where $\tilde{\boldsymbol{X}}=\boldsymbol{W}$ for a mis-specified model and $\tilde{\boldsymbol{X}}=\boldsymbol{X}$ for correctly specified model. Note that $\alpha$-quantiles are calculated either on $\boldsymbol{X}^*$ or $\boldsymbol{W}$ depending which scenario is considered (for the mis-specified model we balance not only on means but also on $\alpha$-quantiles of $\boldsymbol{W}$). Weights $\omega_{1}$ and $\omega_{0}$ are defined as

$$ \omega_{1}(\mathcal{D}, \tilde{\boldsymbol{X}} ; \boldsymbol{\gamma}) =\frac{\mathcal{D}}{p(\tilde{\boldsymbol{X}} ; \boldsymbol{\gamma})} / \mathbb{E}\left[\frac{\mathcal{D}}{p(\tilde{\boldsymbol{X}} ; \boldsymbol{\gamma})}\right] \; \text{ and } \; \omega_{0}(\mathcal{D}, \tilde{\boldsymbol{X}} ; \boldsymbol{\gamma}) =\frac{1-\mathcal{D}}{1-p(\tilde{\boldsymbol{X}} ; \boldsymbol{\gamma})} / \mathbb{E}\left[\frac{1-\mathcal{D}}{1-p(\tilde{\boldsymbol{X}} ; \boldsymbol{\gamma})}\right]. $$

Table (ref) shows results for a mis-specified model based on a Monte Carlo study with 500 replications. Results for correctly specified models are presented in Appendix (ref). In all cases, DPS performs better than CBPS and IPS in terms of both bias and variance, resulting in a more efficient estimator. This pattern is observed particularly for the over-identified DPS with decile constraints, since the simulation study only includes continuous variables. The proposed approach leads to nearly unbiased estimates of QTE for all $\alpha$, especially for the upper part of the distribution.

Table (ref) shows a comparison of the mean and median of the CVM and KS statistics for mis-specified model. In all cases, the proposed method produces more balanced distributions than IPS. Similar pattern is observed for correctly specified model as presented in Appendix with one exception of the IPS method with indicators weights.

Empirical example

In this section we study the effect of 401(k) retirement plans on asset accumulation as discussed in benjamin2003does, belloni2017program, sant2022covariate and others. We use the dataset sample of 9,910 households from the 1991 SIPP; the treatment assignment follows according to plan eligibility and the outcomes of interest are net financial assets and total wealth. For the purpose of estimation we use the same approach as sant2022covariate i.e. for the (instrument) propensity score, we estimate a logistic specification and use all two-way interactions between income, log-income, age, family size, years of education, dummies for home ownership, marital status, two-earner status, defined benefit pension status, and individual retirement account participation status.

In the study we compare just-identified CBPS, the IPS method involving projection-based weights, as suggested in sant2022covariate, and the proposed DPS, in which just-identified CBPS are extended by the inclusion of $\alpha$-quantiles. In our method we balance means and quantiles of income, log-income, age, family size and years of education. For the first three we balance deciles (10%, 20%,.., 90%) and for the last two we balance quartiles and 90% percentile.

Table (ref) shows point estimates and standard errors (in parentheses) for the effect of 401(k) eligibility (participation among compliers) on net financial assets and total wealth. We present estimators for ATE and QTE ($\alpha=0.1,0.25,0.5,0.75,0.90$). We also report two measures of covariate distributional imbalance for each estimator.

The results show that the proposed approach is more efficient than the IPS method. For both outcomes, the standard errors of both ATE and QTE are smaller (with the exception of QTE(0.25) for net financial assets). In particular, the standard errors for total wealth are smaller for the ATE and the upper part of the distribution. The point estimates of the DPS method are close to those of the IPS and CBPS, with the largest difference in QTE(0.50) for total wealth (6,943 for DPS and 7,419 for IPS and CBPS).

When we compare covariate distribution imbalance metrics we get different results for KS and CVM. While the CVM statistic for the proposed method is smaller than that for the IPS and CPS method (0.53 vs 0.57 and 0.56) the KS statistic results for the DPS method are larger than for IPS but smaller than those for CBPS. This may be because the range of propensity scores for the DPS method is between 0.001153 to 0.912971 whereas it is between 0.008307 and 0.868753 for IPS.

As can be seen, the proposed approach yields more similar results to the method proposed by sant2022covariate with lower standard errors. Furthermore, the computational time to fit the parameters based on the IPS method is several hours, while the DPS method takes only a few seconds (on MacBook Air M2 16 GB RAM with 8 cores).

Summary

In this paper we have proposed a simple method for balancing distributions based on the theory of calibration estimators for quantiles. The proposed methods are flexible and allow the researcher to balance an arbitrary number of quantiles that may vary according to pre-treatment variables. In particular, if the researcher is interested in estimating a particular $\alpha$ quantile of treatment effects, they can focus on balancing only these $\alpha$ quantiles for continuous variables.

Furthermore, the proposed methods perform well for linear and especially for nonlinear and misspecified models. In the two simulation studies, we show that DEB and DPS reduce bias and RMSE for both average and quantile treatment effects. The DEB method is comparable to kernel entropy balancing, but significantly faster and less complicated. DPS outperformed the recently proposed integrated propensity score method. The proposed methods are computationally simple and can be implemented in existing statistical software (e.g. Stata, Python). For the purpose of this study, we developed the jointCalib package, which allows the user to run DEB with different distance functions, such as raking, logit, hyperbolic sinus or empirical likelihood. The DPS method is based on the CBPS package and uses its estimation techniques to balance means and quantiles. For implementation for other methods we suggest to use the Weightit package and methods that allow to specify quantile parameter.

The main limitation of the proposed approach is the uniqueness of the quantiles. For example, one may be interested in balancing the first and second deciles, while these values may be exactly the same in the treatment group (e.g. equal to 0). In such cases, the researcher has to carefully study the distribution of the pre-treatment variables in the control and treatment groups in order to select appropriate quantiles. This problem is also related to the selection of $\alpha$-quantiles, which could be solved with an appropriate penalty, but this requires further research.

Further work may involve adapting this approach to high-dimensional settings and variable selection (which should be selected first: variables or $\alpha$-quantiles?), comparison with double/debiased machine learning framework of Chernozhukov2018 or synthetic control methods by modifying the set of control variables to account for quantiles rather than means. This approach may be an attractive alternative to chen_distributional_2020 and gunsilius_distributional_2023.