Structure Matters: Brain Graph Augmentation via Learnable Edge Masking for Data-efficient Psychiatric Diagnosis

2025-09-11

AI Summary

Limited annotated brain network data hinder the accuracy and interpretability of psychiatric disorder diagnosis. Existing self-supervised methods suffer from structural distortion during graph augmentation, leading to the loss of clinically relevant connectivity patterns. To address this, we propose SAM-BG—a novel framework that introduces a learnable edge masking mechanism for the first time, enabling automatic identification and preservation of critical structural priors during pretraining, thereby achieving semantics-preserving graph augmentation. Integrated with graph neural networks, SAM-BG performs structure-aware representation learning in two stages under few-shot settings. Evaluated on two real-world neuroimaging datasets, SAM-BG significantly outperforms state-of-the-art baselines, improving diagnostic accuracy by up to 8.2% in ultra-low-shot regimes (≤10 samples per class). Furthermore, interpretability analysis uncovers cross-regional functional connectivity patterns that align with established clinical evidence, enhancing biological plausibility and diagnostic transparency.

Technology Category

• Machine Learning: Graph-based Machine Learning
• Data Mining & Knowledge Management: Graph Mining, Social Network Analysis & Community
• Knowledge Representation and Reasoning: Diagnosis and Abductive Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding

Abstract

The limited availability of labeled brain network data makes it challenging to achieve accurate and interpretable psychiatric diagnoses. While self-supervised learning (SSL) offers a promising solution, existing methods often rely on augmentation strategies that can disrupt crucial structural semantics in brain graphs. To address this, we propose SAM-BG, a two-stage framework for learning brain graph representations with structural semantic preservation. In the pre-training stage, an edge masker is trained on a small labeled subset to capture key structural semantics. In the SSL stage, the extracted structural priors guide a structure-aware augmentation process, enabling the model to learn more semantically meaningful and robust representations. Experiments on two real-world psychiatric datasets demonstrate that SAM-BG outperforms state-of-the-art methods, particularly in small-labeled data settings, and uncovers clinically relevant connectivity patterns that enhance interpretability. Our code is available at https://github.com/mjliu99/SAM-BG.

Problem

• Addresses limited labeled brain network data for psychiatric diagnosis
• Proposes structure-aware augmentation to preserve brain graph semantics
• Enhances interpretability by identifying clinically relevant connectivity patterns

Innovation

• Edge masker learns structural semantics
• Structure-aware augmentation preserves brain graphs
• Two-stage framework enhances interpretability