Bayesian nonparametric mixtures of categorical directed graphs for heterogeneous causal inference

Estimating cause-and-effect relations between variables is a pervasive issue in many applied domains and primarily medical science. Typically in this setting, interest lies in measuring the (direct or indirect) effect of a therapy on the progression of a disease. Our methodology is motivated by a dataset of patients diagnosed with breast cancer and treated with different oncological therapies. In this context, the protein Human Epidermal growth factor Receptor 2 (HER2) has been identified as one of the main responsibles of tumor progression and growth. Recent studies have shown that therapies targeting HER2 have a strong antitumor effect, improving the overall and progression-free survival. These therapies are commonly based on both monoclonal antibodies, such as trastuzumab, as well as anticancer drugs, in particular antracyclines; see Slamon et al. (1987) and Katzorke et al. (2013). However, they can cause cardiotoxicity as a side effect, with consequent heart failure and loss of left ventricular contractile function (Dempke et al., 2023; Bowles et al., 2012). Establishing the (causal) effect of anti-HER2 therapies on cardiotoxicity is therefore of key importance for the administration of appropriate anticancer treatments, and to develop strategies for preventing and detecting cardiotoxiciy in high-risk patients. In addition, there exist several factors, such as advanced age, hypertension, valvulopathy and arrhythmia, that predispose to cardiotoxicity; see in particular Dempsey et al. (2021), Lotrionte et al. (2013) and references therein. Accordingly, in a related causal-effect analysis one should account for all those clinical features that may act as risk factors in the occurrence of cardiotoxicity.

When several variables are entertained, as in the framework above, one should account for possible interactions/dependencies between them in order to provide a coherent quantification of causal effects. Typically however, such dependence structure is unknown, or can be partially drawn only based on experts’ knowledge and one need to learn it from the data. Graphical models based on Directed Acyclic Graphs (DAGs) offer a powerful tool for this structure learning task. Importantly to our purposes, DAGs allow to properly define the causal effect on a target variable of interest induced by a hypothetical intervention on another variable in the system. Additionally, learning such causal effect can be achieved from observational data alone, under suitable causal assumptions on the data generating process (Pearl, 2000). An important issue in this general framework is however represented by heterogeneity, which implies that causal effects may vary across individuals, as the consequence of an existing, yet unknown, clustering structure in the population. Causal-inference methodologies accounting for heterogeneity can provide a more reliable quantification of treatment effects across patients, leading to personalized strategies for the administration of therapies; see in particular Ma et al. (2015) for an overview.

Available methods for clustering multivariate (categorical) data include distance- and model-based approaches. Among the first, k-modes (Huang, 1998) is the most popular methodology, which is based on a modified version of the k-means algorithm (MacQueen, 1967) and implements a dissimilarity measure between modes of categorical variables computed across clusters. On the other side, poLCA (polytomous variable Latent Class Analysis) (Linzer and Lewis, 2011) is a model-based method which applies to a collection of categorical random variables. In the model, each mixture component corresponds to a multivariate categorical distribution built under the assumption of independence between marginal distributions and for a known number of components in the mixture. Importantly however, none of these methods accounts for possible dependence relationships between variables in the underlying multivariate statistical model, which instead represents a peculiar feature of our methodology. Model-based clustering methods have been extensively developed in the Bayesian literature from both a finite and infinite mixture-model perspective; see in particular Frühwirth-Schnatter et al. (2021), Argiento and De Iorio (2022) and references therein for a review and connections between the two approaches. In a multivariate Gaussian framework, infinite-mixture models based on a Dirichlet Process (DP) prior are considered by Rodríguez et al. (2011) and Castelletti and Consonni (2023) for clustering and structure learning of undirected and directed graphs respectively. Recently, Argiento et al. (2022) proposed a finite-mixture model specifically designed for multivariate unordered categorical data. This is based on a newly-introduced class of Hamming distributions which is assumed for each categorical variable marginally, and from which a joint distribution over $q$ variables is built under the assumption of (local) independence. Finally, Malsiner-Walli et al. (2024) proposed a two-layer mixture model which also allows for associations among categorical variables within each mixture-component. Such dependencies arise from a second-layer mixture, assumed within each component of the main mixture model, rather than a multivariate model with allied dependence parameter, which is instead a distinctive feature of our method for causal discovery and inference.

