SwiFT: Swin 4D fMRI Transformer
Authors: Peter Yongho Kim
Junbeom Kwon
Sunghwan Joo
Sangyoon Bae
Donggyu Lee
Yoonho Jung
Shinjae Yoo
Jiook Cha
Taesup Moon
Venue: NeurIPS 2023 poster (2023)
Keywords: fMRI
Swin Transformer
4D
neuroscience
TL;DR: We propose SwiFT, the first 4D Swin Transformer architecture that can learn from raw fMRI signals in an end-to-end manner.

Abstract: 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.

SwiFT: Swin 4D fMRI Transformer
Peter Yongho Kim* Seoul National University
Junbeom Kwon' Seoul National University
Sunghwan Joo SungKyunKwan University
Sangyoon Bac Seoul National University
Donggyu Lee SungKyunKwan University
Yoonho Jung Scoul National University
Shinjae Yoo Brookhaven National Lab
Jiook Cha Seoul National University
Taesup Moon Seoul National University
Abstract
Modeling spatiotemporal brain dynamics from high-dimensional data, such as functional Magnetic Resonance Imaging (fMRI), is formidable task in neuro- science. 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-efficiee 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 con- sistently outperforms recent state-of-the-ard 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. Addition- ally, we employ an explainable AI method 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- nigh-dimensiooal fMRI. http://gtthbbboom//aassmccccooo Traasccomeettom//SwiiFF
manner. SwiFT achieves this by
Project page: https://github.co
Introduction
The human brain is a dynamic system characterized as extensive (feedback) network generating complex spatiotemporal dynamics of its activity. Analyzing such dynamics and linking them to normative adaptive cognition and behaviors, well as their maladaptive forms in a discase condition of paramount importance in neuroscience and medicine. However, the fields have been hampered by lack of predictive power owing 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 digh-dimessional brain imaging data and learning rich representations of brain function will help close this gap and advance toward precision neuroscience.
Equual Contribution
"Co-CCoooea co--crresponding author
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 temporal sequence of 31) images (a stack of 2D slices) of Blood Oxygenation Level Dependent (BOLD) signals with time resolution ranging from 0.5 to 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 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, 7J. 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, 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 [MRI Transformer (SwiFT), a 4D extension of the Swin Transformer [19] architecture, which can jointly learn the spatiotemporal representations of the brain's 's intrinsic activity directly from nigh-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 basclines mentioned above. Furthermore, 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 construct large-scale foundation models for fMRI, akin to those utilized in several other application domains. Finally, 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.
2 Related Work
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 treat the brain network as 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-bar based model that learns the individual connectivity strengths between each ROL. Namely, they apply an orthonormal clustering readout operation for functional
2

---

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 [MRI inputs Existing DNNs that take raw 4D fMRI as imput 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 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 reconstruuttion-based pre-training steps. To learn spatial features before the downstream task, TFF concatenates decoder layers after the Transformer and minimizes three reconstruction--enn
losses; LI. perceptual, and intensity losses. Nguyen al. [18] suggested pre-training method to predict the types of cognitive tasks performed during an IMRI 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 natural language pro- cessing [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 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, 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 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 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 correlatinn map of two CNN-extracted 21) 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 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 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 H XW x xx11, which consists of a length T sequence of fMRI volumes (H X WXX x DD) with single channel. During the initial patch partitioning step, the input fMRI data is partitioned into Tx H p x F x fa patches with pa 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 p3 are transformed into C-dimensional tokens. This transformation effectively maps the spatially neighboring pixels within a patch onto token.
3

---

SwiFT
Global Attention Black
MLP
MLP
Stage
43
Hf
40
ABC
Patch Merging Embed.
Layer Norm
Layer Norm
gir1
Stage 2.3
Swin Transformer Block
#
4DW-MSA
4DSW MSA
Patch Merging & Pos. Embed.
Layer Norm
Layer Norm
Swin Transformer Black
Stage
Tx
f W
(b) Successive 4D Swin Transformer Blocks
Embed. Pos. Embed.
Patch Partitioning
VSW-AS
W
4D IMRI image
4D local window for selt anencion
A token
(a) Overall architecture of SwiFT
(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 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 2 x neighboring patches into single patch for each time frame. During the
a group of patch merging operation, the spatial dimensions are reduced by half, and the channel size (C) is doubled as compensation. Thus, in Figure la, the numbers P2, Ps, C2 and C3 are 12, 24, 20, and 4C, respectively.
As a general example, consider tensor with arbitrary dimensions Tx H' xW' DCC before passing through the patch merging layer. During patch merging, this tensor is reshaped into a new tensor with dimensions Tx H² W' are concatenated along the channel dimension. Then, each 8C" channel in the resulting tensor projected onto a 2C' dimensional space by applying a single fully connected layer, resulting Tx 7 x x T x2C" 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 well-incorporated via the windowed attention.
2 2 x D' x x8C", where 2x2x2 spatially neighboring patches

---

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 extended to 4D windows; given Tx H' x W x D' input tokens, the tokens are partitioned by a
window of size PxM x x., resulting in f x HH' M x [W" x D M
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, 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; Px M M x M windows obtained from the W-MSA layer, we shift the windows of the successive layer by (
M M M tokens. 2
2
The detailed operations of our 4D W-MSA and SW-MSA are shown in Figure Ic. In this example, the applied size of input tokens and the windows are H'x x DD' = 4 x 8 x 8 x 8 and PxMxMXXX XM = 2 x 4 x 4, respectively. Then, by following the window partitioning methods described above, the numbers of grouped windows in W-MSA and SW-MSA become 2 x 2 x 2 x 2 = 16 and 3 x 3 81, respectively. Such separately applied window self-attention essential in effectively extracting spatiotemporal feature representation from the 4D fMRI data. Although the number of windows increases in SW-MSA, A, 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 Ib, are computed as the following:
2 4DW-MSA(LN(3/-)) + 2111, z MML((L(2E)) 2
gl+1 = 4DSW-MSA(LN( + ++1 = MLP(LN () + 2!+1,
in which W-MSA, IN, and MIP denote the 4D (Shifted) Window Multi-head Self-Attention, Layer Normalization, and Multi-Layer Perceptron module, respectively. Moreover, á and 2¹ de- note the output features of the 4D(S)WMMSA module and the following MLP module for block 1,
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 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 T'x H' W' x D' x C", we define spatial and temporal positional embedding tensors with dimensions of x H' x xD' and Tx1x1xIxXC respectively. These tensors are then added to the input tensor using broadcasting. The effect of this switch is investigated in Appendix C.I, and overall absolute positional embedding seems to be the superior method of choice.
3.2 Self-supervised Pre-training
Our proposed
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 two-step learning approach for 4D fMRL. 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 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 (Lic) is a type of contrastive-baz 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 multi-layer perceptron (MLP) head. To clearly define this loss function, we denote the feature representation as Si.p. where E {1,..., B} refers to the subject index for given batch size B, and p refers to the fMRI sub-sequence index (either I or 2). Since we
5

---

sampling two sub- sequences for each of the B subjects, this amounts to total of 2B representations. For each subject i, the anchor, positive, and negative representations are set as fú,1 f(2 and the remaining 2B 2 feature representations, respectively. Using this setup, the instance contrastive loss Cic = log
for subject is denoted and defined as
1 f1.2)
B D (M.,,,, f3.1) + M [22))
in which h denotes the exponential of the cosine similarity between two vectors, and denotes an indicator function that equals one 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 (LLL) determines both positive and negative pairs within single subject. Namely, positive representation is derived from the same IMRI sub-sequence but is applied with different random augmentations. On the other hand, negative representations come from distinct fMRI sub- sequeeeses 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 fi.p. but the range of p is changed to {1, 2, .... .N. where N is the number of fMRI sub-sequences from single subject. We denote an fMRI sub-sequence of the same subject and IMRI sub-sequence p but different random augmentation as ft.p- Since we are sampling two differently augmented versions for each of the N sequeequensss, this amounts to total of 2.N representations, Using this setup, the local-local temporal contrastive loss for subject is denoted as Li, and defined
N
AlippFfip s== (att)) (h(Fep. fig) + h(S Pia)))
LLL
log
in which denotes the exponential of the cosine similarity between two vectors.
ha
SwIFT
SWIFT
Subject
SWiET
Bubject
fa
Subjecti
SwiFT
SWIFT
Temporal
Anchor
Posttyv
Negative Input seq.
Temporal
Figure 2: Illustration of two different contrastive losses for the pre-training of SwiFT.
4 Experiments
4.1 Experimental Settings
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 fMR1 preprocessing, we used a well-established pipeline, iMRIprep [42, 43], which includes reducing the bias field, skull-stripping alignment to structural image, and spatial normalization 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].
We also used the resting-state state fMRI of 1,084 healthy young adults (age 28.80 = 3.70 years, 54.4% female) from the Human Connectome Project (HCP) (S 1200 data) [46, 47]. and 5,935 middle and old aged adults (age 54.971 7.53 years) from the UK Biobank (UKB) [48]. We used preprocessed data provided by Human Connectome Project [49] and UK Biobank [26, 27). which follows the fMRI volume pipeline, including reducing the bias field, skull-stripping, cross- registration, and spatial normalization standard space.

---

For each of the 4D fMRI volumes, we globally normalized brain images over the four dimensions except for the background regions. Then we filled the background with the minimal voxel intensity value. To easily divide the volume into patches for SwiFT, we changed the 3DD volume into shape of 96 x 96 x 96 by cropping and padding on the background. To evaluate the performances of ROI-based models, following the preprocessing steps of [12] as closely as possible, we applied the HCP MMP1 atlas [50] to each (MRI volume to obtain the time series data for each ROI. Subsequently, we processed this ROI series to generate functional connectivities, which involves computing the Pearson correlation coefficient to construct a correlation matrix. The correlation matrix is then Fisher Transformed, serving as the input for the ROI-based models in Section 4.1.2.
To evaluate our models, we constructed three random splits with a ratio of (train: validation: test) (0.7: 0.15: 0.15) and reported the average performances across the three splits.
Targets We chose the sex [5]], age [52], and cognitive intelligence (NIH Toolbox [53] for ABCD, HCP datasets, and "fluid intelligence" for UKB dataset) of each subject as the prediction target for our models. These targets are significant since the relationship between the brain and these targets represents fundamental brain-biology and brain-cognition association. Also, the capability to predict these outcomes can prove the model's capability process fMRI volumes, possibly leading to the prediction of clinical outcomes of debilitating brain disorders, such as Alzheimer's 's disease, schizophrenia, autism, and bipolar disorder [54, 55, 56, 57]. For these reasons, predicting these outcomes from brain imaging has been an important benchmark task in recent computational neuroscience [12, 8, 58].
The regression targets (age. intelligence) were z-normalized to bring stable training regardless of the range of the target variable. Since the age has a unit (e.g.. years months), when reporting the performance metrics, we transformed the z-scaled age back to its original scale of months or years.
For the binary classification task, balanced accuracy and AUC (Area Under ROC Curve) were used to evaluate model performances. For the regression tasks, Mean Squared Error (MSE) and Mcan Absolute Error (MAB) were used evaluate model performances.
4.1.1 Implementation Details
For SwiFT, we use the same architecture across all of our experiments, using the architecture corresponding to the Swin Transformer-variant from previous work [19, 23] with channel number of C = 36. The numbers of layers are fixed to {L1, L2. L3. L1} = (2, 6, 2} which corresponds to model with three stages with 2, 2, 6 consecutive 4D Swin Transformer blocks for each stage and final stage with 2 consecutive global attention Transformer blocks. In the cases of 4D(S)W-MSA set P = M 4. The final output of the model is obtained by applying global average pooling layer on the output feature map of Stage 4, followed by an MLP head. For training, the Binary Cross Entropy (BCE) loss was used for the binary classification task, and the Mean Squared Error (MSE) loss was used for regression tasks.
For the ABCD dataset, input training images randomly augmented to prevent the model from overfitting. The augmentations include affine transformation, adding Gaussian noise, and Gaussian smoothing, and they were also applied for the contrastive pre-training in Section 4.3.
Instead of inputting the entire fMRI volume of subject, we divided the volume into 20-frame sub- sequences and used them as the input, mainly due to memory constraints. However, we believe this choice may benefit the model, as discussed in Appendix C.2 in detail. For training, each individual sub-sequence was treated backpropagated for each sub-sequence. For inference, the logits from the sub-sequences of each particular subject were averaged, yielding a single output for each subject.
data point, meaning the appropriate loss function was calculated and
Computational complexity The computational complexities of a single global attention Trans- former block (denoted as MSA & MLP) and a 4D Swin Transformer block (denoted as W-MSA & MLP) for input with a dimension of T D'''D''XX can be calculated as
SMMS & MLP) = 1 + C' ?(W-MSA & MLP) = 12N + NP',
in which the number of tokens N = T''WD'' D'. .'... In practice, on Stage I of SwiFT, setting the values used for the experiments C² = 36, T = 20, H' = W = D' = 16, P = M 4, the two terms 12 and NPM are balanced with the second term only being 1.19 times the first term.
7

---

Table 1: Performance comparison to baselines classification and regression tasks
ABCD
Method
Sex
Intelligence
Sex
Age (year)
Inselligence
Sex
Age (year)
Intelligence
ACC
AUC
MSE
MAF
ACC
ALIC
MSF
MAF
MSF
MAF
ACC
AUC
MSF
MAR
MSR
MAE
XGBoost
76.1
0.977
75.5
14.3
79.3
87.6
0.866
BrainNeiCNN[301
80.1
87.9
0.969
0.767
77.1
84.9
12.6
2.97
0.984
0.805
86.8
93.8
42.7
5.36
1,00]
0.800
VanillaTF1121
77.4
85.1
0.961
0.761
77.9
85.2
12.5
2.95
0.987
0.812
87.0
95.1
41.4
5.26
0.999
0.799
12]
79.1
88.9
0.955
0.767
81.0
88.0
12.8
2.98
1.001
0.830
87.0
94.8
39.6
5.17
0.998
0.798
0.783
SwiFT Couns)
79.3
87.8
0,932
0.756
929
98.0
8.6
2,36
0,903
0,786
97.7
99.8
18.2
3,40
0.392
0.796
0.78
0.85
From scratch
From suntich
0.84
Pre-trained
0.77
Presriied
0.83
MAE
MAE
0.82
0.76
0.81
0.8
0.79
0.75
0.78
0.74
0 2
6 8 10
12 14
16
18
20
5
Training epoch
Training epoch
(a) Results HCP
(b) Results for ABCD
Figure 3: Effect of UKB pre-training (a) HCP and (b) ABCD intelligence prediction tasks.
Compared to this, with global attention the C term becomes 379 times larger than the 12NC term, taking up most of the computation budget and creating bottleneck. For successive stages, N reduced by a factor of 8, and C" is increased by a factor of 2, resulting in Stage being the most computationally expensive.
4.1.2 Baselines
ROI-based models We used ROI-based deep learning methods as baseline models, which are listed as BrainNetCNN [30], VanillaTF [12], and Brain Network Transformer (BNT) [12]. These models utilize functional connectivity data as input, which is computed using temporal correlations (Pearson correlation) of every pair of brain regions. To evaluate such methods, followed the hyper-parameter and implementations of these three models from [12]. In addition, we also employed XGBoost (eXtreme Gradient Boosting) [59] in conjunction with the features described in [13] to compare a traditional machine learning model with that of deep learning-base based models. We used the flattened upper triangular correlation matrix as the input for XGBoost.
TFF (a two-step method) As mentioned in Section 2, TFF [17] consists of 3D CNNs reduce the dimensionality of fMRT volumes, which are then passed to a transformer encoder layer for temporal processing. has been reported that the model achieves state-of-the-art performances in sex classification and age regression in HCP datasets compared to other deep neural networks [11, 60]. However, due the separate processing of spatial and temporal information, TFF may lose important spatio-temporal information and suffer from performance degradation. Furthermore, since TFF also accepts the fMRI volume its input, due memory constraints, we also used the technique of dividing the volume into 20-frame uub-sequences described in Section 4.1.1.
4.2 Classification and Regression Results
In Table we compared the performance of SwiFT, which was trained from scratch as in Section 4.1.1, against various baselines mentioned above on sex classification and age, intelligence regression tasks. On the sex classification task, SwiFT outperforms all of the baselines on the HCP and UKB datasets while showing competitive results against the best bas-based models on the ABCD dataset. On the regression tasks, we observe SwiFT outperforms all baselines significantly for the age prediction tasks, although all models still have large room for improvement the UKB intelligence prediction task. We believe our quantitative results clearly underscore the effectiveness of the end-to-end training of our SwiFT model.

---

PCO/POU
mPFC
S. Mar
PUC
IMFC
Oco
Inf Temp
Opc
(a) Children (male)
(b) Young adults (male)
(c) Elderlies (male)
Figure 4: Interpretation maps with Integrated Gradients (IG) for classification. (Sagittal plane)
(a) ABCD (b) HCP (c) UKB
4.3 Effects of Pre-training on Downstream Tasks
To demonstrate the effectiveness of contrastive pre-training described in Section 3.2, SwiFT pre- trained on a larger dataset (UKB) was fine-tuned on smaller dataset (HCP), and comparable-sized dataset (ABCD) for the intelligence prediction task, which has room for improvement compared to sex and age prediction tasks. The model was pre-trained using the combination of the instance contrastive loss function (Lic) and the local-local temporal contrastive loss function (Ct.)) such that the training objective is to minimize the sum (Lic + LI.L). The feature representations used in the loss calculation were obtained the same manner as other tasks; by applying a global average pooling layer on the output feature map of Stage 4, followed by an MLP head. Figure 3 depicts the model's performance for each training epoch during the fine-tuning process. The performances are averaged over three splits. The results from model trained from scratch (Section 4.2) are also shown for comparison. In HCP, the pre-trained model consistently performs better during the early stage of the training compared to the model trained from scratch. In ABCD, on the other hand, the pre-trained model exhibits dramatic drops in performance at the early stage of training and gradually attains a better performance the later stage of fine-tuning This initial worse performance might result from the sub-optimal training rate. Overall, believe results suggest that the contrastive pre-training of SwiFT on larger dataset has a promising effect on improving downstream performance, and we plan to investigate more extensively this topic in future research.
for fine-tuning, such as the sub-optimal initial learning
4.4 Interpretation Results
Using an Integrated Gradient with Smoothgrad sQuare (IG-SQ) implemented Captum framework [15, 61, 62], we identified the brain regions showing high explanatory power on the sex classification task. We acquired 4D IG-SQ maps from test sets and filtered out incorrectly predicted samples. To find the spatial patterns of the brain showing explanatory power across subjects, normalized the IG-SQ maps, smoothing the maps with a Gaussian filter. Then, we averaged the maps over time dimensions and across subjects.
Figure 4 shows the brain regions contributing to successful sex classification. These include, in children (ABCD), the medial prefrontal gyrus (mPFC), posterior cingulate cortex (PCC), precuneus (PCu), and parietal gyrus (the default mode network). In young adults (HCP), similar brain regions observed with broader and higher intensity in mPFC, while showing unique brain regions such as the thalamus and insular cortex. middle and old-aged adults (UKB), we acquired the highest IG values in the inferior temporal gyrus and medial orbitofrontal cortex. Of note, the results confirm those regions implicated in prior brain sex difference literature [63, 64, 65, 66].
4.5 Model Efficiency
Table 2: Efficiency of 4D fMRI Transformers
Method
# Param.
FLOPs
Throughput (samples/sec)
TFF
729M
40.72G
53.60
SwiFT
4.64M
2.62G
104.46
In Table 2, we compared the computational efficiency of SwiFT against TFF by using dummy data with random numbers. Each sample consists of 20 volumes with the shape of 96 96 x 96.

---

0.81
0.93
2.7
3.52
0.83
MAE (HCP)
0.91
2.68
8.49
MAE (HCP)
0.82
MAE (UKB)
0.89
MAE (UKB)
MAE (HCP)
2.06
3.46
MAE (UKB)
0.81
0.87
2.64
3.43
0.8
0.85
2.02
3.4
0.79
0.83
0.78
0.81
2.6
MAE (HCP)
3.37
0.77
0.79
2.58
MAE (UKB)
3.34
0.76
0.77
2.56
3.31
4
8
12
16
20
24
28
32
8
12
16
20
24
28
32
Number of time frames imput volume
Number frames of input volume
(a) Results for intelligence prediction.
(b) Results for age prediction.
Figure 5: Effect of the number of time frames of the input fMRI volume on (a) intelligence prediction
and (b) age prediction tasks.
Floating point operations (FLOPs) were used to estimate the amounts of multiply-ade computations required for processing the 4D fMRI volumes. The throughput (samples/s) denotes the number of 20-volume samples processed per second, calculated using a single NVIDIA A 100 GPU. For an accurate measurement, the throughput was measured using synchronized timing with initial GPU warmup step and was repeated 100 times. The results show that SwiFT has 158.4 times fewer parameters, requires 15.5 times fewer multi-add operations, and processes input data 1.94 times faster than TFF. This result shows SwiFT is much more computationally efficient compared to TFF while also attaining better performance as given in Table
4.6 Effect of Input Time Sequence Length
To justify our use of 4D model and investigate the consequences of dividing the input fMRI volume into sub-sequences, we evaluated the effect of the input fMRI volume's sequence length (number of time frames). To ensure a fair comparison, we kept the model architecture and typer-parameters, such the local window size, constant. In Figure 5, we compared the model performance by input sequence lengths (4, 8, 16. 32 time frames) on intelligence and age prediction tasks, respectively. In the intelligence prediction task which 0.8 (z-score) was only slightly better than the variance of one longer sequence lengths (16 and 32) led to better results in both young adults (HCP) and elders (UKB). Also, in age prediction, for young adults, we found the performance peaking at 16 time frames. However, in predicting the age of elders, longer sequence lengths (16, 32) resulted in poorer performance in age prediction. Namely, the findings of the former three cases showed that the longer sequence lengths enabled better learning of temporal dynamics needed to predict intelligence in young adults and elders and age in young adults. However, the last case of the age prediction task in elders showed that the benefit might not be generalizable to all the cases.
challenging task considering the MAE of about
These findings point out that the optimal value of the input time sequence length may vary depending on the given task and dataset, and our choice to employ constant-sized sub-sequences may have proven beneficial in light of this variability. A more detailed analysis can be found in Appendix C.2.
5 Concluding Remarks
Investigating the spatiotemporal dynamics of the human brain poses a formidable challenge owing to the lack of powerful analytics permitting individual-level prediction of cognitive or clinical outcomes such as psychiatric or neurological diseases. In this study, we present SwiFT, an efficient Transformer model designed for nigh-dimensional 4D brain functional MRI, aimed at learning these spatiotemporal dynamics and predicting biological and cognitive outcomes. Throughout various tasks, our method consistently outperforms the Kate-of-the-aa ROI-based and two-step based methods the domain. Furthermore, compared to TFF, a Transformer-basee uses significantly less memory and training time. Lastly, our IG-SQ interpretation results show the feasibility of interpreting the spatial patterns of the functional representations from SwiFT that contribute to a given task. We believe the simple and effective end-to-end learning capability of our SwiFT has the potential to significantly contribute to enhancing the scalability of spatiotemporal learning of fMRI in both computational and clinical neuroscience research. For future work, we plan to more extensively test the effectiveness of pre- pre-trainin SwiFT on large-scale datasets.
baseline that takes raw fMRI as input, SwiFT
10

---

6
Acknowledgements
This work was supported by the U.S. Department of Energy (DOE), Office of Science (SC), Advanced Scientific Computing Research program under award DE-SC-0012700 and used resources of the National Energy Research Scientific Computing Center, a DOE Office Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02- 05CH1 1231 using NERSC award ASCR-FRCAP0023081 Also, this research used resources of the Argonne Leadership Computing Facility, which is U.S. Department of Energy Office of Science User Facility supported under Contract DE-AC02-06CH1 1357.
This work was also supported in part by the National Research Foundation of Korea (NRF) grant funded by the Korean government (MSIT) [No. 202IRICICI005003 202 2021M335D220002447999 2021RIA3C20008884)
by Pinteeriii Researchers Program through Seoul National University (No. 200-20230058), and by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) N0.22021-0-0133 No. No.2022-0-00959) Taesup Moon is also supported in part by ASRI / INMC Seoul National University.
No.2022-0-00113
This research has been conducted using data from UK Biobank. major biomedical database (https: / /wwww ukbi obank. ac. uk). Our project ID of the UK Biobank is 32575. Data used in the preparation of this article were obtained from the Adolescent Brain Cognitive Development (ABCD) Study (https: / abodstudyy org), held in the NIMH Data Archive (NDA). This is a multisite, longitudinal study designed to recruit more than 10,000 children aged 9-10 and follow them over 10 years into early adulthood. The ABCD Study@ is supported by the National Institutes of Health and additional federal partners under award numbers U0IDA041048, U0IDA050989, U01DA05 1016, U01DA04 1022, 1010 N01DA051037, U01DA041117 U01DA041028, U01DA001 1134 UO1DA050988, U01DA05 1039, U01DA041 156, U01DA041025 U01DA041 120, U01DA05103K, U01DA041 148, U01DA041093, U01DA041089, U24DA04 123, U24DA041147. A full list of supporters is available at (https:/ / abcdstudy. org/ -padeerer-p htmm). A listing of participating sites and a complete listing of the study investigators can be found at (https:/ /awp//aassttt/ udy. wrggconsoree ium wttrrrtuimmmmmmmmeeeee www.mrre///eeee.. /). ABCD consortium investigators designed and implemented the study and/or provided data but did not necessarily participate the analysis or writing of this report. This manuscript reflects the views of the authors and may not reflect the opinions or views of the NIH or ABCD consortium investigators. The ABCD data repository grows and changes over time. The ABCD data used in this report came from http: / /xx. www.:ooooo......... 1555511510450 503209. DOIs can be found at https: : / //nda . www.nihhh.g gov /aw.//ac////////////o.o//oo/oo./
U01DA050987, UO1DA041 1174, UO1DA041 106,
54 eases .ttml. Data were provided [in part] by the Human Connectome Project, WU-Minn Consortium (Principal Investigators: David Van Essen and Kamil Ugurbil; IU54MH091657) funded by the 16 NIH Institutes and Centers that support the NIH Blueprint for Neuroscience Research; and by the McDonnell Center for Systems Neuroscience at Washington University.
References
[1] Nicola K Dinsdale, Emma Bluemke, Vaanathi Sundaresan, Mark Jenkinson. Stephen M Smith, and Ana IL Namburete. Challenges for machine learning in clinical translation of big data imaging studies. Neuron, 2022.
[2] Ming Song, Yuan Zhou, Jun Li, Yong Liu, Lixia Tian, Chunshui Yu, and Tianzi Jiang. Brain spontaneous functional connectivity and intelligence. Neuroimage, 4183):168--116,, July 2008.
[3] Emily S Finn, Xilin Shen, Dustin Scheinost, Monica D Rosenberg. Jessica Huang, Marvin M Chun, Xenophon Papademetris, and Todd Constable. Functional connectome fingerprinting: identifying individuals using patterns of brain connectivity. Nature Neuroscience, 18(11):1664- 1671, November 2015.
[4] Nicolai Franzmeier, Julia Neitzel, Anna Rubinski, Ruben Smith, Olof Strandberg, Rik Os- senkoppele, Oskar Hansson, and Michael Ewers. Functional brain architecture is associated
11

---

with the rate of tau accumulation in alzheimer's disease. Nature Communications, 11(1):347, 2020.
[5] Mary-Ellen Lynall, Danielle S Bassett, Robert Kerwin, Peter McKenna, Manfred Kitzbichler, Ulrich Muller, and Ed Bullmore. Functional connectivity and brain networks in schizophrenia. Journal of Neuroscience, 30(28):9477-949, 2010.
[6] Andrew T Drysdale, Logan Grosenick, Jonathan Downar, Katharine Dunlop, Farrokh Man- souri, Yue Meng. Robert N Fetcho, Benjamin Zebley, Desmond J Oathes, Amit Etkin, Alan F Schatzberg, Keith Sudheimer, Jennifer Keller, Helen S Mayberg, Faith M Gunning. George Alexopoulos, Michael D Fox, Alvaro Pascual-Leong Henning U Voss, B Casey, Mare J Dubin, and Conor Liston. Resting-state connectivity biomarkers define neurophysiological subtypes of depression. Nature Medicine, 23(1):28-33) 38, January 2017.
[7] Jeanette C Mostert, Elena Shumskaya, Maarten Mennes, A Marten H Onnink, Martine Hoogman, Cornelis C Kan, Alejandro Arias Vasquez, Jan Buitelaar, Barbara Franke, and David G Norris. Characterising resting-state functional connectivity in large sample of adults with adhd.
Progress in Neur--Pychopharmacooog and Biological Psychiatry, 67:82-91, 11882299,222 2016.
[8] Byung- Moon- Kim, Jong Chul Ye, and Jae-Jin Kim. Learning dynamic graph representation of brain connectome with spatio-temporal attention. Advances in Neural Information Processing Systems, 34:4314 4327, 2021.
[9] Nima Asadi, Ingrid R Olson, and Zoran Obradovic. A transformer model for learning spatio- temporal contextual representation in fmri data. Network Neuroscience, pages 1-41, 2022.
[10] Sandra Vieira, Walter H L Pinaya, and Andrea Mechelli. Using deep learning to investigate the neuroimaging correlates of psychiatric and neurological disorders: Methods and applications. Neuroscience Biobekavioral Reviews, 74(Pt 15,5 March 2017.
[11] Soham Gadgil, Qingyu Zhao, Adolf Pfefferbaum, Edith V Sullivan, Ehsan Adeli, and Kilian M Pohl. Spatio-temporal graph convolution for resting-state fmri analysis. In Medical Image Computing and Computer Assisted InterventionMMICAAI 2020: 23rd International Conference, Lima, Peru, October 8-8, 2020, Proceedings, Part VII 23, pages 528-538. Springer, 2020.
[12] Xuan Kan, Wei Dai, Hejie Cui, Zilong Zhang, Ying Guo, and Carl Yang. Brain network transformer. In Advances in Neural Information Processing Systems, 2022.
[13] Jonathan D Power, Alexander L Cohen, Steven M Nelson, Gagan S Wig. Kelly Anne Barnes, Jessica A Church, Alecia C Vogel, Timothy O Laumann, Fran M Miezin, Bradley L Schlaggar, al. Functional network organization of the human brain. Neuron, 72(4):665-678, 2011.
[14] Lisa D Nickerson, Stephen M Smith, Döst Öngür, and Christian F Beckmann. Using dual regression to investigate network shape and amplitude in functional connectivity analyses. Frontiers in Neuroscience, 11:115, 2017.
[15] Md Mahfuzur Rahman, Usman Mahmood, Noah Lewis, Harshvardhan Gazula, Alex Fedorov, Zening Fu, Vince D Calhoun, and Sergey M Plis. Interpreting models interpreting brain dynamics. Scientific Reports, (2(1):12023, July 2022.
[16] Wei Li, Xuefeng Lin, and Xi Chen. Detecting alzheimer's disease based on 4d finri: An exploration under deep learning framework. Neurocomputing, 388:280-222 287, 2020.
[17] Itzik Malkiel, Gony Rosenman, Lior Wolf, and Talma Hendler. Self-supervised transformers for fmri representation. In International Conference on Medical Imaging with Deep Learning. pages 993- PMLR, 2022.
[18] Sam Nguyen, Brenda Ng, Alan D) Kaplan, and Priyadip Ray. Attend and decode: 4d fmri task state decoding using attention models. In Machine Learning for Health, pages 267-279. PMLR, 2020.
[19] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012-10022, 2021.
12

---

[20] Yucheng Tang, Dong Yang, Wenqi Li, Holger R Roth, Bennett Landman, Daguang Xu, Vishwesh Nath, and Ali Hatamizadeh. Self-supervised pre-training of swin transformers for 3d medical image analysis. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pages 20730-20740, 2022.
[21] Anurag Arab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid. Vivit: A video vision transformer. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 6836-6846, 2021.
[22] Gedas Bertasius, Heng Wang, and Lorenzo Torresani. Is space-time attention all you need for video understanding? In International Conference on Machine Learning, volume 2, page 2021.
[23] Ze Liu, Jia Ning, Yue Cao, Yixuan Wei, Zheng Zhang, Stephen Lin, and Han Hu, Video swin transformer. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 3202-3211, 2022.
[24] David C Van Essen, Stephen M Smith, Deanna M Barch, Timothy EJ Behrens, Essa Yacoub, Kamil Ugurbil, Wu-Minn HCP Consortium, et al. The wu-minn human connectome project: an overview. Neuroimage, 80:62-79, 2013.
[25] Betty Jo Casey, Tarig Cannonier, May I Conley, Alexandra O Cohen, Deanna M Barch, Mary M Heitzeg, Mary E Soules, Theresa Teslovich, Danielle V Dellarco, Hugh Garavan, al. The adolescent brain cognitive development (abed) study: imaging acquisition across 21 sites. Developmental Cognitive Neuroscience, 32:43-54, 2018.
[26] Karla L Miller, Fidel Alfaro-Almagro, Neal K Bangerter, David L Thomas, Essa Yacoub, Junqian Xu, Andreas Bartsch, Saad Jbabdi, Stamatios N Sotiropoulos, Jesper L R Andersson, Ludovica Griffanti,, Gwenaëlle Douaud, Thomas W Okell, Peter Weale, Iulius Dragonu, Steve Garratt, Sarah Hudson, Rory Collins, Mark Jenkinson, Paul M Matthews, and Stephen M Smith. Multimodal population brain imaging in the UK biobank prospective epidemiological study. Nature Neuroscience, 19(11):1523- 15333,666,3336,,,,, November 2016.
[27] Fidel Alfaro-Almagro, Mark Jenkinson, Neal K Bangerter, Jesper L Andersson, Ludovica Grif- fanti, Gwenaëlle Douaud, Stamatios N Sotiropoulos, Saad Jbabdi, Moises Fernandeez Emmanuel Vallee, Diego Vidaurre, Matthew Webster, Paul McCarthy, Christopher Rorden, Alessandro Daducci, Daniel C Alexander, Hui Zhang, Iulius Dragonu, Paul M Matthews, Karla I. Miller, and Stephen M Smith. Image processing and quality control for the first 10,000 brain imaging datasets from UK biobank. Neuroimage, 166:400-424. 424, February 2018.
[28] Stephen M Smith, Karla L Miller, Gholamreza Salimi-Khorsine Khorshidi, Matthew Webster, Christian Beckmann, Thomas E Nichols, Joseph D Ramsey, and Mark W Woolrich. Network modelling methods for fmri. Neuroimage, (4(2))875-891 2011.
[29] Ed Bullmore and Olaf Sporns. Complex brain networks: graph theoretical analysis of structural and functional systems. Nature Reviews Neuroscience, 10(3): 18--198, 2009.
[30] Jeremy Kawahara, Colin J Brown, Steven P Miller, Brian G Booth, Vann Chau, Ruth E Grunau, Jill G Zwicker, and Ghassan Hamarnhh. Brainnetcnn: Convolutional neural networks for brain networks; towards predicting neurodevelopment. Neurolmage, 146: 1038-1049, 2017.
[31] Xin Deng, Jiahao Zhang, Rui Liu, and Ke Liu. Classifying asd based on time-series fmri using transformer. Computers in Biology and Medicine, 151:106320, 2022.
[32] Hasan Atakan Bedel, Irmak Şivgin, Onat Dalmaz, Salman UI Hassan Dar, and Tolga Çukur. Bolt: Fused window transformers for fmri time series analysis. arXiv preprint arXiv:2205.1157. 2022.
[33] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
13

---

[34] Jacob Devlin, Ming- Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXi::1810.04800 04805 2018.
[35] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11922 2020.
[36] Ali Hatamizadeh, Vishwesh Nath, Yucheng Tang, Dong Yang. Holger R Roth, and Daguang Xu. Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images. In Brainlesion:
Glioma, Multiple Sclerosis, Stroke and Traumatic Bruin Injuries: 7th International Workshop, BrainLes 2021, Held in Conjunction with MICCAI 2021, Virtual Event, September
27, 2021, Revised Selected Papers, Part I, pages 272-284. Springer, 2022.
[37] Sunghwan Hong, Seokju Cho, Jisu Nam, Stephen Lin, and Seungryong Kim. Cost aggregation with 4d convolutional swin transformer for few-shot segmentation. In Computer Vision-ECC ECCV 2022: 17th European Conference, Tel Aviv, Israel, October 23-27, 2022, Proceedings, Purt XXIX, pages 108-126. Springer, 2022.
[38] Christoph Feichtenhofer. X3d: Expanding architectures for efficient video recognition. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 203-213, 2020.
[39] Zhaofan Qiu, Ting Yao, and Tao Mei. Learning spatio-temporal representation with pseudo-3d residual networks. In proceedings of the IEEE International Conference on Computer Vision, pages 5533-5541, 2017.
[40] Ishan Dave, Rohit Gupta, Mamshad Nayeem Rizve, and Mubarak Shah. Telr: Temporal contrastive learning for video representation. Computer Vision and Image Understanding. 219:103406, 2022.
[41] Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the Thirteenth International Conference on Artificial Intelligence and Statistics, pages 297-304. JMLR Workshop and Conference Proceedings, 2010.
[42] O Esteban, C. Markiewicz, RW Blair, C Moodie, AI Isik, AE Aliaga, and KJ Gorgolewski. Fmriprep: a robust preprocessing pipeline for functional mri. biorxiv, 306951, 2018.
[43] Oscar Esteban, Rastko Ciric, Karolina Fine, Ross W Blair, Christopher J Markiewiez, Craig Moodie, James D Kent, Mathias Goncalves, Elizabeth DuPre, Daniel EP Gomez, et al. Analysis of task-based functional mri data preprocessed with fmriprep. Nature Protocols, 15(7):2186 2202, 2020.
[44] Vladimir Fonov, Alan C Evans, Kelly Botteron, C Robert Almli, Robert C McKinstry, D Louis Collins, Brain Development Cooperative Group, et Unbiased average age-appropriate atlases for pediatric studies. Neuroimage, 54(1):313-33 2011.
[45] Yashar Behzadi, Khaled Restom, Joy Liau, and Thomas T Liu. A component based noise correction method (compeor) for bold and perfusion based fmri. Neuroimage, 37(1)990-101 2007.
[46] Matthew F Glasser, Stamatios N Sotiropoulos. Anthony Wilson, Timothy S Coalson, Bruce Fischl, Jesper L. Andersson, Junqian Xu, Saad Jbabdi, Matthew Webster, Jonathan R Polimeni, David C Van Essen, and Mark Jenkinson. The minimal preprocessing pipelines for the human connectome project. Neuroimage, 80:105-124, October 2013.
[47] Stephen M Smith, Christian F Beckmann, Jesper Andersson, Edward J Auerbach, Janine Bi- jsterbosch, Gwenaëlle Douaud, Eugene Duff, David Feinberg, Ludovica Griffanti, Michael Harms, Michael Kelly, Timothy Laumann, Karla L. Miller, Steen Moeller, Steve Petersen, Jonathan Power, Gholamreza Salimi-Khorshin Khorshidi, Abraham 7. Snyder, An T Vu, Mark Wool- rich, Junqian Xu, Essa Yacoub, Kamil Uğurbil, David C Van Essen, Matthew F Glasser, and
14

---

WU-Minn HCP Consortium. Resting-state fMRI in the human connectome project. Neuroimage, 80:144-168, October 2013.
[48] Cathie Sudlow, John Gallacher, Naomi Allen, Valerie Beral, Paul Burton, John Danesh, Paul Downey, Paul Elliott, Jane Green, Martin Landray, Bette Liu, Paul Matthews, Giok Ong. Jill Pell, Alan Silman, Alan Young. Tim Sprosen, Tim Peakman, and Rory Collins. UK biobank: an open access resource for identifying the causes of a wide range of complex diseases of middle and old age. PLOS Medicine, 12 1001779, March 2015.
[49] HCP WU-Minn. 1200 subjects data release reference manual. URL https:///www humanconnec- tome. org, 2017.
[50] Matthew F Glasser, Timothy S Coalson, Emma C Robinson, Carl D Hacker, John Harwell, Essa Yacoub, Kamil Ugurbil, Jesper Andersson, Christian F Beckmann, Mark Jenkinson, et al. A multi-modal parcellation human cerebral cortex. Nature, 536(7615):171-178 2016.
[51] Rebecca M Shansky and Catherine S Woolley. Considering sex as biological variable will be valuable for neuroscience research. Journal of Neuroscience, 36(47): 2016.
[52] James H Cole, Riccardo E. Marioni, Sarah E Harris, and lan J Deary. Brain age and other bodily 'ages": implications for neuropsychiatry. Molecular Psychiatry, 241
[53] Richard C Gershon, David Cella, Nathan A Fox, Richard Havlik, Hugh C Hendrie, and Molly V Wagster. Assessment of neurological and behavioural function: the nih toolbox. The Lancer Neurology, 9()):138-339 2010.
[54] Peter R Millar, Brian A Gordon, Patrick H Luckett, Tammic LS Benzinger, Carlos Cruchaga, Anne M Fagan, Jason Hassenstab, Richard Perrin, Suzanne E Schindler, Ricardo F Allegri, al. Multimodal brain age estimates relate to alzheimer disease biomarkers and cognition in early stages: a cross-sectional observational study. Elife, 12:e81865 2023.
[55] Kaustubh Supekar, Carlo de Los Angeles, Srikanth Ryali, Kaidi Cao, Tengyu Ma, and Vinod Menon. Deep learning identifies robust gender differences in functional brain organization and their dissociable links clinical symptoms in autism. The British Journal of Psychiatry, 220(4):202-209) 2022.
[56] Teresa M Karrer, Danielle S Bassett, Birgit Derntl, Oliver Gruber, André Aleman, Renaud Jardri, Angela R Laird, Peter Fox, Simon Eickhoff, Olivier Grisel, et al. Brain-based ranking of cognitive domains to predict schizophrenia. Human Brain Mapping, 40(15):448744500 2019.
[57] Jaroslav Rokicki, Thomas Wolfers, Wibeke Nordhøy, Natalia Tesli, Daniel Quintana, Dag Alnæs, Genevieve Richard, Ann-Marie G de Lange, Martina J Lund, Linn Norbom, et al. Multimodal imaging improves brain age prediction and reveals distinct abnormalities patients with psychiatric and neurological disorders. Human brain mapping, 42(6): 1726, 2021.
[58] Nico UF Dosenbach, Binyam Nardos, Alexander L Cohen, Damien A Fair, Jonathan D Power, Jessica A Church, Steven M Nelson, Gagan S Wig, Alecia C Vogel, Christina N Lessov- Schlaggar, et al. Prediction of individual brain maturity using fini. Science, 329(5997):1358- 1361, 2010.
[59] Tianqi Chen and Carlos Guestrin. Xgboost: A scalable tree boosting system. In Proceedings of the 22nd ACM SIGKDD International Conference Knowledge Discovery and Data Mining, pages 785-794, 2016.
[60] Atif Riaz, Muhammad Asad, SM Masudur Rahman Al Arif, Eduardo Alonso, Danai Dima, Philip Corr, and Greg Slabaugh. Deep fmri: An end-to-end deep network for classification of fmri data. In 2018 IEEE 15th International Symposium Biomedical Imaging (ISBI 2018), pages 1422. IEEE, 2018
[61] Julius Adebayo, Justin Gilmer, Ian Goodfellow, and Been Kim. Local explanation methods for deep neural networks lack sensitivity to parameter values. arXiv preprint arXiv:1810.00 1810.03307 2018.
15

---

[62] Narine Kokhlikyan, Vivek Miglani, Miguel Martin, Edward Wang, Bilal Alsallakh, Jonathan Reynolds, Alexander Melnikov, Natalia Kliushkina, Carlos Araya, Siqi Yan, and Orion Reblitz- Richardson. Captum: A unified and generic model interpretability library for pytorch, 2020.
[63] Tim Koscik, Dan O'Leary, David Moser, Nancy C Andreasen, and Peg Nopoulos. Sex differences in parietal lobe morphology: relationship to mental rotation performance. Brain and Cognition, (013):45-4559 2009.
[64] Joel Salinas, Elizabeth D Mills, Amy L Conrad, Timothy Koscik, Nancy C Andreasen, and Peg Nopoulos. Sex differences in parietal lobe structure and development. Gender Medicine, 911):44-55, 2012.
[65] Paul M Macey, Nicholas S Rieken, Rajesh Kumar, Jennifer A Ogren, Holly Middlekauff, Paula Wu, Mary A Woo, and Ronald M Harper. Sex differences in insular cortex gyri responses the valsalva maneuver. Frontiers in Neurology, page 87, 2016.
[66] Susanne Weis, Kaustubh R Patil, Felix Hoffstaedter, Alessandra Nostro, BT Thomas Yeo, and Simon B Eickhoff. Sex classification by resting state brain connectivity. Cerebral Cortex, 30(2):824-33,, 2020.
[67] Alexander B. Jung, Kentaro Wada, Jon Crall, Satoshi Tanaka, Jake Graving, Christoph Reinders, Sarthak Yadav, Joy Banerjee, Gábor Vecsei, Adam Kraft, Zheng Rui, Jirka Borovec, Christian Vallentin, Semen Zhydenko, Kilian Pfeiffer, Ben Cook, Ismael Fernández, François-Michh De Rainville, Chi-Hung Weng, Abner Ayala-Acevedo, Raphael Meudec, Matias Laporte, et al. imgaug. https:://// /ttg//gttt/////////// com/ ale ju/ imgaug, 2020. Online; accessed 01-Feb-2020.
16

---

A Performance Comparison with Standard Deviation
Here, we detail the standard deviation of the results posted in Table 1 (manuscript), which compares the performance of SwiFT against other baseline models various downstream tasks, All of the experiments were performed using three pre-determined random data splits for each dataset, which was shared across all models for a fair comparison. The following Tables (A.1, A.2, A.3) show the performance posted in Table (manuscript) along with the standard deviation among the three splits. From the tables, we conclude that SwiFT outperforms all bascline models above the margin of variability at the following tasks: ABCD intelligence, HCP sex, HCP age, HCP intelligence, UKB sex, and UKB age prediction tasks. Note that the age prediction task was not carried out the ABCD dataset due to the narrow age range to 11 years, making it hard to obtain meaningful results from the experiments.
Table A.I: Performance of various models on the ABCD dataset with standard deviation
Dataset: ABCD
Method
Sex
Intelligenn litelligence
ACC
AUC
MSE
MAE
XGBoost
69.5
76.7
0.977+0.087
0.770+0.016
BrainNetCNN [30]
80.1 40.09
87.9
0,969 40.042
0.76740.016
Vanilla [12]
77.4
85,
0,961
050
0.764 025
BNT [12]
79.1
88.9
0.955 +0.058
0.767 40.022
TFF [17]
13
80.2
0.968 +0. 024
0.768 0.768 +0.000 009
SwiFT (ours)
79.3
87.8
0.932 ::0 017
0.756 008
Table A.2: Performance of various models on the HCP dataset with standard deviation
Dataset: HCP
Method
Sex
Age
Intelligence
ACC
AUC
MSE
MAE
MSE
MAE
XGBoost
68.5+3.03
75.5+3.14
14.3+1.61
3.1220.165
0,991
0.813+0.032
BrAinNetCNN [30]
77.1
84.942.28
12.6
2.97
0.984
0.805
VanillaTF [12]
77.9
85.240.69
12.5
2.95
0.987 100039
0.81240.014
BNT [12]
81.0+3.11
88.043.10
12.8
2.98
1.001
+: 00.
0.83020.014
TFF [17]
92.5+1.12
97.5+1.77
13.8
3.11 +0.200
0.953 40.074
0.795-+0.028
SwiFT (ours)
92.9+1..1
98.0
8.640.57
2.36
0.903
0.786=0.030
Table A.3: Performance of various models on the UKB dataset with standard deviation
Dataset: UKB
Method
Age
Intelligence
ACC
AUC
MSE
MAE
MSE
MAE
XGBoost
79.5
87.640.94
48.8
5.85
1.055
0.81640.078
BrainNetCNN [30]
86.8
93.840.31
42.7+0.17
5.36
1.001 10.141
0.800
VanillaTF [12]
87.0
41.4
5.26
0.999 104
0.799
BNT [12]
87. $7.0+1.10
94.8+.0.46
39.6g1.07
5.17
2..998
0.798
TFF [17]
96.8+0.25
99.54
42.1
5.10
0..
0.783
SwiFT (ours)
97.7+0.31
99.840.04
18.2.0..4
3.40=0.08s
0.992+0.105
0.79640.014
B Implementation Details
SwiFT To obtain the results detailed in Section 4.2 (manuscript), we trained SwiFT from scratch using the following configuration:
Optimizer: AdamW using cosine decay learning rate scheduler with linear warm-up (around 5% of total iterations)
Learning rate: After the warm-up, an initial learning rate of 10-5 for classification tasks and 5 x 10-5 for regression tasks on the HCP dataset and 10 for regression tasks on the
ABCD dataset.
Mini-batt batch size: Mini-batch of size 8
17

---

Epochs: 10 epochs of training for the sex classification and intelligence prediction tasks, and maximum of 30 epochs for the age prediction task.
After training, we selected the model instances with the highest validation AUC or lowest validation MSE to report the performance the test dataset. The model trained using four NVIDIA A GPUs using the distributed data-parallel (DDP) strategy provided by Pytorch Lightning, with single training session typically lasting from 4 to 30 hours depending the dataset and whether data augmentation was used during training. The model can be trained on an NVIDIA A5000 GPU with 24GB memory without problem.
To obtain the results detailed Section 4.3. (manuscript), the model was pre- pre-trained using the following:
Optimizer: A optimizer using a cosine decay learning rate scheduler with 2000 steps of linear warm-up
Learning rate: After warm-up, an initial learning rate of 10-5
Mini-batch size: 3
Epochs: six epochs of training
After pre- pre-tra training, we fine-tuned the model without a learning rate scheduler, using 1/10 of the initial learning rate used for nom-scratch training.
We used the same data splits to compare our baseline 4D Transformer, TFF, with SwiFT. We followed the number of attention heads (16) and embedding size (2,640) proposed by [17]. To alleviate over-fitting, we applied data augmentation methods to the brain images, such as Gaussian blur and additive Gaussian noise implemented by imgaug[67]. Since TFF requires more computational resources than SwiFT to run the codes, at least 8 hours of training were required using 2 nodes with A100 GPUs
TFF
We trained TFF with the following training setup:
Optimizer: Adam optimizer using cosine decay learning rate scheduler with linear warm-up by 5% of total iterations
Learning rate: After warm-up, an initial learning rate of 10-4
Mini-batch size: 32
Epochs: 10 epochs of training
ROI-based models The four ROI-based model baselines, XGBoost [59], BrainNetCNN [30], VanillaTF [12], and Brain Network Transformer [12] were reproduced for our experiments. The reproductions were based on the official code of [12]. However, since the preprocessing codes for the ABCD dataset were not provided by [12], we followed the preprocessing steps described in [12], potentially causing some differences. This obscurity in the preprocessing step is suspected to be one of the reasons for the slight performance gap the BNT model between our experiments and the results posted in the original paper [12), despite utilizing the same ABCD dataset.
We trained BrainNeiCNN, VanillaTE, and Brain Network Transformer models with the following
setup:
Optimizer: Adam optimizer using cosine decay learning rate scheduler
Learning rate: Learning rate of 5 x 10-5
Mini-batch size: Mini-batch of size 16
Epochs: 200 epochs of training
We used grid search for hyper-parameter tuning of XGBoost, adjusting the maximum depth and minimal child weight, gamma, learning rate, and colsample by tree. addition, we conducted 5-fold cross-validation. Hyperparameters are tuned with the following setup:
Maximum depth: Chosen between and
Minimal child weight: Chosen between and 7
Gamma: Chosen between 0.0 and 0.4
Learning rate: Chosen between 0.05 and 0.3
Colsample by tree: Chosen between 0.6 and 0.9
Software Version
The major software used for experiments are as the following:
18

---

Table A.4: Performance comparison of SwiFT for different positional embedding methods.
Sex
Age
Intelligence
Dataset
Method
ACC
AUC
MSE
MAE
MSE
MAE
Relative
89.8
95.941.21
9.0.2.56
2.44 20.16
0.90840.009
4.771
HCP
Absolute
92.941.61
98.0+1.79
8.6
2.36410.11
0.903
0.786
Relative
80.2+1166
88,9 20.26
0.93640.029
0.761 #0.013
ABCD
N/A
Absolute
79.3
87.8
0.902
0.756
Relative
97.5+0.10
99.8+0.05
19.4
3.5340.071 3.40
1.019+0.083
0.807
UKB
Absolute
97.7
99.8+0.04
18.2
0.992
0.796400.04
Table A.5: Efficiency comparison of SwiFT for different positional embedding methods.
Method
Param.
FLOPs
Throughput
Relative
4.66M
2.62G
94.17
Absolute
4.64M
2.62G
104.16 16
python 3.10.4
pytorch 1.12.1
pytorch-lightning 1.6.5
monai 1.1.0
neptume-client 0.16.4
scipy 1.8.1
torchvision 0.13.1
torchaudio 0.12.1
The full dependency list released alongside the code.
Additional Experiments
C.1
Comparison of Positional Embedding Methods
We discuss the impact of the switch from a relative positional bias scheme used in most Swin Transformer variants [19, 23, 20] to an absolute positional embedding scheme, as detailed in the paragraph "4D absolute positional embedding" of Section 3.1 (manuscript). We implemented the 4D relative positional bias scheme by extending the 3D relative positional bias described in [23]. Given window with dimensions of PxM X x M, the 4D relative positional bias B I M X AM xM2 for each self-attention head is integrated as
Attention(Q2 Q, K, v) SoftMax(QK//VLL BYV,
(I)
where Q K,V PM are the query, key, value matrices and C' is the channel number. A parameterized bias matrix B E (2P ()
is used to calculate the values in B since there are only 2P or 2M - possible position differences for every axis. separate parameterized bias matrix is used for each attention head at each layer of the Transformer.
Table A.4 shows the overall performance comparison of SwiFT between the absolute positional embedding scheme and the relative positional bias scheme for the HCP, ABCD, and UKB datasets. Comparing the performance of both methods, we found that the absolute positional embedding scheme shows better performance in most cases. Note that the age prediction task on the ABCD dataset was not tested due to the same reasons detailed in Section A.
Additionally, we compared the efficiency of the two methods in Table A.5 through the number of parameters, the number of FLOPs per forward pass, and the throughput. Throughput measures how many fMRI sub-sequences of length 20 the model processes per second during inference on single A100 GPU. The absolute positional embedding scheme more memory efficient as it only requires parameters at the beginning of each stage. In addition, the absolute positional embedding scheme also more computationally efficient as it does not require the 4D relative positional bias B to be reconstructed during each self-attention computation, resulting in a 9.6% throughput improvement. Overall, we conclude that the absolute positional embedding scheme is appropriate for our tasks compared to the relative positional bias scheme.
19

---

100
86
99.5
85.5
HCP, UKB AUC (%
99
85
ABCD ABCD AUC (%))
98.5
84.5
98
84
97.5
83.5
97
UKB AUC
ABCD AUC
96.5
82.5
HCP AUC
96
82
4
8
12
16
20
24
28
32
Number of time frames of input volume
Figure A.1: Effect the number of time frames of the input 1MRI volume on sex classification.
C.2 Effect of Input Time Sequence Length for Sex Classification
In our setup, instead of inputting the entire fMRI sequence of subject all at once, we have opted to divide it into 20-frame sub-sequences and use the divided sub-sequences as the input to our model. This was mainly due to memory constraints and to standardize the number of input tokens to our model. We have investigated the effect of changing the length of this input sub-sequence in section 4.6. Here, we present our experiments on the sex classification task (Figure A.1) and elaborate as to why the optimal length seems to change for each dataset and task.
The SwiFT model was trained with varying number of input fMRI volumes, with 20 volumes being the default configuration used for the previous experiments. We adjusted the mini-batch size to keep the total number of training iterations per epoch constant. All the performances in the figure averaged performances from three pre-determined data splits. We note that the training data augmentation on the ABCD dataset was not applied in this experiment to keep the training environment consistent compared other datasets, and thus the model has lower performance compared to the results posted in Section 4.2 (manuscript).
The performance of the classification task on the HCP and UKB dataset is already saturated near 100% AUC, and thus we observed that changing the number of input time frames has a small or negligible effect the performance, ranging within dataset, the performance of the model has a larger fluctuation. The performance peaked at around 20 input frames, which is the default number of input volumes used for the other experiments.
standard deviation. However, for the ABCD
Inputting longer time sequence into the model would allow the model to attend longer sequences once, at the expense of bloating the model with more parameters and causing it to overfit. Therefore a setting where we are dealing with a limited number of training data, using shorter sub-sequence has possibility to be beneficial to the model's generalization capabilities. We also note that the total amount of information used for our predictions remains the same regardless of the input time sequence length. the averaged logits from all sub- seb-eequences of the subject used for the prediction. Additionally, the datasets used for this work are resting-state fMRI scans, where the image does not change drastically over time. In the future, it would be interesting investigate the effect of input time sequence length on more temporally dynamic datasets, such as task-based fMRI scans.
C.3 Time-windor Analysis
In accordance with the details provided in Section 4.1.1, the SwiFT model was trained by processing individual 20-frame time windows (sub-sequences) of fMRI data sequentially, while sliding the time windows encompass the entire fMRI dataset. For inference purposes, the predictions obtained from the fMRI time windows were aggregated by averaging the logit values of each subject. This averaged value then utilized as the final prediction for the respective subject. To ensure that the predictions
20

---

of time windows from the same subjects are homogeneous and the final predictions are not decided by a few noisy time windows, we verified how many subjects exhibited distinct predictions among their time windows using the sex classification task of ABCD, HCP, and UKB datasets. Each subject of ABCD, HCP, and UKB has 18, 60, and 24 time windows.
(a) UKB
(b) HCP
(c) ABCD
969
88.3%
74.3%
40
14
20
Number of subjects
01
30
3.2%
7.2%
10
04
Percentage of corrected time windows
Figure A.2: Inner-subject accuracy of sex classification.
As seen the histogram in Figure A.2, over 90% of subjects showed identical predictions among the time windows (0.992 for UKB, 0.907 for HCP, and 0.927 for ABCD). Of note, 1.0 the x-axis denotes that the predictions from the time windows of the subjects are perfectly correct, while 0.0 means none of the time windows exhibited correct predictions. This suggests that the predictions of time windows are homogeneous and that the final predictions each subject are not biased toward a few time windows.
D SwiFT's Ability to Capture Long-range Interactions
Owing to the local nature of the attention mechanism, there could be concerns about SwiFT not being able to capture long-range interactions between spatially/temporally separated tokens. However, there some built-in mechanisms that could learn about long-rang range communications between tokens.
On the spatial dimensions, the Swin Transformer model can cover large ranges on the later parts of its layers through the patch merging step. On our configuration, after the initial patch embedding step where 6
216 neighboring voxels are embedded into a patch, a x 4 window can cover range equivalent 24 x 24 x 24 voxels from our 96 x 96 x 96 input. However, after the patch merging step where 8 neighboring tokens are merged, same window can cover a range equivalent to 48 x 48 x 48 voxels. After another patch merging step, stage 3, this grows to 96 x 96 x 96 which is already the same size as our entire input. From this point on, single window covers the entire span of our input in the spatial dimension, thanks to the patch merging step.
On the temporal dimension, where patch merging does not occur on the current version of the model, we needed another method that can ensure long-range interaction between tokens. The shifting window allows a token interact with tokens that are about 10 time frames away from our setup, but to allow longer range interactions we made the model do full global attention on the final stage as mentioned in Section 3.1 (manuscript). Thus, model processes the entire temporal range in the final stage, allowing it to introduce long-range temporal interactions in the end. We acknowledge that there could be more effective methods to allow long-range interactions, and would be an exciting problem to tackle when are dealing with more temporally dynamic datasets.
21

---

E
Loss Curves
Here we detail the model's performance the validation set during the training procedure in order to provide reference for our training procedure. Note that the validation MSE spikes for ABCD intelligence prediction are most likely due to the restarts of the learning rate scheduler.
54X
ABCD
MCP
UKB
8.000
6795
2.44
ABR0C
0204
6252
epoch
Validation AUROC per training epoch for sex classification
HCP
UKB
0.000
0.000
Mae
0.200
(normatzed]
0340
0.600
0.50
6:03
2012
epoch
(b) Validation MSE per training epoch for age prediction
Inteligence
ADGD
HGP
URB
wilawaan Porne Redd
3532
4400
0832
2.00
epoch
(c) Validation MSE per training epoch for intelligence prediction
Figure A.3: Performance on the validation set during the training procedure.
F
Limitations
The biggest limitation of research is that the advantages of self-supervised pre-training evident in some tasks. The pre-training was found to foster early convergence during fine-tuning on some downstream tasks, such as HCP and ABCD intelligence prediction, but some other cases, the performance gain from pre-training was not substantial. This result can be attributed to the distinct age range between UK Biobank used for pre-training and the other two datasets for downstream tasks, ABCD and HCP. In addition, the scanner effect of fMRI, which confounds the biological and cognitive features, can also hamper knowledge transfer. Otherwise, the performance increase might be limited because we have already reached the upper bounds for the sex and age prediction tasks. Previously, few studies targeted a transfer learning of 4D fMRI between different data sources. We suggest that various sources of the dataset should be included during pre-training in supervised or unsupervised ways. Though the observed performance increase from transfer learning was limited, established the foundation for large-scaling pre-training for 4D fMRL. Consequently, optimizing
22

---

self-supervised pre- training methodology for fMRI with larger amounts of datasets are promising future research topic.
G Licences
We used existing assets follows:
TFF [17]: hutp://gihhhkoommmmonoeemmmaa/a/ No explicit license.
Brain Network Transformer [12]: http://ghii...ooo/////////////at
MIT License.
Swin Transformer http://gghhooccoommmmmoooss////t MIT License.
MonAl: http://jihnnccmo//peccccx/oooooo Apache License.
TCLR [40]: http://gghtlbbb.ommm/DaaSSSSAAAR No explicit license
23


Reviews:
[{"review_id": "bBLnqgE9rg", "reviewer": "Reviewer_bLJJ", "summary": "This paper introduces SwiFT, a novel Swin Transformer-based architecture for analyzing high-dimensional functional brain MRI data. SwiFT learns end-to-end to process spatiotemporal brain functional data, and it achieves state-of-the-art performance on several tasks, including sex classification, age prediction, and cognitive intelligence prediction. The authors also demonstrate the feasibility of applying the pre-train and fine-tune framework to SwiFT, empowering researchers to construct large-scale foundation models for fMRI. ", "questions": "- Could the authors provide more details on the computational resources required to train and evaluate SwiFT? This would help to better understand the practical feasibility of using SwiFT in real-world applications.\n- Can the authors provide more details on the implementation of the 4D window multi-head self-attention mechanism and absolute positional embeddings used in SwiFT? This would help to better understand the specifics of SwiFT's architecture and how it differs from other Transformer-based models.\n- Could the authors provide a more in-depth discussion of the limitations of SwiFT, such as its sensitivity to noise in the data or its ability to generalize to new datasets?", "limitations": "Consider the potential impact of SwiFT on the broader field of neuroscience and medicine, such as its potential to facilitate scalable learning of functional brain imaging. This would help to better understand the potential benefits and drawbacks of using SwiFT in research and clinical settings.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper propose Swin Transformer-based architecture that can learn brain dynamics directly from 4D functional brain MRI data in an end-to-end fashion, encouraging researchers to construct foundation models for fMRI.\n- The authors provide a clear and concise introduction to the problem of analyzing high-dimensional functional brain MRI data and explain the motivation for developing SwiFT. \n- Results reveal that SwiFT consistently outperforms recent state-of-the-art models. The authors also conduct ablation studies to substantiate the modeling choices and present interpretation results using IG-SQ for SwiFT's predictions.\n- SwiFT has the potential to facilitate scalable learning of functional brain imaging in neuroscience research by reducing the hurdles associated with applying Transformer models to high-dimensional fMRI.\n- The authors also provide interpretability of SwiFT's predictions, which can help researchers better understand the brain's spatiotemporal dynamics. ", "weaknesses": "- The paper compares SwiFT with limited recent state-of-the-art models,. It would be interesting to see how SwiFT compares with more types of models, such as CNNs or GNNs, which are commonly used in fMRI analysis.\n- While the paper presents interpretation results using IG-SQ for SwiFT's predictions, it would be helpful to have a more in-depth discussion and how they can be used to better understand the brain's spatiotemporal dynamics.\n- It would be helpful to have a more in-depth discussion of the limitations of SwiFT, such as its sensitivity to noise in the data or its ability to generalize to new datasets.\n- The paper does not provide open-source code for SwiFT, which could limit its adoption and reproducibility. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed analysis of our work and for the insightful suggestions. Here we address some of the questions and concerns raised:\n\n> *...It would be interesting to see how SwiFT compares with more types of models...*\n\n-  Thank you for providing us with some pointers. To faithfully evaluate the effectiveness of our model, we have already reproduced multiple models as baselines. One of such models is BNT ([12] of our paper), the current state-of-the-art model to the best of our knowledge. BNT improves upon the traditional GNN-based approaches for fMRI analysis and represents the cutting-edge method of similar GNN-based approaches. \n\n> *While the paper presents interpretation results ...*\n- Thank you for your inquiry. Currently, we employ a method of calculating IG-SQ for each sub-sequence and subsequently averaging along the time axis. Through this approach, we have identified which regions, on average, exhibit the highest explanatory power for gender differences.\n- For these IG-SQ values to hold further significance, conducting additional analyses on their relationship with brain connectivity is necessary. The brain regions showing high IG-SQ values in our findings, such as the thalamus, insular cortex, and inferior temporal gyrus, are primarily hubs for integrating various sensory information and exhibit connectivity with several other brain regions. A crucial research question is investigating whether this high connectivity contributes to the elevated IG-SQ values. \n- Furthermore, exploring how well IG-SQ values can account for individual differences is worthwhile. IG-SQ values serve as a more direct measure of sex difference. Each individual possesses unique IG-SQ values, which might also represent individual differences. Consequently, investigating the predictive capacity of Interpretation maps for mental disorders characterized by significant gender differences, such as depression or ASD, is a crucial research avenue in the future.\n\n> *The paper does not provide open-source code ...*\n\n-  We apologize that we had to provide our code only as a file in the supplementary material due to the anonymous review process. We would be happy to release our open-source code alongside the final draft of our paper.\n\n> *Could the authors provide more details on the computational resources ...*\n\n-  Please refer to the general rebuttal for the additional details.\n\n> *Can the authors provide more details on the implementation...*\n- Thank you for the comment. We supplement our implementation details with the following:\n- The 4D window multi-head self-attention (4DW-MSA) mechanism is an extension of its 2D variant from the Swin Transformer. The core idea of the windowed attention mechanism is that tokens only attend to tokens that are spatially or temporally adjacent to them. This would allow the Transformer to achieve lower computation and memory complexity.\n- To be specific, we divide the entire 4D space into smaller sub-spaces (windows), and tokens within a sub-space only attend to tokens within the same sub-space. Additionally, after each layer the windows are shifted by half of its width to allow the tokens to exchange information with new adjacent tokens. \n- The absolute positional embeddings, although not commonly used, were introduced in our work to improve the model's efficiency. The absolute positional embeddings have a simple premise: we add a (learnable) embedding vector to each of the tokens based on their position. This would be adding a unique vector for each of the 81,260 possible positions for our configuration. As a design choice, we separate the time dimension and space dimension and employ 20 time-embedding vectors and $16 \\times 16 \\times 16 = 4,096$ space-embedding vectors and add them separately. \n\n    \n> *a more in-depth discussion of the limitations of SwiFT...*\n- A limitation of SwiFT is its inability to process entire sequence as a whole, necessitating the division into sub-sequences. We theorize it may not be an entirely negative aspect, however it would be difficult to test out processing entire sequences due to memory constraints. For a more comprehensive explanation, please refer to the concept on the general rebuttal.\n- For limitations associated with generalizing to new datasets, please refer to the response provided below.\n\n> *Consider the potential impact of SwiFT on the broader field ...*\n\n-\tThank you for your interest in the potential impact of SwiFT. We expect SwiFT to become a foundation model for functional brain imaging that can solve various problems in neuroscience and medicine. In recent years, it has been a trend to train large-scale models and then perform new tasks through few-shot learning. However, functional imaging has not had such an approach due to the lack of models and computational resources for large-scale training, and we believe that SwiFT can be an effective solution.\n-\tIn research and clinical settings, SwiFT can be used to solve specific, well-defined problems. For instance, SwiFT could be directly employed for the diagnosis and prognosis of diseases, and the features learned by SwiFT could be harnessed for subtyping heterogeneous psychiatric disorders. However, owing to the intrinsic nature of transformers, the efficacy of SwiFT might not be fully realized with smaller datasets, and it becomes necessary to fine-tune the model using weights pre-trained on larger datasets.\n-\tIn our work it has been observed that transfer learning exhibits restricted enhancements in specific tasks. However, this limitation can be due to conducting pre-training solely on a single type of data. We anticipate that broadening the scope of pre-training to encompass diverse datasets, such as task fMRI, will augment the effectiveness of transfer learning. This stands as a crucial avenue for future inquiries."}, {"review_id": "zqUAIE0aFw", "reviewer": "Reviewer_pTkQ", "summary": "* The authors seek to develop an approach for modeling spatiotemporal brain dynamics, as measured with resting state functional MRI. For this, they have extended a transformer-based architecture to handle time-varying 3D scan data. \n* They seek to predict subject traits and attributes like age, sex, and cognitive scores from the learned representations using three publicly available datasets. \n* They have succeeded in outperforming state-of-the-art models in these tasks while reducing the computational complexity and increasing throughput in model evaluations. \n* Additionally, they have demonstrated the usefulness of pre-training such models on downstream tasks. \n* Lastly, the relationships captured by the model between the subject attributes and the brain regions are consistent with prior work.\n\n", "questions": "* Could the authors clarify if the parameters for z-scoring, i.e., mean and variance, were learned over the whole data or only on the training data (& applied to validation and test data)?\n* I would expect that using the whole fMRI time series for any task prediction would provide better results than using sub-sequences.  Could the authors provide more intuition about the seemingly counterintuitive results obtained in Section 4.6, especially in UKB (e.g., age prediction)?\n* While we commend that the authors have shared their code, it would be helpful if they could also share their trained models (especially due to the high compute power required to generate these results). It would help the scientific community at large and also be in line with the spirit of Section 4.3.", "limitations": "* Explanation methods like IG are generally brittle. It would be helpful if one more explanation method could confirm the results in Section 4.4\n* While a shifting window enables efficient feature extraction, it limits model expressivity. It would be useful to know how to extract features that span large spatial/temporal ranges. This could be particularly relevant for task-based fMRI data since signals of varying timescales could affect behavior in such paradigms.\n* Authors have mentioned that they divided the fMRI data into sub-sequences due to memory constraints. Moreover, as the authors have used four Nvidia A100s (a total of 160GB of GPU memory with NvLink support), they are already on the higher end of compute power. It would be helpful to understand what those memory constraints were and suggestions, if any, on tackling them, as this will have implications in replicating and/or extending their work.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "* The problem statement is well-motivated, and the limitations of prior works (i.e., ROI-based and two-step approaches) are detailed.\n* The results were replicated across three datasets and various tasks, showcasing this architecture's robustness.\n* The authors show the learned features are general enough to help downstream tasks through pre-training.\n* The authors have performed computational complexity analysis. Because their model has almost double the throughput compared to the second-best-performing model, it paves the way to use these models in real-time settings.\n* Regions of the brain used by the model (as in Section 4.4) to predict age align with previous work, which increases confidence that the model is learning meaningful features.\n* The authors have submitted their full code, which assists replicability and helps the scientific community extend their work.", "weaknesses": "* The novelty of the architecture vis-a-vis latest advances in transformer architectures ([23], [34], [35]) seems limited. \n* The authors aimed to learn representations for brain dynamics but focussed only on predicting traits (age, sex etc), because they used only resting state scans. Dynamic attributes such as task-level performance measures -- e.g., reaction times and accuracies in HCP tasks), cognitive load in a working memory task etc -- are relevant to establish the generality of the findings. The model should be tested on these tasks also.\n* While the authors have detailed the approach for the two contrastive losses if they could justify the choice of positive and negative samples through domain knowledge or established methods, it would help increase the confidence in the proposed approach.\n* The authors have mentioned performing ablation studies to \"substantiate their modeling choices.\" I expect that this includes studying the effect of hyperparameters like the number of layers, channel size, window size, etc., given that this is a novel architecture. Details of such studies should have been mentioned.\n* Additionally, it would be helpful to see the loss curves of models for different datasets and targets. It will help understand the choice of a low number of epochs for training various models.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for raising some important questions and providing us with helpful feedback. Here is our response to the questions and concerns raised:\n\n>*The novelty of the architecture ...*\n\n- The novelty of our architecture may seem to be limited, however there are several points that we claim to be novel.\n- As previous works (e.g., [23] in the manuscript) have extended the Swin Transformer to accept 3D inputs, we extend it one step further to directly process 4D inputs. This comes with its own set of challenges, such as how to deal with the patch merging step combining larger number of tokens or the memory and computation constraints due to the added dimension.\n- VAT ([35] of our paper) technically utilizes the Swin Transformer to process 4D inputs, although the Swin Transformer module was used to process pair of intermediate 2D feature maps instead of a \"true\" 4D data. The Swin Transformer is also missing some key features such as the patch merging step and the hierarchical structure. Due to these distinctions we believe that our introduction of a fully-fledged 4D Swin Transformer architecture can still be considered as a novel contribution.\n\n>*...The model should be tested on these tasks also.*\n- Thank you for the suggestion. As briefly mentioned in the general rebuttal regarding long-range interactions, we also believe this to be an exciting future work. We expect SwiFT to perform well on these datasets compared to ROI-based models since SwiFT does not \"delete\" the temporal information during the ROI preprocessing step.\n\n>*...approach for the two contrastive losses...*\n\n-  We will add the following to the final version of the paper.\n- Our choices are based on an established method called TCLR ([38] of our paper). We adapt two loss functions from here to help SwiFT achieve a better understanding of temporal dynamics by a) **distinguishing fMRI scans from different subjects** and b) **distinguishing fMRI scans from different timestamps from the same subject**.\n- The instance contrastive loss accomplishes a) by considering fMRI scans from the same subject as positive pairs and fMRI scans from a different subject as negative pairs. \n- The local-local contrastive loss accomplishes b) by considering different augmentations of the same scan as positive pairs and fMRI scans from different timestamps as negative pairs.\n- We hope this clarifies the use of our contrastive loss. \n\n> *... performing ablation studies to \"substantiate their modeling choices.\" ...*\n\n- Although we had originally intended the phrase \"substantiate our modeling choices\" to refer to our qualitative design choices such as the switch to the absolute position embedding scheme, we understand the reviewer's concerns. \n- Here are some of the test results on the HCP sex classification task. Tweaking the hyperparameters had surprisingly little effect, and after toying on a larger ABCD dataset, we settled on the architecture mentioned in the paper while balancing the model's performance, efficiency, and memory usage.\n\n| Temporal Window Size $P$ |4| 2 | 6 | 4 | 4 |\n|-|-|-|-|-|-|\n| **Channel Number $C$**|**36**|**36**|**36**|**24**|**48**|\n|Accuracy (%)|$92.9\\pm1.51$|$93.2\\pm2.18$|$93.2\\pm3.22$|$94.1\\pm2.41$|$92.7\\pm1.76$|\n|AUC (%)|$98.0\\pm1.79$|$97.9\\pm1.40$|$98.0\\pm1.65$|$98.6\\pm0.6$|$97.2\\pm1.1$|\n\n>*Additionally, it would be helpful to see the loss curves of models...*\n\n- Thank you for the suggestion, and we have included the loss curves in the general rebuttal PDF (figure 1). We would be happy to share the details in our final version of the paper.\n\n>*Could the authors clarify if the parameters for z-scoring...*\n\n- We originally utilized the mean and variance from the whole data to normalize the regression target labels, but we acknowledge the err in our methodology and modified our normalization to only use the mean and variance from the training data for each split. \n- Even though in practice this should not impact the test results since the same normalization method was used for all baseline models, and the change of the mean and variance due to this switch is minimal (table 1 of PDF), we trained our model using the correct normalization method and posted the training curve to make sure the training process was not impacted due to this change (figure 2 of PDF). We will make sure to update the regression results using the correct normalization method in the final version of our paper.\n\n> *I would expect that using the whole fMRI time series ...*\n\n- We agree that in an ideal scenario where we have access to a very large dataset the reviewer's expectation should be valid, however we believe that the input time sequence length should be considered as a hyperparameter in more limited settings. Please refer to the general rebuttal response for our opinion regarding the matter.\n\n> *...it would be helpful if they could also share their trained models...*\n\n- We agree with the sentiment and would be happy to release our trained models and open-source code alongside the final draft of our paper.\n\n>*Explanation methods like IG are generally brittle...*\n\n- The interpretation method we used for our analysis, which is Integrated gradient with Smoothgrad sQuare (IG-SQ), is known to be more robust [1] thanks to noise tunneling and the squaring of estimates. We believed this interpretation method to be the most appropriate for our model, and we would be happy to try other methods if there are any suggestions.\n\n> *... a shifting window ... limits model expressivity...*\n\n- We understand the concerns, and please refer to the general rebuttal regarding our analysis of the matter.\n\n> *Authors have mentioned that they divided the fMRI data into sub-sequences...*\n\n- We apologize for the confusion. Please refer to our general rebuttal discussing our computing setup.\n\n[1] Hooker, Sara, et al. \"A benchmark for interpretability methods in deep neural networks.\" Advances in neural information processing systems 32 (2019).\n"}, {"review_id": "VSlQtLpMa6", "reviewer": "Reviewer_JwtU", "summary": "The paper focuses on the expansion of the Swin Transformer to a 4D version, enabling it to be trained on 4D fMRI data in an end-to-end manner. Specifically, the proposed method commences by constructing an absolute positional embedding layer across spatially neighboring patches, prior to implementing deep 4D Swin Transformer blocks. To facilitate interaction between the windows, a shifted window MSA technique is employed. The paper also explores the self-supervised pre-training of the method, leveraging instance contrastive loss and local-local temporal contrastive loss. Ultimately, the manuscript presents an evaluation of the method on classification and regression tasks, and conducts an ablation study on the model's efficiency in comparison to the TFF method.", "questions": "The method's evaluation is currently confined to classification and regression. What about what other real-world applications (that might be feasible with the proposed method)? For instance, could there be potential for extending the method to tasks like brain segmentation or image reconstruction?\n\nCould you possibly expand on the concept of integrating 206 spatially neighboring patches into a token? Might it not be more efficient to consider alternative techniques such as tokenization based on brain anatomical regions? \n\nAs for the issue of sex classification, the method either failed to surpass baseline methodologies or the difference in accuracy was negligible. Could you delve deeper into potential reasons for this constraint within the study?", "limitations": "Yes, the paper covers some limitations of the work.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The manuscript is coherently presented, and the concept of adapting the Swin Transformer for representation learning using 4D fMRI data is interesting.\n \nThe paper's advancement into self-supervised learning, employing both local-local and instance contrastive learning, effectively demonstrates that the proposed SwiFT model exhibits robust generalization capabilities in the context of downstream tasks.", "weaknesses": "Swin Transformers primarily restrict self-attention computations to specific sub-windows, which potentially curtails the model's ability to capture information from brain regions that are spatially distant from one another. \n\nThe computational cost associated with the proposed method is substantial. Instead of merely contrasting it with TFF, wouldn't a more comprehensive ablation study on the computational cost, compared to other state-of-the-art methods, offer a broader perspective? \n\nThe method's experimental evaluation appears restricted (please refer to the questions section for further details).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and providing us with helpful suggestions. Here is our response to some of the questions and concerns raised:\n\n> *Swin Transformers primarily restrict self-attention computations to specific sub-windows...*\n\n- We agree that it is important to address the issue. Please refer to the general rebuttal for our analysis on the matter.\n\n> *...wouldn't a more comprehensive ablation study on the computational cost ... offer a broader perspective?*\n\n- We could compare the computational cost our model against other baseline ROI-based models, although the comparison would be heavily favored towards the ROI-based models. If we compare the computation cost of the model given the input from a subject, the ROI-based models would have to only process a $180 \\times 180$ input while SwiFT and TFF has to directly process a $96 \\times 96 \\times 96 \\times 1200$ input. We could take account of the extra one-time preprocessing step that ROI-based models require, although even at that point it would be unfair to directly compare the computation cost of those two types of models. \n\n> *... What about what other real-world applications...*\n\n- Thank you for the good suggestion. SwiFT can be extended to solve many important psychological and neuroscientific problems. Let us introduce three feasible applications for SwiFT:\n    -\tPredicting various task-related brain activity from resting-state functional connectivity is a well-established task that has received much attention [1]. The task holds a substantial potential value, particularly for patients or children who encounter challenges performing complicated tasks within fMRI. With SwiFT demonstrating exceptional predictive performance, we anticipate an improved ability to forecast task-related brain activity using the raw resting-state BOLD signal in comparison to functional connectivity methods.\n    -\tSwiFT can also be utilized to execute functional segmentation. Previous researchers have extracted spatially independent components from fMRI based on multivariate decomposition methods such as independent component analysis (ICA). Coherent functional brain networks can be discovered based on the relationship between these spatial components. By analyzing the brain dynamics present in resting-state fMRI in a non-linear manner, SwiFT is expected to generate higher quality component maps than existing methods.\n    -\tSwiFT can also be extended to brain decoding tasks, where information about what a subject sees or hears is reconstructed from the fMRI. Recent decoding studies have shown that it is possible to predict fMRI activity levels in specific brain regions, such as the visual cortex or inferior temporal gyrus, from the features of words learned by a large language model. By utilizing task fMRIs, we can extend the scope of brain decoding to the entire brain and understand how whole brain regions relate to the visual cortex and external stimuli.\n \n> *... consider alternative techniques such as tokenization based on brain anatomical regions?*\n\n-\tThank you for the intriguing suggestion. To clarify, during the initial patch embedding step, 6×6×6=216 neighboring voxels within a patch are embedded into a 36-dimensional token, which is then used as the input for the Transformer. Incorporating information from the brain anatomical regions could enhance the model's ability to effectively learn the brain's structure. The ROI-based tokenization scheme could also help us reduce the number of tokens and make our model more efficient.\n-\tHowever, the reason for our current approach is that ROI-based tokenization based on a specific population-based template image, such as an atlas, may not reflect the individual specificity of the subject, such as age or race. If we ignore these characteristics and define ROI based on a population-based atlas, we may introduce bias into model training. However, if we differentiate atlases based on the demographic attributes of the data, we become constrained to training distinct models for varying atlases. This may limit the ability to generalize data from multiple ages and races in a single model. Therefore, we adopted a simple yet consistent method for SwiFT, which may also be one argument for using a end-to-end model that does not rely on hand-crafted features.\n\n> *As for the issue of sex classification, the method either failed to surpass baseline methodologies...*\n\n- While for each dataset there are baseline models that show comparable results against our model, on a whole there is no baseline model that consistently performs on all three datasets. The ROI-based models perform well on the ABCD dataset but underperforms on the HCP and UKB dataset. The TFF model performs well on the HCP and UKB dataset but underperforms on the ABCD dataset.\n- For the HCP and UKB dataset, the best performing baseline model (TFF) already performs extremely well, with 0.980 and 0.998 AUC for the HCP and UKB dataset respectively. This does not leave much room for improvement, and may be the reason why our model was only able to have minor gains compared to TFF.\n- As for the ABCD dataset, the ROI-based baseline models (BrainNetCNN, VanillaTF, BNT) have showed comparable or slightly better performance compared to SwiFT, although they have underperformed in other datasets. This may be the result of our choice of brain atlas for these models was appropriate for the ABCD dataset but not for the other datasets. We argue that this is one inherent problem of relying on such hand-crafted features, since we have to make sure the application of such feature is appropriate on a case-by-case basis.\n\n[1] Tavor, I. et al. (2016). Task-free MRI predicts individual differences in brain activity during task performance. Science, 352(6282), 216-220.\n\n[2] Hatamizadeh, A. et al. (2021). Swin unetr: Swin transformers for semantic segmentation of brain tumors in mri images. In International MICCAI Brainlesion Workshop (pp. 272-284)."}, {"review_id": "KD5OroVKAE", "reviewer": "Reviewer_ksE2", "summary": "This paper presents SwiFT (Swin 4D fMRI Transformer), a novel Swin Transformer architecture designed for modeling spatiotemporal brain dynamics from high-dimensional 4D functional MRI data. The architecture incorporates a 4D window multi-head self-attention mechanism and absolute positional embeddings, making memory usage and computation efficient. SwiFT outperforms recent state-of-the-art models in several tasks, including predicting sex, age, and cognitive intelligence, based on the evaluation on multiple largest-scale human functional brain imaging datasets. The paper further demonstrates the feasibility of self-supervised pre-training of SwiFT using contrastive loss for improved performance on downstream tasks, marking the first end-to-end learning application of Swin Transformer architecture on dimensional spatiotemporal brain functional data.", "questions": "1. Which ROI system is used to preprocess data for these ROI-based methods, like BNT?\n2. How do you choose the number of layers?", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The ability to utilize pretraining technology and leverage large datasets to aid small datasets is a considerable strength of this paper. This is particularly beneficial for fMRI analysis, where large public datasets and smaller private datasets are common. The paper's approach could help mitigate the challenges of small sample sizes.\n2. The paper marks an important advancement by directly applying deep learning models to fMRI data. This approach could unify various preprocessing pipelines and simplify the analysis process, which is a significant step forward in the field.\n3. The experiments were carried out on three large fMRI datasets, which adds credibility and robustness to the results. By testing their approach on different datasets, the authors ensured that their findings were not limited to a specific dataset, thereby improving the generalizability of the results.", "weaknesses": "1. The paper lacks a proper discussion of its limitations. Understanding the constraints of the presented approach is important for future research and application of the study's findings. The authors should consider addressing potential limitations, caveats, and assumptions made in their methodology to provide a more comprehensive view of the work.\n2. In Section 4.6, the authors had better thoroughly explain why different patterns emerge when the Input Time Sequence Length varies among different tasks. This lack of in-depth discussion and analysis might hinder the reader's understanding of the method's behavior under different conditions. Therefore, it would be beneficial to elucidate these differences.\n3. The discussion in Section 4.4 seems insufficient in terms of relating the study's findings to previous literature. A more detailed comparison with past studies would provide readers with a better understanding of the novelty and contribution of this work. This could include a more explicit discussion of how their findings support previous studies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and raising some important questions. Here is our response:\n    \n>*The paper lacks a proper discussion of its limitations...*\n- Although our model has demonstrated high performance and efficiency compared to existing models, it still presents certain limitations for neuroscientists aiming to apply it to their specific subjects. The fMRI data utilized in this experiment ranges from a minimum volume of 383 (586 megabytes) to 1200 (1.3 gigabytes) per subject. While SwiFT significantly reduces the number of parameters and enhances computational speed compared to the existing fMRI Transformer, training a model on such data necessitates more than 24 gigabytes of GPU resources and storage space for thousands of fMRI images. This can pose a significant challenge for researchers with limited computing resources.\n- Our study is based on a sliding window approach and learning is performed on sub-sequences. This only offers a limited description of its capability to handle long-term dynamics. Processing entire fMRI volumes, which can amount to several gigabytes for multiple subjects, is unfeasible considering limitations in GPU resources. When using a sliding window, the model primarily focuses on the local temporal dynamics of the fMRI, which restricts its ability to learn long-term temporal patterns. Hence, when reporting performance, it is essential to conduct additional validation on the uniformity of the distribution of logit values within sub-sequences for each subject.\n\n> *In Section 4.6, the authors had better thoroughly explain why different patterns emerge ...*\n\n- We apologize for not elaborating as to why this behavior might have emerged. Please refer to the general response for our analysis on the matter.\n\n> *The discussion in Section 4.4 seems insufficient...*\n- Across all age groups, we observed brain regions associated with the default mode network, such as the medial prefrontal gyrus (mPFC), posterior cingulate cortex (PCC), precuneus (PCu), and parietal gyrus. Sex difference in these regions were implicated in the literature ([1], [2], [3]). These regions are known to exhibit concurrent activation when not engaged in specific tasks, where strong functional connection exists between each other.\n- Notably, we found different brain contributors to sex classification across different age groups. Unique regions were observed in young adults (HCP) and older-aged adults (UKB). In youth (HCP), the observation centered on the thalamus and insular cortex. On the other hand, among the older adults (UKB), the focus shifted to the inferior temporal gyrus (ITG) and medial orbitofrontal cortex (mOFC). These regions serve as hubs for various cognitive functions, encompassing multisensory integration (Thalamus), emotional processing (insular cortex), higher-level visual processing (ITG), and decision-making (mOFC), where diverse sensory information is synthesized. This result is in line with the previous findings reporting the sex differences in those regions ([4], [5], [6]). \n- These findings suggest that sex difference is prominently evident in regions characterized by robust connectivity with other areas and involved in substantial information exchange. Moreover, we found unique brain contributors to sex across different age groups. \n\n>*Which ROI system is used to preprocess data for these ROI-based methods, like BNT?*\n\n- As mentioned in **Datasets** of section 4.1 on page 6, we have utilized the HCP MMP1 atlas to preprocess our data for ROI-based methods.\n\n>*How do you choose the number of layers?*\n\n- During our earlier phase of testing, starting with the number of layers suggested by the first Swin Transformer paper ([19] of our paper), we have toyed with different configurations by adding or subtracting some layers, or even omitting the last stage. However we discovered that the original number of layers performed the best in general and decided to keep it while shifting our focus to other hyperparameters. As for the effect of other hyperparameters, please refer to our response to the question from reviewer pTkQ.\n\n[1] Ficek-Tani, B., Horien, C., Ju, S., Xu, W., Li, N., Lacadie, C., ... & Fredericks, C. (2023). Sex differences in default mode network connectivity in healthy aging adults. Cerebral Cortex, 33(10), 6139-6151.\n    \n[2] Ernst, M., Benson, B., Artiges, E., Gorka, A. X., Lemaitre, H., Lago, T., ... & Martinot, J. L. (2019). Pubertal maturation and sex effects on the default-mode network connectivity implicated in mood dysregulation. Translational psychiatry, 9(1), 103.\n\n[3] Weis, S., Patil, K. R., Hoffstaedter, F., Nostro, A., Yeo, B. T., & Eickhoff, S. B. (2020). Sex classification by resting state brain connectivity. Cerebral cortex, 30(2), 824-835.\n\n[4] Wu, X., Lu, X., Zhang, H., Bi, Y., Gu, R., Kong, Y., & Hu, L. (2023). Sex difference in trait empathy is encoded in the human anterior insula. Cerebral Cortex, 33(9), 5055-5065.\n    \n[5] Leming, M., & Suckling, J. (2021). Deep learning for sex classification in resting-state and task functional brain networks from the UK Biobank. NeuroImage, 241, 118409.\n    \n[6] Weis, S., Patil, K. R., Hoffstaedter, F., Nostro, A., Yeo, B. T., & Eickhoff, S. B. (2020). Sex classification by resting state brain connectivity. Cerebral cortex, 30(2), 824-835.\n"}, {"review_id": "PKV3UXst3v", "reviewer": "Reviewer_st1h", "summary": "This paper extends Swin Transformer to 4D brain functional MRI data. Unlike existing deep learning methods in the neuroscience field that either use an ROI-based method (taking as inputs functional connectivity) or a two-step method (3D spatial encoding followed by temporal encoding), the proposed model SwiFT takes as inputs 4D fMRI data and can be trained in an end-to-end manner. This also allows pre-training using contrastive losses. Experiments on multiple large human functional brain imaging datasets indicate the effectiveness of SwiFT compared to prior ROI-based and two-stage methods.", "questions": "1. In Section 3.2 “Instance contrastive loss”, what’s a clip? Please define it.\n2. In page 6 lines 230-231, the authors mention that the average performance across 3 splits are reported. But in Appendix page 2 lines 26-27, the authors state that they report the performance on the test set. Please clarify. Also, is the train-validation-test data split patient-wise (i.e., different splits have different patients)? If not, the tasks are much easier and the results may be over-saturated.  \n3. Ablation studies: What’s the effect of the contrastive losses? Also, the ablation of absolute vs relative position biases is hidden in the Appendix, which should be mentioned in section 3.1 “4D absolute positional embedding”.", "limitations": "Limited performance grain from pre-training is discussed in Appendix. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Originality: this paper is the first attempt at extending Swin Transformer to 4D fMRI data. \n2. Methods are technically sound. Experiments are carefully designed to support the authors’ claims. \n3. While the building blocks of SwiFT exist in the literature, I consider this paper a significant contribution because 1) it extends Swin Transformer to 4D fMRI data and 2) it addresses a technically challenging task of modeling fMRI data in an end-to-end manner. ", "weaknesses": "Due to memory constraint, the fMRI sequence needs to be divided into shorter sub-sequences and the model predictions for sub-sequences are aggregated through averaging. In my opinion, this is a limitation of the method and should be discussed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review and giving us some helpful pointers. Here is our response to the questions and concerns:\n\n>*Due to memory constraint, the fMRI sequence needs to be divided into shorter sub-sequences and the model predictions for sub-sequences are aggregated through averaging. In my opinion, this is a limitation of the method and should be discussed.*\n\n- We understand your concerns regarding the issue, however we believe that there to be a trade-off for using sub-sequences. Please refer to the general rebuttal regarding the matter.\n\n> *In Section 3.2 “Instance contrastive loss”, what’s a clip? Please define it.*\n\n- We apologize for not stating it clearly. A *clip* refers to an fMRI sub-sequence that acts as an input for our model. Our goal for the contrastive pre-training process is to train the model to differentiate clips (inputs) that come from different subjects or different time frames. We would be happy to clarify the term for the final version of our paper.\n\n> *In page 6 lines 230-231, the authors mention that the average performance across 3 splits are reported. But in Appendix page 2 lines 26-27, the authors state that they report the performance on the test set. Please clarify. Also, is the train-validation-test data split patient-wise (i.e., different splits have different patients)? If not, the tasks are much easier and the results may be over-saturated.*\n- Thank you for the suggestion. To clarify our setup, we have created 3 randomly generated splits that separate our subjects to training, validation, and test sets with a 70:15:15 ratio. For example, if we had 1000 subjects in a dataset, for each split, 700 of them would be used for training, 150 would be used for validation, and 150 would be used for testing. For each split separately, the model was trained on subjects from the training set and after each epoch evaluated on the validation set, the best model iteration was chosen based on the performance on the validation set, and the chosen iteration was evaluated on the test set. The evaluation results (on the test set) from the 3 splits were averaged and reported on our paper.\n\n>*Ablation studies: What’s the effect of the contrastive losses? Also, the ablation of absolute vs relative position biases is hidden in the Appendix, which should be mentioned in section 3.1 “4D absolute positional embedding”.*\n\n- Regarding the effect of the contrastive losses, we have conducted experiments comparing a SwiFT model trained from scratch and a contrastive pre-trained SwiFT model in section 4.3 and appendix C.3. \n- Thank you for the suggestion and we would mention the positional embedding experiment in section 3.1 of the final version of our paper."}]

Author rebuttal:

We thank the reviewers for taking their time to review our work and provide constructive feedback that would help us improve our work. Here are our responses to common questions and concerns raised by multiple reviewers.

### Our computational resources (for reviewer bLJJ, pTkQ)

- We have used four NVIDIA A100 GPUs with a distributed data-parallel (DDP) strategy. Using multiple GPUs speeds up the training process by processing subjects in parallel. However, the number of samples (and the amount of memory each GPU can use) each GPU can process does not change as the number of GPUs increases.
- For training, on a single NVIDIA A100 GPU, it takes about 40 minutes for each epoch to train on the HCP dataset,composed of 1,084 subjects each with \~1,200 volumes of ($96\times96\times96$) fMRI scans, totaling \~6 hours to complete training (e.g., sex classification). The GPU profiling tool showed that the GPU was using 17.6 GB of memory during our typical training session with a mini-batch size of 8. Although we have mainly used a GPU with a 40GB memory, we have confirmed that NVIDIA A5000 GPU with a 24GB memory is enough to train our model. Using a DDP strategy with 4 GPUs speeds up the training by a factor of 3\~4, provided that the input data can be loaded from the storage into the GPUs fast enough.
- For evaluation, using the throughput measured on a single NVIDIA A100 GPU (Table 2), the model can process the 1200 frames of a single subject in 0.57 seconds. With a lower-powered A5000 GPU we could still process a subject under a second.
- We would be happpy to add these additional details to the final version of our paper.

### Swin Transformer's ability to capture long-range interactions (for reviewer pTkQ, JwtU)
- As reviewers pTkQ and JwtU have pointed out, owing to the local nature of the attention mechanism, our model might not capture long-range interactions between spatially/temporally separated tokens. We understand the reviewer's concerns, but we believe our model have some built-in mechanisms to learn about long-range communication between tokens. Here we would like to explain how our model may learn about long-range relationships.
- On the spatial dimensions, the Swin Transformer model can cover large ranges on the later parts of its layers through the patch merging step. On our configuration, after the initial patch embedding step where $6 \times 6 \times 6 = 216$ neighboring voxels are embedded into a patch, a $4 \times 4 \times 4$ window can cover a range equivalent to $24 \times 24 \times 24$ voxels from our $96 \times 96 \times 96$ input. However, after the patch merging step where 8 neighboring tokens are merged, the same window can cover a range equivalent to $48 \times 48 \times 48$ voxels. After another patch merging step, on stage 3, this grows to $96 \times 96 \times 96$ which is already the same size as our entire input. From this point on, a single window covers the entire span of our input in the spatial dimension, thanks to the patch merging step.
- On the temporal dimension, where patch merging does not occur on the current version of the model, we needed another method that can ensure long-range interation between tokens. The shifting window allows a token to interact with tokens that are ~10 time frames away with our setup, but to allow longer range interactions we made the model do a full global attention on the final stage as mentioned on section 3.1. Thus, the model processes the entire temporal range on the final stage, allowing it to introduce long-range temporal interactions in the end. We acknowledge that there could be more effective methods to allow long-range interactions, and it would be an exciting problem to tackle when we are dealing with more temporally dynamic datasets as reviewer pTkQ have metioned.

### Limitations regarding the utilization of sub-sequences (section 4.6) (for all reviewers)
- In our setup, instead of inputting the entire fMRI sequence of a subject all at once, we have opted to divide it into 20-frame sub-sequence and use the divided sub-sequences as an input to our model. This was mainly due to memory constraints and to regularize the number of input tokens to our model. We have investigated the effect of changing the length of this input sub-sequence in section 4.6. Here, we would like to elaborate as to why the optimal length seems to change for each dataset and task.
- Inputting a longer time sequence into the model would allow the model to attend to longer sequences at once, at the expense of bloating the model with more parameters and causing it to overfit. Therefore in a setting where we are dealing with a limited number of training data, using a shorter sub-sequece has a possibility to be beneficial to the model's generalization capabilities. 
- We also note that the total amount of information used for our predictions remains the same regardless of the input time sequence length, as the averaged logits from all sub-sequences of the subject is used for the prediction. 
- The datasets used for this work are resting-state fMRI scans, where the image does not change drastically over time. In the future, it would be interesting to investigate the effect of input time sequence length on more temporally dynamic datasets, such as task-based fMRI scans.


Meta-review: This paper introduces SwiFT, a novel Swin Transformer-based architecture for analyzing high-dimensional functional brain MRI data. The reviewers unanimously support the acceptance of the paper, based on that novel application with end-to-end learning approach, as well as strong empirical results. However, there are some limitations pointed out by the reviewers such as limited comparison to some other model types like CNNs and GNNs, high computational resources required for training. Overall, it is a good paper.
Decision: Accept (poster)