Covariate-Adjusted Deep Causal Learning for Heterogeneous Panel Data Models

Covariate-Adjusted Deep Causal Learning for Heterogeneous Panel Data Models

\newcounter{appendixcounter} \refstepcounter{appendixcounter} \refstepcounter{appendixcounter} \refstepcounter{appendixcounter} \refstepcounter{appendixcounter}

\pagestyle{plain}

\setstretch{1.2}

Keywords: causal panel data models, counterfactual estimation, heterogeneous treatment effects, matrix completion, missing not at random, multi-output autoencoders, nonlinear factor models

\setstretch{1.5}

Introduction

Causal inference in panel data settings has attracted growing interest in recent years, with broad applications across various fields such as economics clarke2024synthetic, life science helske2024estimating, political science imai2021use, and social science imbens2024causal. Causal panel data analysis is inherently a missing data problem due to the fundamental problem of causal inference: for each unit and time period, we could only observe either the treated or the untreated outcome, but never both. The panel data structure brings unique challenges to causal analysis. Different units may get exposed to the treatment at various time periods, forming a block-wise non-random missing pattern athey2022design. The presence of both temporal and cross-sectional dependencies could bias the estimation if not properly addressed sun2021estimating, and unit-specific heterogeneity can confound causal interpretation millimet2023fixed, especially if additional covariates are present.

To estimate causal effects, one natural approach is to first impute the missing counterfactual, and then examine the difference between the imputed counterfactuals and the actual observed outcomes. athey2021matrix formulated the imputation of missing counterfactuals as a matrix completion task with nuclear norm minimization. It is important to note that standard matrix completion approaches are not directly suitable for causal imputation choi2024matrix. Most conventional matrix completion methods operate under the assumption that missing entries are missing at random, which enables estimation of missing values without explicitly modeling the missingness mechanism. However, in causal panel data settings where different units receive treatment at different times, missing entries correspond to counterfactual outcomes that are systematically unobserved due to treatment assignments, invalidating the missing-at-random assumption. As a result, applying standard matrix completion blindly without properly accounting for this structured, non-random missingness can lead to biased estimation agarwal2023causal.

In recent years, a growing body of work has focused on developing causal matrix completion methods that handle missing-not-at-random panel designs. Besides the optimization-based causal matrix completion athey2021matrix,choi2024matrix, another line of research leverages factor models bai2003,bai2009panel to capture the latent low-rank structure of observed entries. Most factor-model-based causal matrix completion methods in the literature xu2017generalized,bai2021matrix,agarwal2023causal,xiong2023large,yan2024entrywise posit a linear factor structure. Linear factor analysis helps capture cross-sectional and temporal dependencies to some extent han2024cp,han2024tensor,yu2024power, but is often limited in capturing complex nonlinear relationships, motivating the generalization to nonlinear factor models yalcin2001nonlinear.

Another key aspect in panel data analysis is to account for unit-specific heterogeneity millimet2023fixed,semenova2023inference. Most relevant works in the literature still focus on the overall average treatment effects across units athey2021matrix,athey2025identification, as opposed to the unit-specific heterogeneous treatment effects. Pooled estimates may mask important variations across units. Units in panel data often differ in characteristics and respond to the treatment differently wager2018estimation,nandy2023generalized, which may bias the estimates of treatment effects, motivating a critical need to estimate unit-specific heterogeneous treatment effects and adjust for the covariate effects if possible.

In this paper, we develop Covariate-Adjusted DEep CAusal Learning (CoDEAL) for panel data models, a deep-learning-based causal learning method for estimating heterogeneous treatment effects. CoDEAL captures covariate effects with a deep neural network regression, and models the covariate-adjusted outcomes via a nonlinear factor structure. Our contributions can be summarized as follows.

Related Work. (a) An overview of causal learning methods in panel data models. One traditional approach to estimate treatment effects in causal panel models is the difference-in-differences imai2021use,athey2022design,wing2024designing that compares the average changes in outcomes over time between the treated and untreated units under a parallel trend assumption. Alternatively, one can estimate the causal effects by first imputing the missing counterfactual outcomes, and then examining the difference between the imputed counterfactuals and the actual observed outcomes. Following this strand, two approaches have emerged in recent causal panel literature: the uncounfoundedness-based approach (also known as horizontal regression) rosenbaum1983central,imbens2015causal and the synthetic control method (also known as vertical regression) abadie2003economic,abadie2010synthetic,abadie2015comparative,doudchenko2016balancing,abadie2021using. The close connections between the difference-in-differences, unconfoundedness-based methods, synthetic control approaches, and causal matrix completion estimators were revealed by athey2021matrix. Despite their distinct appearances, athey2021matrix showed that the linear versions of all four estimators can be characterized as solutions to the same optimization problem with the exact same objective function, but subject to different restrictions on parameters of that optimization. We refer readers to arkhangelsky2024causal for a comprehensive review of past and recent progress on causal panel models.

