Cause-Aware Heterogeneous Graph Learning

• The paper introduces a cause-aware heterogeneous graph learning architecture that integrates causal inference to enhance comorbidity risk prediction.
• It leverages causal nonnegative latent factorization and relation-specific message passing to fuse and assess causal signals across diverse node and edge types.
• By incorporating counterfactual interventions and a composite loss function, the framework improves interpretability and robustness in predictive modeling.

A cause-aware heterogeneous graph learning architecture integrates principles of causal inference with heterogeneous graph neural network design to address the challenge of distinguishing causation from spurious correlation in relational data. In the context of comorbidity risk prediction for chronic obstructive pulmonary disease (COPD), such an architecture explicitly encodes different node and edge types—e.g., patients and diseases—and fuses causal representations via tailored message passing, counterfactual interventions, and specialized objective functions. The method aims to enhance the detection of genuine causal relationships among entities, thereby improving the interpretability and robustness of predictive models (Zhou et al., 22 Dec 2025).

1. Heterogeneous Graph Definition and Construction

The approach begins with the construction of a heterogeneous comorbidity graph, defined as $G = (V, E, T_v, T_E, R)$, where:

• $V = P \cup D$ is the set of nodes, partitioned into patients ($P$) and diseases ($D$).
• $T_v$ enumerates node types ("Patient", "Disease").
• $E$ is the edge set, partitioned by relation set $R$.
• $T_E$ defines edge types, specifically:
  • $r_1$: Patient–Patient relations (e.g., similarity or contact)
  • $r_2$: Disease–Disease relations (e.g., comorbidity, shared etiology)
  • $r_3$: Patient–Disease relations (direct comorbid links)

Each relation $r \in R$ induces a sparse adjacency matrix $A^{(r)} \in \mathbb{R}^{|V| \times |V|}$. The collection $\{A^{(r)}\}_{r \in R}$ can be indexed as a relation-type incidence tensor. Initial node features are provided as follows: for patients, 68-dimensional laboratory test vectors, often incomplete, are embedded using causal non-negative latent factorization (see Section 2); for diseases, one-hot or learned embeddings are used.

2. Model Architecture: Preprocessing and Causal Graph Neural Inference

The cause-aware framework consists of two major components:

A. Causal Nonnegative Latent Factorization (CSINLF) for Preprocessing:

• Patient lab measurements (matrix $M$) are factorized into $U \in \mathbb{R}^{n_p \times k}$ and $V \in \mathbb{R}^{n_d \times k}$, optimizing:

$$\min_{U, V} \| M - UV^\top \|_F^2 + \gamma \cdot \Omega_{\text{causal}}(U, V)$$

where $\Omega_{\text{causal}}$ regularizes towards known causal constraints (e.g., disease $\rightarrow$ lab value). Latent vector $U_i$ for patient $i$ is used as initial graph input.

B. Heterogeneous Graph Neural Network with Causal Attention and Counterfactual Reasoning:

For each layer $\ell=0,\ldots,L-1$:

1. Relation-Specific Message Passing: Relation-wise message aggregation,

$$h_v^{(\ell+1, \text{base})} = \sigma \left( \sum_{r \in R} \sum_{u \in N_r(v)} W_r^{(\ell)} h_u^{(\ell)} + b_r^{(\ell)} \right)$$

2. Pairwise Causal Strength Estimation: Computation of $CS_{u \rightarrow v}$,

$$CS_{u \rightarrow v} = \sigma \left( W_2 \,\text{ReLU}(W_1 [h_u^{(\ell)} \Vert h_v^{(\ell)}] + b_1 ) + b_2 \right)$$

3. Causal Attention: Computation of edge-wise attention score $a_{u \rightarrow v}$,

$$a_{u \rightarrow v} = \sigma \left( W_4 \,\text{ReLU}( W_3 [ h_u^{(\ell)} \Vert h_v^{(\ell)} \Vert CS_{u \rightarrow v}] + b_3 ) + b_4 \right)$$

4. Do-Intervention (Counterfactual Feature Transformation):

$$h_u^{do} = W_6\,\text{ReLU}( W_5\,\text{ReLU}( W_6 h_u^{(\ell)} + b_5) + b_6 )$$

1. Causal Message Passing: Message $$m_{u \rightarrow v} = a_{u \rightarrow v} \cdot h_u^{do}$$
2. Causal Aggregation: Sum of received causal messages,

$$h_v^{(\ell+1, \text{causal})} = \sum_{u \in N(v)} m_{u \rightarrow v}$$

1. Feature Fusion: Additive combination of base and causal outputs,

$$h_v^{(\ell+1)} = h_v^{(\ell+1, \text{base})} + h_v^{(\ell+1, \text{causal})}$$

