Learnable Community-Aware Transformer for Brain Connectome Analysis with Token Clustering

Yanting Yang,Beidi Zhao,Zhuohao Ni,Yize Zhao,Xiaoxiao Li

2024-03-13

Abstract:Neuroscientific research has revealed that the complex brain network can be organized into distinct functional communities, each characterized by a cohesive group of regions of interest (ROIs) with strong interconnections. These communities play a crucial role in comprehending the functional organization of the brain and its implications for neurological conditions, including Autism Spectrum Disorder (ASD) and biological differences, such as in gender. Traditional models have been constrained by the necessity of predefined community clusters, limiting their flexibility and adaptability in deciphering the brain's functional organization. Furthermore, these models were restricted by a fixed number of communities, hindering their ability to accurately represent the brain's dynamic nature. In this study, we present a token clustering brain transformer-based model ($\texttt{TC-BrainTF}$) for joint community clustering and classification. Our approach proposes a novel token clustering (TC) module based on the transformer architecture, which utilizes learnable prompt tokens with orthogonal loss where each ROI embedding is projected onto the prompt embedding space, effectively clustering ROIs into communities and reducing the dimensions of the node representation via merging with communities. Our results demonstrate that our learnable community-aware model $\texttt{TC-BrainTF}$ offers improved accuracy in identifying ASD and classifying genders through rigorous testing on ABIDE and HCP datasets. Additionally, the qualitative analysis on $\texttt{TC-BrainTF}$ has demonstrated the effectiveness of the designed TC module and its relevance to neuroscience interpretations.

Neurons and Cognition,Machine Learning,Image and Video Processing

What problem does this paper attempt to address?

The paper attempts to address the problem of how to use dynamic clustering methods to identify functional communities in brain connectome analysis and improve the accuracy of Autism Spectrum Disorder (ASD) and gender classification. Specifically, traditional models are limited by predefined functional community clusters, lacking flexibility and adaptability, and unable to accurately capture the dynamic nature of brain functional organization. This paper proposes a Transformer architecture based on learnable prompt tokens—TC-BrainTF. This model, through a novel Token Clustering (TC) module, can dynamically learn ROI embeddings specific to functional communities without the need for predefined community structures. This approach not only improves the model's performance on ASD and gender classification tasks but also enhances the interpretability of neuroscience research findings. Experimental results show that TC-BrainTF outperforms existing methods on both the ABIDE and HCP datasets, particularly excelling in terms of AUROC, accuracy, and specificity.