Gleason grading is the standard method pathologists use to assess prostate cancer aggressiveness by examining tissue patterns under a microscope. The grade directly drives treatment decisions, from active surveillance of low-grade cancer to aggressive surgery or radiation for high-grade disease.
The problem is that Gleason grading is subjective. Studies have documented significant inter-observer variability, meaning different pathologists examining the same biopsy sample frequently assign different grades. This inconsistency can lead to over-treatment of indolent cancers or under-treatment of aggressive ones.
Whole-slide images of prostate biopsies contain enormous amounts of visual data, and human experts must mentally segment and classify tissue regions across an entire slide. This process is time-consuming and cognitively demanding, particularly in high-volume pathology settings where fatigue and time pressure can affect accuracy.
Deep learning models that can automatically segment prostate tissue into Gleason grade regions offer a path toward faster, more consistent grading. However, most existing models are opaque black boxes, which limits pathologist trust and clinical adoption. This paper addresses both the accuracy and the explainability challenges simultaneously.
The proposed TAH U-Net (Transformer-Attention Hybrid U-Net) is a segmentation architecture that combines the strengths of two fundamentally different AI approaches. Convolutional neural networks (CNNs) are excellent at detecting fine-grained local texture features in tissue images. Vision Transformers (ViT) excel at capturing long-range spatial relationships across an entire image, something CNNs inherently struggle with.
The encoder, which extracts features from the input image, uses a ResNet50 CNN backbone followed by a ViT block. Early encoding layers use ResNet50's convolutional filters to capture local histological features like nuclear shape and gland architecture. Later encoding layers use the ViT's self-attention mechanism to understand how these local features relate to each other across larger tissue regions.
The decoder reconstructs a pixel-level segmentation map from these features. TAH U-Net replaces the standard skip connections in classic U-Net with attention-guided skip connections. These attention gates selectively amplify features from encoder layers that are most relevant to the current decoding step, while suppressing irrelevant background tissue signals.
The training process uses a multi-stage guided loss function that applies both binary cross-entropy and Dice loss at multiple decoder stages simultaneously, not just at the final output. This forces the model to learn useful representations at intermediate depths, improving gradient flow and reducing the risk that early layers learn poorly during training.
The model was trained and evaluated on SICAPv2, a publicly available dataset containing approximately 10,000 histopathological whole-slide image patches from 1,500 prostate cancer patients. Each image patch is annotated at the pixel level with one of four tissue classes: non-cancerous tissue and Gleason Grades 3, 4, and 5, where higher grades indicate more aggressive cancer.
The dataset was split into training, validation, and test sets following a patient-level split, meaning all image patches from a given patient appear in only one set. This prevents data leakage where the model could implicitly learn a patient's tissue characteristics from the training set and then perform artificially well on their other patches in the test set.
Training used standard augmentation techniques including random flips, rotations, and color jitter to increase effective dataset size and improve generalization. The model was compared against several established baselines including classic U-Net, UNETR, SwinUNETR, TransUNet, and SegFormer, all trained under identical conditions on the same dataset.
An ablation study tested the contribution of each TAH U-Net component by training variants with the ViT block removed, the attention gates replaced with standard skip connections, or the multi-stage guided loss replaced with a standard single-output loss. This allows the contribution of each individual design choice to be isolated and measured.
TAH U-Net achieved a Dice score of 0.852 on the validation set and 0.824 on the test set, where the Dice score measures the overlap between predicted and ground-truth tissue segmentations on a scale from 0 to 1. These results represent a substantial improvement over all comparison architectures tested under the same conditions.
The nearest competitor, SwinUNETR, a transformer-based architecture using shifted window attention, achieved a Dice score of 0.705. Classic U-Net scored only 0.612, UNETR scored 0.684, and TransUNet scored 0.698. The gap between TAH U-Net and the next best model exceeds 0.12 Dice points, which is a meaningful clinical margin in tissue segmentation accuracy.
Beyond Dice score, TAH U-Net achieved precision of 76.44%, recall of 95.66%, and F1 score of 94.33% on the test set. The high recall figure is particularly important in cancer detection: it means the model misses very few cancerous tissue regions, prioritizing sensitivity over the alternative of leaving cancer unidentified.
The Intersection over Union (IoU) metric, another standard measure of segmentation quality, reached 0.745 on validation and 0.732 on the test set. These results were consistent across Gleason grade classes, indicating the model generalizes well rather than excelling only on the most common grade category.
The ablation study confirmed that all three core architectural innovations contribute meaningfully to performance. Removing the ViT block from the encoder and relying on ResNet50 alone caused the Dice score to drop by approximately 0.06 points, confirming that transformer-based long-range context modeling adds information that convolutional layers cannot provide.
Replacing the attention-guided skip connections with standard skip connections caused a smaller but still significant performance drop. This finding confirms that filtering which encoder features are passed to the decoder, rather than passing everything indiscriminately, helps the decoder focus on diagnostically relevant tissue characteristics.
Removing the multi-stage guided loss and training only on the final decoder output caused another measurable drop in performance, confirming that supervising intermediate decoder stages during training improves the quality of learned representations throughout the entire network, not just at the output layer.
The full model with all three components achieved the highest performance, and no single component was redundant. The combined design captures local texture through ResNet50, spatial context through ViT attention, selective feature routing through attention gates, and robust multi-level training through guided loss, each component addressing a distinct limitation of simpler architectures.
Beyond performance, the paper addresses the clinical adoption barrier of model opacity using explainable AI (XAI) techniques. Three distinct methods were applied: LIME (Local Interpretable Model-agnostic Explanations), Occlusion Sensitivity mapping, and Partial Dependence Analysis (PDP).
LIME explains individual predictions by fitting a simpler interpretable model around each specific input image. It identifies which image regions most influenced the prediction by measuring how the output changes when those regions are masked. In this application, LIME achieved a fidelity score of 0.89, meaning its explanations closely approximate the actual model behavior, and a comprehensiveness score of 0.21, indicating focused rather than diffuse explanations.
Occlusion Sensitivity systematically blocks different parts of the input image and measures how much each blocked region changes the predicted segmentation. High sensitivity in a region indicates the model depends heavily on that area. The resulting sensitivity maps showed that TAH U-Net focuses on glandular architecture and nuclear features, which are precisely the tissue characteristics pathologists use to assign Gleason grades, validating that the model learned clinically meaningful patterns.
Partial Dependence Analysis examines how specific imaging features such as texture gradients and color intensity distributions influence model outputs across many images. This analysis confirmed that the model's behavior tracks known histopathological criteria rather than exploiting irrelevant staining artifacts or background patterns, an important safety check for clinical deployment.
The combination of strong segmentation accuracy and XAI transparency positions TAH U-Net as a practical tool for computer-aided pathology. Rather than replacing pathologists, the model could serve as a second-reader system that pre-segments tissue regions, flags areas of high-grade cancer, and highlights the specific tissue features driving its assessment for the pathologist to review.
The high recall of 95.66% is particularly relevant clinically, as it means the system is unlikely to miss significant cancerous regions. In practice, a system that consistently identifies all grade-4 and grade-5 regions, even if it occasionally over-segments adjacent grade-3 tissue, would be valuable in reducing the risk of diagnostic miss in high-volume clinical settings.
The XAI validation adds a layer of trustworthiness that is missing from most published models. Demonstrating that the model attends to the same tissue features a human pathologist would inspect is a necessary condition for regulatory approval and clinical trust. The LIME fidelity of 0.89 and the occlusion maps showing focus on glandular architecture directly address the explainability requirements being discussed in AI medical device regulation.
Remaining challenges include validating the model across staining protocols from different laboratories, testing on whole-slide images rather than extracted patches, and clinical studies comparing pathologist accuracy with and without the AI assistance. Extending the approach to other cancer types with grading systems, such as bladder or cervical cancer, is a natural next step.