Prostate MRI model improves tasks through continual pretraining
Continual pretraining Prostate MRI ViT Encoder for task alignment
Computer Vision and Pattern RecognitionMachine Learning
Summary
Medical images like prostate MRIs need computer programs that can understand them for different tasks, like finding organs or abnormalities. The paper’s authors discovered that continuing to train a vision model pretrained on general images specifically on prostate MRI data helped the model better prepare for certain tasks. This continual training aligned the model’s internal understanding with the needs of segmentation tasks, improving performance on some tests. They also analyzed how the model’s internal features changed during this process to explain why the improvements happened.
What this means in practice
- •For medical image analysts: Improve prostate MRI image segmentation models by continually pretraining with task-focused adaptation to enhance accuracy and performance.
- •For medical ai developers: Develop more task-aligned vision transformer encoders through continual pretraining to boost downstream task outcomes in medical imaging.
Authors
Bilal A. Sidiqi, Yipei Wang, Shaheer U. Saeed, Shiqi Huang, Daniel C. Alexander, Yipeng Hu
Abstract
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.