The literature on heterogeneous causal inference has grown extensively in the last years, particularly under the potential outcome framework (Rubin, 2005). Assuming the existence of latent sub-groups of individuals in the population, this issue is addressed through the definition and estimation of group-specific causal effects, known as Conditional Average Treatment Effects (CATEs). Machine learning methods are employed to identify the underlying clustering structure by stratification of subjects based on the levels of available covariates; see Dominici et al. (2021) for a review, Athey and Imbens (2016), Hahn et al. (2020) and Bargagli Stoffi et al. (2022), the latter proposing a method which is specifically designed to handle imperfect compliance. Recent methods also aim at improving interpretability of causal results. An instance in this direction is the Causal Rule Ensamble (CRE) method (Lee et al., 2021), which adopts multiple trees to identify patterns of heterogeneity in the data and to ensure stability in sub-group identification. Other approaches to heterogeneous causal inference using counterfactuals are based on Bayesian nonparametric methods; see for instance Linero and Antonelli (2022) and references therein. Among these, Zorzetto et al. (2024) employ a Dependent Probit Stick-Breaking mixture model to simultaneously impute the missing outcomes (counterfactuals), and to identify mutually exclusive groups, thus allowing to estimate causal effects in the presence of population heterogeneity. Still in a potential outcome framework, Roy et al. (2016) adopt marginal structural models and implement a dependent Dirichlet Process (DP) prior for the evaluation of heterogeneous causal effects of treatments on survival outcomes; Oganisian et al. (2021) instead consider a DP mixture of zero inflated regression models for pathological data exhibiting excesses of zeros. DP priors for clustering and heterogeneous causal inference are also employed by Castelletti and Consonni (2023) in a multivariate framework based on Gaussian graphical models where causal effects are defined and estimated according to do-calculus theory (Pearl, 2000).

We propose a Bayesian methodology based on a infinite mixture of categorical DAGs for causal discovery and causal effect estimation in the presence of heterogeneous data. Our model allows for the presence of latent sub-groups of individual/patients in the sample, each characterized by a possibly different causal structure and battery of related causal-effect parameters. Specifically, we assume that the multivariate distribution of the observables belongs to a Dirichlet Process (DP) mixture of categorical DAG models. Each mixture component reflects a factorization of the sampling distribution satisfying a set of conditional independencies imposed by the DAG. Under the latter, causal effects between variables are then defined according to do-calculus. With regard to the DP prior, we define a baseline measure over the space of priors on $(\mathcal{D},\bm{\theta})$ where $\mathcal{D}$ is a DAG and $\bm{\theta}$ the parameter of a categorical DAG model. We employ a constructive procedure based on local and global parameter independence to assign priors to DAG parameters, providing closed-form expressions for both the prior and posterior predictive of DAGs, as well as for the posterior distribution of DAG-parameters. We then leverage these results to develop a computational scheme for posterior inference of our DP model. When applied to breast cancer data, our methodology ultimately allows to quantify treatment effects of assigned therapies w.r.t. the occurrence of cardiotoxicity at subject-specific levels, thus leading to a more reliable decision process for the development of personalized therapies.

The rest of the paper is organized as follows. In Section 2 we provide some background material on DAGs and causal effects within a categorical modelling framework. In Section 3 we introduce our mixture model based on a DP prior, for which we detail the construction of the baseline mixing measure over the space of DAGs and allied parameters. We then describe a Markov Chain Monte Carlo (MCMC) strategy for posterior inference in Section 4. In Section 5 we evaluate our methodology relative to the tasks of clustering and causal discovery through extensive simulation studies, which include comparisons with alternative state-of-the-art methods. Section 6 is devoted to the analysis of breast cancer data and includes our causal-effect analysis to evaluate heterogeneous side effects of anti-HER2 therapies with respect to the occurrence of cardiotoxicity. We finally provide a discussion to our methodology in Section 7, together with possible future developments. Some technical results, including the computation of prior and posterior predictive distributions required by our posterior sampler, are reported in the Supplementary Material.