(b) Causal matrix completion and deep matrix completion. A surge of interest in the recent past has led to new research on causal matrix completion methods capable of addressing missing-not-at-random panel data, e.g., optimization-based methods athey2021matrix,choi2024matrix and factor-model-based methods xu2017generalized,bai2021matrix,agarwal2023causal,xiong2023large,yan2024entrywise. There has been a handful of deep-learning-based methods for matrix completion, e.g., fan2018matrix,radhakrishnan2022simple,xiu2024deep,fan2024neuron. However, most of them assume missing-at-random, and therefore, not directly suitable for causal imputation in panel data models.

Methodology

Problem Setup

Heterogeneous Treatment Effects \footnote{Here, the “heterogeneous treatment effect” refers to the heterogeneity across units. } in Causal Panel Data Models. We begin by setting up the causal learning problem in panel data. Consider a panel data setting with $N$ units over $T$ periods. Let $[n]$ denote the set $\{1, 2, \ldots, n\}$. For each unit $i \in [N]$ and time point $t \in [T]$, there are two potential outcomes, $Y_{it}(1)$ and $Y_{it}(0)$, representing the outcomes the $i$-th unit would experience at time $t$ under treatment and control, respectively. The treatment assignment is captured by a binary indicator variable $W_{it} \in \{0, 1\}$, where $W_{it} = 1$ indicates that the unit $i$ receives the treatment at time $t$ and $W_{it} = 0$ otherwise. Suppose we also have access to some unit-specific covariates $\mathrm{\bf X}_i \in \mathbb{R}^P$ that could potentially affect the observations $\{Y_{it}, t \in [T]\}$. We denote the matrix forms of the observed outcomes, treatment indicators, and covariates by $\mathrm{\bf Y} = (Y_{it}) \in \mathbb{R}^{N \times T}$, $\boldsymbol{\Omega} = (W_{it}) \in \mathbb{R}^{N \times T}$, and $\mathrm{\bf X} = (\mathrm{\bf X}_1, \ldots, \mathrm{\bf X}_N)^\top \in \mathbb{R}^{N \times P}$, respectively. Let $\mathrm{\bf Y}(1) = (Y_{it}(1))\in \mathbb{R}^{N \times T}$ and $\mathrm{\bf Y}(0) = (Y_{it}(0))\in \mathbb{R}^{N \times T}$ be the matrices of potential outcomes of the treated and untreated, respectively. Our aim\footnote{Aligning with relevant works in the literature athey2021matrix,athey2025identification, we focus on the unit-specific average treatment effects on the treated (ATT) as our primary estimand of interest to demonstrate our proposed model and algorithm. The proposed method can be also be used to many other estimands, such as the overall (non-unit-specific) ATT $\mathbb{E}_{(i,t)} [Y_{it}(1) - Y_{it}(0) \mid W_{it}=1]$, the unit-specific average treatment effects (ATE) $ \mathbb{E}_{t} [Y_{it}(1) - Y_{it}(0)]$, and the overall (non-unit-specific) ATE $\mathbb{E}_{(i,t)} [Y_{it}(1) - Y_{it}(0)]$. } in this work is to estimate the unit-specific (potentially heterogeneous) average treatment effects on the treated (ATT), $\tau_i^\star = \mathbb{E}_{t} [Y_{it}(1) - Y_{it}(0) \mid W_{it} = 1]$ for each treated unit $i$.

Staggered Adoption Design. This work focuses on staggered adoption design athey2022design, where treatments can be rolled out at different times across various units and are irreversible once initiated. We assume there is at least one never-treated unit in the panel. Details about staggered design are in Appendix (ref).

Causal Assumptions. To ensure the validity of the estimators' interpretation, we adopt the following assumptions throughout the analysis. (i) (Stable Unit Treatment Value Assumption (SUTVA).) The potential outcomes for a given unit at a particular time depend solely on the treatment administered to that unit at that time, and are unaffected by treatments assigned to other units or at other time points. (ii) (Static Treatment Effects.) The treatment effects are non-dynamic, i.e., do not change over time.

Challenges and Limitations of Existing Methods. The fundamental problem of causal inference implies that $Y_{it}(0)$ and $Y_{it}(1)$ cannot be observed simultaneously. The panel data forms two partially observed data matrices $\mathrm{\bf Y}(0)$ and $\mathrm{\bf Y}(1)$, where the missing pattern is determined according to the treatment assignments. To estimate the treatment effects, one natural idea is to estimate the unobserved counterfactuals, in other words, to impute the missing potential outcomes. Since our focus is the unit-specific ATT and $\{Y_{it}(1): W_{it}=1\}$ is observed, it suffices to estimate $\{Y_{it}(0): W_{it}=1\}$ to impute missing elements in $\mathrm{\bf Y}(0)$. Let $\{\widehat{Y}_{it}(0): W_{it}=1\}$ denote imputed values. The unit-specific ATT can then be estimated by $\widehat\tau_i = \sum_{t: W_{it} = 1} [Y_{it}(1)-\widehat{Y}_{it}(0)]/\sum_{t} W_{it}$ for each treated unit $i$.

