English

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

Machine Learning 2025-09-25 v4 Artificial Intelligence

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.

Keywords

Cite

@article{arxiv.2509.09744,
  title  = {Structure Matters: Brain Graph Augmentation via Learnable Edge Masking for Data-efficient Psychiatric Diagnosis},
  author = {Mujie Liu and Chenze Wang and Liping Chen and Nguyen Linh Dan Le and Niharika Tewari and Ting Dang and Jiangang Ma and Feng Xia},
  journal= {arXiv preprint arXiv:2509.09744},
  year   = {2025}
}
R2 v1 2026-07-01T05:32:35.197Z