SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table A.3: Performance of various models on the UKB dataset with standard deviation
Method	Dataset: UKB
Sex	Age	Intelligence
ACC	AUC	MSE	MAE	MSE	MAE
XGBoost	79.5±1.28	87.6±0.94	48.8±1.01	5.85±0.046	1.055±0.199	0.816±0.078
BrainNetCNN [30]	86.8±0.19	93.8±0.31	42.7±0.17	5.36±0.113	1.001±0.141	0.800±0.060
VanillaTF [12]	87.0±1.31	95.1±0.37	41.4±1.16	5.26±0.142	0.999±0.144	0.799±0.059
BNT [12]	87.0±1.10	94.8±0.46	39.6±1.07	5.17±0.092	0.998±0.139	0.798±0.058
TFF [17]	96.8±0.25	99.5±0.06	42.1±4.80	5.10±0.331	0.997±0.123	0.783±0.046
SwiFT (ours)	97.7±0.31	99.8±0.04	18.2±0.94	3.40±0.083	0.992±0.105	0.796±0.044
SwiFT To obtain the results detailed in Section 4.2 (manuscript), we trained SwiFT from scratch using the following configuration:
- Optimizer: AdamW using a cosine decay learning rate scheduler with a 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 \times 10^{-5}$  for regression tasks on the HCP dataset and  $10^{-5}$  for regression tasks on the ABCD dataset.
- Mini-batch size: Mini-batch of size 8
- Epochs: 10 epochs of training for the sex classification and intelligence prediction tasks, and a 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 on the test dataset. The model was trained using four NVIDIA A100 GPUs using the distributed data-parallel (DDP) strategy provided by Pytorch Lightning, with a single training session typically lasting from 4 to 30 hours depending on the dataset and whether data augmentation was used during training. The model can be trained on an NVIDIA A5000 GPU with 24GB memory without a problem.
To obtain the results detailed in Section 4.3. (manuscript), the model was pre-trained using the following:
- Optimizer: AdamW 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-training, we fine-tuned the model without a learning rate scheduler, using 1/10 of the initial learning rate used for from-scratch training.
TFF 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 4 A100 GPUs.
We trained TFF with the following training setup:
- Optimizer: Adam optimizer using a cosine decay learning rate scheduler with a 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 of the BNT model between our experiments and the results posted in the original paper [12], despite utilizing the same ABCD dataset.
We trained BrainNetCNN, VanillaTF, and Brain Network Transformer models with the following setup:
- Optimizer: Adam optimizer using a cosine decay learning rate scheduler
- Learning rate: Learning rate of  $5 \times 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. In addition, we conducted 5-fold cross-validation. Hyperparameters are tuned with the following setup:
Maximum depth: Chosen between 3 and 6
- Minimal child weight: Chosen between 1 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 our experiments are as the following:
Table A.4: Performance comparison of SwiFT for different positional embedding methods.
Dataset	Method	Sex	Age	Intelligence
ACC	AUC	MSE	MAE	MSE	MAE
HCP	Relative	89.8±1.87	95.9±1.21	9.0±0.56	2.44±0.116	0.908±0.009	0.775±0.011
Absolute	92.9±1.51	98.0±1.79	8.6±0.57	2.36±0.114	0.903±0.077	0.786±0.030
ABCD	Relative	80.2±1.65	88.9±0.26	N/A	0.936±0.029	0.761±0.013
Absolute	79.3±1.29	87.8±1.31	0.932±0.017	0.756±0.009
UKB	Relative	97.5±0.10	99.8±0.05	19.4±0.53	3.53±0.071	1.019±0.083	0.807±0.035
Absolute	97.7±0.31	99.8±0.04	18.2±0.94	3.40±0.083	0.992±0.105	0.796±0.044
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
- python 3.10.4
pytorch 1.12.1
pytorch-lightning 1.6.5
monai 1.1.0
- neptune-client 0.16.4
scipy 1.8.1
torchvision 0.13.1
- torchaudio 0.12.1
The full dependency list is released alongside the code.
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 a window with dimensions of  $P \times M \times M \times M$ , the 4D relative positional bias  $B \in \mathbb{R}^{P^2 \times M^2 \times M^2 \times M^2}$  for each self-attention head is integrated as
$$
\operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {S o f t M a x} \left(Q K ^ {T} / \sqrt {d} + B\right) V, \tag {1}
$$
where  $Q, K, V \in \mathbb{R}^{PM^3 \times C'}$  are the query, key, value matrices and  $C'$  is the channel number. A parameterized bias matrix  $\hat{B} \in \mathbb{R}^{(2P - 1) \times (2M - 1) \times (2M - 1) \times (2M - 1)}$  is used to calculate the values in  $B$  since there are only  $2P - 1$  or  $2M - 1$  possible position differences for every axis. A 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 a 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 a single A100 GPU. The absolute positional embedding scheme is more memory efficient as it only requires parameters at the beginning of each stage. In addition, the absolute positional embedding scheme is 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.
Figure A.1: Effect of the number of time frames of the input fMRI 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 a 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 a 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 are 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 to other datasets, and thus the model has a lower performance compared to the results posted in Section 4.2 (manuscript).
The performance of the sex 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 on the performance, ranging within one standard deviation. However, for the ABCD 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.
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-sequence 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 are 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 to investigate the effect of input time sequence length on more temporally dynamic datasets, such as task-based fMRI scans.
C.3 Time-window 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 to 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 was then utilized as the final prediction for the respective subject. To ensure that the predictions
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.
Figure A.2: Inner-subject accuracy of sex classification.
As seen in 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 on 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 of 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 are some built-in mechanisms that could learn about long-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 \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 interaction between tokens. The shifting window allows a token to interact with tokens that are about 10 time frames away from our setup, but to allow longer range interactions we made the model do a full global attention on the final stage as mentioned in Section 3.1 (manuscript). Thus, the 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 it would be an exciting problem to tackle when we are dealing with more temporally dynamic datasets.
E Loss Curves
Here we detail the model's performance on the validation set during the training procedure in order to provide a 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.
(a) Validation AUROC per training epoch for sex classification
(b) Validation MSE per training epoch for age prediction
(c) Validation MSE per training epoch for intelligence prediction
Figure A.3: Performance on the validation set during the training procedure.
The biggest limitation of our research is that the advantages of self-supervised pre-training were not 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 in 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, we established the foundation for large-scaling pre-training for 4D fMRI. Consequently, optimizing
self-supervised pre-training methodology for fMRI with larger amounts of datasets are promising future research topic.
G Licences
We used existing assets as follows:
- TFF [17]: https://github.com/GonyRosenman/TFF, No explicit license.
- Brain Network Transformer [12]: https://github.com/Wayfear/BrainNetworkTransformer, MIT License.
- Swin Transformer [19]: https://github.com/microsoft/Swin-Transformer, MIT License.
- MonAI : https://github.com/Project-MONAI/MONAI, Apache License.
- TCLR [40]: https://github.com/DAVEISHAN/TCLR, No explicit license