After $L$ layers, a patient’s embedding $h_i^{(L)}$ is the input for a multilayer perceptron (MLP) or linear-softmax classifier to yield predicted risk scores $\hat{y}_i$.

3. Integrated Causal Inference Mechanisms

The architecture adopts a structural causal model (SCM) perspective:

• Variables: Node representations $h$, edge-level causal strengths $CS$, and attentions $a$ are considered endogenous variables governed by the model's structural equations.
• Treatment Variable: The do-intervention $h_u^{do}$ (Eq.(5)) operationalizes the "treatment" by actively modifying source node features, thereby blocking confounding from the graph context.
• Outcome: The predicted risk $\hat{y}_i$ for each patient node post-aggregation.
• Confounders: Baseline lab features and graph-neighborhood messages that, unadjusted, may introduce non-causal associations.

For each patient, two outputs are constructed: the factual prediction $\hat{y}_i = f(h_i^{(L)})$ and a counterfactual $\hat{y}_i^{do}=f(h_i^{do,(L)})$ simulating intervened features. The difference between these is used to quantify and adjust for estimated causal effects. While no formal do-calculus is applied, the inclusion of explicit intervention operations and their integration into the optimization criteria aligns the method with SCM-based learning.

4. Objective Function and Optimization

The CHGRL method minimizes a composite loss:

$$L = L_{\text{cls}} + L_{\text{cf}} + \alpha L_{\text{cr}}$$

where:

• Classification Loss $L_{\text{cls}}$: Cross-entropy over patient labels, enforcing predictive accuracy.

$$L_{\text{cls}} = -\sum_{i \in P} \big( y_i \log \hat{y}_i + (1 - y_i) \log(1 - \hat{y}_i) \big)$$

• Counterfactual Reasoning Loss $L_{\text{cf}}$: Penalizes disparities between factual and intervention-based predictions:

$$L_{\text{cf}} = \mathbb{E}_{i \in P} [( \hat{y}_i - \hat{y}_i^{do} )^2 ]$$

• Causal Regularization Loss $L_{\text{cr}}$: Enforces consistency of latent representations with causal constraints, arising from $\Omega_{\text{causal}}$ on factorized embeddings and, in general, penalizes spurious edges or encourages similarity between $P(Y \mid do(T))$ and $P(Y \mid T)$.

This multi-term loss enforces that the learned representations retain predictive power, maintain counterfactual (causal) consistency, and adhere to known mechanistic knowledge.

5. Workflow Summary: Layerwise and Loss Function Overview

Step	Key Formula (Eq.)	Description
Relation-specific MP	(2): $h_v^{(\ell+1,\text{base})}$	Aggregation per edge type
Causal strength	(3): $CS_{u \rightarrow v}$	Neural pairwise causal estimator
Causal attention	(4): $a_{u \rightarrow v}$	Learnable edge weighting using $CS$
Do-intervention	(5): $h_u^{do}$	Counterfactual transformation on $u$
Causal message	(6): $m_{u \rightarrow v}$	Attention-weighted causal message
Aggregation	(7): $h_v^{(\ell+1, \text{causal})}$	Sum of incoming messages at $v$
Feature fusion	(8): $h_v^{(\ell+1)}$	Combine base and causal features
Full loss	(9): $L$	Overall objective (classification, counterfactual, regularizer)

In this architecture, each component targets a fundamental aspect of cause-aware learning: initial embeddings reflect feasible causal structures (via CSINLF); edge propagation selectively weights paths based on causal salience; counterfactual operations ensure that the system focuses on genuine interventions, and the composite loss penalizes non-causal or confounded explanations.

6. Relevance, Significance, and Relation to Prior Work

The cause-aware heterogeneous graph learning architecture, as instantiated in CHGRL (Zhou et al., 22 Dec 2025), addresses the limitations of standard GNNs in medical prediction tasks where causal distinction is critical. By integrating edge-type specific propagation, explicit modeling of pairwise causal influences, and counterfactual interventions, it provides a unified framework for causal discovery and risk prediction. The inclusion of the counterfactual loss and causal regularization differentiates this method from correlation-focused approaches.

A plausible implication is the extensibility of the approach to other domains where graph-structured confounding is prevalent. While CHGRL does not enumerate meta-paths, its multi-relation aggregation and explicit causal weighting generalize the meta-path attention paradigm by making each edge a candidate for causal interpretation. The structure ensures that the model is robust to spurious correlations often encountered in heterogeneous biomedical and social graphs. This suggests utility in early disease detection settings, but also motivates future exploration of formal SCM analysis and integration with broader families of causal GNNs.