Diffusion-Guided Pretraining for Brain Graph Foundation Models

TL;DR

This paper introduces a diffusion-guided pretraining framework for brain graph models that preserves meaningful connectivity and captures global structure, leading to improved transferability across diverse neuroimaging datasets.

Contribution

It proposes a novel diffusion-based approach for structure-aware augmentation and global reconstruction in brain graph pretraining, addressing limitations of existing methods.

Findings

Consistent performance improvements across multiple neuroimaging datasets.

Effective preservation of brain graph semantics during augmentation.

Enhanced global structural understanding in learned representations.

Abstract

With the growing interest in foundation models for brain signals, graph-based pretraining has emerged as a promising paradigm for learning transferable representations from connectome data. However, existing contrastive and masked autoencoder methods typically rely on naive random dropping or masking for augmentation, which is ill-suited for brain graphs and hypergraphs as it disrupts semantically meaningful connectivity patterns. Moreover, commonly used graph-level readout and reconstruction schemes fail to capture global structural information, limiting the robustness of learned representations. In this work, we propose a unified diffusion-based pretraining framework that addresses both limitations. First, diffusion is designed to guide structure-aware dropping and masking strategies, preserving brain graph semantics while maintaining effective pretraining diversity. Second, diffusion enables topology-aware graph-level readout and node-level global reconstruction by allowing graph embeddings and masked nodes to aggregate information from globally related regions. Extensive experiments across multiple neuroimaging datasets with over 25,000 subjects and 60,000 scans involving various mental disorders and brain atlases demonstrate consistent performance improvements.