There are three main limitations of existing methods that this work aims to overcome in the task of modeling unit-specific ATT and imputing missing potential outcomes. First, the vast majority of matrix completion methods rely on the assumption of missing at random, an assumption that is typically violated in staggered adoption designs, thereby leading to biased estimates. Second, existing matrix-completion-based causal panel models mostly focus on the overall (non-unit-specific) ATT, without adequately accounting for unit-level heterogeneity. Third, many existing methods often lack the capacity to disentangle the causal impact of treatment from confounding covariate effects and to model the underlying complex nonlinear dependencies, limiting both the precision of counterfactual predictions and the interpretability of causal analysis.

Covariate-Adjusted Deep Causal Learning (CoDEAL) for Panel Data Models

To tackle the above challenges, we propose the CoDEAL method that (i) leverages the staggered adoption structure to fully exploit both cross-sectional and temporal dependencies for more accurate imputation of missing potential outcomes, (ii) models the unit-specific heterogeneity explicitly to accommodate variation in treatment effects across units, and (iii) involves flexible model structures and employs powerful neural network architectures to efficiently capture complex nonlinear relationships.

For each unit $i \in [N]$ at time $t \in [T]$, we model the observed outcome by

where $\mathrm{\bf F}^\star_t \in \mathbb{R}^K$ is a $K$-dimension vector of latent factors, $\phi^\star_{i}(\cdot): \mathbb{R}^K \rightarrow \mathbb{R}$ is a potentially nonlinear factor loading function, $g^\star_t(\cdot): \mathbb{R}^P \rightarrow \mathbb{R}$ is a potentially nonlinear function that captures how covariates influence the outcome, and $ \mathds{1}\{\cdot\}$ is an indicator function, and $\varepsilon_{it}$ is the idiosyncratic noise. CoDEAL employs a nonlinear factor model yalcin2001nonlinear to characterize an underlying low-rank structure of the covariate-adjusted outcomes. By letting $\phi^\star_{i}(\cdot)$ take different forms, CoDEAL unifies a class of factor-model-based causal matrix completion approaches, e.g., xu2017generalized,bai2021matrix,agarwal2023causal,xiong2023large,yan2024entrywise. As shown by (ref), CoDEAL accommodates heterogeneous but non-dynamic treatment effects. The covariate effects can be heterogeneous across units and time-varying over periods. Here, we allow $N$ and $T$ to be large and diverging, while $P$ and $K$ are fixed and small.

For clarity, we first present the learning procedures using the simplest form of a $2\times 2$ staggered adoption (also known as the four-block design) in Section (ref), and generalize the algorithm to the more general staggered adoption design in Section (ref). More details about the four-block and staggered adoption design as well as their graphical illustrations are in Appendix (ref). A graphical illustration of CoDEAL in four-block design is provided in Figure (ref). The complete CoDEAL method in general staggered adoption design is summarized in Algorithm (ref).

CoDEAL in Four-Block Design

We begin with the four-block design, the simplest form of the staggered adoption and often regarded as foundational building blocks for more general staggered adoption designs. The four-block design is a $2 \times 2$ staggered adoption design, where a subset of $N_1$ units never receive the treatment while the remaining $N_2 = N-N_1$ units are exposed to the irreversible treatment starting at time $t=T_1+1$. Without loss of generality, we rearrange the panel such that the first $N_1$ units ($i=1,\ldots, N_1$) are untreated. We call the time periods $t\in [T_1]$ as pre-treatment periods and $t\in [T]\backslash [T_1]$ as post-treatment periods. Let $T_2 = T-T_1$ denote the time duration of the post-treatment periods. The four-block design naturally decomposes the panel into four blocks:

where $\mathrm{\bf Y}_\mathcal{A} \in \mathbb{R}^{N_1\times T_1}$, $\mathrm{\bf Y}_\mathcal{B} \in \mathbb{R}^{N_1\times T_2}$, $\mathrm{\bf Y}_\mathcal{C} \in \mathbb{R}^{N_2\times T_1}$, and $\mathrm{\bf Y}_\mathcal{D} \in \mathbb{R}^{N_2\times T_2}$. By construction, $\mathrm{\bf Y}_\mathcal{A}$, $\mathrm{\bf Y}_\mathcal{B}$, $\mathrm{\bf Y}_\mathcal{C}$ are observed outcomes of the untreated with $Y_{it} = Y_{it}(0)$ for $(i,t)\in \mathcal{A} \cup \mathcal{B} \cup \mathcal{C}$, while $\mathrm{\bf Y}_\mathcal{D}$ are observed outcomes of the treated with $Y_{it} = Y_{it}(1)$ for $(i,t)\in \mathcal{D}$, leaving a missing block in $\mathrm{\bf Y}(0)$. To estimate $\tau_i^\star$, it suffices to impute the missing values of $\mathrm{\bf Y}_{\mathcal{D}}(0)$. Denote the imputed block as $\widehat{\mathrm{\bf Y}}_{\mathcal{D}}(0)$. The unit-specific ATT can be then estimated by $\widehat{\tau}_i = \frac{1}{T_2} \sum_{t: (i,t)\in\mathcal{D}} [Y_{it}(1) -\widehat{Y}_{it}(0)] \text{ for } i \in [N]\backslash [N_1]$.

