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

arXiv: https://arxiv.org/abs/2509.09744 · HTML: https://mlsanity.com/paper/2509.09744/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 11 of its 34 references that are in the map.
- First posted on arXiv 11 Sep 2025.

## Builds on (most cited first, 11 of 11)

- [Graph Contrastive Learning with Augmentations](https://mlsanity.com/paper/2010.13902/) (2020)
- [The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables](https://mlsanity.com/paper/1611.00712/) (2016)
- [Graph Contrastive Learning with Adaptive Augmentation](https://mlsanity.com/paper/2010.14945/) (2020)
- [Graph Self-Supervised Learning: A Survey](https://mlsanity.com/paper/2103.00111/) (2021)
- [Self-Supervised Learning of Graph Neural Networks: A Unified Review](https://mlsanity.com/paper/2102.10757/) (2021)
- [From Canonical Correlation Analysis to Self-supervised Graph Neural Networks](https://mlsanity.com/paper/2106.12484/) (2021)
- [BrainGB: A Benchmark for Brain Network Analysis with Graph Neural Networks](https://mlsanity.com/paper/2204.07054/) (2022)
- [Interpretable Graph Neural Networks for Connectome-Based Brain Disorder Analysis](https://mlsanity.com/paper/2207.00813/) (2022)
- [Recognizing Predictive Substructures with Subgraph Information Bottleneck](https://mlsanity.com/paper/2103.11155/) (2021)
- [Learning on Multimodal Graphs: A Survey](https://mlsanity.com/paper/2402.05322/) (2024)
- [MA-GCL: Model Augmentation Tricks for Graph Contrastive Learning](https://mlsanity.com/paper/2212.07035/) (2022)

## Abstract (arXiv)

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.
