Spatio-temporal directed acyclic graph learning with attention mechanisms on brain functional time series and connectivity

Abstract

We develop a deep learning framework, spatio-temporal directed acyclic graph with attention mechanisms (ST-DAG-Att), to predict cognition and disease using functional magnetic resonance imaging (fMRI). This ST-DAG-Att framework comprises of two neural networks, (1) spatio-temporal graph convolutional network (ST-graph-conv) to learn the spatial and temporal information of functional time series at multiple temporal and spatial graph scales, where the graph is represented by the brain functional network, the spatial convolution is over the space of this graph, and the temporal convolution is over the time dimension; (2) functional connectivity convolutional network (FC-conv) to learn functional connectivity features, where the functional connectivity is derived from embedded multi-scale fMRI time series and the convolutional operation is applied along both edge and node dimensions of the brain functional network. This framework also consists of an attention component, i.e., functional connectivity-based spatial attention (FC-SAtt), that generates a spatial attention map through learning the local dependency among high-level features of functional connectivity and emphasizing meaningful brain regions. Moreover, both the ST-graph-conv and FC-conv networks are designed as feed-forward models structured as directed acyclic graphs (DAGs). Our experiments employ two large-scale datasets, Adolescent Brain Cognitive Development (ABCD, n=7693) and Open Access Series of Imaging Study-3 (OASIS-3, n=1786). Our results show that the ST-DAG-Att model is generalizable from cognition prediction to age prediction. It is robust to independent samples obtained from different sites of the ABCD study. It outperforms the existing machine learning techniques, including support vector regression (SVR), elastic net's mixture with random forest, spatio-temporal graph convolution, and BrainNetCNN.