Efficient Fine-Tuning for Brain Foundation Model: Insights into Scaling and Resource Optimization
Poster No:
1126
Submission Type:
Abstract Submission
Authors:
Ruilin Li1, Zijian Dong2, Thomas Yeo3, Juan Helen Zhou1
Institutions:
1National University of Singapore, Singapore, Singapore, 2National University of Singapore, Singapore, N/A, 3Centre for Sleep and Cognition, National University of Singapore, Singapore, Singapore
First Author:
Co-Author(s):
Introduction:
Recent breakthroughs in brain foundation models have revolutionized the analysis of brain activity. However, as these models grow in complexity and capability, a critical challenge emerges: understanding their scaling behavior during adaptation to downstream tasks. In this work, we investigate this scaling behavior during the fine-tuning stage by exploring two critical dimensions: dataset size and the number of tunable layers. Using a state-of-the-art brain foundation model, we conduct experiments varying the amount of downstream task data and the number of trainable layers to understand their impact on model performance. Our analysis reveals that the model mainly benefits most from fine-tuning higher-level transformer features and can achieve optimal performance using 60%-80% of the training data. Our findings contribute to the understanding of scaling laws in brain foundation models and offer practical implications for model adaptation.
Methods:
Our experiments leverage Brain-JEPA (Dong & Li et al., 2024), a state-of-the-art foundation model for fMRI analysis that employs Joint-Embedding Predictive Architecture (JEPA). Brain-JEPA incorporates two key innovations: Brain Gradient Positioning, which enables functional Region of Interest (ROI) localization, and Spatiotemporal Masking, which captures complex patterns across both spatial and temporal dimensions of brain activity. The model's architecture is specifically designed to extract and predict abstract representations from fMRI data, providing a robust foundation for analyzing neural activity patterns.
We utilize the 86M-parameter version of Brain-JEPA, pretrained on 80% of the fMRI data from 40,162 participants in the UK Biobank dataset (Miller et al., 2016), with the remaining 20% reserved for held-out evaluation. We fine-tune the pretrained model using varying proportions (20%, 40%, 60%, 80%, and 100%) of the held-out data. On the other hand, we investigate layer-wise adaptation by fine-tuning different numbers of layers (0-12) while keeping the remaining layers frozen, using the complete held-out dataset. We perform demographic prediction (age and sex) for our downstream evaluation. We conducted all experiments using 5 random seeds and reported the average results.
The data was preprocessed following the pipeline established by Di Biase et al. The fMRI data underwent corrections for EPI susceptibility, gradient distortion, and motion, followed by grand-mean intensity normalization and high-pass temporal filtering. An ICA-FIX procedure was then applied to remove remaining artifacts.
We utilize the 86M-parameter version of Brain-JEPA, pretrained on 80% of the fMRI data from 40,162 participants in the UK Biobank dataset (Miller et al., 2016), with the remaining 20% reserved for held-out evaluation. We fine-tune the pretrained model using varying proportions (20%, 40%, 60%, 80%, and 100%) of the held-out data. On the other hand, we investigate layer-wise adaptation by fine-tuning different numbers of layers (0-12) while keeping the remaining layers frozen, using the complete held-out dataset. We perform demographic prediction (age and sex) for our downstream evaluation. We conducted all experiments using 5 random seeds and reported the average results.
The data was preprocessed following the pipeline established by Di Biase et al. The fMRI data underwent corrections for EPI susceptibility, gradient distortion, and motion, followed by grand-mean intensity normalization and high-pass temporal filtering. An ICA-FIX procedure was then applied to remove remaining artifacts.
Results:
In the analysis of fine-tuning with different portions of downstream dataset, we observe a clear improvement in model performance as the portion size increases, as shown in Figure 1. The results indicate that performance gains are substantial up to 60% of the data, after which the improvement becomes marginal. This suggests that for the tasks under consideration, utilizing 60%-80% of the data is sufficient to effectively capture the necessary task-specific information, minimizing the need for additional data beyond this range.
When freezing varying numbers of transformer blocks, the results reveal a convergence in performance once five or more blocks are frozen, as shown in Figure 2. This highlights that the model primarily relies on adjustments to the higher-level task-specific features of the transformer architecture for these tasks. Freezing the lower-level layers does not significantly impact the results, underscoring the minimal role of fine-tuning low-level features for task-specific adaptation.
When freezing varying numbers of transformer blocks, the results reveal a convergence in performance once five or more blocks are frozen, as shown in Figure 2. This highlights that the model primarily relies on adjustments to the higher-level task-specific features of the transformer architecture for these tasks. Freezing the lower-level layers does not significantly impact the results, underscoring the minimal role of fine-tuning low-level features for task-specific adaptation.
·Figure 1. Performance of sex and age prediction with different portions of the training dataset.
·Figure 2. Performance of sex and age prediction by fine-tuning varying numbers of transformer blocks from the upper to lower levels of the model.
Conclusions:
This work demonstrates that efficient utilization of resources is achievable for demographic prediction by focusing on higher-level feature adjustments and using a moderate portion of the available data. This insight reduces computational demands while maintaining task performance, providing a practical guideline for future applications of the model to similar scenarios.
Modeling and Analysis Methods:
Classification and Predictive Modeling 1
Methods Development
Task-Independent and Resting-State Analysis
Other Methods 2
Keywords:
FUNCTIONAL MRI
Machine Learning
Other - brain foundation model
1|2Indicates the priority used for review
Abstract Information
By submitting your proposal, you grant permission for the Organization for Human Brain Mapping (OHBM) to distribute your work in any format, including video, audio print and electronic text through OHBM OnDemand, social media channels, the OHBM website, or other electronic publications and media.
The Open Science Special Interest Group (OSSIG) is introducing a reproducibility challenge for OHBM 2025. This new initiative aims to enhance the reproducibility of scientific results and foster collaborations between labs. Teams will consist of a “source” party and a “reproducing” party, and will be evaluated on the success of their replication, the openness of the source work, and additional deliverables. Click here for more information. Propose your OHBM abstract(s) as source work for future OHBM meetings by selecting one of the following options:
Please indicate below if your study was a "resting state" or "task-activation” study.
Healthy subjects only or patients (note that patient studies may also involve healthy subjects):
Was this research conducted in the United States?
Were any human subjects research approved by the relevant Institutional Review Board or ethics panel? NOTE: Any human subjects studies without IRB approval will be automatically rejected.
Were any animal research approved by the relevant IACUC or other animal research panel? NOTE: Any animal studies without IACUC approval will be automatically rejected.
Please indicate which methods were used in your research:
Provide references using APA citation style.
2. Dong, Z., Li, R., Wu, Y., Nguyen, T. T., Chong, J. S. X., Ji, F., ... & Zhou, J. H. Brain-JEPA: Brain Dynamics Foundation Model with Gradient Positioning and Spatiotemporal Masking. In The Thirty-eighth Annual Conference on Neural Information Processing Systems.
3. Miller, K. L., Alfaro-Almagro, F., Bangerter, N. K., Thomas, D. L., Yacoub, E., Xu, J., ... & Smith, S. M. (2016). Multimodal population brain imaging in the UK Biobank prospective epidemiological study. Nature neuroscience, 19(11), 1523-1536.