SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Peter Yongho Kim*
Junbeom Kwon*
Sunghwan Joo
SungKyunKwan University
Sangyoon Bae
Donggyu Lee
Yoonho Jung
Shinjae Yoo
Brookhaven National Lab
Jiook Cha
Taesup Moon
Modeling spatiotemporal brain dynamics from high-dimensional data, such as functional Magnetic Resonance Imaging (fMRI), is a formidable task in neuroscience. Existing approaches for fMRI analysis utilize hand-crafted features, but the process of feature extraction risks losing essential information in fMRI scans. To address this challenge, we present SwiFT (Swin 4D fMRI Transformer), a Swin Transformer architecture that can learn brain dynamics directly from fMRI volumes in a memory and computation-efficient manner. SwiFT achieves this by implementing a 4D window multi-head self-attention mechanism and absolute positional embeddings. We evaluate SwiFT using multiple large-scale resting-state fMRI datasets, including the Human Connectome Project (HCP), Adolescent Brain Cognitive Development (ABCD), and UK Biobank (UKB) datasets, to predict sex, age, and cognitive intelligence. Our experimental outcomes reveal that SwiFT consistently outperforms recent state-of-the-art models. Furthermore, by leveraging its end-to-end learning capability, we show that contrastive loss-based self-supervised pre-training of SwiFT can enhance performance on downstream tasks. Additionally, we employ an explainable AI method to identify the brain regions associated with sex classification. To our knowledge, SwiFT is the first Swin Transformer architecture to process dimensional spatiotemporal brain functional data in an end-to-end fashion. Our work holds substantial potential in facilitating scalable learning of functional brain imaging in neuroscience research by reducing the hurdles associated with applying Transformer models to high-dimensional fMRI.
Project page: https://github.com/Transconnectome/SwIFT
The human brain is a dynamic system characterized as an extensive (feedback) network generating complex spatiotemporal dynamics of its activity. Analyzing such dynamics and linking them to normative adaptive cognition and behaviors, as well as their maladaptive forms in a disease condition is of paramount importance in neuroscience and medicine. However, the fields have been hampered by a lack of predictive power owing to the gap between the complexity of the brain network and the contrasting simplicity of brain imaging analytics [1]. Hence, developing an accurate predictive model using high-dimensional brain imaging data and learning rich representations of brain function will help close this gap and advance toward precision neuroscience.
† Co-corresponding author
Among many brain imaging modalities, functional Magnetic Resonance Imaging (fMRI) provides a unique opportunity to model the spatiotemporal patterns of brain electrochemical activity. Functional MRI captures a temporal sequence of 3D images (a stack of 2D slices) of Blood Oxygenation Level Dependent (BOLD) signals with time resolution ranging from 0.5 to 3 seconds, rendering in 4D data. It has accelerated the discovery of detailed functional anatomy of the human brain and significantly contributed to the understanding of brain diseases, such as Alzheimer's disease and psychiatric disorders. To better process the fMRI data, statistical methods were conventionally used to estimate the spatiotemporal patterns related to cognition [2, 3] and brain diseases [4, 5, 6, 7]. More recently, deep neural networks (DNN) have been applied to fMRI to investigate the nonlinear relationship of brain dynamics with human cognition and behaviors [8, 9, 10, 11, 12].
Regarding the DNN-based approaches, researchers have broadly pursued two distinct lines of work. The first approach is the so-called ROI-based method, in which the high-dimensional fMRI data (with around 300,000 voxels) is clustered into the temporal sequence of hundreds of pre-defined brain regions (ROIs) using anatomical segmentation [13] or statistical clustering [14] before learning a predictive DNN model. Despite the computational efficiency, these pre-processing approaches may be prone to losing information important to capture subtle variability across the individual brains [15]. The second DNN-based approach is the two-step method, in which the raw fMRI data is used as input, and specialized architectures for spatial and temporal domains are separately used. Namely, for learning spatial features, convolutional neural networks (CNNs) are used, and for temporal, long short-term memory (LSTM) [16] or Transformers [17, 18] are used. This separation of DNN architecture necessarily results in two-step learning for better computational and memory efficiency. However, it could also limit the capability of capturing comprehensive information among brain regions across the temporal dimension. Therefore, a critical unresolved issue is whether an efficient, end-to-end DNN that utilizes raw 4D fMRI input can be formulated to better model and learn the brain dynamics compared to previous approaches.
To that end, we propose Swin 4D fMRI Transformer (SwiFT), a 4D extension of the Swin Transformer [19] architecture, which can jointly learn the spatiotemporal representations of the brain's intrinsic activity directly from high-dimensional fMRI in an end-to-end fashion. The main gist of our method is to employ the 4D local window attention structure, which makes SwiFT readily applicable to process large-scale, high-dimensional 4D data with low computational complexity. We note that while 3D variants of Swin Transformers have been proposed before [20, 21, 22, 23] to take video or medical image inputs, to the best of our knowledge, this is the first work to extend Swin Transformer to take 4D data input and to apply it to the fMRI data.
Our experimental results show that the end-to-end learning capability of SwiFT unlocks its potential to effectively learn complex spatiotemporal patterns in fMRI. Specifically, we evaluate SwiFT's performance on three representative fMRI benchmarks: the Human Connectome Project (HCP) [24], the Adolescent Brain Cognitive Development (ABCD) [25], and the UK Biobank (UKB) [26, 27]. Across various classification and regression tasks, including sex classification and age/intelligence prediction, SwiFT significantly outperforms the recent baselines mentioned above. Furthermore, we also demonstrate that applying the widely successful "pre-train and fine-tune" framework to SwiFT would be feasible. We believe this capability has the potential to empower researchers to construct large-scale foundation models for fMRI, akin to those utilized in several other application domains. Finally, to provide a comprehensive analysis, we present the interpretation results using IG-SQ for SwiFT's predictions and conduct ablation studies to substantiate our modeling choices.
ROI-based methods To analyze the brain network, researchers typically reduce the sequences of brain volumes into low-dimensional multivariate time series by aggregating voxel intensities in specific regions of interest (ROI) based on standardized brain atlases, considering pairwise correlations between the time series of each ROI as functional connectivity [28, 29]. The majority of the DNNs, such as graph neural networks (GNN), were designed to treat the brain network as a graph, considering each ROI as node and pairwise correlations between the ROIs as edges. For instance, BrainNetCNN, consisting of multiple graph convolutional filters, was proposed to model various levels of topological interactions in structural [30] and functional connectivity [12]. Kan et al. [12] proposed a Transformer-based model that learns the individual connectivity strengths between each ROI. Namely, they apply an orthonormal clustering readout operation for functional
connectivity to locate functional clusters related to specific human behaviors, and those clusters serve as informative embeddings for predicting psychiatric and biological outcomes. Some recent studies proposed methods to capture spatiotemporal dynamics directly from extracted fMRI time series, utilizing Transformer by separating spatial and temporal attention units [31] and focusing on local representations with fused window multi-head self-attention (FW-MSA) [32].
Two-step methods for 4D fMRI inputs Existing DNNs that take raw 4D fMRI as input typically process spatial and temporal information separately. Li et al. [16] integrated 3D convolutional neural networks (CNNs) to extract spatial embeddings in each 3D volume of a 4D fMRI and then feed the spatial embeddings to an LSTM for temporal encoding. Transformer Framework for fMRI (TFF) by Malkiel et al. [17] replaces the LSTM with a Transformer to extract temporal features and propose reconstruction-based pre-training steps. To learn spatial features before the downstream task, TFF concatenates decoder layers after the Transformer and minimizes three reconstruction-based losses; L1, perceptual, and intensity losses. Nguyen et al. [18] suggested a pre-training method to predict the types of cognitive tasks performed during an fMRI scanning. A 3D CNN encoder is then trained to predict the target variable from fMRI volumes. Furthermore, the pre-trained encoder learns temporal features by attaching multi-head self-attention layers. The aforementioned models may have issues of unstable training of spatiotemporal data, which involves multiple training steps and large memory usage. These limitations may result in sub-optimal model computation and learning capability.
Transformers for vision tasks Following the success of Transformers in natural language processing [33, 34], many works applied the Transformer architecture in vision tasks. One of the major challenges for such an extension is the computation complexity scaling quadratically with respect to the number of tokens. Namely, as images have a much larger number of tokens (pixels), rendering the direct pixel-to-token use of Transformers infeasible in most cases. Dosovitskiy et al. [35] tackled this issue with ViT, converting image patches into tokens rather than treating each individual pixel as a token, but it still did not resolve the quadratic scaling of the self-attention layers. Liu et al. [19] suggested the Swin Transformer, a model that reduces the computational complexity to be linear in the number of tokens by running self-attention only within a local window instead of running it globally. Along with the local windowed attention, Swin Transformer also introduces shifting windows to add cross-window connections and patch merging (downsampling) steps to produce hierarchical representations. Swin UNETR [36] extended the Swin Transformer to handle 3D images and demonstrated its utility by coupling a Swin Transformer encoder with CNN-based decoders. VAT [37], a 4D Convolutional Swin Transformer, extended the Swin Transformer model to accept a 4D correlation map of two CNN-extracted 2D image features, utilizing the model for cost aggregation. While these works show that the Swin Transformer model can be extended to handle more spatial dimensions, Video Swin Transformer [23] shows that the Swin Transformer model can also be extended to the temporal dimension by processing 3D video inputs. Overall, these works present the feasibility of applying Swin Transformers to higher spatiotemporal dimensions, but to the best of our knowledge, such a method has yet to be applied to 4D functional brain imaging.
3 Main Method
3.1 Swin 4D fMRI Transformer (SwiFT)
Overall architecture In line with the recent studies [23, 20], which introduce 3D extensions to enhance the capabilities of the Swin Transformer [19], we propose an advancement in the architecture to add an additional temporal dimension, thereby enabling its application to 4D data.
The overall architecture of our model is depicted in Figure 1a. The SwiFT architecture utilized in our study consists of four distinct stages. Each stage is constructed through the implementation of patch merging, with linear embedding employed in the case of Stage 1. Additionally, positional embedding is incorporated, and multiple (Swin) Transformer blocks are applied repeatedly within the stages. The model processes an input fMRI data of size  $T \times H \times W \times D \times 1$ , which consists of a length  $T$  sequence of fMRI volumes  $(H \times W \times D)$  with a single channel. During the initial patch partitioning step, the input fMRI data is partitioned into  $T \times \frac{H}{P} \times \frac{W}{P} \times \frac{D}{P}$  patches with  $P^3$  voxels. In this study,  $H$ ,  $W$ , and  $D$  are 96, and the initial patch size  $P$  is 6. During the linear embedding process, patches with size  $P^3$  are transformed into  $C$ -dimensional tokens. This transformation effectively maps the spatially neighboring pixels within a patch onto a token.
(a) Overall architecture of SwiFT
(b) Successive 4D Swin Transformer Blocks
(c) 4DW-MSA and 4DSW-MSA
Figure 1: Figures depicting the structure of SwiFT and its components.
Next, following an absolute positional embedding layer, multiple layers of 4D Swin Transformer blocks are applied on the embedded patches, forming Stage 1 of SwiFT. Starting from Stage 2, a patch merging layer at the beginning of each stage reduces the number of tokens by merging 8 spatially neighboring patches. After the patch merging layer, an absolute positional embedding layer followed by multiple layers of 4D Swin Transformer blocks is applied, forming Stage 2 and onward. In the final stage, Stage 4, the 4D Swin Transformer blocks are replaced by global attention Transformer blocks which carry out global attention instead of local window attention. This computationally expensive global attention is made possible by the significant reduction in the number of tokens achieved through the patch merging steps executed in the preceding stages. Global attention Transformer blocks allow each token to globally attend to all other tokens rather than being restricted to the tokens within the local window.
Patch merging Following prior works [23, 20, 38, 39], the patch merging step is only performed for the three spatial dimensions  $(H,W,D)$  and not for the temporal dimension  $(T)$ , thereby merging a group of  $8 = 2 \times 2 \times 2$  neighboring patches into a single patch for each time frame. During the patch merging operation, the spatial dimensions are reduced by half, and the channel size  $(C)$  is doubled as compensation. Thus, in Figure 1a, the numbers  $P_{2}, P_{3}, C_{2}$ , and  $C_{3}$  are 12, 24, 2C, and 4C, respectively.
As a general example, consider a tensor with arbitrary dimensions  $T \times H' \times W' \times D' \times C'$  before passing through the patch merging layer. During patch merging, this tensor is reshaped into a new tensor with dimensions  $T \times \frac{H'}{2} \times \frac{W'}{2} \times \frac{D'}{2} \times 8C'$ , where  $2 \times 2 \times 2$  spatially neighboring patches are concatenated along the channel dimension. Then, each  $8C'$  channel in the resulting tensor is projected onto a  $2C'$  dimensional space by applying a single fully connected layer, resulting  $T \times \frac{H'}{2} \times \frac{W'}{2} \times \frac{D'}{2} \times 2C'$  in total. The process of patch merging facilitates the hierarchical feature-extraction structure of SwiFT and reduces the computational complexity of the subsequent layers. We clarify that while the patch merging is operated only on the spatial dimensions, the temporal information is still well-incorporated via the windowed attention.
4D window multi-head self-attention The core of the Swin Transformer model is the window multi-head self-attention (W-MSA) layer, which allows the model to process a large number of tokens by limiting self-attention only within a local window. In SwiFT, the 3D window mechanism is extended to 4D windows; given  $T \times H' \times W' \times D'$  input tokens, the tokens are partitioned by a window of size  $P \times M \times M \times M$ , resulting in  $\lceil \frac{T}{P} \rceil \times \lceil \frac{H'}{M} \rceil \times \lceil \frac{W'}{M} \rceil \times \lceil \frac{D'}{M} \rceil$  non-overlapping windows.
However, simply stacking multiple window self-attention layers would be undesirable since there would be no crosstalk across different windows. To that end, a shifted window multi-head self-attention (SW-MSA) layer enables cross-window connections. Namely, we extend the 3D shifted window mechanism to 4D shifted windows as well; in  $P \times M \times M \times M$  windows obtained from the W-MSA layer, we shift the windows of the successive layer by  $\left( \frac{P}{2}, \frac{M}{2}, \frac{M}{2}, \frac{M}{2} \right)$  tokens.
The detailed operations of our 4D W-MSA and SW-MSA are shown in Figure 1c. In this example, the applied size of input tokens and the windows are  $T \times H' \times W' \times D' = 4 \times 8 \times 8 \times 8$  and  $P \times M \times M \times M = 2 \times 4 \times 4 \times 4$ , respectively. Then, by following the window partitioning methods described above, the numbers of grouped windows in W-MSA and SW-MSA become  $2 \times 2 \times 2 \times 2 = 16$  and  $3 \times 3 \times 3 \times 3 = 81$ , respectively. Such separately applied window self-attention is essential in effectively extracting spatiotemporal feature representation from the 4D fMRI data. Although the number of windows increases in SW-MSA, the actual computation cost remains similar by leveraging the cyclic-shifting batch computation proposed in [19].
Combining the W-MSA layer and the SW-MSA layer, two successive 4D Swin Transformer blocks, as shown in Figure 1b, are computed as the following:
$$
\begin{array}{l} \hat {z} ^ {l} = 4 \mathrm {D W - M S A} (\mathrm {L N} (z ^ {l - 1})) + z ^ {l - 1}, \quad z ^ {l} = \operatorname {M L P} (\mathrm {L N} (\hat {z} ^ {l})) + \hat {z} ^ {l} \\ \hat {z} ^ {l + 1} = 4 \operatorname {D S W - M S A} (\operatorname {L N} (z ^ {l})) + z ^ {l}, \quad z ^ {l + 1} = \operatorname {M L P} (\operatorname {L N} (\hat {z} ^ {l + 1})) + \hat {z} ^ {l + 1}, \\ \end{array}
$$
in which 4D(S)W-MSA, LN, and MLP denote the 4D (Shifted) Window Multi-head Self-Attention, Layer Normalization, and Multi-Layer Perceptron module, respectively. Moreover,  $\hat{z}^l$  and  $z^l$  denote the output features of the 4D(S)W-MSA module and the following MLP module for block  $l$ , respectively.
4D absolute positional embedding Even though previous models utilize relative position biases to encode positional information [19, 23], we have opted instead for an absolute position embedding scheme for SwiFT. While absolute positional embeddings are more computationally expensive for low-dimensional data [19], since we are dealing with much larger scale 4D data, absolute positional embeddings become more cost-effective than relative positional bias. To that end, we add a learnable embedding at the beginning of each stage of the Transformer right after the patch merging step. In line with [22], we separately add positional embeddings for the spatial and temporal dimensions. Specifically, given an input tensor with dimensions of  $T \times H' \times W' \times D' \times C'$ , we define spatial and temporal positional embedding tensors with dimensions of  $1 \times H' \times W' \times D' \times C'$  and  $T \times 1 \times 1 \times 1 \times C'$ , respectively. These tensors are then added to the input tensor using broadcasting. The effect of this switch is investigated in Appendix C.1, and overall absolute positional embedding seems to be the superior method of choice.
Our proposed end-to-end model structure allows efficient self-supervised pre-training of SwiFT, which can then be fine-tuned for specific tasks. This unique capability sets our method apart from other approaches in Section 2 relying on ROI-based brain data or a two-step learning approach for 4D fMRI. We achieve this by using two different contrastive loss-based pre-training objectives adapted from [40]. Figure 2 depicts the positive and negative pairs for the two loss functions, where the InfoNCE [41] loss of the pairs is calculated for the final loss function.
Instance contrastive loss To allow the model to distinguish fMRI scans that come from different subjects, the instance contrastive loss  $(\mathcal{L}_{IC})$  is a type of contrastive-based loss function that considers a representation to be positive if two distinct fMRI sub-sequences come from the same subject and negative if they come from different subjects. The feature representation passes through three layers: SwiFT, global average pooling, and a multi-layer perceptron (MLP) head. To clearly define this loss function, we denote the feature representation as  $f_{i,p}$ , where  $i \in \{1, \dots, B\}$  refers to the subject index for a given batch size  $B$ , and  $p$  refers to the fMRI sub-sequence index (either 1 or 2). Since we are
sampling two sub-sequences for each of the  $B$  subjects, this amounts to a total of  $2B$  representations. For each subject  $i$ , the anchor, positive, and negative representations are set as  $f_{i,1}$ ,  $f_{i,2}$ , and the remaining  $2B - 2$  feature representations, respectively. Using this setup, the instance contrastive loss for subject  $i$  is denoted and defined as
$$
\mathcal {L} _ {I C} ^ {i} = - \log \frac {h \left(f _ {i , 1} , f _ {i , 2}\right)}{\sum_ {j = 1} ^ {B} \left[ \mathbb {1} _ {[ j \neq i ]} \left(h \left(f _ {i , 1} , f _ {j , 1}\right) + h \left(f _ {i , 1} , f _ {j , 2}\right)\right) \right]},
$$
in which  $h$  denotes the exponential of the cosine similarity between two vectors, and  $\mathbb{1}$  denotes an indicator function that equals one if the condition is true and zero otherwise.
Local-local temporal contrastive loss To allow the model to distinguish fMRI scans that come from different timestamps within the same subject, the local-local temporal contrastive loss  $(\mathcal{L}_{LL})$  determines both positive and negative pairs within a single subject. Namely, a positive representation is derived from the same fMRI sub-sequence but is applied with different random augmentations. On the other hand, negative representations come from distinct fMRI sub-sequences from the same subjects. The feature representations are obtained in the same manner as the instance contrastive loss. To clearly define this loss function, we use the same notation of feature representation as  $f_{i,p}$ , but the range of  $p$  is changed to  $\{1,2,\dots,N\}$ , where  $N$  is the number of fMRI sub-sequences from a single subject. We denote an fMRI sub-sequence of the same subject  $i$  and fMRI sub-sequence  $p$  but different random augmentation as  $\tilde{f}_{i,p}$ . Since we are sampling two differently augmented versions for each of the  $N$  sub-sequences, this amounts to a total of  $2N$  representations. Using this setup, the local-local temporal contrastive loss for subject  $i$  is denoted as  $\mathcal{L}_{LL}^{i}$  and defined as
$$
\mathcal {L} _ {L L} ^ {i} = - \sum_ {p = 1} ^ {N} \log \frac {h (f _ {i , p} , \tilde {f} _ {i , p})}{\sum_ {q = 1} ^ {N} [ \mathbb {1} _ {[ q \neq p ]} (h (f _ {i , p} , f _ {i , q}) + h (f _ {i , p} , \tilde {f} _ {i , q})) ]},
$$
in which  $h$  denotes the exponential of the cosine similarity between two vectors.
Figure 2: Illustration of two different contrastive losses for the pre-training of SwiFT.
Datasets The Adolescent Brain Cognitive Development (ABCD) study is a longitudinal, multi-site investigation of brain development and related behavioral outcomes in children [25]. The dataset is open to the scientific community but requires authorization. After quality control, we used the resting state fMRI of 9,128 children (age = 118.95 ± 7.46 months, 52.4% female) from release 2. For fMRI preprocessing, we used a well-established pipeline, fMRIprep [42, 43], which includes reducing the bias field, skull-stripping, alignment to structural image, and spatial normalization to standard space for a pediatric brain [44]. After fMRIprep, we applied low pass filtering, head movement correction, and artifact removal, regressing out signals from non-grey matters (aCompcor) [45].