Heterogeneous Treatment Effects in Panel Data

Causal inference using panel data is often approached through the synthetic control framework, where treated units are compared to a synthetic control group constructed from untreated units [1, 2]. However, for these methods, the treatment is restricted to a block. Matrix completion approaches offer a more flexible alternative, allowing for general treatment patterns [4, 7, 35]. The treated entries are viewed as missing entries, and low-rank matrix completion is utilized to estimate counterfactual outcomes for the treated entries. In most of the matrix completion literature, the set of missing or treated entries are assumed to be generated randomly [12]. However, this assumption is not realistic in our panel data setting. Other works explore matrix completion under deterministic sampling [10, 17, 19, 22, 24], and our work extends the methodology from [17].

Several methods for estimating heterogeneous treatment effects use supervised learning techniques to fit models for the outcome or treatment assignment mechanisms [14, 15, 18, 23, 27]. Although these methods are intended for cases where observations are drawn i.i.d. from one or multiple predefined distributions, they could still be applied to panel data. In contrast, our approach is specifically designed for panel data with underlying structure. By avoiding complex machine learning methods, our approach not only provides accurate results for panel data but also offers interpretable estimates.

The first step of our method involves grouping data into clusters that vary in their treatment effects using a regression tree. As in [8], we greedily search for good splits, iteratively improving the tree’s representation of the true heterogeneous treatment effects. However, unlike supervised learning where the splitting criterion is directly computable, treatment effects are not readily observed. Therefore, we estimate the splitting criterion, an approach used in [5, 6, 29, 33]. While methods like causal forest only use information within a given leaf to estimate the treatment effect for observations in that leaf [33], our method leverages global information from the panel structure to determine the average treatment effect of each cluster, enabling us to achieve high accuracy with just one tree.

The second step involves estimating the average treatment effect in each cluster. For this step, we utilize the de-biased convex estimator introduced by [17]. This estimator is applicable to panel data with general treatment patterns but is specifically designed for estimating average, rather than heterogeneous, treatment effects. It uses a convex estimator, similar to those used in [20, 34], and introduces a novel de-biasing technique that improves upon the initial estimates. Additionally, [17] builds on previous works [11, 13] to provide theoretical guarantees for recovering the average treatment effect with provably minimal assumptions on the intervention patterns. The results of [17] are specific to the case where there is only one treatment. We build on their proof techniques to show convergence when the number of treatments is $O(\log n)$, and also nest this technique into the first step of our algorithm to estimate heterogeneous effects.

For any matrix $A$, $\left\lVert A\right\rVert$ denotes the spectral norm, $\left\lVert A\right\rVert_{\mathrm{F}}$ denotes the Frobenius norm, $\left\lVert A\right\rVert_{\star}$ denotes the nuclear norm, and $\|A\|_{\text{sum}}$ denotes the sum of the absolute values of the entries in $A$. For any vector $a$, $\left\lVert a\right\rVert$ denotes its Euclidean norm, and $\left\lVert a\right\rVert_{\infty}$ denotes the maximum absolute value of its components. We use $\circ$ to denote element-wise matrix multiplication. For a matrix subspace $\mathbf{T}$, let $P_{{\mathbf{T}}^{\perp}}(\cdot)$ denote the projection operator onto ${\mathbf{T}}^{\perp}$, the orthogonal space of $\mathbf{T}$. We define the sum of two subspaces $\mathbf{T}_{1}$ and $\mathbf{T}_{2}$ as $\mathbf{T}_{1}+\mathbf{T}_{2}=\left\{T_{1}+T_{2}\mid T_{1}\in\mathbf{T}_{1},T_{2}\in\mathbf{T}_{2}\right\}$. We write $a\lesssim b$ whenever $a\leq c\cdot b$, for some fixed constant $c$; $a\gtrsim b$ is defined similarly.

Initially, the entire dataset with $n\times T$ entries is in one single cluster. This cluster serves as the starting point for the iterative partitioning process, where we split the data into $\ell$ more granular clusters for each treatment that better represent the true heterogeneous treatment effects. We present our method for clustering observations in Algorithm 1. Below, we discuss its two main steps.

After partitioning all observations into $\ell$ clusters for each of the $q$ treatments, the next step in our algorithm is to estimate the average treatment effect within each cluster. This is done using an extension of the de-biased convex estimator introduced in [17]. We first obtain an estimate of the average treatment effects by solving a convex optimization problem. This initial estimate is then refined through a de-biasing step.

