Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
84,916 characters · 19 sections · 42 citation commands
MALTS: Matching After Learning to Stretch
\editor{Russ Greiner}
\allowdisplaybreaks
Matching methods are used throughout the social and health sciences to make causal conclusions where access to randomized trials is scarce but observational data are widely available. Matching methods construct sets of similar individuals, some of whom select into treatment and some of whom select into control, allowing for direct comparison of outcomes between the samples from these populations. These methods are particularly interpretable since they allow fine-grained troubleshooting of the data. For instance, examining a matched group of patients through chart review of their medical data and doctors' notes may allow an analyst to determine whether the matched groups are indeed trustworthy, and if not, determine what other factors should be included in the analysis. Having high-quality matches also allows the user to estimate nonlinear treatment effects with lower bias than parametric approaches.
As a concrete example of the importance of match group quality, Table (ref) presents a series of matched groups from the Lalonde dataset lalonde, dehejia_wahba_nonexp. A simple visual inspection of the matched groups produced by standard-bearer methods like propensity score matching and prognostic score matching reveals that the units being considered similar by these methods are not similar on underlying covariates. On the other hand, the matches generated by our proposed method are qualitatively (and quantitatively) better. The quality of the matches is our main consideration in this work.
Typically, matching methods place units that are close together into the same matched group, where closeness is measured in terms of a pre-defined distance (e.g., exact, coarsened exact, Euclidean, etc.), while maintaining balance constraints between treatment and control units. Despite its merits, this classical paradigm has flaws, namely that it relies heavily on a prespecified distance metric. The distance metric cannot be determined without an understanding of the importance of the variables; for instance, the quality of matches for any prespecified distance that weighs all covariates equally will degrade as the number of irrelevant covariates increases. This is true irrespective of the matching methodology employed. This issue has previously been referred to as the toenail problem wang2017flame, DiengEtAl2018, where the inclusion of irrelevant covariates (like “toenail length”) with nonzero weights can worsen the metric for matching. A related concern is that the covariates may be scaled differently, where a given distance along one covariate has a different impact than the same distance along a different covariate; in this case, if the scaling or weights on the covariates are chosen poorly, the total distance metric can inadvertently be determined by less relevant covariates, again leading to lower quality matches.
Ideally, the distance metric would focus on important covariates that significantly contribute to the outcome, so that after matching, treatment effect estimates computed using the matched groups would be accurate. If the researcher knows how to choose the distance metric so that it yields accurate treatment effect estimates, it would solve the problem. However, there is no reason to believe that this is achievable in complex high-dimensional data settings. Producing high dimensional functions to characterize data is a task at which humans are not naturally adept.
In this work, we propose a framework for matching where an interpretable distance measure between matched units is learned from a training set. As long as the distance metric generalizes from the training set to the full sample, we are able to compute high-quality matches and accurate estimates of conditional average treatment effects (CATEs) within the matched groups. One can use any form of distance metric to train, and in this work, we focus on exact matching for discrete variables and generalized Mahalanobis distances for continuous variables. By definition, the generalized Mahalanobis distance is determined by a matrix. If the matrix is diagonal, the distance calculation represents a stretch for each covariate. Irrelevant covariates will be compressed so that their values are always effectively zero. Highly relevant covariates will be stretched so that for two units to be considered a match, they must have very similar values for those covariates. In this way, diagonal matrices lead to very interpretable distance metrics. If the Mahalanobis distance matrix is not constrained to be diagonal, then it induces a stretch and rotation, leading to more flexible but less interpretable notions of distance.
The new framework is called Learning-to-Match, and the algorithm introduced in this work is called Matching After Learning to Stretch (MALTS). Figure (ref) shows the main steps of MALTS, which are: divide the data into training and estimation sets, learn the distance metric on the training set, use the learned distance metric to perform nearest neigbhor matching on estimation set, and use those matched groups to estimate conditional average treatment effects (CATEs). We tested MALTS against several other matching methods in simulation studies (Section (ref)), where ground truth CATEs are known. In these experiments, MALTS consistently achieves substantially better results than other matching methods including Genmatch, propensity score matching, and prognostic score matching for estimating CATEs. Even though our method is heavily constrained to produce interpretable matches, it performs at the same level as non-matching methods that are designed to fit extremely flexible but uninterpretable models directly to the response surface.
In Section (ref), we introduce the learning-to-match framework and show that under a choice of smooth distance metric (Definition (ref)) we can estimate conditional average treatment effects accurately with high probability. Section (ref) discusses MALTS' optimization set up and training procedure that learns a smooth distance metric. In Section (ref), we prove that the distance metric learned by MALTS is multi-robust (Definition (ref)) and generalizable (Definition (ref)). Thus, the distance metric estimated by MALTS' framework facilitates the correct estimates of CATEs under SUTVA and positivity assumptions.
Since the 1970's, the causal inference literature on matching methods has been concentrated on dimension reduction techniques rubin1973matching,rubin1973use,rubin1976multivariate,cochran1973controlling. In this literature, the leading approach for dimension reduction uses the propensity score, which is the conditional probability of treatment given covariate information. Propensity score methods are designed for calculating average treatment effects (as opposed to conditional average treatment effects) and do not produce exact or almost-exact matches. When treatment is binary, they project data onto one dimension, and closeness of units in propensity score does not imply their closeness in covariate space. As a result, the matches cannot directly be used for estimating heterogeneous treatment effects.
Other causal inference methods have been studied in the literature gu1993, imbens2004nonparametric, but almost all of them suffer from at least one of four possible problems: using a black box model that is uninterpretable (i.e., almost all machine learning methods), having a distance metric that is predefined (rather than learned), computational inefficiency, or not being applicable to CATE estimation (as we discussed with propensity scores). These issues cause the vast majority of matching methods to be ineffective in producing high quality interpretable CATE estimates. Regression methods can be used for CATE estimation, but only when the regression method is correctly specified -- or in the case of doubly robust estimation farrell2015robust, either the propensity model or the outcome model needs to be correctly specified. Machine learning approaches generalize regression approaches and can create models that are extremely flexible and predict outcomes accurately for both treatment and control groups hill2011bayesian,10.1111/ectj.12097,hahn2017bayesian. However, complicated regression methods lose the interpretability inherent to almost-exact matches and are difficult to troubleshoot and trust. In practice, MALTS performs similarly to (or better than) several machine learning methods in our experiments, despite being restricted to interpretable almost-exact matches with an interpretable distance metric.
A flexible setup for producing high-quality matches is provided by the optimal matching literature rosenbaum2016imposing. These are built on network flow algorithms and integer programming to produce matches that are constrained in user-defined ways zubizarreta2012using,zubizarreta2014matching,keele2014optimal,resa2016evaluation,pmlr-v54-kallus17a,MorucciEtAl2022. In all of these approaches, the user defines the distance metric (rather than learning it from data), potentially leading to poor quality matched groups. An alternative to optimal matching is coarsened exact matching iacus2012causal, an approach that requires users to specify explicit bins for all covariates on which to construct matches. This requires users to know in advance that the outcomes are insensitive to movements within many high-dimensional bins, which is essentially equivalent to the user knowing the answer to the problem we investigate in this work. Large amounts of user choice to define these bins can also lead to unintentional user bias. By learning the stretching rather than asking the user to define it as in CEM, this bias is potentially reduced.
zhao2004 and imbens2004nonparametric discuss the choice of distance metric for matching. The approach by zhao2004 depends on the correlations between treatment choice, outcome and covariates. However, this approach assumes a model for the relationship between the outcome and covariates, or the treatment choice and covariates. Hence, under model misspecification, the estimator may not be consistent. MALTS learns a distance metric without any model assumptions.
The present work builds on work of wang2017flame, DiengEtAl2018 where a discrete distance metric is learned by considering the prediction quality of the covariate sets. That work does not pertain to continuous covariates, whereas ours does. There is substantial work on learning distance metrics goldberger2005neighbourhood,weinberger2006distance,weinberger2009distance, where the goal is to learn a distance metric in latent space to separate different classes of data in supervised learning, often with a margin. This is different from our goal of matching for causal inference, but some of our proofs were inspired by this work in supervised learning.
A sister work, developed in parallel, is that of AHB, which learns adaptively-sized hyperboxes as matched groups. MALTS was previously used on the ACIC 2018 Causal Inference Challenge Data harsh2019acic. An extension of MALTS for multi-level treatments has been used to study the effect of seizures on the discharge status of critical ill patients parikh2022interpretable.
Within this framework, we perform treatment effect estimation using following three stages: 1) learning a distance metric, 2) matching samples, and 3) estimating CATEs.
We denote the $p$ dimensional covariate vector space as $\mathcal{X}\subset \mathbb{R}^p$ and the unidimensional outcome space by $\mathcal{Y} \subset \mathbb{R}$. Let $\mathcal{T}$ be a finite label set of treatment indicators (in this paper we consider only the binary case). Let $\mathcal{Z}=\mathcal{X}\times\mathcal{Y}\times\mathcal{T}$ such that $z=(\mathbf{x},y,t)\in \mathcal{Z}$ means that $\mathbf{x}\in\mathcal{X}$, $y\in\mathcal{Y}$ and $t\in\mathcal{T}$. Let $\mu$ be an unknown probability distribution over $\mathcal{Z}$ such that $\forall z \in \mathcal{Z}, ~ \mu(z)>0$. We assume that $\mathcal{X}$ is a compact convex space with respect to $\|\cdot\|_2$, thus there exists a constant $\mathbf{C}_x$ such that $\|\mathbf{x}\|_2\leq\mathbf{C}_x$. Also, $|y|\leq\mathbf{C}_y$. A distance metric is a symmetric, positive definite function with two arguments from $\mathcal{X}$ such that $\mathbf{d}: \mathcal{X}\times\mathcal{X} \to \mathbb{R}^+$. A distance metric must obey the triangle inequality. Let $\mathcal{S}_n$ denote a set of $n$ observed units $\{s_1,...,s_n\}$ drawn i.i.d$.$ from $\mu$ such that $\forall i, ~s_i \in \mathcal{Z}$. We parameterize $\mathbf{d}$ with parameter $\mathcal{M}(\cdot)$, explicitly calling it $\mathbf{d}_\mathcal{M}$, and let $\mathcal{M}(\mathcal{S}_n)$ denote the parameter learned using $\textsc{MALTS}$ methodology which is described in Section (ref). For ease of notation, we will denote the observed sample of treated units as $\mathcal{S}^{(T)}_n := \{s^{(T)}_i = (\mathbf{x}_i,y_i,t_i)~|~ s^{(T)}_i \in \mathcal{S}_n \text{ and } t_i = T \}$ and the observed sample of control units as $\mathcal{S}^{(C)}_n := \{s^{(C)}_i = (\mathbf{x}_i,y_i,t_i)~|~ s^{(C)}_i \in \mathcal{S}_n \text{ and } t_i = C \}$.
We assume no unobserved confounders and standard ignorability assumptions, i.e., $\forall i,~ (Y_i^{(T)},Y_i^{(C)})~ \rotatebox[origin=c]{90}{$\models$} ~T_i ~|~ (X_i=\mathbf{x}_i)$ Rubin2005 where $Y_i^{(T)}$ and $Y_i^{(C)}$ are potential outcomes for unit $i$ under treatments $(T)$ and $(C)$ respectively, $T_i$ is unit $i$'s treatment choice and $X_i$ corresponds to the vector of covariates for unit $i$. For each individual unit $s_i = (\mathbf{x}_i,y_i,t_i) \in \mathcal{Z}$ we define its conditional average treatment effect (or individualized treatment effect) as the difference of potential outcomes of unit $i$ under the treatment and control, $\tau(\mathbf{x}_i) =$ $\mathbbm{E}\left[Y_i^{(T)} - Y_i^{(C)} | X_i=\mathbf{x}_i \right]$ $=\mathbbm{E}\left[Y_i^{(T)}|X_i=\mathbf{x}_i \right] - \mathbbm{E}\left[Y_i^{(C)} | X_i=\mathbf{x}_i \right]$. We use the $\widehat{Y}^{(t)}_{\mathbf{x}_i}$ to refer to the estimated conditional average potential outcome, $\mathbbm{E}\left[Y_i^{(t)} | X_i=\mathbf{x}_i \right]$, for treatment $t\in\mathcal{T}$ and covariate level $\mathbf{x}_i\in\mathcal{X}$. $\widehat{\tau}(\mathbf{x}_i)$ refers to the estimated conditional average treatment effect for covariate value $\mathbf{x}_i$.
Our goal is to minimize the expected loss between estimated treatment effects $\widehat{\tau}(\mathbf{x})$ and true treatment effects $\tau(\mathbf{x})$ across target population $\mu(z)$ (this can either be a finite or super-population).
Let the population expected loss be:
For a finite random i.i.d$.$ sample $\{s_i=(\mathbf{x}_i,y_i,t_i)\}^n_{i=1}$ from the distribution $\mu$, the finite sample version of the average loss can be written as
However, we do not observed true values of $\mathbb{E}\left[ Y_i^{(T)} | X_i = \mathbf{x}_i \right]$ and $\mathbb{E}\left[ Y_i^{(C)} | X_i = \mathbf{x}_i \right]$.
Instead, we could estimate the upper bound of sample average loss as
Here, we use can $y_i$ for $t_i=1$ as the unbiased estimate of $\mathbb{E}\left[ Y_i^{(T)} | X_i = \mathbf{x}_i \right]$ and similar for $t_i=0$.
For a unit $s_i$, we estimate the conditional average potential outcomes, $\hat{Y}^{(T)}_{\mathbf{x}_i}$ and $\hat{Y}^{(C)}_{\mathbf{x}_i}$, using the treated and control units' outcomes in the constructed matched group using the observed data. The matched group $\mathrm{MG}$ of unit $s_i$ for treatment $t'$ under the distance metric $\mathbf{d}_{\mathcal{M}}$ on covariate space is defined as a set of $K$ nearest neighbors of $s_i$ from set $\mathcal{S}^{(t')}_n = \{ s_k | t_k = t', s_k\in\mathcal{S}_n \}$.
We allow reuse of units in multiple matched groups. Thus for a chosen estimator $\phi$,
where $K$ is the size of the matched group $\mathrm{MG}(s_i,\mathbf{d}_{\mathcal{M}},\mathcal{S}^{(t')}_n,K)$. A simple example of $\phi$ is the mean estimator, i.e. $\phi\left(\mathrm{MG}(s_i,\mathbf{d}_{\mathcal{M}},\mathcal{S}^{(t')}_n,K)\right) = \frac{1}{K}\sum_{k\in\mathrm{MG}(s_i,\mathbf{d}_{\mathcal{M}},\mathcal{S}^{(t')}_n,K)} y_k$. However, one can choose the estimator to be a weighted mean, linear regression or a non-parametric model like random-forest, within the matched group.
Our framework performs honest causal inference by learning a distance metric from a separate training set of data (not the estimation data considered in the averages above), and we denote this training set by $\mathcal{S}_{tr}$. To learn $\mathbf{d}_{\mathcal{M}}$, we minimize the following:
where $\hat{Y}^{(C)}_{\mathbf{x}_i}$ and $\hat{Y}^{(T)}_{\mathbf{x}_i}$ are defined by Equations ((ref)) and ((ref)) including its dependence on the $\mathrm{distance}$ $\mathbf{d}_{\mathcal{M}}$, which is parameterized by $\mathcal{M}$, using the training data to create matched groups.
Once $\mathcal{M}(\mathcal{S}_{tr})$ is learned from the training set, it is used for matching (and estimation) on the estimation data.
In this subsection, we discuss that if a distance metric is a smooth distance metric, then we can estimate the individualized treatment effect using a finite sample with high probability. First, let us define a smooth distance metric.
The concept of the smooth distance metric is analogous to commonly assumed Lipschitz continuity in the matching literature abadie2006. Note that because the range of $Y$ is bounded, there always exists a choice of the function $\delta_{\mathcal{M}}(\cdot)$ such that a distance metric $d_\mathcal{M}$ is smooth. This choice of $\delta_{\mathcal{M}}(\cdot)$ controls the quality of inference from the matching as we see in Theorem 1 below.
Theorem (ref) directly follows from Lemma (ref) in the (ref) which proves that for all $t' \in \mathcal{T}$ and $\mathbf{x}\in\mathcal{X}$, we can estimate average conditional potential outcomes, $\mathbb{E}[Y^{(t')}|X=\mathbf{x}]$, correctly with high probability using nearest neighbor matching under any smooth distance metric, and Lemma (ref) in (ref) which proves that estimating average conditional potential outcomes correctly with high probability leads to estimating CATEs, $\tau$, correctly with high probability.
Our setup and Definition (ref) are similar to one described by kara2017data. Our result in Lemma (ref) proves the consistency for a uniform weighted nearest neighbor estimator where the weights are probability weights. The result is in congruence with consistency results by stone1977consistent and jiang2019non; those works handled the special case where the weights are uniform probability weights instead of any probability weights.
Note that matching using any type of stretch norm that induces a smooth distance metric, including Mahalanobis distance (or its special case with an identity covariance matrix, the $L_2$ distance), to adjust for confounding produces consistent estimates of average treatment effects. Prognostic score hansen2008prognostic and other approaches that induce a smooth distance metric also produce consistent estimates of ATE.
MALTS performs weighted nearest neighbors matching, where the weights for the nearest neighbors can be learned by minimizing the following objective. This objective is simply the loss of the in-sample nearest neighbor estimator:
where $\textrm{Reg}(\cdot)$ is a regularization function. We let $\widetilde{W}_{i,l}$ be a function of $\mathbf{d}_{\mathcal{M}}(\mathbf{x}_i,\mathbf{x}_l)$. For example, the $\widetilde{W}_{i,l}$ can encode whether $l$ belongs to $i$'s $K$-nearest neighbors. Alternatively, they can encode soft $\textrm{KNN}$ weights where $\widetilde{W}_{i,l} \propto e^{-\mathbf{d}_{\mathcal{M}}(\mathbf{x}_i,\mathbf{x}_l)}$. Thus, the intuition is to learn $\mathbf{W}$ such that the in-sample nearest-neighbors estimator is as accurate as possible.
As a reminder of our notation, we consider distance metric $\mathbf{d}_{\mathcal{M}}$ parameterized by a set of parameters $\mathcal{M}$. We use Euclidean distances for continuous covariates, namely distances of the form $\|\mathcal{M} \mathbf{x}_a -\mathcal{M} \mathbf{x}_b\|_2$ where $\mathcal{M}$ encodes the orientation of the data. In the past, $\mathcal{M}$ has been hard-coded rather than learned; an example in the causal inference literature is the classical Mahalanobis distance ($\mathcal{M}$ is fixed as the inverse covariance matrix for the observed covariates). This approach has been demonstrated to perform well in settings where all covariates are observed and the inferential target is the average treatment effect stuart2010matching. We are interested instead in individualized treatment effects, and just as the choice of Euclidean norm in Mahalanobis distance matching depends on the estimand of interest, the stretch metric needs to be amended for this new estimand. We propose learning the parameters of a distance metric, $\mathcal{M}$, directly from the observed data rather than setting it beforehand. The parameters of distance metric $\mathcal{M}$ can be learned such that $\mathbf{W}$ minimizes the objective function on the training set.
In our framework, we can define “approximate closeness" differently for discrete covariates if desired. For continuous covariates, MALTS uses Euclidean distance, which is also a reasonable metric to use for binary data stuart2010matching; however, there are benefits to using other metrics, such as weighted Hamming distances, for comparison among sets of binary covariates. To accommodate a combination of Euclidean and Hamming distances, we parameterize our distance metric in terms of two components: one is a learned weighted Euclidean distance for continuous covariates while the other is a learned weighted Hamming distance for discrete covariates as in the FLAME and DAME algorithms wang2017flame, DiengEtAl2018. These components are separately parameterized by matrices $\mathcal{M}_c$ and $\mathcal{M}_d$ respectively, $\mathcal{M} = \left[\mathcal{M}_c, \mathcal{M}_d\right]$ (here $c$ indicates “continuous,” and $d$ indicates “discrete”). Let $a = (a_c,a_d)$ and $b = (b_c,b_d)$ be the covariates for two individuals split into continuous and discrete pairs respectively.
\noindentOperationalizing Equation (ref): To perform the step called “Distance Metric Learning” in Figure (ref) we propose the following form for the distance metric: $$\mathbf{d}_\mathcal{M}(a,b) = d_{\mathcal{M}_c}(a_c,b_c) + d_{\mathcal{M}_d}(a_d,b_d) \text{, where}$$ $$d_{\mathcal{M}_c}(a_c,b_c) = \|\mathcal{M}_{c}a_c - \mathcal{M}_{c}b_c\|_{2}, \hspace{0.25 cm}d_{\mathcal{M}_d}(a_d,b_d) = \sum_{j=0}^{|a_d|} \mathcal{M}_d^{(j,j)} \mathbbm{1}[a_d^{(j)}\neq b_d^{(j)}], $$ and $\mathbbm{1}[A]$ is the indicator that event $A$ occurred. We thus perform learned Hamming distance matching on the discrete covariates and learned-Mahalanobis-distance matching for continuous covariates.
MALTS performs an “honest” causal inference by splitting the observed sample dataset $\mathcal{S}_n$ into a training set $\mathcal{S}_{tr}$ (not for matching) and an estimation set $\mathcal{S}_{est}$ (for matching). We learn $\mathcal{M}(\mathcal{S}_{tr})$ using the training set $\mathcal{S}_{tr}$ such that in Equation (ref), $\widetilde{W}_{i,l} = \frac{e^{-\mathbf{d}_\mathcal{M}(\mathbf{x}_i,\mathbf{x}_l)}}{\sum_{s_k \in \mathcal{S}^{(t_i)}_{tr}} e^{-\mathbf{d}_\mathcal{M}(\mathbf{x}_i,\mathbf{x}_k)}}$ and $Reg(\widetilde{W})=\|\mathcal{M}\|_{\mathcal{F}}$ which defines MALTS' main implemented optimization problem:
where $\|\cdot\|_{\mathcal{F}}$ is the Frobenius norm of the matrix, and:
\noindentMatching and Estimation: To perform the step called “Nearest Neighbor Matching,” which produces “Matched Groups” that are used to estimate “CATEs” in Figure (ref), we use the learned distance metric $\mathcal{M}(\mathcal{S}_{tr})$. To estimate conditional average treatment effects (CATEs) for each unit in the estimation set, we use its nearest neighbors from the same estimation set. Specifically, for any given unit $s$ in the estimation set, we construct a K-nearest neighbor matched group $\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}(\mathcal{S}_{tr})},\mathcal{S}_{est},K)$ using control set $\mathcal{S}^{(C)}_{est}$ and treatment set $\mathcal{S}^{(T)}_{est}$. For a choice of estimator $\phi$, the estimated CATE for a treated unit $s = (\mathbf{x}_s,y_s,t_s=t')$ is calculated as follows: $$\hat{\tau}(\mathbf{x}) = \phi \left(\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}(\mathcal{S}_{tr})},\mathcal{S}^{(T)}_{est},K)\right) - \phi \left(\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}(\mathcal{S}_{tr})},\mathcal{S}^{(C)}_{est},K)\right).$$
A simple example of $\phi$ is the empirical mean, i.e., $$\phi\left(\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}},\mathcal{S}^{(t)}_n,K)\right)=\frac{1}{K}\sum_{k\in\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}},\mathcal{S}^{(t)}_n,K)}y_k.$$ However, one can choose the estimator to be a weighted mean, linear regression or a non-parametric model like Random Forest. Particular choices of $\phi$ can also play a role in bias-adjustment to improve the matching estimator of the ATE as in abadie2011bias.
For $\phi\left(\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}},\mathcal{S}_n,K)\right) = \sum_{k\in\mathrm{MG}(s,\mathbf{d}_{\mathcal{M}},\mathcal{S}_n,K)} \widetilde{W}_{k}y_k$, if $\widetilde{W}_{k}$ is chosen to be proportional to $e^{\mathbf{d}_{\mathcal{M}}(\mathbf{x},\mathbf{x}_k)}$, then it leads to multi-robust (defined shortly) and generalizable CATE estimates via soft KNN (as shown in Theorem (ref) and Theorem (ref) below), while letting $\widetilde{W}_k$ be proportional to $\mathbbm{1}\left[s_k\in\textrm{KNN}^{\mathcal{S}^{(C)}_{est}}_{\mathcal{M}(\mathcal{S}_{tr})}\right]$ produces interpretable matched groups.
\noindentHyperparameter choice: MALTS has four main hyperparameters: 1) K, which is the number of nearest neighbors used to estimate the counterfactual, which can be chosen by cross-validation. 2) $n$, the size of training set, i.e., the size of the split on the left of Figure (ref). This can be chosen based on the amount of data relative to the number of features, though typically we choose it to be 10% of the data. 3) The maximum allowed diameter or caliper to prune bad matched groups. If the matches have a larger diameter, the matches are not tight and we may not be able to trust their estimates. The maximum diameter can be chosen by domain knowledge; the user defines how far apart points can be to make the matched group less interpretable. 4) The number of repeats refers to the number of times we shuffle the data and re-partition it for MALTS’ training and estimation procedure. A larger number of repeats of the whole process helps with smoothing out the estimates over different train/test splits.
In this section, we show that the MALTS framework correctly estimates the distance metric, facilitating correct estimates of CATEs under SUTVA and a positivity assumption. After basic definitions, and after showing that the learned distance metric and objective are bounded, we introduce and define the concepts of multi-robustness and generalizability of the learned distance metric. Multi-robustness implies that for any possible pair of points the empirical average loss is not far away from the population average loss. Theorem (ref) proves that the distance metric learned by the MALTS algorithm is multi-robust. We use these results along with the error bound shown in Lemma (ref), to show that MALTS' distance metric is generalizable, i.e., the population average loss and the empirical average loss on the observed data for the learned distance metric are close with high probability. Lastly, we show that MALTS' distance metric is asymptotically generalizable and that the empirical average loss approaches the population average loss as the size of the dataset goes to infinity.
\noindentBasic definitions of empirical loss and population loss. First, we define a pairwise loss for $s_i$ and $s_l$ so that it is only finite for treatment-treatment or control-control matched pairs,
This loss is high for pairs of points that are close (i.e., with small $\mathbf{d}_\mathcal{M}(\mathbf{x}_i,\mathbf{x}_l)$) when the outcomes $y_i$ and $y_l$ values are very different. Further, we define an empirical average pairwise loss over finite sample $\mathcal{S}_n$ of size $n$ as
and define an average loss over population $\mathcal{Z}$ as
\noindentThe search space over distance metrics is bounded. We show a basic result about the optimization-based approach we take to learn the distance metric. Specifically, we show that the learned distance metric will be in a bounded region of search space.
Now, because the learned $\mathcal{M}(\mathcal{S}_{tr})$ on the set $\mathcal{S}_{tr}$ is the distance metric that minimizes the given objective function, we know that the following inequality is true, which states that the learned parameter has a lower training objective than that of the trivial parameter $\mathbf{0}$:
Denoting the right hand side of the inequality by $g_0$ we note that we can limit our search space over distance metrics $\mathcal{M}$ that satisfy the following inequality:
\noindentThe objective function terms are bounded. The objective terms $\Delta_{\mathcal{S}_{tr}}^{(C)}$ and $\Delta_{\mathcal{S}_{tr}}^{(T)}$ (defined in Equation (ref)) for learning the distance metric are also bounded, although it is not that easy to see this directly because their denominators are somewhat complicated, involving a sum over exponential terms. Here, we point out that because the learned distance metric is bounded, the objective's terms ($\Delta_{\mathcal{S}_{tr}}^{(C)}$ and $\Delta_{\mathcal{S}_{tr}}^{(T)}$) are also bounded. Specifically, their upper bound is proportional to the empirical average pairwise losses $L_{emp}(\mathcal{M},\mathcal{S}^{(C)}_{tr})$ and $L_{emp}(\mathcal{M},\mathcal{S}^{(T)}_{tr})$, defined above. Further, in Theorem (ref), we show that for $t'\in\{T,C\}$ the empirical average loss $L_{emp}(\mathcal{M},\mathcal{S}^{(t')}_{tr})$ is close to population average pairwise loss $L_{pop}(\mathcal{M}(\mathcal{S}_n),\mathcal{Z}^{(t')})$ with high probability. Following Equations (ref) and (ref) and Theorem (ref), the objective terms $\Delta_{\mathcal{S}_{tr}}^{(C)}(\mathcal{M})$ and $\Delta_{\mathcal{S}_{tr}}^{(T)}(\mathcal{M})$ are upper-bounded by a term proportional to the population average pairwise loss with high probability.
We know that:
Together, the two previous lines imply:
Similarly for the treatment units, we have
Now, we define a few concepts important for our results including covering number, multi-robustness, and generalizability. The following definitions and results closely align with the theoretical guarantees of distance metric learning algorithms in bellet2015robustness and xu2012robustness. Our work extends these results to learn a distance metric for causal inference.
Note that $\mathbf{N}(\gamma,\mathcal{V},\mathbf{d})$ is finite if $\mathcal{U}$ is compact.
Intuitively, robustness means that for any possible unit in the support, the loss is not far away from the loss of nearby units in the training set, should some training units exist nearby. bellet2015robustness,xu2012robustness
Intuitively, multi-robustness means that for any possible pair of points from any two partitions of $\mathcal{X}$, the empirical average loss over training points is not far away from the population average loss. As the training procedure aims at minimizing the total loss, we can safely say that a multi-robust method will not perform poorly out of sample.
Given these definitions, we first show that the distance metric learned using MALTS is robust in Theorem (ref) and we extend the argument to show that it is also generalizable in Theorem (ref).
Proof (Theorem (ref)). Given $\mathcal{Z}=\mathcal{X}\times\mathcal{Y}\times\mathcal{T}$, we consider the following definition of a minimum sized $\gamma$-cover $\hat{\mathcal{V}}$ of the set $\mathcal{X}$ under the distance metrix $\|\cdot\|_2$: Partition the set into $K$ disjoint subsets $\textbf{C}_\gamma = \{C_i\}_{i=1}^{K}$ such that $K$ is the $\gamma$-covering-number of $\mathcal{X}$ under $\|\cdot\|_2$ (which is exactly equal to $|\hat{\mathcal{V}}|$) where each $C_i$ is contained in the $\gamma$-neighborhood of each $\hat{v}_i\in\hat{\mathcal{V}}$ and each $C_i$ contains at least one control and one treated sample. Note that if $\mathcal{X}$ is a compact convex set, then such a cover and the corresponding packing $\textbf{C}_\gamma$ exists and $K=|\textbf{C}_\gamma|$ is finite.
For any arbitrary $C_i$ and $C_l$ in $\textbf{C}_\gamma$, consider the empirical average loss for all training units $s_i \in C_i$ and $s_l \in C_l$ with treatment $t'$:
and the expected loss for units $Z_i$ and $Z_l$:
Let $f$ be a function of the set of independent random variables such that $$f(s_1,\dots,s_{|C^{(t')}_i|},s_{|C^{(t')}_l|+1},\dots,s_{|C^{(t')}_i|+|C^{(t')}_l|}) = \frac{1}{|C^{(t')}_i\|C^{(t')}_l|} \sum_{j=1}^{C^{(t')}_i}\sum_{i=C^{(t')}_i+1}^{C^{(t')}_l} loss[\mathcal{M}(\mathcal{S}_n),s_i,s_l].$$ Thus, $f(s_1,\dots,s_{|C^{(t')}_i|},s_{|C^{(t')}_l|+1},\dots,s_{|C^{(t')}_i|+|C^{(t')}_l|}) = \widehat{\overline{loss}}\left[\mathcal{M}(\mathcal{S}_n),C^{(t')}_i,C^{(t')}_l\right]$.
Now, let $\rho^{(t')}_{\gamma}$ be the density of the $\gamma$-cover for treatment $t'$, defined as the number of units with treatment $t'$ in the smallest partition set $\rho^{(t')}_{\gamma} = \min_i |C^{(t')}_i|$ and $B = \max_{z_1,z_2} loss(\mathcal{M}(\mathcal{S}_n),z_1,z_2)$. Now, we show that $f(\cdot)$ has bounded difference. Without loss of generality, consider an index $j \leq |C^{(t')}_i|$, then
Similarly, for any $j > |C^{(t')}_i| $, $$|f(s_1,\dots,s_{j},\dots,s_{|C^{(t')}_i|+|C^{(t')}_l|}) - f(s_1,\dots,s'_{j},\dots,s_{|C^{(t')}_i|+|C^{(t')}_l|})| \leq \frac{2B}{\rho^{(t')}_{\gamma}}.$$.
As $f()$ is a function of independent $|C^{(t')}_i| + |C^{(t')}_l|$ random variables, by McDiarmid's inequality:
\ensuremath{\blacksquare}
We will need the following lemma to prove Theorem (ref). The lemma provides a bound for a particular treatment assignment, while the theorem sums over all treatment assignments.
Now that we have theoretically proven the functionality of MALTS, we will next discuss and compare MALTS performance with other methods on different datasets.
In this section, we discuss and compare the performance of MALTS with other competing methods on a few different simulation setups with continuous covariates, discrete covariates and mixed (continuous and discrete) covariates. Lastly, we demonstrate MALTS performance for estimating ATE on LaLonde's NSW and PSID-2 data samples lalonde,dehejia_wahba_nonexp.
MALTS performs an $\eta$-fold honest causal inference procedure with the estimator $\phi$ inside each matched group being linear regression. We split the observed samples $\mathcal{S}_n$ into $\eta$ equal parts such that the ratio of treated to control units in each part is similar. For each fold, we use one of the $\eta$ partitions as the training set $\mathcal{S}_{tr}$ (not used for matching) and the rest of the $\eta-1$ partitions as the estimation set $\mathcal{S}_{est}$. Using the output from each of the $\eta$ folds, we calculate the estimated CATE for each unit (averaged across folds), estimated distance metric (averaged across folds) and a weighted unified matched group for each unit $s_i \in \mathcal{S}_n$. The weight of each matched unit $s_k$ corresponds to the number of times a particular unit $s_k$ was in the matched group of unit $s_i$ across the $\eta-1$ constructed matched groups. Here, $\eta$ was chosen to be 5 in our experiments.
For interpretability, we let $\mathcal{M}_c$ be a diagonal matrix, which allows stretches of the continuous covariates. (Note that $\mathcal{M}_d$, which is the stretch matrix over discrete covariates, is always set to be diagonal.) This way, the magnitude of an entry in $\mathcal{M}_c$ or $\mathcal{M}_d$ provides the relative importance of the indicated covariate for the causal inference problem.
We further analyzed strategies for variance estimation for MALTS in Section (ref), and performance under limited overlap between the covariates distribution of treated and control groups, and sensitivity to unobserved confounding. Detailed results are shown in (ref).
The main results of these experiments are that MALTS' performance is on par with existing state-of-the-art methods for causal inference, including black box methods. MALTS tends to have fairly consistent performance, even if the training set is fairly small or the number of irrelevant covariates is large. Further, MALTS provides interpretable distance metrics and matched groups that black box machine learning methods do not provide.
In this subsection we describe the data generation process (DGP) used in the simulation experiments. We use two main data-generation processes: The first DGP has a linear baseline with linear and quadratic treatment effects while the second DGP is the extension of Friedman's function introduced to test performance of prediction algorithms of friedman1991multivariate. This second DGP, also termed as Friedman's DGP, has a scaled cosinusoidal treatment effect.
This simulation includes both linear and quadratic terms. Let $\mathbf{x}_{i,p} = \{\mathbf{x}_{i,p_c}, \mathbf{x}_{i,p_d}\}$ be a $p$-dimensional covariate vector composed of $|p_c|$ continuous covariates and $|p_d|$ discrete ones. There are $k = k_{c} \cup k_{d}$ relevant covariates and the rest of the dimensions are irrelevant. Here, $p_c,k_c,p_d,$ and $k_d$ refer to the the subsets of indices of the covariates: all continuous, relevant continuous, all discrete, and relevant discrete, respectively. $\mathbf{x}_{i,k_c}$ and $\mathbf{x}_{i,k_d}$ refer to the vectors of relevant continuous and discrete covariates respectively. $\mathbf{x}_{i,k}$ refers to all $|k|$ relevant covariates. $\kappa_c \subseteq k_c$ is the set of continuous covariates and $\kappa_d \subseteq k_d$ is the set of discrete which are relevant in determining the treatment choice. The potential outcomes and treatment assignment are determined as follows:
Here expit$(z) = \exp(z)/(1+\exp(z))$. The variance of $\epsilon_{i,\textrm{treat}}$ determines how much confounding and overlap there is in the dataset: higher values of the variance make the dataset look like a randomized experiment with good overlap, while very small values of the variance lead to poor overlap and a very hard to analyze observational study. We explore these issues in detail in (ref).
The data generation process of friedman1991multivariate was first proposed to assess the performance of prediction methods. We augmented Friedman's simulation setup to evaluate causal inference methods. The potential outcome under control is Friedman's function as provided by friedman1991multivariate and Chipman10bart:bayesian. The expected treatment effect we study is equal to the cosine of the product of the first two covariates scaled by the third covariate.
We use the data-generation process described in Section (ref) to generate 2500 units with no discrete covariates, 15 important continuous covariates and 25 irrelevant continuous covariates. Further, we set the parameters for the DGP as follows: $\mu = 1$, $\Sigma = 1.5 \textbf{I}$, $\psi=0.5$, $\sigma^2 = 1$ and $\kappa_c = \{0,1\}$. We estimate CATE for each unit using matching methods like propensity score matching, prognostic score matching and genetic matching, and non-matching (uninterpretable) methods like causal forest and BART. Figure (ref) shows the performance of these methods. MALTS' performance is on par with existing state-of-the-art non-matching methods and outperforms all other matching methods for continuous covariates in the quadratic data generation process.
We use the data-generation process described in Section (ref) to generate 2500 units with no continuous covariates, 15 important discrete covariates and 10 irrelevant discrete covariates. Further, we set the parameters of the DGP as follows: $\sigma^2 = 1$, $c = 2$ and $\kappa_d = \{0,1\}$. We used the weighted Hamming distance metric for this experiment.
Figure (ref) shows the performance comparison, again showing that MALTS' performance is on par with existing state-of-the-art non-matching methods; it also performs better than FLAME (a state-of-the-art matching method for discrete data) as it is able to provide additional smoothing in this relatively small-$n$ setting. Hence, MALTS performs well for discrete covariates in the quadratic data generation process.
We use the data-generation process used for experiments on continous and discrete covariates (described in Section (ref)) to generate 2500 units with 5 relevant continuous covariates, 15 relevant discrete covariates, 10 irrelevant continuous and 10 irrelevant discrete covariates. We used the same set of parameters for the DGP as the previous two experiments. Similar to the previous two experiments, Figure (ref) shows that MALTS performs on par with the state-of-the-art non-matching methods and outperforms all matching methods that can handle mixed covariates for the quadratic data generation process.
We studied the performance of various causal inference methods to estimate CATEs as the number of covariates ($p$) changes, keeping the number of relevant covariates ($|k|$) constant and equal to $8$. We simulated the data using the DGP described in Section (ref). The number of units is constant ($n=2048$) while the number of covariates ($p$) changes from 8 to 256. The performance of MALTS is on-par with or better than other causal inference methods as the number of irrelevant covariates increases (see Figure (ref)). This indicates that MALTS can be used to help reduce the effects of the curse of dimensionality.
We studied the change in CATE estimation error-rates as the number of units in a dataset increases. We simulated the data using the DGP described in Section (ref), keeping the number of covariates constant and equal to $20$ (all of them are relevant in outcome determination). We changed the number of units from $2^8$ to $2^{12}$. MALTS' performance is on-par with or better than BART and the error-rate is significantly lower than that of other causal inference methods (see Figure (ref)).
We further compare MALTS and other flexible methods' performance on data generated using the process described in Section (ref). This DGP is particularly interesting because the potential outcomes are highly non-linear functions with trigonometric expressions.
As shown in Figure (ref), we observe that MALTS performs on par with Causal Forest while BART's error-rate is significantly higher (worse) than MALTS, for the Friedman's data generation process.
We use the DGP described in Section (ref) with 2 relevant continuous covariates and no irrelevant covariates for the coverage study. Further, we set the parameters to the DGP as follows: $\mu = 1$, $\Sigma = 1.5 \textbf{I}$, $\psi=0.5$, and $c = 2$. We selected 9 reference points in a grid from the covariate space as shown in Figure (ref)(b) and conducted an experiment that considered these reference points, over 100 repetitions. We compared coverage for CATEs estimated using MALTS for different values of the variance, ranging from 1.0 to 4.0, for noise term $\epsilon_0$ and $\epsilon_1$ in the potential outcomes function.
Variance estimation is notoriously hard in matching problems, even for overall quantities such as the average treatment effect abadie2006. We consider both a conservative variance estimator wang2017flame and estimators that sacrifice some interpretability for better coverage. Specifically, we consider the CATEs estimated using MALTS and study how well an uninterpretable method can predict those estimates to obtain a variance estimate. We use the predictive variance from gradient boosting regression, from gaussian process regression and from Bayesian ridge regression on the covariates, where we estimated CATEs and quantify variance of each CATE estimate.
Based on Figure (ref)(a), the coverage for each the nine points of interest is between 0.85 and 1 for most values of the variance using any of the three variance estimation approaches.
The LaLonde data pertain to the National Support Work Demonstration (NSW) temporary employment program and its effect on income level of the participants lalonde. This dataset is frequently used as a benchmark for the performance of methods for observational causal inference. We employ the male sub-sample from the NSW in our analysis as well as the PSID-2 control sample of male household-heads under age 55 who did not classify themselves as retired in 1975 and who were not working when surveyed in the spring of 1976 dehejia_wahba_nonexp. The outcome variable for both experimental and observational analyses is earnings in 1978 and the considered variables are age, education, whether a respondent is Black, is Hispanic, is married, has a degree, and their earnings in 1975. Previously, it has been demonstrated that almost any adjustment during the analysis of the experimental and observational variants of these data (both by modeling the outcome and by modeling the treatment variable) can lead to extreme bias in the estimate of average treatment effects lalonde.
Performance results: Tables (ref) and (ref) present the average treatment effect estimates based on MALTS, state-of-the-art modeling methods, and matching methods. MALTS (after appropriately pruning low-quality matched groups) is able to achieve accurate ATE estimation on both experimental and observational datasets.
Figure (ref) illustrates how the matched groups were pruned. There was a clear visual separation between high-quality matched groups, which had low diameters, and low-quality matched groups, with larger diameters.
Model Interpretability: One difference between MALTS and the other methods is that its solution can be described concisely: MALTS produces a total of seven numbers that define the distance metric on the LaLonde data. The distribution of the learned distance metric values across folds is shown in Figure (ref)(a). Once the researcher has these seven numbers, along with the value of $k$ in $k$-nearest neighbors used to train MALTS, they know precisely which units should be matched. In contrast, causal forest and BART require a model whose size depends on the number of trees, where each tree is several levels deep--in this case, 2000 trees and 150 trees, respectively.
Interpretability of Matched Groups: To examine the interpretability of MALTS' matched groups, we present two of the matched groups from MALTS for the observational Lalonde dataset in Table (ref), corresponding to two “query” individuals in the dataset. Query 1 is a 22 year old with no income in 1975. MALTS was able to construct a tight matched group for this individual (both in control and in treatment). In contrast, Query 2 is a 42-year-old high-income individual without a degree, which is an extremely unlikely scenario, leading to a matched group with a very large diameter, which should probably not be used during analysis. Such granular analysis is not possible for regression methods like BART and matching methods like prognostic score or propensity score matching.
This further highlights the troubleshooting capabilities of interpretable matching methods: by identifying units that are poorly matched, we know exactly which units to study in more detail. In this case, it is possible that the “degree” field might have a data error, which means it would be better not to match this unit and to potentially follow up on the veracity of responses to the survey.
This paper introduces the MALTS algorithm, which learns a distance metric on the covariate space for use with matching. The learned metric stretches important covariates and compresses irrelevant covariates for outcome prediction in order to produce high-quality matches. Unlike other methods, MALTS can handle a large number of irrelevant covariates by compressing them to the point where they are effectively eliminated, which helps handle the curse of dimensionality. Unlike black-box machine learning methods, MALTS produces interpretable matched groups and returns the stretch matrix on covariates for counterfactual prediction. The stretch matrix is chosen here to be diagonal, so that it can be represented using only a few “stretch” numbers that determine the importance of each covariate in determining the matched groups.
Whereas deep neural networks mainly show improvements over other methods for problems that do not have natural data representations (computer vision, speech, etc.), we conjecture that the stretch/almost-exact match combination should suffice for most datasets. A natural extension, however, is to use neural networks to learn a flexible distance metric in a latent space, thus allowing us to match on medical records, images, and text documents. This will allow us to incorporate complex data structures by introducing a flexible learning framework (e.g., interpretable neural networks) for coding the data. That is, we can redefine the distance metric via
where $\omega_\mathcal{M}$ is a summary of relevant data features learned using a complex modeling framework.
In the future, the MALTS framework could be extended to deal with missing covariates, and can be adapted to instrumental variables.
\acks{We gratefully acknowledge funding from the National Science Foundation under grants III 1703431, CCF 1934964, IIS 2130250, IIS 2147061 (with Amazon), and CAREER DMS 2046880, and the National Institute of Health under grants NIDA DA054994 and R01EB025021. We also acknowledge funding from an Amazon Graduate fellowship.}