Prostate cancer is the second most common cancer in men worldwide. After a suspicious PSA reading or biopsy, one of the central clinical questions is whether the cancer is low-grade (LG) -- slow-growing and potentially manageable with active surveillance -- or high-grade (HG) -- aggressive and requiring prompt treatment. Getting this distinction right directly determines whether a patient undergoes radical treatment or is spared its side effects.
Multiparametric MRI (mpMRI) combined with the PI-RADS scoring system is now the standard imaging tool for prostate cancer. Radiologists evaluate T2-weighted (T2w), diffusion-weighted, and dynamic contrast-enhanced images to score suspicious lesions from 1 to 5. However, significant inter-reader variability remains, and not all centers have access to expert prostate MRI radiologists.
Deep learning offers a path toward automated, objective grading of prostate cancer on MRI. However, most published models either work in 2D (losing volumetric context), use only a single architecture, or require multiple imaging sequences like T2w plus ADC (apparent diffusion coefficient) maps. A model that works from T2-weighted images alone in full 3D would be particularly valuable given that T2w is universally acquired.
This study from the University of Siena and the EU ProCAncer-I project explores whether 3D Vision Transformers (ViTs) -- the self-attention-based architecture that transformed natural image recognition -- can classify prostate cancer aggressiveness from 3D T2w MRI, and whether combining multiple ViTs in a stacking ensemble can improve performance beyond what any single model achieves.
A Vision Transformer (ViT) divides an image into small fixed-size patches and processes them as a sequence using the self-attention mechanism originally developed for natural language processing. Unlike convolutional neural networks (CNNs), which detect features through local filters, ViTs capture long-range relationships across all patches simultaneously. This makes them particularly powerful for detecting spatial patterns that span large regions of an image.
To handle 3D MRI volumes, the team trained a 3D ViT from scratch (not pretrained on natural images), dividing each volumetric lesion patch into 3D sub-volumes. Each sub-volume becomes a token, and the self-attention mechanism learns which combinations of regions across the full 3D volume are most informative for distinguishing LG from HG cancer. The final classification head is a single fully connected (MLP) layer applied to the class token output.
The team conducted a systematic grid search over architecture parameters: the embedding dimension (d), the MLP hidden size (D), the number of transformer layers (L), and the number of attention heads (k). This produced 18 distinct base model configurations, each trained and evaluated independently. Configurations varied embedding sizes from 32 to 64, attention heads from 4 to 8, and transformer depth from 4 to 8 layers.
The ProstateX-2 challenge dataset provided 112 biopsy-confirmed lesions from 99 patients: 77 low-grade (ISUP grade group 1 and 2) and 35 high-grade (ISUP grade group 3 and above). Each lesion was represented by a 5-slice axial stack centered on the lesion center, cropped to 128 by 128 pixels. Data augmentation was applied to address class imbalance, generating additional synthetic high-grade examples through rotation, flipping, and intensity perturbation.
A stacking ensemble combines the outputs of multiple independently trained base models using a separate meta-learner. In this study, the team concatenated the feature vectors extracted from the second-to-last layer of two or three base 3D ViTs, then fed the combined representation into a single fully connected output layer. This allows the ensemble to learn which combination of base model representations is most predictive, rather than simply averaging outputs.
Starting from 18 base models, the team generated all possible 2-model and 3-model combinations, creating 966 candidate ensembles. Each was evaluated using 5-fold cross-validation on the training set. The top-performing ensemble (configurations 5, 9, and 11) was then re-trained 100 times on bootstrapped 100-sample subsets of the full training data, and all 100 models were evaluated on a completely separate held-out test set.
The study introduced two novel evaluation metrics designed to better reflect clinical requirements. Confident Specificity (CSP) measures the fraction of correct negative (LG) predictions among those made with high confidence, defined as an output probability below 0.3. Confident Sensitivity (CSE) measures the same for correct positive (HG) predictions made with high confidence, defined as probability above 0.7. These metrics reward not just correctness but certainty, filtering out predictions the model was unsure about.
Performance was also measured using standard metrics including AUROC (overall discrimination between LG and HG), AUPRC (precision-recall performance specifically for the HG class, which is clinically the more important class to correctly identify), and the Brier Score (BS) as a calibration metric measuring the mean squared difference between predicted probabilities and true outcomes.
The best-performing ensemble (combining base model configurations 5, 9, and 11) achieved a median AUROC of 0.89 (95% CI: 0.61 to 1.0) and a median AUPRC of 0.87 (95% CI: 0.57 to 1.0) on the held-out test set. The best single base model (configuration 5) achieved AUROC 0.86 but only AUPRC 0.65. The ensemble thus outperformed the base model in identifying high-grade lesions by 22 percentage points in AUPRC, a difference confirmed as statistically significant by Wilcoxon signed-rank test (p less than 0.001).
The ensemble's balanced accuracy was 0.75, with specificity of 0.83 and sensitivity of 0.67. The Confident Specificity (CSP) equaled the classical specificity at 0.83, meaning all correct low-grade predictions were made with high confidence. However, the Confident Sensitivity (CSE) was 0.30, reflecting that many correct high-grade predictions were made with low confidence, indicating the ensemble was less certain when classifying aggressive lesions.
Comparing with the prior literature, the ensemble outperformed all published models using T2w images alone: Jensen et al. achieved AUROC 0.83 with KNN radiomics; Yuan et al. achieved AUROC 0.81 with CNN transfer learning; Bertelli et al. achieved AUROC 0.88 with attention CNN. The ensemble AUROC of 0.89 equaled or exceeded these benchmarks, and was comparable to multimodal approaches using T2w plus ADC (Yuan et al. achieved 0.90 using both sequences).
For the best base model, the AUROC and AUPRC differences versus the ensemble were not significant for AUROC (p equals 0.73), but the AUPRC difference was highly significant (p less than 0.001). This pattern confirms that the ensemble's key advantage is in correctly identifying high-grade cancer specifically, not in general discrimination performance.
A key finding was that the ensemble model's improvement in AUPRC for high-grade lesions came at the cost of reduced calibration for that same class. The ensemble had a higher Brier Score for positive class (BSPC) compared to the base model, meaning its probability estimates for HG predictions were less well-calibrated. In clinical practice, a model that correctly flags HG cancer but assigns an uncertain probability of 0.55 is less actionable than one that assigns 0.90.
At the same time, the ensemble was better calibrated and more confident for low-grade predictions. Its Brier Score for negative class (BSNC) was lower than the base model, and CSP equaled the clinical specificity, meaning every correct benign classification was made with a probability of 0.3 or lower. This asymmetry suggests the ensemble is better suited for ruling out high-grade cancer (high-confidence negative calls) than for confirming it.
An analysis of architecture parameters revealed that models with 8 attention heads and embedding size 32 (configurations 5 and 8) and models with 4 attention heads and embedding size 64 (configurations 4, 7, and 16) appeared most frequently in the top-10 ensemble combinations. Notably, MLP hidden size (parameter d) had minimal impact on performance, suggesting that adding more parameters to the feedforward layers beyond a certain threshold yields diminishing returns.
The stacking ensemble principle proved broadly effective: combining weaker base models into ensembles consistently produced better accuracy than any individual model. This aligns with the general machine learning principle that model diversity -- achieved here through architecture variation -- allows ensemble components to compensate for each other's weaknesses, even when no individual component is particularly strong.
The authors acknowledge three main limitations. First, no hyperparameter optimization was performed during training -- learning rate, batch size, and other training parameters were fixed throughout. Systematic tuning of these values could potentially improve performance further. Second, the ensemble search was limited to 2- and 3-model combinations; larger ensembles might yield additional gains in diversity and performance.
Third, and most substantially, the models were trained on the ProstateX-2 dataset, which at the time of the study was one of the few publicly available prostate MRI datasets with more than 100 lesions. With only 112 lesions (77 LG, 35 HG), the cohort is small by deep learning standards. The class imbalance between LG and HG cases limits the model's ability to learn robust HG-specific features without data augmentation.
The authors frame this work as a preliminary investigation that will serve as the foundation for future research using larger datasets being collected through the EU Horizon 2020 ProCAncer-I project and the Tuscany Region PAR FAS NAVIGATOR project. These initiatives are assembling multi-center prostate MRI datasets that will enable retraining and validation at a much larger scale.
Despite these limitations, the study demonstrates that 3D Vision Transformer ensembles yield state-of-the-art performance for prostate cancer grading from T2w images alone, opening a path toward automated MRI reading tools that do not require contrast injection or multiple imaging sequences, and that could be deployed in clinical settings where access to expert prostate MRI radiologists is limited.