Continual pretraining Prostate MRI ViT Encoder for task alignment
Authors: Bilal A. Sidiqi, Yipei Wang, Shaheer U. Saeed, Shiqi Huang, Daniel C. Alexander, Yipeng Hu
Organizations: Department of Medical Physics and Biomedical Engineering, University College London, London, UK · Hawkes Institute, University College London, London, UK · Department of Computer Science, University College London, London, UK
Building foundation models for MRI has emerged as a promising direction for advancing medical image analysis and improving accessibility of task specific models. Existing approaches primarily focus on effective pretraining strategies and evaluating their performance across diverse downstream tasks. While these methods show performance gains, limited attention has been given to understanding task alignment between pretraining and downstream objectives. Unlike language models, where the two stages are naturally aligned, such alignment is less evident in vision based applications. Motivated by this gap, we explore continual pretraining. Specifically, continually pretraining a pretrained ViT encoder using Low-Rank Adaptation (LoRA) on a prostate segmentation task to enhance task alignment and downstream performance. We further perform layer-wise representational analysis using Centered Kernel Alignment (CKA) to assess representational shifts and their correlation with task outcomes. Our experiments show improved performance in 3 of 6 downstream tasks and comparable results in the others. Our Feature analysis reveals that, for segmentation tasks, continually pretrained encoder features exhibit greater similarity to those of the finetuned encoder than the original pretrained encoder, correlating with observed performance gains.
Figures & tables
Figure 1: Overview of the proposed continual pretraining framework, integrating LoRA within a pretrained ViT encoder to continually pretrain the encoder for better task alignment and downstream performance.
Dataset
Data Split
MRI Input
Output
PROMIS
(396,57,113)
T2W, DWI, ADC
Lesion Mask
Anatomy
(412,60,117)
T2W
Prostate Zones Mask
RISK
(718,103,207)
T2W, DWI, ADC
PIRADS Score (2–5)
UCL Set
(1386,178,387)
T2W
Prostate Volume
Table 1: Overview of datasets, data splits, and corresponding input–output configurations.
PIRADS Classification (RISK Data)
AUC
AUC ≥ 3
AUC ≥ 4
Avg Std
ViT Random Weights
58.81
69.98
62.82
2.51
ProFound ViT OMW
61.36
74.76
64.03
2.56
ProFound ViT Ours
60.94
79.92
64.00
2.63
Zone Classification (PROMIS Data)
ViT Random Weights
64.41
–
–
1.97
Table 2: Results for classification tasks.
Prostate Volume Regression (UCL Data)
MSE
Std
ViT Random Weights
0.038
0.028
ProFound ViT OMW
0.013
0.009
ProFound ViT Ours
0.015
0.012
Table 3: Results for regression task.
Lesion Segmentation (PROMIS Data)
Avg Dice
Class Dice
Std
ViT Random Weights
0.216
–
0.186
ProFound ViT OMW
0.248
–
0.203
ProFound ViT Ours
0.253
–
0.206
Zone Segmentation (Anatomy Data)
Avg Dice
Class Dice
Std
Table 4: Results for segmentation tasks.
T2 MRI Super-Resolution (PROMIS Data)
SSIM
Std
ViT Random Weights
0.900
0.009
ProFound ViT OMW
0.931
0.005
ProFound ViT Ours
0.930
0.005
Table 5: Results for super-resolution task.
Encoder Pair ( Z1−Z2 )
L3
L5
L7
L11
Avg.
Pretrained–Continual
0.997
0.996
0.997
0.995
0.996
CKA Similarity for Lesion Segmentation
Pretrained–Finetuned
0.918
0.915
0.893
0.340
0.767
Continual–Finetuned
0.915
0.910
0.890
0.344
0.765
CKA Similarity for PIRADS Classification
Pretrained–Finetuned
1.000
1.000
1.000
0.993
0.998
Table 6: Layer-wise and average CKA similarity between encoder pairs. “Pretrained” denotes the encoder with pretrained weights; “Continual” indicates the pretrained encoder further trained with LoRA on the prostate segmentation task; “Finetuned” refers to the encoder finetuned on downstream tasks (lesion segmentation or PIRADS classification).
UCL Hawkes Institute, University College London, London, UK · Department of Medical Physics and Biomedical Engineering, University College London, London, UK · School of Artificial Intelligence, Beijing University of Posts and Telecommunications, Beijing, China +11
Department of Electronic Technology and Communications, Universidad Autónoma de Madrid, Madrid, Spain · Department of Computing, Imperial College London, London, United Kingdom · Department of Brain Sciences, Imperial College London, London, United Kingdom +3