Modeling Covariate Effects. To account for the potentially nonlinear, heterogeneous, and time-varying effects brought by covariates, we consider a fully connected DNN for $g_t^\star(\cdot)$. We construct our model using a fully connected DNN with ReLU activation, due to its strong empirical performance. We refer to this architecture as a deep ReLU network. Let $L\in\mathbb N$ denote the network depth and let $\mathrm{\bf d} = (d_1, \ldots, d_{L+1}) \in \mathbb{N}^{L+1}$ specify the layer widths. A deep ReLU network is a function mapping $\mathbb{R}^{d_0}$ to $\mathbb{R}^{d_{L+1}}$ and takes the form

where each $\mathcal{L}_\ell(z) = M_\ell z + b_\ell$ is an affine transformation with weight matrix $M_\ell \in \mathbb{R}^{d_\ell \times d_{\ell-1}}$ and bias vector $b_\ell \in \mathbb{R}^{d_\ell}$, and $\bar{\sigma}_\ell : \mathbb{R}^{d_\ell} \to \mathbb{R}^{d_\ell}$ applies the ReLU activation function entrywise. For simplicity, we refer to both $M_\ell$ and $b_\ell$ as the weights of the deep ReLU network.

The covariate effects can be estimated via

for certain functional class $\mathcal{G}_{P}^{1}(L,\mathrm{\bf d}, B, C)$ specified in Appendix (ref). We define the covariate-adjusted observed outcomes for unit $i \in [N]$ at time $t \in [T]$ as $\widetilde{Y}_{it} = Y_{it} - \widehat{g}_t(\mathrm{\bf X}_i).$

Estimating Counterfactuals by Matrix Completion via a Multi-Output AE. With covariate-adjusted outcomes, it remains to model the nonlinear factor structure $\phi^\star_i(\mathrm{\bf F}_t^\star)$ in (ref). Inspired by xiu2024deep, we consider tackling $\phi^\star_i(\mathrm{\bf F}_t^\star)$ using a multi-output AE.

The architecture of the multi-output autoencoder (AE) consists of a single DNN encoder and multiple independent DNN decoders. The encoder $\rho(\cdot) \in \mathcal{G}^{K_1}_{N}(L_1, \mathrm{\bf d}_1, B,C)$ maps the input to a $K_1$-dimensional bottleneck representation. For each of the $N$ output components, a distinct DNN decoder is used; specifically, the decoder for the $i$-th output is given by $\phi_i(\cdot) \in \mathcal{G}^{1}_{K_1}(L_2, \mathrm{\bf d}_2, B, C)$ for $i \in [N]$. Consequently, the $i$-th output of the multi-output AE for an input $x$ is given by $\phi_i \circ \rho(x)$. We formally define the function class of the multi-output AE as:

where $K_1\ge K$. Further details on this function class are provided in Appendix (ref).

With the multi-output AE, the encoder $\rho^\star(\cdot)$ effectively approximates the latent factors, facilitating the factor-model-based matrix completion. The unit-specific decoders $\phi_i^\star(\cdot)$, $i \in [N]$, are better equipped to capture the heterogeneity across units, in contrast to a traditional single-output AE which limits $\phi_1^\star = \ldots =\phi_N^\star$. DNNs involved in the multi-output AE are estimated by

The counterfactuals $\mathrm{\bf Y}_{\mathcal{D}}(0)$ can then be imputed by $\widehat{Y}_{it}(0) = \widehat{\phi}_i (\widehat\rho(\widetilde{\mathrm{\bf Y}}_{\cdot t})) + \widehat{g}_t(\mathrm{\bf X}_i)$ for any $(i,t) \in \mathcal{D}$. The complete steps of imputing counterfactuals in the four-block design are summarized as a standalone algorithm in Algorithm (ref) to facilitate further generalizations to staggered design.

CoDEAL in Staggered Adoption Design