Consider a Directed Acyclic Graph (DAG) $\mathcal{D}=(V,E)$ with set of nodes $V=\{1,\dots,q\}$ and set of directed edges $E\subseteq V\times V$. For a given $\mathcal{D}$, if $(u,v)\in E$, we say that $u$ is a parent of $v$ and let $\mathrm{pa}_{\mathcal{D}}(v)$ be the set of all parents of $v$ in $\mathcal{D}$. Moreover, we let $\mathrm{fa}_{\mathcal{D}}(v):=v\cup\mathrm{pa}_{\mathcal{D}}(v)$ be the family of node $v$ in the DAG. Consider now a collection of random variables $X=(X_{1},\dots,X_{q})$, such as clinical features that can be measured on patients, and binary categorical variables indicating the administration of a therapy and the absence/presence of a disease. In the following, we assume that each $X_{j}$, $j\in V$, is categorical with set of levels $\mathcal{X}_{j}$ and let $x_{j}\in\mathcal{X}_{j}$ be one of its levels. Accordingly, $X\in{\mathcal{X}}:=\times_{j\in V}\mathcal{X}_{j}$, whose generic element is $x\in{\mathcal{X}}$. In addition, if for any $S\subset V$ we let $X_{S}=(X_{j},j\in S)$, then $X_{S}\in{\mathcal{X}}:=\times_{j\in S}\mathcal{X}_{j}$, with $x_{S}\in{\mathcal{X}}_{S}$. Under $\mathcal{D}$, the joint probability $p(x)=\text{Pr}(X_{1}=x_{1},\dots,X_{q}=x_{q})$ admits the factorization

$$p(x)=\prod_{j=1}^{q}\Pr(X_{j}=x_{j}\,|\,X_{\mathrm{pa}(j)}=x_{\mathrm{pa}(j)}).$$

(1)

For the remainder of this section we omit DAG $\mathcal{D}$ from our notation and reason conditionally on a fixed DAG. Let now $\theta_{s}^{S}=\Pr(X_{S}=s)$, $s\in{\mathcal{X}}_{S}$, be a marginal probability for variables in $S\subseteq V$. Moreover, let $\theta^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s}=\Pr(X_{j}=m\,|\,X_{\mathrm{pa}(j)}=s)$ be a conditional probability for $X_{j}$ given configuration (level) $s$ of $X_{\mathrm{pa}(j)}$, with $m\in{\mathcal{X}}_{j},s\in{\mathcal{X}}_{\mathrm{pa}(j)}$. Consider $n$ observations from $X$, $\bm{x}^{(1)},\dots,\bm{x}^{(n)}$, where each $\bm{x}^{(i)}=(x^{(i)}_{1},\dots,x^{(i)}_{q})^{\top}$, and $\bm{x}^{(i)}\in{\mathcal{X}}$, $i=1,\dots,n$. Also, let $\bm{x}_{S}^{(i)}$ be the sub-vector of $\bm{x}^{(i)}$ with components indexed by $S\subset V$. If we collect the $\bm{x}^{(i)}$’s into an $(n,q)$ data matrix $\bm{X}$, then the likelihood function can be written as

	$\displaystyle p(\bm{X}\,|\,\bm{\theta})$	$\displaystyle=\prod_{i=1}^{n}\left\{\prod_{x\in{\mathcal{X}}}\left\{p(\bm{x}^{(i)}\,|\,\bm{\theta})\right\}^{\mathbbm{1}\{\bm{x}^{(i)}=x\}}\right\}$		(2)
		$\displaystyle=\prod_{j=1}^{q}\left\{\prod_{s\in{\mathcal{X}}_{\mathrm{pa}(j)}}\left\{\prod_{m\in\mathcal{X}_{j}}\left\{\theta^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s}\right\}^{n^{\mathrm{fa}(j)}_{(m,s)}}\right\}\right\},$

