Task-guided Generative Adversarial Networks for Synthesizing and Augmenting Structural Connectivity Matrices for Connectivity-Based Prediction

Abstract

Recent machine learning techniques have improved connectome-based predictions by modelling complex dependencies between brain connectivity and cognitive traits. However, they typically require large datasets that are costly and time-consuming to collect. To address this, we propose Task-guided GAN II, a novel data augmentation method that uses generative adversarial networks (GANs) to expand sample sizes in connectome-based prediction tasks. Our method incorporates a task-guided branch within the Wasserstein GAN framework, specifically designed to synthesize structural connectivity matrices and improve prediction accuracy by capturing task-relevant features. We evaluated Task-guided GAN II on the prediction of fluid intelligence using the NIMH Health Research Volunteer Dataset. Results showed that data augmentation improved prediction accuracy. To further assess whether augmentation can substitute for increasing real-world sample sizes, we conducted additional validation using the HCP WU-Minn S1200 dataset. Task-guided GAN II improved prediction performance with limited real data, with gains of up to twofold augmentation observed. However, excessive augmentation did not result in further improvements, suggesting that augmentation complements, but does not fully replace, real data augmentation. These results suggest that Task-guided GAN II is a promising tool for harnessing small datasets in human connectomics research, improving predictive modelling where large-scale data collection is impractical.