The extension to staggered adoption design is a natural generalization of the four-block design. Under the staggered adoption design, we rearrange the panels by sorting the $N$ units based on the time at which they began receiving the treatment in a descending order (with never-treated units at the top, followed by the latest adopters, and then progressively earlier adopters towards the bottom). This sorting procedure yields a block structure in the treatment indicator matrix $\boldsymbol{\Omega}$. By grouping the units based on their adoption timing and segmenting the overall time periods into relevant intervals, we partition the indicator matrix $\boldsymbol{\Omega}$ (as well as the outcome matrix $\mathrm{\bf Y}$) into submatrices $\boldsymbol{\Omega}_{(\xi, \eta)} \in \mathbb{R}^{N_{\xi} \times T_{\eta}}$ and $\mathrm{\bf Y}_{(\xi, \eta)} \in \mathbb{R}^{N_{\xi} \times T_{\eta}}$ for $\xi=1,\ldots, r$ and $\eta =1,\ldots, r$, where $\sum_{\xi=1}^{r} N_{\xi} =N $ and $\sum_{\eta=1}^{r} T_{\eta} = T $. Examples of block partitions in the staggered design are included in Appendix (ref). Let $\mathcal{I}_{tr} = \{(\xi,\eta): \boldsymbol{\Omega}_{(\xi, \eta)} = 1\}$ denote the index set of the treated blocks. By the construction of these partitions, the indices in $\mathcal{I}_{tr}$ satisfy $\xi+\eta > r+1$.

To estimate the unit-specific ATT, it suffices to estimate the counterfactuals of $\mathrm{\bf Y}_{(\xi,\eta)}$ for all $(\xi,\eta) \in \mathcal{I}_{tr}$. The imputation of $\mathrm{\bf Y}_{(\xi,\eta)}(0)$ can be reduced to a four-block design problem yan2024entrywise, therefore the method proposed in the previous Section (ref) can be directly applied. Specifically, to estimate $\mathrm{\bf Y}_{(\xi_0, \eta_0)}(0)$ for some $(\xi_0, \eta_0) \in \mathcal{I}_{tr}$ , we construct a four-block data matrix $ \mathrm{\bf Y}^{(\xi_0,\eta_0)}$ as follows

where $\mathcal{A}_{(\xi_0,\eta_0)} = \{(\xi,\eta): 1 \leq \xi \leq k_1, 1 \leq \eta \leq k_2 \}$, $\mathcal{B}_{(\xi_0,\eta_0)} = \{(\xi,\eta):1\leq \xi \leq k_1, k_2+1 \leq \eta \leq \eta_0 \}$, $\mathcal{C}_{(\xi_0,\eta_0)} = \{(\xi,\eta): k_1+1 \leq \xi \leq \xi_0, 1 \leq j \leq k_2 \}$, and $\mathcal{D}_{(\xi_0,\eta_0)} = \{(\xi,\eta): k_1+1 \leq \xi \leq \xi_0, k_2+1 \leq j \leq \eta_0 \}$ with \( k_1 = r + 1 - \eta_0 \) and \( k_2 = r + 1 - \xi_0 \). The associated treatment indicator matrix $\boldsymbol{\Omega}^{(\xi_0,\eta_0)}$ is defined analogously and the associated covariate matrix $\mathrm{\bf X}^{(\xi_0)}$ can be obtained by retaining the units included in $\mathrm{\bf Y}^{(\xi_0,\eta_0)}$. The missing potential outcomes $\mathrm{\bf Y}_{\mathcal{D}_{(\xi_0,\eta_0)}}(0)$ can be directly estimated using Algorithm (ref). By construction, $\mathrm{\bf Y}_{\mathcal{D}_{(\xi_0,\eta_0)}} \supseteq \mathrm{\bf Y}_{(\xi_0,\eta_0)}$. In the case when $\mathrm{\bf Y}_{\mathcal{D}_{(\xi_0,\eta_0)}} \supset \mathrm{\bf Y}_{(\xi_0,\eta_0)}$, we treat all the blocks of $\mathrm{\bf Y}_{\mathcal{D}_{(\xi_0,\eta_0)}}$ as missing and impute their values, but only extract $\widehat{\mathrm{\bf Y}}_{(\xi_0,\eta_0)}(0)$ from the imputed $\widehat{\mathrm{\bf Y}}_{\mathcal{D}_{(\xi_0,\eta_0)}}(0)$; see Figure (ref) for examples.

With the four-block construction in (ref) and the CoDEAL method in four-block design introduced in Section (ref), we obtain the imputed values of $\mathrm{\bf Y}_{(\xi, \eta)}(0)$ for all $(\xi,\eta) \in \mathcal{I}_{tr}$. The unit-specific ATT can then be estimated by $\widehat\tau_i = \sum_{t: W_{it} = 1} [Y_{it}(1)-\widehat{Y}_{it}(0)]/\sum_{t} W_{it}$ for each treated unit $i$. The complete algorithm of CoDEAL in the staggered adoption design is summarized as Algorithm (ref).

Theoretical Properties