now emphasizing the dependence on the DAG-parameter $\bm{\theta}$ (corresponding to the collection of conditional probabilities in the equation) and where $n^{\mathrm{fa}(j)}_{(m,s)}=\sum_{i=1}^{n}\mathbbm{1}\left\{\bm{x}^{(i)}_{\mathrm{fa}(j)}=(m,s)\right\}$ is the number of observations for which $X_{\mathrm{fa}(j)}=(m,s)$. See also Castelletti et al. (2024) for further notation on categorical DAG models.

For a given collection of random variables whose multivariate distribution factorizes according to a DAG, we now focus on the causal effect of an intervention on $X_{h}$, $h\in V$, on a response variable of interest, say $X_{j}:=Y$, $j\neq h$. In practice, such an intervention corresponds to assigning a treatment to an individual, equivalently fixing $X_{h}=\tilde{x}$, where $X_{h}$ is typically an exposure of $Y$, and this action can be denoted using Pearl’s do-operator $\textnormal{do}(X_{h}=\tilde{x})$ (Pearl, 2003). This implies a change in the observational distribution (1), leading to the so-called post-intervention distribution

$$p(x\,|\,\textnormal{do}(X_{h}=\tilde{x}))=\begin{cases}\prod\limits_{j\neq h}p\big{(}X_{j}=x_{j}\,|\,X_{\mathrm{pa}(j)}=x_{\mathrm{pa}(j)}\big{)}&\textrm{if}\ X_{h}=\tilde{x}\\ \,\,0&\textrm{otherwise}.\end{cases}$$

(3)

Assuming for simplicity that both $X_{h}$ and $Y$ are binary variables with levels in $\{0,1\}$, the causal effect of $\textnormal{do}\{X_{h}=\tilde{x}\}$ on $Y$ can be defined as

$$c_{y,h}\vcentcolon=\mathbb{E}\big{[}Y\,|\,\textnormal{do}(X_{h}=1)\big{]}-\mathbb{E}\big{[}Y\,|\,\textnormal{do}(X_{h}=0)\big{]};$$

(4)

see Pearl (2003). More in general, if $X_{h}$ is polytomous with levels labeled as $\{0,1,\dots,L\}$, one can define a battery of causal effects by considering $X_{h}=l$, for each $l=1,\dots,L$ in the first expectation of Equation (4). According to the definition above, $c_{y,h}$ involves a (marginal) post-intervention distribution of $Y$. However, because of (3), the latter can be expressed in terms of observational distributions, simply by conditioning and then marginalizing w.r.t. a valid adjustment set $Z\subset X$; see Pearl (2003). A common choice for such an adjustment set is $Z=X_{\mathrm{pa}(h)}$, namely the parents of $X_{h}$, leading to

	$\displaystyle c_{y,h}=\sum_{s\in{\mathcal{X}}_{\mathrm{pa}(h)}}\mathbb{E}\big{(}Y\,|$	$\displaystyle X_{h}=1,X_{\mathrm{pa}(h)}=s\big{)}\Pr\big{(}X_{\mathrm{pa}(h)}=s\big{)}$		(5)
	$\displaystyle-$	$\displaystyle\sum_{s\in{\mathcal{X}}_{\mathrm{pa}(h)}}\mathbb{E}\big{(}Y\,|\,X_{h}=0,X_{\mathrm{pa}(h)}=s\big{)}\Pr\big{(}X_{\mathrm{pa}(h)}=s\big{)};$

see in particular Theorem 3.2.3 in Pearl (2000). Under model (2), the causal effect in (5) can be expressed as a function of the DAG parameter $\bm{\theta}$ as

