Developing interpretable attention-based generative models for neuroimaging data

The prevailing scientific consensus posits that cognitive functions of the brain are underpinned by networks of structured, synchronous neural activities with distinct spatiotemporal dynamics. The development of generative models, such as the Time-Delay Embedded Hidden Markov Model [1] and Dynamic Network Modes [2], has been pivotal in identifying these dynamic functional brain networks.

However, these models still face challenges in representing long-range temporal dependencies and inadequate global context perception that arise from their sequential nature [3, 4]. Recently, the introduction of a transformer [5] has offered a potential solution to these issues, leveraging its ability to represent sequences using non-directional, contextually adaptive attention patterns.

While the efficacy of attention-based methods in neural data decoding has been recognized in brain-computer interfaces [6-9] and computer vision [10, 11] studies, their "black box" nature makes interpreting their inner workings difficult. Hence, the objectives of this project are twofold: 1. Develop an attention-based generative model for decoding neuroimaging data.

2. Interpret its learned representations to elucidate the brain's spatiotemporal dynamics.

The initial focus of our project is on developing a transformer model with multi-head self-attention networks capable of predicting and generating MEG data.1 To interpret the model and its internal representations of brain networks, we will employ ablation scoring [12] and permute feature importance methods [13, 14]. These methodologies will be instrumental in identifying the contributions of distinct sub-architectures and input variables to the model predictions.

This phase will include an in-depth inspection of the model's dynamic machinery, utilizing secondary analyses such as clustering algorithms applied to time-varying attention matrices. Additionally, we aim to expand the applicability of our model to EEG and fMRI data, potentially utilizing simultaneous multimodal datasets to evaluate the consistency of network estimations across different neuroimaging modalities.

Upon achieving a satisfactory interpretation of the model, we intend to fine-tune it for targeted applications in neuropsychiatric disorders. We will begin with the New Therapeutics in Alzheimer's Disease (NTAD) dataset [15], where we plan to investigate the dynamics of brain

networks in Alzheimer's disease and explain these with spatial and temporal features of the attention matrices. We will then evaluate the potential of these features in early-stage disease prediction, examining both their predictive power and diagnostic accuracy. If time permits, we aspire to extend our analysis to additional clinical datasets and cognitive tasks.

Our proposed model is designed to quantitatively analyze how the brain processes information across different scales of time, space, and dimensionality, thereby revealing previously unobserved latent information from various data modalities and bringing novel insights to the field of neuroimaging. The successful application of the HMM in clinical studies has already enhanced our understanding of various neuropsychiatric conditions, including multiple sclerosis, schizophrenia, Alzheimer's disease, and Parkinson's disease [16-19].

Our developed method is expected to deepen our understanding of how various modalities record unique neural information. This, in turn, will aid in establishing a biological mechanism-based nosology for neuropsychiatric disorders, fulfilling the growing demand for more objective and personalized medical approaches.