Assumption (ref) is a standard assumption in the literature of nonparametric regression. Note that assuming $\mathrm{\bf X}_i$ has compact support is equivalent to assuming that the support is $[0,1]^P$ via centering and scaling. The smoothness condition is standard in the nonparametric regression literature (cf. tsybakov2003introduction, Chapter 1), as it governs the complexity of the underlying function class. A formal definition is provided in Appendix (ref). Assumption (ref) generalizes the classical pervasiveness condition (with $n = N$) from linear factor models bai2003. It effectively enables approximate recovery of the latent factors from the observed input: when $x$ denotes the latent factors, condition (ref) ensures the existence of a reconstruction map $\rho^\star \circ M^\star$ that approximately inverts the encoding $\phi^\star(x)$ and recovers $x$ from the noiseless data. To illustrate, consider the linear factor model setup in bai2003, where $\phi^\star(x) = A x$ with bounded loadings satisfying $\|A\|_\infty \lesssim 1$ and $\|A\|_0 \asymp N$ indicating that each factor affects nearly all observed variables. Assumption (ref) is satisfied by setting $M^\star = N^{-1} A^\top$ and $\rho^\star(x) = N (A^\top A)^{-1} x$.

The following theorem establishes the convergence of the estimated counterfactuals under a four-block design. The result can be readily extended to general staggered adoption designs.

Additional theoretical properties, technical conditions and proofs are detailed in Appendix (ref).

Simulation Studies

In this section, we use extensive simulation studies to demonstrate the finite-sample performance of the proposed CoDEAL\footnote{Our code for implementing CoDEAL is available in the supplementary material.}.

Benchmark Methods

We consider five benchmark methods, including a single-output AE (Single AE) (detailed below), neuron-enhanced AE (AEMC-NE) fan2024neuron, matrix completion with nuclear norm minimization (MC-NNM) athey2021matrix, vertical regression (Vert-Reg) doudchenko2016balancing, and difference-in-differences (DiD) imai2021use.

Additionally, we evaluate the algorithm's performance in terms of the covariate effect removal. To this end, we compare the performance of CoDEAL using different methods to remove the covariate effects, specifically, DNN-based removal (as illustrated in Section (ref)), and linear-regression(LR)-based removal (detailed below).

We evaluate the algorithm's performance using two evaluation metrics: the mean absolute error $MAE=(\sum_{(i, t): W_{it} = 1}|Y_{it}(1)-\widehat{{Y}}_{it}(0)-\tau_{i}^\star|)/(\sum_{(i,t)} W_{it})$ and the mean squared error $MSE=(\sum_{(i,t): W_{it} = 1}(Y_{it}(1)-\widehat{{Y}}_{it}(0)-\tau_{i}^\star)^2)/(\sum_{(i,t)} W_{it})$.

{Matrix Completion by Single-Output AE (as a benchmark method).} Our proposed CoDEAL utilizes a multi-output autoencoder (as illustrated in Figure (ref)) to recover nonlinear latent factors while explicitly modeling the unit-specific heterogeneity through multiple decoders. The multi‐output autoencoder uses a shared encoder network \(\rho(\cdot)\) to map each column \(\widetilde{\mathbf{Y}}_{\!\cdot t}\) into a \(K\)-dimensional latent space, followed by \(N\) separate decoder networks \(\{\phi_1,\dots,\phi_N\}\) to reconstruct each unit’s value. To demonstrate CoDEAL's ability to account for the unit-specific heterogeneity, we include an additional benchmark that replaces the multi-output autoencoder with a single-output autoencoder in the CoDEAL algorithm. This benchmark is referred to as the Single AE in later discussions. A single-output autoencoder uses the same encoder \(\rho(\cdot)\) but only one decoder \(\phi(\cdot)\) (i.e. $\phi=\phi_1=\dots=\phi_N$) for all units. By design, the common decoder makes the single-output AE inherently unable to model the unit-specific heterogeneity effectively. For implementation, DNNs in the single-output AE are estimated by

{Covariate Effect Removal by Linear Regression (as a benchmark method).} Our proposed CoDEAL uses a DNN to remove the potentially nonlinear, heterogeneous, and time-varying effects brought by covariates. To evaluate the algorithm's performance in terms of the covariate effect removal, we include a linear-regression(LR)-based removal as a benchmark. The LR-based removal assumes that covariates have a linear impact on the outcome. With covariates $\mathbf{X}_i\in\mathbb{R}^{P}$, $i\in [N]$, we denote the effect of covariates on the outcome by a coefficient matrix $\mathrm{\bf H} = [\mathrm{\bf H}_1, \ldots, \mathrm{\bf H}_T] \in\mathbb{R}^{P\times T}$, and we estimate $\mathrm{\bf H}_t \in \mathbb{R}^P$ separately for each period. For each pre-treatment time $t\in\{1,\dots,T_1\}$, we estimate the covariate effect vector $\widehat{\mathrm{\bf H}}_t\in\mathbb{R}^{P}$ using OLS on all units: $\widehat{\mathrm{\bf H}}_t = \arg\min_{\mathrm{\bf H}_t} \sum_{i=1}^{N}\left(Y_{it}-\mathrm{\bf X}_i^\top\mathrm{\bf H}_t\right)^2.$ For each post-treatment period $t\in\{T_1+1,\dots,T\}$, we use only control units to estimate the covariate effect: $\widehat{\mathrm{\bf H}}_t = \arg\min_{\mathrm{\bf H}_t} \sum_{i: W_{it}=0}\left(Y_{it}-\mathrm{\bf X}_i^\top\mathrm{\bf H}_t\right)^2.$ With the estimated $\widehat{\mathrm{\bf H}}_t, t\in [T]$, we obtain the LR-based covariate-adjusted observed outcomes by $\widetilde{Y}_{it} = Y_{it} - \mathrm{\bf X}_i^\top\widehat{\mathrm{\bf H}}_t$ for each $i \in [N]$ and $t \in [T]$.