$\displaystyle\gamma_{y,h}(\bm{\theta})=\sum_{s\in\mathcal{X}_{\mathrm{pa}(h)}}\left\{\left(\theta^{Y\,|\,\mathrm{fa}(h)}_{1\,|\,(1,s)}-\theta^{Y\,|\,\mathrm{fa}(h)}_{1|(0,s)}\right)\theta^{\mathrm{pa}(h)}_{s}\right\}.$

(6)

Conditionally on DAG $\mathcal{D}$, we first assign a prior $p(\bm{\theta}\,|\,\mathcal{D})$. To this end, consider for each node $j\in\{1,\dots,q\}$ and $s\in{\mathcal{X}}_{\mathrm{pa}(j)}$ the parameter $\big{(}\theta^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s},m\in\mathcal{X}_{j}\big{)}:=\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}$ corresponding to a $|\mathcal{X}_{j}|$-dimensional vector collecting conditional probabilities for variable $X_{j}$, given a configuration $s$ of its parents $X_{\mathrm{pa}(j)}$. We assign to each $\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}$ a Dirichlet prior with hyper-parameter $\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}=\big{(}a^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s},m\in\mathcal{X}_{j}\big{)}$, written as $\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}\sim\textnormal{Dirichlet}(\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s})$, whose p.d.f. is

$\displaystyle p\big{(}\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}\big{)}=h\big{(}\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}\big{)}\prod_{m\in\mathcal{X}_{j}}\left\{\theta^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s}\right\}^{a^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s}-1},$

(10)

and where $h\big{(}\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}\big{)}$ is the prior normalizing constant. Let now $\bm{\theta}^{j\,|\,\mathrm{pa}(j)}=\big{(}\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s},s\in{\mathcal{X}}_{\mathrm{pa}(j)}\big{)}$. By assuming global and local parameter independence (Geiger and Heckerman, 1997), respectively $\bot\bot_{j}\bm{\theta}^{j\,|\,\mathrm{pa}(j)}$ and $\bot\bot_{s}\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}$, a joint prior on $\bm{\theta}=\big{\{}\bm{\theta}^{j\,|\,\mathrm{pa}(j)},j\in V\big{\}}$ can be written as

$\displaystyle p(\bm{\theta})=\prod_{j=1}^{q}\left\{\prod_{s\in{\mathcal{X}}_{\mathrm{pa}(j)}}p\big{(}\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}\big{)}\right\}.$

(11)

In what follows we implement the default choice $a^{j\,|\,\mathrm{pa}(j)}_{m\,|\,s}=a/|{\mathcal{X}}_{\mathrm{fa}(j)}|$, $a>0$, leading to the Bayesian Dirichlet Equivalent uniform (BDEu) score (Heckerman et al., 1995), which guarantees that Markov equivalent DAGs are assigned the same marginal likelihood; see also Castelletti et al. (2024).

Also notice that the resulting prior is conjugate with the likelihood (2) since, for generic dataset $\bm{X}$, $\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}\,|\,\bm{X}\sim\textnormal{Dirichlet}\big{(}\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}+\bm{n}_{s}^{\mathrm{fa}(j)}\big{)}$ with $\bm{n}_{s}^{\mathrm{fa}(j)}=\big{(}n_{(m,s)}^{\mathrm{fa}(j)},m\in\mathcal{X}_{j}\big{)}$. Accordingly, the posterior of $\bm{\theta}$ is

$\displaystyle p(\bm{\theta}\,|\,\bm{X})=\prod_{j=1}^{q}\left\{\prod_{s\in{\mathcal{X}}_{\mathrm{pa}(j)}}p\big{(}\bm{\theta}^{j\,|\,\mathrm{pa}(j)}_{s}\,|\,\bm{X}\big{)}\right\},$

(12)

with each term corresponding to a Dirichlet p.d.f., so that direct sampling from $p(\bm{\theta}\,|\,\bm{X})$ is possible. Finally, under the same prior, a marginal (i.e. integrated w.r.t. to $\bm{\theta}$) likelihood $m(\bm{X}\,|\,\mathcal{D})=\int p(\bm{X}\,|\,\bm{\theta},\mathcal{D})p(\bm{\theta}\,|\,\mathcal{D})\,d\bm{\theta}$ is available and admits the factorization

