Causal Invariance-aware Augmentation for Brain Graph Contrastive Learning

Minqi Yu, Jinduo Liu, Junzhong Ji

TL;DR: We propose a Causally Invariance-aware Augmentation for brain Graph Contrastive Learning, called CIA-GCL.

Abstract: Deep models are increasingly used to analyze brain graphs for the diagnosis and understanding of brain diseases. However, due to the multi-site data aggregation and individual differences, brain graph datasets exhibit widespread distribution shifts, which impair the model’s generalization ability to the test set, thereby limiting the performance of existing methods. To address these issues, we propose a Causally Invariance-aware Augmentation for brain Graph Contrastive Learning, called CIA-GCL. This method first generates a brain graph by extracting node features based on the topological structure. Then, a learnable brain invariant subgraph is identified based on a causal decoupling approach to capture the maximum label-related invariant information with invariant learning. Around this invariant subgraph, we design a novel invariance-aware augmentation strategy to generate meaningful augmented samples for graph contrast learning. Finally, the extracted invariant subgraph is utilized for brain disease classification, effectively mitigating distribution shifts while also identifying critical local graph structures, enhancing the model’s interpretability. Experiments on three real-world brain disease datasets demonstrate that our method achieves state-of-the-art performance, effectively generalizes to multi-site brain datasets, and provides certain interpretability.

Lay Summary:

Brain graph analysis is a key tool for computer-aided diagnosis and understanding brain diseases, but identifying the commonalities of patients across multi-site data remains difficult.
To tackle this challenge, we aim to extract local structures that remain highly similar across complex environments to better distinguish patients. Building on this, we propose CIA-GCL，a Causally Invariance-aware Augmentation framework for brain graph contrastive learning.

Specifically, we recognize the importance of key brain areas and connections. The brain is a highly intricate system, like a finely tuned instrument, where even the disruption of certain crucial areas or pathways can throw off the entire system.
We design an invariant subgraph extractor to capture critical brain regions and use a brain-specific augmentation strategy to simulate diverse environments.
We extract key distinguishing biomarkers by mining subgraphs inward and externally enforcing their consistency across different environments.

These biomarkers not only help in differentiating patients but also provide valuable insights into the abnormal regions associated with brain diseases.

Link To Code: https://github. com/qinsheng1900/CIA-GCL

Primary Area: Applications->Health / Medicine

Keywords: Graph contrastive learning, Invariant learning, fMRI analysis, Neurodevelopmental disorders

Submission Number: 3026