Data Generating Process and Simulation Configurations

We perform extensive simulation studies under various linear and nonlinear generating mechanisms, with different setups of unit numbers $N$, time periods $T$, matrix ranks $K$, and treatment assignments. Specifically, we consider the following configurations 1-4:

For the generation of factor effects and covariate effects, we first generate covariates $\mathrm{\bf X} = (X_{it}) \in \mathbb{R}^{N\times P}$, factors $\mathrm{\bf F} = (F_{it}) \in \mathbb{R}^{T\times K}$, factor loadings $\boldsymbol{\Lambda} = (\Lambda_{it})\in \mathbb{R}^{N\times K}$ with $X_{it},F_{it}, \Lambda_{it}\overset{\text{i.i.d.}}{\sim}N(0,1)$. With these building blocks, we consider various linear and nonlinear generating mechanisms to generate factor effects $\phi_i^\star(\mathrm{\bf F}^\star_{t})$ and covariate effects $g^\star_t(\mathbf{X}_i)$. Specifically, we experiment with 4 variants for generating factor effects and 6 variants for generating covariate effects, as listed in Table (ref). The unit-specific causal effect $\tau_i^\star \overset{i.i.d}{\sim}N(12,5) $ is designed to be added to treated units in treated periods. The observed outcome matrix $\mathrm{\bf Y} = (Y_{it}) \in \mathbb{R}^{N \times T}$ is then obtained by $Y_{it}= Y^\star_{it} + \tau^\star_i \cdot \mathcal{I}\{W_{it}=1\} + \varepsilon_{it}$ where $W_{it}$ are determined based on the treatment assignments according to the simulation configuration, and the idiosyncratic error $\varepsilon_{it}$ follows $\varepsilon_{it}\overset{i.i.d}{\sim}\mathcal{N}(0,0.5^2)$.

Evaluation on the Performance of CoDEAL in the Removal of Covariate Effects

To evaluate the algorithm’s performance in terms of the covariate effect removal, we compare the performance of CoDEAL using different methods, specifically, DNN (as illustrated in Section (ref)) and LR (detailed in Section (ref)), to remove the covariate effects. Six scenarios of covariate effects are considered as detailed in Table (ref). We also include an oracle benchmark, when no covariate effects are present and no covariate-adjusted operations are needed.

Table (ref) and Table (ref) report the comparisons of MAE and MSE using the two covariate effect removal methods in Config.1 and Config.2, respectively. Under each configuration, we consider both linear factor effects and nonlinear Sine factor effects, with all six types of covariate effects.

As shown in Tables (ref) and (ref), LR-based removal and DNN-based removal behave very differently depending on the underlying true types of covariate effects. In cases where the covariate effect is strictly linear, LR often outperforms NN by a small margin, indicating that a simple linear adjustment suffices for purely linear effects. In contrast, whenever the covariate effects enter nonlinearly (e.g.\ tanh, log, polynomial, or ReLU), the proposed DNN covariate removal achieves substantially lower errors, with MAE reduced by 10–30 % and MSE by similar amounts relative to LR. These results confirm that LR is adequate for linear covariate effects but that the DNN approach is essential to capture complex, nonlinear covariate–outcome relationships.

Evaluation on the Performance of CoDEAL in Causal Matrix Completion

Table (ref) reports the comparisons of prediction accuracy between the proposed CoDEAL and five benchmark methods (detailed in Section (ref)) in Config.1-4. Under each configuration, we consider all four types of factor effects (listed in Table (ref)), with linear covariate effects and nonlinear tanh covariate effects.

CoDEAL in Four-Block Design. As shown in Config.1 and Config.2 in Table (ref), CoDEAL consistently achieves the lowest MAE and MSE across nearly all settings of factor effects and covariate effects, and achieves substantial improvements in MAE and MSE relative to all competing methods. Under the simplest linear factor model, CoDEAL reduces MAE by approximately 5–10% and MSE by about 10–20% compared to MC-NNM, the second-best performing method, and significantly outperforms other baseline methods by even larger margins. For scenarios involving nonlinear factor models, CoDEAL continues to show superior performance, typically surpassing MC-NNM by approximately 5–20% reduction in MAE and 10–20% reduction in MSE. Overall, these findings clearly demonstrate that CoDEAL’s integration of DNN-based covariate adjustment and deep autoencoder-based latent factor extraction is robust and highly effective across diverse linear and nonlinear data-generating settings.