$$m(\bm{X}\,|\,\mathcal{D})=\prod_{j=1}^{q}m(\bm{X}_{j}\,|\,\bm{X}_{\mathrm{pa}(j)}),$$

(13)

with

$$m\big{(}\bm{X}_{j}\,|\,\bm{X}_{\mathrm{pa}(j)}\big{)}=\prod_{s\in{\mathcal{X}}_{\mathrm{pa}(j)}}\frac{h\big{(}\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}\big{)}}{h\big{(}\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}+\bm{n}^{\mathrm{fa}(j)}_{s}\big{)}}$$

(14)

and where $h\big{(}\bm{a}^{j\,|\,\mathrm{pa}(j)}_{s}+\bm{n}^{\mathrm{fa}(j)}_{s}\big{)}$ is the posterior normalizing constant. See also the Supplementary Material (Section 1) for full details.

Let $\mathcal{S}_{q}$ be the (discrete) space of all DAGs with $q$ nodes. Additionally, we can restrict $\mathcal{S}_{q}$ to a subset of DAGs satisfying some structural constraints, typically edge orientations that can be postulated in advance based on the specific real-data problem. As an instance, in our application to breast cancer data we regard age as an exogenous variable and forbid any incoming edge to it from other variables; conversely, we regard the occurrence of cardiotoxic side effect as a response variable and accordingly forbid any outgoing edge from it. Each DAG $\mathcal{D}=(V,E)$ in $\mathcal{S}_{q}$ can be represented through a 0-1 adjacency matrix $\bm{A}^{\mathcal{D}}$, whose $(u,v)$-element $\bm{A}^{\mathcal{D}}_{u,v}=1$ if $(u,v)\in E$, $0$ otherwise. Additionally, let $\bm{S}^{\mathcal{D}}$, be the adjacency matrix of the skeleton of $\mathcal{D}$, namely the undirected graph obtained from $\mathcal{D}$ by disregarding edges orientations. We assign for each $u>v$, $\bm{S}^{\mathcal{D}}_{u,v}\,|\,\pi\overset{\textnormal{{iid}}}{\sim}\textnormal{Ber}(\pi)$ where $\pi$ is a prior probability of edge inclusion. We then assume hierarchically $\pi\sim\textnormal{Beta}(a,b)$, leading to the integrated prior on DAG $\mathcal{D}$

$$p(\mathcal{D})\propto p\big{(}\bm{S}^{\mathcal{D}}\big{)}=\frac{\Gamma(a+b)}{\Gamma(a)\Gamma(b)}\cdot\frac{\Gamma\big{(}|\bm{S}^{\mathcal{D}}|+a\big{)}\Gamma\big{(}q(q-1)/2-|\bm{S}^{\mathcal{D}}|+b\big{)}}{\Gamma\big{(}q(q-1)/2+a+b\big{)}},$$

where $|\bm{S}^{\mathcal{D}}|$ the number of non-null elements in $\bm{S}^{\mathcal{D}}$, corresponding to the number of edges in $\mathcal{D}$, and $q(q-1)/2$ is the maximum number of edges in a DAG having $q$ nodes. Sampling from the baseline over DAGs, as required by our MCMC sampler (Section 4), is possible through an acceptance-rejection algorithm over the space $\mathcal{S}_{q}$; see also the Supplementary Material (Section 2).

The structure of our baseline measure (Section 3.1) is such that we can integrate out the DAG parameter $\bm{\theta}$, which allows for the implementation of a collapsed sampler approximating the marginal posterior of $\{\xi_{i}\}_{i=1}^{n},\alpha,\{\mathcal{D}_{k}\}_{k=1}^{K},K$. The resulting scheme has a Gibbs-sampling structure as Algorithm 2 in Neal (2000) and implements the following steps.