SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We also used the resting-state fMRI of 1,084 healthy young adults (age = 28.80 ± 3.70 years, 54.4% female) from the Human Connectome Project (HCP) (S1200 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-modality registration, and spatial normalization to 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 3D volume into a shape of  $96 \times 96 \times 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 fMRI 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 [51], 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 a fundamental brain-biology and brain-cognition association. Also, the capability to predict these outcomes can prove the model's capability to process fMRI volumes, possibly leading to the prediction of clinical outcomes of debilitating brain disorders, such as Alzheimer'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 or 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 Mean Absolute Error (MAE) were used to 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 a channel number of  $C = 36$ . The numbers of layers are fixed to  $\{L_1, L_2, L_3, L_4\} = \{2, 2, 6, 2\}$  which corresponds to a model with three stages with 2, 2, 6 consecutive 4D Swin Transformer blocks for each stage and a final stage with 2 consecutive global attention Transformer blocks. In the cases of 4D(S)W-MSA, we set  $P = M = 4$ . The final output of the model is obtained by applying a 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 were 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 a subject, we divided the volume into 20-frame subsequences 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 as a data point, meaning the appropriate loss function was calculated and 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.
Computational complexity The computational complexities of a single global attention Transformer block (denoted as MSA & MLP) and a 4D Swin Transformer block (denoted as W-MSA & MLP) for input with a dimension of  $T \times H' \times W' \times D' \times C'$  can be calculated as
$$
\Omega (\mathrm {M S A} \& \mathrm {M L P}) = 1 2 N C ^ {\prime 2} + 2 N ^ {2} C ^ {\prime} \quad \Omega (\mathrm {W - M S A} \& \mathrm {M L P}) = 1 2 N C ^ {\prime 2} + 2 P M ^ {3} N C ^ {\prime},
$$
in which the number of tokens  $N = TH'W'D'$ . In practice, on Stage 1 of SwiFT, setting the values used for the experiments  $C' = 36$ ,  $T = 20$ ,  $H' = W' = D' = 16$ ,  $P = M = 4$ , the two terms  $12NC'^2$  and  $2PM^3 NC'$  are balanced with the second term only being 1.19 times the first term.
Table 1: Performance comparison to baselines on classification and regression tasks
Method	ABCD	HCP	UKB
Sex	Intelligence	Sex	Age (year)	Intelligence	Sex	Age (year)	Intelligence
ACC	AUC	MSE	MAE	ACC	AUC	MSE	MAE	MSE	MAE	ACC	AUC	MSE	MAE	MSE	MAE
XGBoost	69.5	76.7	0.977	0.770	68.5	75.5	14.3	3.12	0.991	0.813	79.5	87.6	48.8	5.85	1.055	0.816
BrainNetCNN[30]	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.001	0.800
VanillaTF[12]	77.4	85.1	0.961	0.764	77.9	85.2	12.5	2.95	0.987	0.812	87.0	95.1	41.4	5.26	0.999	0.799
BNT[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
TFF[17]	73.8	80.2	0.968	0.768	92.5	97.5	13.8	3.11	0.953	0.795	96.8	99.5	42.1	5.10	0.997	0.783
SwiFT (ours)	79.3	87.8	0.932	0.756	92.9	98.0	8.6	2.36	0.903	0.786	97.7	99.8	18.2	3.40	0.992	0.796
(a) Results for HCP
(b) Results for ABCD
Figure 3: Effect of UKB pre-training on (a) HCP and (b) ABCD intelligence prediction tasks.
Compared to this, with global attention the  $2N^{2}C^{\prime}$  term becomes 379 times larger than the  $12NC^{\prime 2}$  term, taking up most of the computation budget and creating a bottleneck. For successive stages,  $N$  is reduced by a factor of 8, and  $C^\prime$  is increased by a factor of 2, resulting in Stage 1 being the most computationally expensive.
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, we 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-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 to reduce the dimensionality of fMRI volumes, which are then passed to a transformer encoder layer for temporal processing. It 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 to 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 as its input, due to memory constraints, we also used the technique of dividing the volume into 20-frame sub-sequences described in Section 4.1.1.
4.2 Classification and Regression Results
In Table 1, 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 ROI-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 a large room for improvement on the UKB intelligence prediction task. We believe our quantitative results clearly underscore the effectiveness of the end-to-end training of our SwiFT model.
Figure 4: Interpretation maps with Integrated Gradients (IG) for sex 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 pretrained on a larger dataset (UKB) was fine-tuned on a smaller dataset (HCP), and a 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  $(\mathcal{L}_{IC})$  and the local-local temporal contrastive loss function  $(\mathcal{L}_{LL})$  such that the training objective is to minimize the sum  $(\mathcal{L}_{IC} + \mathcal{L}_{LL})$ . The feature representations used in the loss calculation were obtained in 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 a 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 at the later stage of fine-tuning. This initial worse performance might result from the sub-optimal training hyper-parameters for fine-tuning, such as the sub-optimal initial learning rate. Overall, we believe our results suggest that the contrastive pre-training of SwiFT on a larger dataset has a promising effect on improving downstream performance, and we plan to investigate more extensively on this topic in our future research.
4.4 Interpretation Results
Using an Integrated Gradient with Smoothgrad sQuare (IG-SQ) implemented in 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, we 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 were observed with a broader and higher intensity in mPFC, while showing unique brain regions such as the thalamus and insular cortex. In 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 \times 96 \times 96$ .
(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-add 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 A100 GPU. For an accurate measurement, the throughput was measured using synchronized timing with an 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 1.
4.6 Effect of Input Time Sequence Length
To justify our use of a 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 hyper-parameters, such as 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 is a challenging task considering the MAE of about 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.
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.
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 high-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 state-of-the-art ROI-based and two-step based methods in the domain. Furthermore, compared to TFF, a Transformer-based baseline that takes raw fMRI as input, SwiFT 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-training SwiFT on large-scale datasets.
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-0012704 and used resources of the National Energy Research Scientific Computing Center, a DOE Office of Science User Facility supported by the Office of Science of the U.S. Department of Energy under Contract No. DE-AC02-05CH11231 using NERSC award ASCR-ERCAP0023081. Also, this research used resources of the Argonne Leadership Computing Facility, which is a U.S. Department of Energy Office of Science User Facility supported under Contract DE-AC02-06CH11357.
This work was also supported in part by the National Research Foundation of Korea (NRF) grant funded by the Korean government (MSIT) [No. 2021R1C1C1006503, 2021K1A3A1A2103751212, 2021M3E5D2A01022515, RS-2023-00265406, 2021M3E5D2A01024795, 2021R1A2C2007884], by Creative-Pioneering 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) [No.2021-0-01343, No.2021-0-02068, No.2022-0-00113, No.2022-0-00959]. Taesup Moon is also supported in part by ASRI / INMC at Seoul National University.
This research has been conducted using data from UK Biobank, a major biomedical database (https://www.ukbiobank.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://abcdstudy.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 U01DA041048, U01DA050989, U01DA051016, U01DA041022, U01DA051018, U01DA051037, U01DA050987, U01DA041174, U01DA041106, U01DA041117, U01DA041028, U01DA041134, U01DA050988, U01DA051039, U01DA041156, U01DA041025, U01DA041120, U01DA051038, U01DA041148, U01DA041093, U01DA041089, U24DA041123, U24DA041147. A full list of supporters is available at (https://abcdstudy.org/federal-partners.html). A listing of participating sites and a complete listing of the study investigators can be found at (https://abcdstudy.org/consortium_membership/). ABCD consortium investigators designed and implemented the study and/or provided data but did not necessarily participate in 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://dx.doi.org/10.15154/1503209. DOIs can be found at https://nda.nih.gov/abcd/abcd-annual-releases.html. Data were provided [in part] by the Human Connectome Project, WU-Minn Consortium (Principal Investigators: David Van Essen and Kamil Ugurbil; 1U54MH091657) 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.
[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, 41(3):1168-1176, July 2008.
[3] Emily S Finn, Xilin Shen, Dustin Scheinost, Monica D Rosenberg, Jessica Huang, Marvin M Chun, Xenophon Papademetris, and R 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 Ossenkoppele, Oskar Hansson, and Michael Ewers. Functional brain architecture is associated
[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 J Moser, Nancy C Andreasen, and Peg Nopoulos. Sex differences in parietal lobe morphology: relationship to mental rotation performance. *Brain and Cognition*, 69(3):451-459, 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, 9(1):44-55, 2012.
[65] Paul M Macey, Nicholas S Rieken, Rajesh Kumar, Jennifer A Ogren, Holly R Middlekauff, Paula Wu, Mary A Woo, and Ronald M Harper. Sex differences in insular cortex gyri responses to 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–835, 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-Michel De Rainville, Chi-Hung Weng, Abner Ayala-Acevedo, Raphael Meudec, Matias Laporte, et al. imgaug. https://github.com/aleju/imgaug, 2020. Online; accessed 01-Feb-2020.
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 on 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 1 (manuscript) along with the standard deviation among the three splits. From the tables, we can conclude that SwiFT outperforms all baseline 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 on the ABCD dataset due to the narrow age range of 9 to 11 years, making it hard to obtain meaningful results from the experiments.
Table A.1: Performance of various models on the ABCD dataset with standard deviation
Method	Dataset: ABCD
Sex	Intelligence
ACC	AUC	MSE	MAE
XGBoost	69.5±0.59	76.7±0.86	0.977±0.037	0.770±0.016
BrainNetCNN [30]	80.1±0.69	87.9±0.37	0.969±0.042	0.767±0.016
VanillaTF [12]	77.4±2.47	85.1±3.21	0.961±0.050	0.764±0.025
BNT [12]	79.1±0.80	88.9±0.64	0.955±0.058	0.767±0.025
TFF [17]	73.8±1.13	80.2±1.06	0.968±0.024	0.768±0.009
SwiFT (ours)	79.3±1.29	87.8±1.31	0.932±0.017	0.756±0.009
Table A.2: Performance of various models on the HCP dataset with standard deviation
Method	Dataset: HCP
Sex	Age	Intelligence
ACC	AUC	MSE	MAE	MSE	MAE
XGBoost	68.5±3.03	75.5±3.14	14.3±1.61	3.12±0.165	0.991±0.084	0.813±0.032
BrainNetCNN [30]	77.1±2.33	84.9±2.28	12.6±0.74	2.97±0.153	0.984±0.034	0.805±0.016
VanillaTF [12]	77.9±2.08	85.2±0.89	12.5±1.15	2.95±0.182	0.987±0.039	0.812±0.014
BNT [12]	81.0±3.11	88.0±3.10	12.8±0.89	2.98±0.155	1.001±0.009	0.830±0.014
TFF [17]	92.5±1.12	97.5±1.77	13.8±1.58	3.11±0.200	0.953±0.074	0.795±0.028
SwiFT (ours)	92.9±1.51	98.0±1.79	8.6±0.57	2.36±0.114	0.903±0.077	0.786±0.030