An approach to building foundation models for brain image analysis

Abstract

Existing machine learning methods for brain image analysis are mostly based on supervised training. They require large labeled datasets, which can be costly or impossible to obtain. Moreover, the trained models are useful only for the narrow task defined by the labels. In this work, we developed a new method, based on the concept of foundation models, to overcome these limitations. Our model is an attention-based neural network that is trained using a novel self-supervised approach. Specifically, the model is trained to generate brain images in a patch-wise manner, thereby learning the brain structure. To facilitate learning of image details, we propose a new method that encodes high-frequency information using convolutional kernels with random weights. We trained our model on a pool of 10 public datasets. We then applied the model on five independent datasets to perform segmentation, lesion detection, denoising, and brain age estimation. Results showed that the foundation model achieved competitive or better results on all tasks, while significantly reducing the required amount of labeled training data. Our method enables leveraging large unlabeled neuroimaging datasets to effectively address diverse brain image analysis tasks and reduce the time and cost requirements of acquiring labels.

Keywords: brain; deep learning; foundation models; neuroimaging.

Fig. 1.

Proposed method. $\mathbf{F}$: Frequency encoding block, where ${\mathit{f}}_{\mathbf{1}}\text{-}{\mathit{f}}_{\mathit{n}}$ denote different frequency encodings and ${\mathit{P}}_{\mathbf{1}}$ is positional encoding. $\mathbf{M}$: Positional encoding and masking, where ${\mathit{P}}_{\mathbf{2}}$ is positional encoding and $\mathbf{B}$ denotes a Bernoulli random process used to mask the sequence. $\mathbf{T}$: Vision transformer block. $\mathbf{U}$: Unmasking. $\mathbf{R}$: Reshaping.

Fig. 2.

Example test images reconstructed by the proposed method. Left: a T2 image from the dHCP dataset; Right: a T1 image from the HCP dataset.

Fig. 3.

(a) Plots of DSC for our foundation model and competing methods on Task 1. (b) Example results for the foundation model fine-tuned with 15 labeled images. (Green voxels: correct segmentation, blue: over-segmentation, red: under-segmentation.)

Fig. 4.

Example CST, MCP, and OPR segmented with the foundation model fine-tuned on 15 images. Green: correct; blue: false positive; red: false negative. (Higher magnification image in supp. material.)

Fig. 5.

Example dMRI denoising results by the foundation model in Task 3. Left: noisy; right: denoised.(Higher magnification image in supp. material.)