SwiFT: Swin 4D fMRI Transformer

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.

Strengths

1. Originality: this paper is the first attempt at extending Swin Transformer to 4D fMRI data. 2. Methods are technically sound. Experiments are carefully designed to support the authors’ claims. 3. 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.

Questions

1. In Section 3.2 “Instance contrastive loss”, what’s a clip? Please define it. 2. 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. 3. 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”.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Limited performance grain from pre-training is discussed in Appendix.

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.

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. 2. 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. 3. 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. 2. 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. 3. 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.

Questions

1. Which ROI system is used to preprocess data for these ROI-based methods, like BNT? 2. How do you choose the number of layers?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Yes

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.

Strengths

The manuscript is coherently presented, and the concept of adapting the Swin Transformer for representation learning using 4D fMRI data is interesting. The 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. The 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? The method's experimental evaluation appears restricted (please refer to the questions section for further details).

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? Could 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? As 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Yes, the paper covers some limitations of the work.

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. * They seek to predict subject traits and attributes like age, sex, and cognitive scores from the learned representations using three publicly available datasets. * They have succeeded in outperforming state-of-the-art models in these tasks while reducing the computational complexity and increasing throughput in model evaluations. * Additionally, they have demonstrated the usefulness of pre-training such models on downstream tasks. * Lastly, the relationships captured by the model between the subject attributes and the brain regions are consistent with prior work.

Strengths

* The problem statement is well-motivated, and the limitations of prior works (i.e., ROI-based and two-step approaches) are detailed. * The results were replicated across three datasets and various tasks, showcasing this architecture's robustness. * The authors show the learned features are general enough to help downstream tasks through pre-training. * 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. * 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. * 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. * 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. * 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. * 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. * 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.

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)? * 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)? * 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.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

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 * 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. * 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.

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.

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. - 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. - 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. - 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. - 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. - 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. - 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. - The paper does not provide open-source code for SwiFT, which could limit its adoption and reproducibility.

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. - 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. - 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

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.