Source-linked AI summary
Continual pretraining Prostate MRI ViT Encoder for task alignment
Bilal A. Sidiqi, Yipei Wang, Shaheer U. Saeed, Shiqi Huang, Daniel C. Alexander, Yipeng Hu
TL;DR
Medical vision pretraining has not consistently yielded strong downstream performance, and alignment between pretraining and downstream objectives remains less evident than in language models. The paper continually pretrains a prostate MRI ViT encoder with LoRA on segmentation and analyzes feature similarity with CKA. It reports improvements in three of six downstream tasks, with comparable results in the others, and closer feature similarity to fine-tuned encoders for segmentation correlates with gains.
Problem
General-purpose and domain-specific pretraining have not consistently produced strong medical downstream performance, while alignment between pretraining and downstream objectives remains unclear.
Source: §1 · p. 1Method
The authors use LoRA to continually pretrain a pretrained prostate MRI ViT encoder on segmentation and analyze layer-wise representation similarity with CKA.
Source: §1 · pp. 1–2, §2.4 · p. 2Results
Three of six downstream tasks improve and the others show comparable results; for segmentation, continual-pretrained features are more similar to fine-tuned features than pretrained features, a pattern correlated with observed gains.
Source: §1 · p. 1, §4 · p. 4, §3.2 · p. 4Takeaways & Limitations
For segmentation tasks, greater representational similarity between continual-pretrained and fine-tuned encoders correlates with observed performance gains.
Source: §4 · p. 4, §3.2 · p. 4Future work may explore later-layer Mixture-of-Experts for multi-task continual pretraining and alternative representational similarity metrics.
Source: §4 · p. 4
Abstract
from arXiv · showhide
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.
1. INTRODUCTION
Medical vision pretraining has not consistently produced strong downstream performance, and the alignment between pretraining and downstream objectives remains unclear. This work continually pretrains a prostate MRI ViT encoder and examines whether representational shifts correspond to downstream outcomes.
1. INTRODUCTION
- Independent task models are resource intensive, limit practical deployment, and hinder generalisation to out-of-distribution data..Source: §1 · p. 1
- General-purpose and domain-specific pretrained features have not consistently yielded strong medical downstream performance without extensive task-specific fine-tuning.The passage also notes that existing pretraining strategies may lack balanced representations for shared downstream semantics.Source: §1 · p. 1
- The authors continually pretrain an already pretrained encoder to better align its features with downstream tasks, reporting improvements in several in-domain and out-of-domain tasks despite using a smaller dataset than prior work.Continual pretraining further trains an existing encoder; a cited prior study used datasets ranging from 30k to 1.3M samples.Source: §1 · p. 1
- The study uses prostate segmentation for continual pretraining because it is closely related to lesion and anatomical segmentation, PI-RADS classification, and prostate volume regression.Across six downstream task examples, performance improves in three, including two segmentation tasks.Source: §1 · p. 1
- CKA compares representations from pretrained, fine-tuned, and continual-pretrained encoders to assess shifts toward fine-tuned features and their association with downstream performance differences.The reported correlation holds when a measurable representational similarity shift is observed.Source: §1 · pp. 1–2
2. METHOD
The method adapts a pretrained 3D ViT encoder with LoRA during supervised prostate segmentation, then evaluates the encoder on downstream tasks and compares layer-wise representations with CKA.
2.1. Applying LoRA to a Pretrained Encoder
- LoRA adds trainable low-rank updates to the query, key, value, and output projections in each transformer block of a pretrained 3D ViT encoder.The encoder was pretrained with a self-supervised MAE objective; the update uses a scaling factor and low-rank matrices.Source: §2.1 · p. 2
2.2. Continual Pretraining
- During continual pretraining, frozen encoder weights feed selected hidden states through an adapter and UperNet decoder to produce voxel-level prostate segmentation.The decoder fuses reshaped multiscale feature maps, while LoRA parameters are trained to adapt the encoder features.Source: §2.2 · p. 2
2.3. Downstream Tasks
- Downstream evaluation uses task-specific heads, including a two-layer linear head for classification and regression, UperNet for segmentation, and a super-resolution module for low-resolution T2 inputs.The encoder weights are unfrozen for these evaluations, which include in-domain and out-of-domain datasets.Source: §2.3 · p. 2
2.4. Geometric Feature Analysis Metric
- CKA compares layer-wise representations across pretrained, continual-pretrained, and fine-tuned encoders, using layers 3, 5, 7, and 11 for tasks with performance differences of approximately 0.2% or greater.The metric assesses whether feature representations share similar geometric structure.Source: §2.4 · p. 2
3. EXPERIMENTS
Across classification, regression, segmentation, and super-resolution, outcomes vary by task; layer-wise CKA analysis reports closer similarity to fine-tuned features for lesion segmentation. The experiments also compare datasets and input configurations.
3.1. Dataset and Experimental Setup
- The experiments used PROMIS for continual pretraining and Anatomy, RISK, and UCL Set for downstream tasks; inputs were 64 × 256 × 256 except for super-resolution, which used one-quarter depth.Standard preprocessing and augmentation included normalization, center cropping, rotation, flipping, zooming, and intensity variation.Source: §3.1 · pp. 2–3
3.2. Results
- For PIRADS classification, the baseline led overall AUC by approximately 0.4%, while the continual-pretraining model exceeded it by over 5% at AUC≥3; the baseline led zone classification by 2.3% on average.Source: Table 2 · p. 3, §3.2 · p. 3
- Regression showed a marginal 0.2% decrease versus baseline, and the baseline achieved a 0.1% improvement in super-resolution SSIM.Source: §3.2 · p. 3, Table 3 · p. 4, Table 5 · p. 4
- Continual pretraining improved average Dice by approximately 0.5% for lesion segmentation and 0.1% for prostate zone segmentation versus the baseline.The study used PROMIS for lesion segmentation and Anatomy for prostate zone segmentation.Source: §3.2 · p. 3, Table 4 · p. 4
- The pretrained and continual-pretrained encoders remained highly similar across layers, with similarity gradually decreasing in deeper layers.The analysis reports layer-wise and average CKA for layers 3, 5, 7, and 11.Source: §3.2 · p. 4, Table 6 · p. 4
- For lesion segmentation, fine-tuned representations became less similar to both encoder variants with depth, but deepest-layer similarity to the continual encoder increased by approximately 0.4%.For PIRADS classification, the fine-tuned encoder remained nearly identical to both variants, suggesting globally stable representations rather than layer-specific adaptations.Source: §3.2 · p. 4
4. CONCLUSION
The study reports improved or comparable performance across six downstream tasks and finds that continually pretrained features are more similar to finetuned features for segmentation. Future work could explore multi-task continual pretraining and alternative representational similarity metrics.
4. CONCLUSION
- Continual pretraining yielded improved or comparable performance across all six downstream tasks.For segmentation tasks, continually pretrained features were more similar to finetuned encoder features than pretrained encoder features, a similarity that correlated with observed performance improvements.Source: §4 · p. 4
- For classification tasks, representational geometry changed minimally during finetuning, limiting a clear correlation between representation shifts and performance.The authors suggest this may reflect a stronger focus on fine-grained discriminative details.Source: §4 · p. 4
- Future work may explore Mixture-of-Experts in later encoder layers for multi-task continual pretraining and more generalisable task alignment.Alternative representational similarity metrics may also provide deeper insight into continual pretraining's effects on feature representations.Source: §4 · p. 4
5. COMPLIANCE WITH ETHICAL STANDARDS
The study reports that it followed the Declaration of Helsinki principles and received institutional ethical approval.
5. COMPLIANCE WITH ETHICAL STANDARDS
- The study followed the Declaration of Helsinki principles and received approval from the UCL Hospitals Data Trust Committee.Approval reference: 24-04-2024/IRAS ID 299136.Source: §5 · p. 5