Treatment Effects Estimation on Networked Observational Data using Disentangled Variational Graph Autoencoder

From the theoretical analysis in the previous section, we have seen that eliminating unnecessary factors is essential to effectively and accurately estimating the treatment effect. However, in practice, we do not know the mechanism of generating $\mathbf{x}$ from $\mathbf{z}$ and the mechanism of disentangling $\mathbf{z}$ into different disjoint sets. This requires us to propose a method that can learn to disentangle the latent factors $\mathbf{z}$ and estimate ITE through what the model has learned.

Therefore, we aim to infer the posterior distribution $p_{\theta}(\mathbf{z}\mid\mathbf{x},\mathbf{A})$ of latent factors $\mathbf{z}$ through the observed proxy covariates $\mathbf{x}$ and network information $\mathbf{A}$, while disentangling $\mathbf{z}$ into latent instrumental factors $\mathbf{z}_{t}$, confounding factors $\mathbf{z}_{c}$, adjustment factors $\mathbf{z}_{y}$, noise factors $\mathbf{z}_{o}$. Since exact inference is intractable, we use the variational inference framework to approximate the posterior distribution with the tractable distribution. We adopt Variational Graph Autoencoders (VGAEs) to construct our model. Proposed by Kipf and Welling (2016b), VGAEs extend Variational Autoencoders (VAEs) to take into consideration the graph structure in the data. For every observed variable $\mathbf{x}$, VGAEs define a multi-dimensional latent variable $\mathbf{z}$. Moreover, VGAEs rely on an adjacency matrix $\mathbf{A}$, which is utilized by the Graph Neural Network (GNN) in the encoder to enforce the structure of the posterior approximation $q_{\phi}(\mathbf{z}\mid\mathbf{x},\mathbf{A})$. As shown in Fig. 2, we use four separate encoders to approximate the variational posterior $q_{\phi_{t}}(\mathbf{z}_{t}\mid\mathbf{x},\mathbf{A})$, $q_{\phi_{c}}(\mathbf{z}_{c}\mid\mathbf{x},\mathbf{A})$, $q_{\phi_{y}}(\mathbf{z}_{y}\mid\mathbf{x},\mathbf{A})$, $q_{\phi_{o}}(\mathbf{z}_{o}\mid\mathbf{x},\mathbf{A})$, disentangling the latent variable $\mathbf{z}$ into $\mathbf{z}_{t}$, $\mathbf{z}_{c}$, and $\mathbf{z}_{y}$, $\mathbf{z}_{o}$, respectively²²2Our method does not employ $t$ and $y$ as inputs to the encoder as done in (Louizos et al., 2017), because we assume that $t$ and $y$ are generated by the latent factors. So, the inference of latent factors relies solely on $\mathbf{x}$. For additional information, see (Zhang et al., 2021b).. Then, these four latent factors are used by the decoder $p_{\theta}(\mathbf{x}\mid\mathbf{z}_{t},\mathbf{z}_{c},\mathbf{z}_{y},\mathbf{z}_{o})$ to reconstruct $\mathbf{x}$, $t$, and $y$³³3Note that, as shown in Figure 1, we obtain the independence property $\mathbf{x}\!\perp\!\!\!\perp\mathbf{A}\mid\{\mathbf{z}_{t},\mathbf{z}_{c},\mathbf{z}_{y},\mathbf{z}_{o}\}$. Thus, the original VGAE decoder $p_{\theta}(\mathbf{x}\mid\mathbf{z}_{t},\mathbf{z}_{c},\mathbf{z}_{y},\mathbf{z}_{o},\mathbf{A})=p_{\theta}(\mathbf{x}\mid\mathbf{z}_{t},\mathbf{z}_{c},\mathbf{z}_{y},\mathbf{z}_{o})$. The derivation for $t$ and $y$ is similar.. Following the standard VGAE design, we select the prior distributions $p(\mathbf{z}_{t})$, $p(\mathbf{z}_{c})$, , $p(\mathbf{z}_{y})$ and $p(\mathbf{z}_{o})$ as factorized Gaussian distributions:

where $d_{\mathbf{z}_{t}}$, $d_{\mathbf{z}_{c}}$, $d_{\mathbf{z}_{y}}$, and $d_{\mathbf{z}_{o}}$ represent the dimensions of latent instrumental, confounding, adjustment, noise factors, respectively. And $\{\mathbf{z}_{t}\}_{j}$ denotes the $j$-th dimension of $\mathbf{z}_{t}$, and the same applies to $\mathbf{z}_{c}$, $\mathbf{z}_{y}$, and $\mathbf{z}_{o}$.

The probabilistic representation of the generative model for $\mathbf{x}$, $t$, and $y$ is as follows:

where $f_{1}$ to $f_{7}$ are functions parameterized by fully connected neural networks, $\sigma(\cdot)$ represents the logistic function, and $Bern$ refers to the Bernoulli distribution. The distribution of $\mathbf{x}$ should be chosen based on the dataset, and in our case, we approximate it with a Gaussian distribution, as the data we use consists of continuous variables. Similarly, for the continuous outcome variable $y$, we also parameterize it as a Gaussian distribution, where the mean and variance are defined by two separate neural networks defining $p(y\mid t=1,\mathbf{z}_{c},\mathbf{z}_{y})$ and $p(y\mid t=0,\mathbf{z}_{c},\mathbf{z}_{y})$, following the two-headed approach proposed by Shalit et al. (2017).

In the inference model, since we input the network information $\mathbf{A}$ into the encoder, we design the encoder based on the idea of Variational Graph Autoencoders (VGAEs). Specifically, we utilize Graph Convolutional Networks (GCNs) (Kipf and Welling, 2016a) as the encoder to obtain latent factor representations. GCN has been shown to effectively handle non-Euclidean data, such as graph-structured data, across diverse settings. To simplify notation, we describe the message propagation rule using a single GCN layer, as shown below:

where $\mathbf{h}\in\mathbb{R}^{d}$ is the output vector of the GCN, $\mathbf{X}\in\mathbb{R}^{n\times k}$ is the feature matrix of the instances, $(\hat{\mathbf{A}}\mathbf{X})_{\mathbf{x}}$ represents the row of the matrix product corresponding to instance $\mathbf{x}$, $\tilde{\mathbf{A}}=\mathbf{A}+\mathbf{I}_{N}$, $\mathbf{I}_{N}$ is the identity matrix, $\tilde{\mathbf{D}}_{ii}=\sum_{j=1}^{N}\tilde{\mathbf{A}}_{ij}$, and $\mathbf{W}\in\mathbb{R}^{k\times d}$ denotes the parameters of the weight matrix. ${Relu}(\cdot)$ denotes the ReLU activation function. This leads to the following definition of the variational approximation of the posterior distribution for the latent factors:

where $\hat{\boldsymbol{\mu}}_{t}$, $\hat{\boldsymbol{\mu}}_{c}$, $\hat{\boldsymbol{\mu}}_{y}$, $\hat{\boldsymbol{\mu}}_{o}$ and ${\rm{diag}}(\hat{\boldsymbol{\sigma}}_{t}^{2})$, ${\rm{diag}}(\hat{\boldsymbol{\sigma}}_{c}^{2})$, ${\rm{diag}}(\hat{\boldsymbol{\sigma}}_{y}^{2})$, ${\rm{diag}}(\hat{\boldsymbol{\sigma}}_{o}^{2})$ are the means and covariance matrix of the Gaussian distributions, parameterized by the GCN as shown in Equation (7). Additionally, $\hat{\boldsymbol{\mu}}_{t}$ and ${\rm log}\,{\boldsymbol{\sigma}}_{t}^{2}$ are learned from two GCNs that share the training parameters of the first layer, and the same applies to the remaining three pairs.

The latent factors $\mathbf{z}_{t}$ and $\mathbf{z}_{c}$ are associated with the treatment $t$, whereas $\mathbf{z}_{c}$ and $\mathbf{z}_{y}$ are associated with the outcomes $y$, as illustrated in Fig. 1. To ensure that the treatment information is effectively captured by the union of $\mathbf{z}_{t}$ and $\mathbf{z}_{c}$, we add an auxiliary classifier to predict $t$ from the encoder’s output, under the assumption that $\mathbf{z}_{t}$ and $\mathbf{z}_{c}$ can accurately predict $t$. Additionally, $y$ is predicted using two regression networks under different treatments to ensure that the outcome information is captured by the union of $\mathbf{z}_{c}$ and $\mathbf{z}_{y}$, based on the assumption that $\mathbf{z}_{c}$ and $\mathbf{z}_{y}$ can accurately predict $y$. Inspired by related approaches (Zhang et al., 2021b; Liu et al., 2024), the classifier and regression networks are defined as follows:

where $h_{1}$ to $h_{5}$ are functions parameterized by fully connected neural networks, and the distribution settings are similar to those in Equations (5) and (6).

Explicitly enhancing the independence of disentangled latent factors encourages the graph encoder to more effectively capture distinct and mutually independent information associated with each latent factor. In the following, we detail the regularization applied to enforce independence among the latent factors.

The goal of our method is for the encoder to capture disentangled latent factors—namely, $\mathbf{z}_{y}$, $\mathbf{z}_{c}$, $\mathbf{z}_{t}$, and $\mathbf{z}_{o}$—that each contain exclusive information. This requires increasing the statistical independence between these latent factors to further strengthen disentanglement. Given the high dimensionality of the latent factors, using histogram-based measures is infeasible. Therefore, we use the Hilbert-Schmidt Independence Criterion (HSIC) (Gretton et al., 2005) to promote sufficient independence among different latent factors.

Specifically, let $\mathbf{z}_{t,*}$ represent the $d_{\mathbf{z}_{t}}$-dimensional random variable corresponding to the latent factor $\mathbf{z}_{t}$. Consider a measurable, positive definite kernel $\kappa_{t}$ defined over the domain of $\mathbf{z}_{t,*}$, with its associated Reproducing Kernel Hilbert Space (RKHS) denoted by $\mathcal{H}_{t}$. The mapping function $\psi_{t}(\cdot)$ transforms $\mathbf{z}_{t,*}$ into $\mathcal{H}_{t}$ according to the kernel $\kappa_{t}$. Similarly, for $\mathbf{z}_{y}$, $\mathbf{z}_{c}$, and $\mathbf{z}_{o}$, the same definitions apply. Given a pair of latent factors $\mathbf{z}_{t}$ and $\mathbf{z}_{c}$, where $\mathbf{z}_{t,*}$ and $\mathbf{z}_{c,*}$ are jointly sampled from the distribution $p(\mathbf{z}_{t,*},\mathbf{z}_{c,*})$, the cross-covariance operator $\mathcal{C}_{\mathbf{z}_{t,*},\mathbf{z}_{c,*}}$ in the RKHS of $\kappa_{t}$ and $\kappa_{c}$ is defined as:

where $\boldsymbol{\mu}_{\mathbf{z}_{t,*}}=\mathbb{E}(\psi_{t}(\mathbf{z}_{t,*}))$, $\boldsymbol{\mu}_{\mathbf{z}_{c,*}}=\mathbb{E}(\psi_{c}(\mathbf{z}_{c,*}))$. Then, HSIC is defined as follows:

where $\|\cdot\|$ is the Hilbert-Schmidt norm, which generalizes the Frobenius norm on matrices. It is known that for two random variables $\mathbf{z}_{t,*}$ and $\mathbf{z}_{c,*}$ and characteristic kernels $\kappa_{\mathbf{z}_{t,*}}$ and $\kappa_{\mathbf{z}_{c,*}}$, if $\mathbb{E}[\kappa_{\mathbf{z}_{t,*}}({\mathbf{z}_{t,*}},{\mathbf{z}_{t,*}})]<\infty,\mathbb{E}[\kappa_{\mathbf{z}_{c,*}}({\mathbf{z}_{c,*}},{\mathbf{z}_{c,*}})]<\infty$, then ${\rm HSIC}({\mathbf{z}}_{t,*},{\mathbf{z}}_{c,*})=0$ if and only if ${\mathbf{z}}_{t,*}\!\perp\!\!\!\perp{\mathbf{z}}_{c,*}$. In practice, we employ an unbiased estimator ${\rm HSIC}({\mathbf{z}}_{t,*},{\mathbf{z}}_{c,*})$ with $n$ samples (Song et al., 2012), defined as:

where $\tilde{\mathbf{U}}$ and $\tilde{\mathbf{V}}$ denote the Grammer matrices with $\kappa_{\mathbf{z}_{t,*}}$ and $\kappa_{\mathbf{z}_{c,*}}$, respectively, with the diagonal elements set to zero. In our approach, we employ the radial basis function (RBF) kernel. The analysis for other pairs of latent factors follows similarly.

The advantage of using Equation (13) to measure the dependence between different latent factors lies in its ability to capture more complex, nonlinear dependencies by mapping latent factors into the RKHS. The HSIC estimator we employ is unbiased, which is both effective and computationally efficient.

The encoder and decoder parameters can be learned by minimizing the negative evidence lower bound (ELBO), consistent with the standard VGAE (Kipf and Welling, 2016b), where $i$ denotes the $i$-th instance:

The factual loss function for predicting potential outcomes, along with the loss function for predicting treatment assignments, is defined as follows:

We apply pairwise independence constraints to the latent factors $z_{t}$, $z_{c}$, $z_{y}$, and $z_{o}$ in order to improve the statistical independence between disentangled representations. Thus, the HSIC regularizer $\mathcal{L}_{reg}$ is calculated as follows:

As shown in Fig. 1, we observe that $\mathbf{z}_{y}\!\perp\!\!\!\perp t$, implying that $p(\mathbf{z}_{y}\mid t=0)=p(\mathbf{z}_{y}\mid t=1)$. Therefore, following the approach in (Hassanpour and Greiner, 2019), we aim for the learned $\mathbf{z}_{y}$ to exclude any confounding information, ensuring that all confounding factors are captured within $\mathbf{z}_{c}$. This is crucial for ensuring the accuracy of the treatment effect estimation. To quantify the discrepancy between the distributions of $\mathbf{z}_{y}$ for the treatment and control groups, we use the integral probability metric (IPM) (Müller, 1997; Sriperumbudur et al., 2012; Guo et al., 2020c). We define the balanced representation loss as $\mathcal{L}_{disc}$,

We utilize the Wasserstein-1 distance, defined as Sriperumbudur et al. (2012), to calculate Equation (18). Additionally, we employ the effective approximation algorithm proposed by (Cuturi and Doucet, 2014) to calculate the Wasserstein-1 distance and associated gradients about the model parameters for training the TNDVGA.

The following provides a summary of the overall objective function for TNDVGA:

where $\alpha_{t}$, $\alpha_{y}$, $\alpha_{1}$, and $\alpha_{2}$ are non-negative hyperparameters that balance the corresponding terms. The final term, $\lambda{\|\Theta\|}_{2}^{2}$, is applied to all model parameters $\Theta$ to avoid overfitting.

After completing the model training, we can predict the ITEs of new instances based on the observed covariates $\mathbf{x}$. We utilize the encoders $q_{\phi_{c}}(\mathbf{z}_{c}\mid\mathbf{x},\mathbf{A})$ and $q_{\phi_{y}}(\mathbf{z}_{y}\mid\mathbf{x},\mathbf{A})$ to sample the posteriors of confounding factors and risk factors $l$ times, and then use the decoder $p_{\theta_{y}}(y\mid t,\mathbf{z}_{c},\mathbf{z}_{y})$ to compute the predicted outcomes $y$ at different $t$, averaging them to obtain the estimated potential outcomes $y^{1}$ and $y^{0}$. The calculation of ATE can be done by performing the above steps on all test samples and then taking the average.

Inspired by (Hassanpour and Greiner, 2019), we generate synthetic datasets named TNDVGASynth, which follow the structure illustrated in Fig. 1 and the relationships defined in Equations (22)-(25).

where $concat(\cdot,\cdot)$ represents the vector concatenation operation. $a$ is an element of the matrix $\mathbf{A}$; $m_{t},m_{c},m_{y},m_{o}$ represent the dimensions of the latent factors $\mathbf{z}_{t},\mathbf{z}_{c},\mathbf{z}_{y},\mathbf{z}_{o}$, respectively; a scalar $\zeta$ determines the slope of the logistic curve; $\cdot$ denotes the dot product, and $\circ$ signifies the element-wise (Hadamard) product. We considered all feasible datasets generated from the grid defined by $m_{t},m_{c},m_{y},m_{o}\in\{4,8\}$, creating 16 scenarios. For each scenario, we synthesize five datasets using different initial random seeds.

