Hierarchical Vision Transformers for prostate biopsy grading: Towards bridging the generalization gap

Practical deployment of Vision Transformers in computational pathology has largely been constrained by the sheer size of whole-slide images. Transformers faced a similar limitation when applied to long documents, and Hierarchical Transformers were introduced to circumvent it. This work explores the capabilities of Hierarchical Vision Transformers for prostate cancer grading in WSIs and presents a novel technique to combine attention scores smartly across hierarchical transformers. Our best-performing model matches state-of-the-art algorithms with a 0.916 quadratic kappa on the Prostate cANcer graDe Assessment (PANDA) test set. It exhibits superior generalization capacities when evaluated in more diverse clinical settings, achieving a quadratic kappa of 0.877, outperforming existing solutions. These results demonstrate our approach’s robustness and practical applicability, paving the way for its broader adoption in computational pathology and possibly other medical imaging tasks. Our code is publicly available at https://github.com/computationalpathologygroup/hvit . • We show H-ViTs match state-of-the-art prostate cancer grading algorithms when tested on cases from the same center as the training data, while demonstrating stronger generalization to more diverse clinical settings. • We demonstrate the bene t of prostate-speci c pretraining over more generic multi-organ pretraining. • We systematically compare two H-ViT variants and offer concrete guidance on when each is preferable. • We provide an extensive analysis of loss function choices for ordinal classi cation, showcasing the superiority of treating prostate cancer grading as a regression task. • We enhance model interpretability by introducing an innovative approach for combining attention scores across hierarchical transformers which balances task-agnostic and task-speci c contributions.