For simplicity, we define separate treatment matrices for each cluster. Letting for $k:=q\times\ell$, $\tilde{Z}_{1},\dots,\tilde{Z}_{k}$ are defined as follows: $\tilde{Z}_{(i-1)\times\ell+j}=W_{i}\circ C_{j}^{i}$ for $i\in[q],j\in[\ell]$. These are binary matrices that identify observations that are in a given cluster and receive a given treatment. We denote the true average treatment effect of the treated in each cluster $i\in[k]$ as $\tilde{\tau}^{\star}_{i}$, and let $\delta_{i}$ be the associated residual matrices, which represent the extent to which individual treatment effects differ from the average treatment effect. More precisely, for $i\in[k]$,

Finally, we define $\delta$ as the combined residual error, where $\delta:=\sum_{i=1}^{k}\delta_{i}$, and $\hat{E}:=E+\delta$ is the total error.

We normalize all of these treatment matrices to have Frobenius norm 1. This standardization allows the subsequent de-biasing step to be less sensitive to the relative magnitudes of the treatment matrices. Define $Z_{i}:=\tilde{Z}_{i}/\|\tilde{Z}_{i}\|_{\mathrm{F}}$. By scaling down the treatment matrices in this way, our estimates are proportionally scaled up. Similarly, we define a rescaled version of $\tilde{\tau}^{\star}$, given by $\tau^{\star}_{i}:=\tilde{\tau}^{\star}_{i}\|\tilde{Z}_{i}\|_{\mathrm{F}}.$

First, we solve the following convex optimization problem with the normalized treatment matrices to obtain an estimate of $M^{\star}$ and $\tau^{\star}$:

Because we introduced the $m\mathbf{1}^{\top}$ term in the convex formulation, which centers the rows of $\hat{M}$ without penalization from the nuclear norm term, $\hat{M}$ serves as our estimate of $M^{\star}$, projected to have zero-mean rows, while $\hat{m}$ serves as our estimate of the row means of $M^{\star}$. For notational convenience, let $P_{\mathbf{1}}$ denote the projection onto $\left\{\alpha\mathbf{1}^{\top}\mid\alpha\in\mathbb{R}^{n}\right\}$. That is, for any matrix $A\in\mathbb{R}^{n\times T}$,

Hence $P_{{\mathbf{1}}^{\perp}}(M^{\star})$ is the projection of $M^{\star}$ with zero-mean rows. Let $m^{\star}:=\frac{M^{\star}\mathbf{1}}{T}$ be the vector representing the row means of $M^{\star}$. Let $r$ denote the rank of $P_{{\mathbf{1}}^{\perp}}(M^{\star})$,²²2Note that we can show $r\leq\text{rank}(M^{*})+1$, so $P_{{\mathbf{1}}^{\perp}}(M^{\star})$ is a low-rank matrix as well. and let $P_{{\mathbf{1}}^{\perp}}(M^{\star})=U^{\star}\Sigma^{\star}V^{\star\top}$ be its SVD, with $U^{\star}\in\mathbb{R}^{n\times r},\Sigma^{\star}\in\mathbb{R}^{r\times r},V^{\star}\in\mathbb{R}^{T\times r}$. We denote by $\mathbf{T}^{\star}$ the span of the space $\left\{\alpha\mathbf{1}^{\top}\mid\alpha\in\mathbb{R}^{n}\right\}$ and the tangent space of $P_{{\mathbf{1}}^{\perp}}(M^{\star})$ in the manifold of matrices with rank $r$. That is,

The closed form expression of the projection $P_{{\mathbf{T}}^{\star\perp}}(\cdot)$ is given in Lemma B.1.

The following lemma gives a decomposition of $\hat{\tau}-\tau^{\star}$ that suggests a method for de-biasing of $\hat{\tau}$. The difference between our error decomposition and the decomposition presented in [17] stems from the fact that we introduced the $m\mathbf{1}^{\top}$ term in (2).

As noted in [17], since $D^{-1}\Delta^{1}$ depends solely on observed quantities, it can be subtracted from our estimate. Thus, we define our de-biased estimator $\tau^{d}$ as $\tau^{d}_{i}=\left(\hat{\tau}-D^{-1}\Delta^{1}\right)_{i}/\|\tilde{Z}_{i}\|_{\mathrm{F}}$ for each $i\in[k]$. The rescaling by $\|\tilde{Z}_{i}\|_{\mathrm{F}}$ adjusts our estimate to approximate $\tilde{\tau}^{\star}$, rather than $\tau^{\star}.$