Tables 2 and 3 present the experimental results on the BlogCatalog and Flick datasets, respectively. Through a comprehensive analysis of the experimental results, we have the following observations:

• •

  The proposed variational inference framework for ITE estimation, TNDVGA, consistently outperforms state-of-the-art traditional baseline methods, including BART, Causal Forest, CFR, and CEVAE, across different settings on both datasets, as these methods do not account for disentangled latent factors or leverage network information for ITE learning.

• •

  TNDVGA and NetDeconf, along with GIAL, outperform other baseline methods in ITE estimation due to their ability to leverage auxiliary network information to capture the impact of latent factors on ITE estimation. This result suggests that network information helps in learning representations of latent factors, leading to more accurate ITE estimation. Furthermore, TNDVGA also outperforms NetDeconf and GIAL in ITE estimation because it learns representations of four different latent factors, whereas NetDeconf and GIAL only learn representations of latent confounding factors.

• •

  TEDVAE also performs reasonably well in estimating ITE, mainly because its model infers and disentangles three disjoint sets of instrumental, confounding, and risk factors from the observed variables. This also highlights the importance of learning disentangled latent factors for ITE estimation. However, TNDVGA outperforms TEDVAE, as it additionally accounts for latent noise factors and effectively leverages network information, whereas TEDVAE struggles to fully utilize network information to enhance its modeling capabilities.

• •

  TNDVGA demonstrates strong robustness in selecting the latent dimensionality parameter. Considering that we did not explicitly model the generation process of latent factors in these two semi-synthetic real datasets, and that the dataset generation did not include instrumental, risk, or noise factors, TNDVGA still exhibits optimal performance under these conditions. These results indicate that, even in more realistic datasets , TNDVGA can effectively learn latent factors and estimate ITE.

• •

  When the influence of hidden confounders increases (i.e., with a growing $\kappa_{2}$ value), TNDVGA suffers the least in $\sqrt{\epsilon_{PEHE}}$ and $\epsilon_{ATE}$. This is because TNDVGA has the ability to identify patterns of latent confounding factors from the network structure, enabling it to infer ITE more accurately.

First, similar to (Bao et al., 2022), we control the magnitude of selection bias in the dataset by setting the size of the scalar $\zeta$. We compare TNDVGA with TEDVAE and NetConf when the dimensions of the latent factors are (8, 8, 8, 8). As shown in Fig. 3, we observe that as the value of $\zeta$ increases, indicating a rise in selection bias, TNDVGA consistently performs the best. Furthermore, TNDVGA’s performance remains stable and is largely unaffected by variations in selection bias. This demonstrates that TNDVGA exhibits better robustness against selection bias, which is crucial when handling real-world datasets. We have similar results on other synthetic datasets.

Next, we investigate the TNDVGA’s ability to recover the latent components $\mathbf{z}_{t}$, $\mathbf{z}_{c}$, $\mathbf{z}_{y}$, and $\mathbf{z}_{o}$ that are utilized to construct the observed covariates $\mathbf{x}$ and examine the contribution of disentangling these latent factors to the estimation of ITE. To this end, similar to the settings in (Hassanpour and Greiner, 2019; Bao et al., 2022), we compare the performance of TNDVGA when the parameters $d_{\mathbf{z}_{t}}$, $d_{\mathbf{z}_{c}}$, $d_{\mathbf{z}_{y}}$ and $d_{\mathbf{z}_{o}}$ are set to correspond with the true number of latent factors against the performance when one of the latent dimensionality parameters is set to zero. For example, setting $d_{\mathbf{z}_{c}}=0$ forces TNDVGA to ignore the disentanglement of confounding factors. If TNDVGA performs better when considering the disentanglement of all latent factors compared to when any one latent factor is ignored, then it can be concluded that TNDVGA can recover latent factors, and that disentangling latent factors is beneficial for ITE estimation. Fig. 4 displays the radar charts corresponding to each factor. We can clearly see that when TNDVGA considers disentangling the latent factors through non-zero dimensionality parameters, its performance outperforms that when any latent dimension is set to zero.