CoDEAL in Staggered Adoption Design. As Config.3 and Config.4 of Table (ref) demonstrate, CoDEAL achieves substantial percentage gains in MAE and MSE relative to all competing methods across staggered adoption designs. Under Linear factors, CoDEAL CoDEAL {ties MC-NNM within 3% while reduces MAE by 16–35% over other benchmarks}. In staggered design with nonlinear factors, CoDEAL delivers evident advantages. Especially when under ReLU-based factors, CoDEAL delivers its largest margins, with {MAE reductions of 12–40% and MSE reductions of 15–41% versus all competing approaches}. Finally, as $r$ increases from 5 to 10, the missingness pattern grows more complex and extensive, so all methods show some drop in accuracy. However, CoDEAL’s errors only suffer a slight increase, whereas simpler methods like Vert-Reg show a dramatic rise in both MAE and MSE when missingness increases. Above results reemphasize CoDEAL’s consistent superiority and robustness in staggered‐treatment settings.

In summary, CoDEAL achieves performance comparable to MC-NNM in linear settings, while exhibiting substantial improvements over all considered methods in the presence of nonlinear effects. By integrating DNN-based covariate adjustments with deep multi-output-AE–based nonlinear factor analysis, CoDEAL is robustly effective across a wide range of linear and nonlinear scenarios.

Real Data: Oxford COVID-19 Government Response Tracker (OxCGRT)

In this section, we illustrate our methodology using data from the OxCGRT Hale2021. The data is publicly available at \href{https://github.com/OxCGRT}{https://github.com/OxCGRT}. OxCGRT systematically collects panel data on the timing and intensity of 24 different government interventions, such as school closures, travel restrictions, and vaccination mandates, across over 185 countries from January 2020 through 2022.

We examine the impact of two specific policy interventions: (1) the implementation of mandatory vaccination, and (2) the termination of internal travel restrictions (results in Appendix (ref)). Our analysis covers their impacts on COVID-19 confirmed cases and deaths across 64 states and territories in the United States and Canada from April 16, 2021, to December 11, 2021. This time frame aligns with the period following the detection of the Delta variant and precedes the rapid emergence of the Omicron variant in North America. In our analysis, we use four OxCGRT indices (overall government response, containment and health, stringency, and economic support), averaged over the time frame, as covariates, resulting in data dimensions of $N=64$, $T=240$, and $P=4$. The policy is adopted in different states at various times, forming a staggered adoption design (visualized in Figure (ref)).

Figure (ref) plots the comparisons of the total confirmed cases and deaths across policy-in-effect states between observed results (with policy executed) and estimated counterfactuals (if no policy). As shown in Figure (ref), if the mandatory vaccination policy was not implemented, there would be an increase in both confirmed cases and deaths, with a notably larger impact on mortality. This observation aligns with recent research highlighting a greater reduction in mortality than in infection rates following vaccination hernandez2023impact,xu2023systematic,mesleEtAl2024estimated,wu2023long. Although the visual differences in the log-transformed graphs appear modest, converting these back to their original scales reveals substantial impacts: on December 11, 2021, the actual cumulative counts of confirmed cases and deaths for states with mandatory vaccination policies were ${26.57\text{ }million, 0.39\text{ }million}$, respectively. In contrast, our estimates suggest that without these policies, the counts would have risen to ${29.81\text{ }million(\uparrow 12.2\%), 0.52\text{ }million(\uparrow 34.6\%)}$.

Discussion of Limitations and Future Extensions

In this paper, we propose CoDEAL, a deep-learning-based causal matrix completion method and a unifying framework for estimating heterogeneous causal effects in panel data models. One limitation of CoDEAL is that it is only designed for binary treatment settings. Causal panel data with multiple treatments is an open research area with limited existing work in the literature abadie2021using, yet this topic is gradually gaining attention due to its growing relevance in practical applications agarwal2020synthetic,squires2022causal.

To extend CoDEAL to multiple treatment settings, one possible direction is to adopt a similar strategy to that mentioned by agarwal2020synthetic, in which the treatment type is incorporated as an additional data dimension in the analysis. Together with the original panel data structures, the inclusion of treatment dimension forms a three-dimensional tensor. Methodologically, with different choices of factor models and advanced neural network architectures, the proposed CoDEAL framework can be further generalized to accommodate various data types with more complex data structures, such as tensors, images, and networks. For example, CoDEAL can be naturally extended to three-dimensional or higher-order tensors by involving tensor factorization methodologies han2022tensor, chen2024estimation, zhou2025factor,chen2025diffusion, providing a more general approach compared to matrix factor analysis bai2021matrix, luo2022inverse,yu2022nonparametric,yu2024dynamic. That being said, to extend to multiple treatments, how to handle the increased level of missingness brought by the treatment assignments in a higher-dimensional tensor remains challenging and requires careful modeling strategies. We leave the further investigation on multiple treatment settings as future work.