The de-biased $\tau^{d}$ serves as our final estimate. Specifically, if an observation’s covariates $X^{zt}$ belong to the $j$-th cluster, then for each treatment $i\in[q]$, our estimate of $\mathcal{T}^{\star}_{i}(X^{zt})$ is given by $\hat{\mathcal{T}}_{i}(X^{zt})=\tau^{d}_{(i-1)\times\ell+j}$, which represents the average treatment effect of treatment $i$ associated with the cluster that $X^{zt}$ belongs to.

We start by bounding the bias introduced by approximating $\mathcal{T}^{\star}_{i}$ by a piece-wise constant function. To guarantee consistency, we establish constraints on valid splits. We adopt the $\alpha$-regularity condition from Definition 4b of [33]. This condition requires that each split retains at least a fraction $\alpha\in(0,\frac{1}{2})$ of the available training examples and at least one treated observation on each side of the split. We further require that the depth of each leaf be on the order $\log\ell$ and that the trees are fair-split trees. That is, during the tree construction procedure in Algorithm 1, for any given node, if a covariate $j$ has not been used in the splits of its last $\pi p$ parent nodes, for some $\pi>1$, the next split must utilize covariate $j$. If multiple covariates meet this criterion, the choice between them can be made arbitrarily, based on any criterion.

Additionally, we assume that the proportion of treated observations with covariates within a given hyper-rectangle should be approximately proportional to the volume of the hyper-rectangle. This is a ‘coverage condition’ that allows us to accurately estimate the heterogeneous treatment effect on the whole domain using the available observations. For the purposes of this analysis, we assume that all covariates are bounded and normalized to $[0,1]^{p}$.

3.1 is a broader condition than the requirement for the covariates to be i.i.d. across observations, as assumed in [33]. Our assumption only requires that the covariates maintain an even density. Indeed, Lemma A.2 in Appendix A illustrates that while i.i.d. covariates satisfy the assumption with high probability, the assumption additionally accommodates (with high probability) scenarios where covariates are either constant over time and only vary across different units, or constant across different units and only vary over time.

In the following result, we demonstrate that, as the number of observations grows, a regression tree, subject to some regularity constraints, contains leaves that become increasingly homogeneous in terms of treatment effect.

Using the above theorem, we can make the following observation. For any $i\in[q]$, denote by $\tilde{\tau}^{\star,i}$ the regression tree function that is piece-wise constant on each cluster $C_{j}^{i}$ for $j\in[\ell]$, with a value of $\tilde{\tau}^{\star}_{(i-1)\ell+j}$ assigned to each cluster $j$. Then, for all covariate vectors $X$, we have $\left\lvert\mathcal{T}_{i}^{\star}(X)-\tilde{\tau}^{\star,i}(X)\right\rvert\leq\frac{2L_{i}\sqrt{p}}{\ell^{s}}$ because the value assigned to each cluster by $\tilde{\tau}^{\star,i}$ is an average of treatment effects within that cluster. This demonstrates that it is possible to approximate the true treatment effect functions with a regression tree function with arbitrary precision. As the number of leaves $\ell$ in the tree $\mathbb{T}$ increases, the approximation error of the tree decreases polynomially. The theorem requires an upper bound on $\ell$, which grows with $\min(n,T)$ as the dataset expands, due to the definition of the margin or error $M$. For PaCE, we constrain $\ell$ to be on the order of $\log n$ (we require $k=O\left(\log n\right)$ in Theorem 3.5), thereby complying with the upper bound requirement.

We now bound the estimation error, showing that it converges as $\max(n,T)$ increases. For simplicity, we will assume that $n\gtrsim T$ so that we can write all necessary assumptions and statements in terms of $n$ directly. We stress, however, that even without $n\gtrsim T$, the statement of Theorem 3.5 also holds by replacing all occurrences of $n$ in 3.3, 3.4, and Theorem 3.5 with $\max(n,T)$. Proving this only requires changing $n$ to $\max(n,T)$ in the proofs accordingly.

Our main technical contribution is extending the convergence guarantee from [17] to a setting with multiple treatment matrices and a convex formulation that centers the rows of $\hat{M}$ without penalization. We require a set of conditions that extend the assumptions made in [17] to multiple treatment matrices.

