Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

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

TL;DR

The paper tackles the rigidity of predefined brain community structures in functional connectome analysis by introducing TC-BrainTF, a token-clustering transformer that learns dynamic, community-aware ROI embeddings. It employs learnable prompt tokens $P \in \mathbb{R}^{K \times E}$ and a Deep Embedded Clustering (DEC) mechanism with an orthogonal loss $\mathcal{L}_{\text{ortho}}$ to produce a soft assignment $A \in \mathbb{R}^{N \times K}$ and merged token representations, followed by a graph readout for task classification. Key contributions include a novel Token Clustering (TC) module integrated into a transformer encoder, enabling flexible determination of the number of communities $K$ and improved ASD and gender classification on ABIDE and HCP datasets, along with qualitative analyses linking clusters to neuroscience interpretations. The approach removes reliance on atlas-defined communities, enhancing adaptability and interpretability in brain connectome analysis and offering a data-driven perspective on functional organization."

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.

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

TL;DR

The paper tackles the rigidity of predefined brain community structures in functional connectome analysis by introducing TC-BrainTF, a token-clustering transformer that learns dynamic, community-aware ROI embeddings. It employs learnable prompt tokens and a Deep Embedded Clustering (DEC) mechanism with an orthogonal loss to produce a soft assignment and merged token representations, followed by a graph readout for task classification. Key contributions include a novel Token Clustering (TC) module integrated into a transformer encoder, enabling flexible determination of the number of communities and improved ASD and gender classification on ABIDE and HCP datasets, along with qualitative analyses linking clusters to neuroscience interpretations. The approach removes reliance on atlas-defined communities, enhancing adaptability and interpretability in brain connectome analysis and offering a data-driven perspective on functional organization."

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 () 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 offers improved accuracy in identifying ASD and classifying genders through rigorous testing on ABIDE and HCP datasets. Additionally, the qualitative analysis on has demonstrated the effectiveness of the designed TC module and its relevance to neuroscience interpretations.
Paper Structure (10 sections, 4 equations, 3 figures, 2 tables)

This paper contains 10 sections, 4 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Overview
  4. Transformer Encoder Layer
  5. Token Clustering Module
  6. Graph Readout Layer
  7. Experiments
  8. Datasets and Experimental Settings
  9. Quantitative and Qualitative Results
  10. Conclusion

Figures (3)

  • Figure 1: (a) Overview of TC-BrainTF framework. (b) Schematic of the orthogonal loss and orthogonal projection of node feature embeddings onto prompt token embeddings space.
  • Figure 2: Qualitative results for communities clustering. (a) Assignment matrix $A$ depicting the relationship between ROIs. (b) Chord diagram showcasing the top 0.5% strongest connections from the token-token attention scores. (c) The ROIs in salient clusters 0 and 7 (d) Radar chart representing the correlation between functional clusters and cognitive keywords decoded by Neurosynth.
  • Figure S1: Qualitative results for communities clustering of HCP dataset. (a) Assignment Matrix $A$ depicting the relationship between ROIs. (b) Chord diagram showcasing the top 0.5% strongest connections from the token-token attention scores. (c) The ROIs in salient clusters 2 and 5. (d) Radar chart representing the correlation between functional clusters and cognitive keywords decoded by Neurosynth.