Throughout the paper, we define the subspace $\mathbf{Z}:=\text{span}\{Z_{i},i\in[k]\}+\text{span}\{\alpha\mathbf{1}^{\top}\mid\alpha\in\mathbb{R}^{n}\}$. We next define $D^{\star}\in\mathbb{R}^{k\times k}$ and $\Delta^{\star 1}\in\mathbb{R}^{k}$ similarly to $D$ and $\Delta^{1}$ in Lemma 2.1, except we replace all quantities associated with $\hat{M}$ with the corresponding quantity for $M^{\star}$. That is, $D^{\star}_{ij}=\left\langle P_{{\mathbf{T}}^{\star\perp}}(Z_{i}),P_{{\mathbf{T}}^{\star\perp}}(Z_{j})\right\rangle$, and $\Delta^{\star 1}_{i}=\left\langle Z_{i},U^{\star}V^{\star\top}\right\rangle$ for $i,j\in[k]$. Now, we are ready to introduce the assumptions that we need to bound the estimation error.

Item 3.3(a) and Item 3.3(b) are the direct generalization of Assumptions 3(a) and 3(b) from [17] to multiple treatment matrices. Proposition 2 of [17] proves that their Assumptions 3(a) and 3(b) are nearly necessary for identifying $\tau^{\star}$. This is because it is impossible to recover $\tau^{\star}$ if $M^{\star}+\gamma Z$ were also a matrix of rank $r$ for some $\gamma\neq 0$. More generally, if $Z$ lies close to the tangent space of $M^{\star}$ in the space of rank-$r$ matrices, recovering $M^{\star}$ is significantly harder (the signal becomes second order and may be lost in the noise). In our setting with multiple treatment matrices, we further require that there is no collinearity in the treatment matrices; otherwise, we would not be able to distinguish the treatment effect associated with each treatment matrix. Mathematically, this is captured by Item 3.3(c).

The assumptions involving $E$ are mild: they would be satisfied with probability at least $1-e^{-n}$ if $E$ has independent, zero-mean, sub-Gaussian entries with sub-Gaussian norm bounded by $\sigma$ (see Lemma D.10), an assumption that is canonical in matrix completion literature [17]. Both Item 3.4(a) and Item 3.4(b) were also used in [17]. Because $\delta$ is a zero-mean matrix that is zero outside the support of $\sum_{i\in[k]}Z_{i}$, the assumption $\left\lVert\delta\right\rVert=O(\sigma\sqrt{n})$ is mild. In fact, it is satisfied with high probability if the entries of $\delta$ are sub-Gaussian with certain correlation patterns [26]. Furthermore, we note that if there is only one treatment matrix $(k=1)$ or if all treatment matrices are disjoint, then Item 3.4(c) is guaranteed to be satisfied. This is because, for each $i\in[k]$, $\delta_{i}$ is a zero-mean matrix that is zero outside of the support of $Z_{i}$. Consequently, in these scenarios, we simply have $\left\lvert\left\langle Z_{i},\delta\right\rangle\right\rvert=\left\lvert\left\langle Z_{i},\delta_{i}\right\rangle\right\rvert=0$.

We denote by $\sigma_{\max}$ and $\sigma_{\min}$ the largest and smallest singular values of $P_{{\mathbf{1}}^{\perp}}(M^{\star})$ respectively, and we let $\kappa:=\sigma_{\max}/\sigma_{\min}$ be the condition number of $P_{{\mathbf{1}}^{\perp}}(M^{\star})$.

Note that the bound becomes weaker as $\|\tilde{Z}_{j}\|_{\mathrm{F}}$ shrinks. In other words, more treated entries within a cluster results in a more accurate estimate for that cluster, as expected. Our rate in Theorem 3.5 matches, up to $\log n$ and $r$ factors, the optimal error rate obtained in Theorem 1 of [17], which bounds the error in the setting with one treatment matrix. Our bound accumulated extra factors of $\log n$ because we allowed $k=O(\log n)$ treatment matrices, and we have an extra factor of $r$ because we assumed $\left\lVert\delta\right\rVert\lesssim\sigma\sqrt{n}$, instead of $\left\lVert\delta\right\rVert\lesssim\sigma\sqrt{n/r}.$ Proposition 1 from [17] proves that their convergence rate is optimal (up to $\log n$ factors) in the “typical scenario" where $\sigma,\kappa,r=O(1)$ and $\sigma_{\min}=\Theta(n)$. We can conclude the same for our rate in Theorem 3.5.

The convergence rate of PaCE is primarily determined by the slower asymptotic convergence rate of the approximation error. However, our result in Theorem 3.5 is of independent interest, as it extends previous results on low-rank matrix recovery with a deterministic pattern of treated entries to allow for multiple treatment matrices.