HiViT: Hierarchical attention-based Transformer for multi-scale whole slide histopathological image classification

Abstract

The classification of Whole Slide Images (WSIs) remains a challenging task in computer-aided diagnostics. Though Multi-Instance Learning (MIL) and patch-wise modeling have become the mainstream methods in current research, they rely on a strong assumption that different patches are independent and identically distributed. Therefore, the contextual correlations within the images and across all the patches have been neglected, resulting in inferior performance, particularly in global-level tasks like mutation predictions. In this paper, we propose HiViT, a multi-scale WSI classification Transformer model utilizing the proposed Cross-scale Hierarchical Self-Attention (CHSA) mechanism to aggregate patch features among different scales and capture context information from the entire image with a reasonable memory cost. The CHSA mechanism is developed based on the multiple magnification scannings of WSIs and restricts cross-scale self-attention computation within local-constrained windows to aggregate patch features among different scales and capture context information from the entire image while getting rid of the infeasible memory cost of vanilla Transformers, and finally achieves the first fully trainable Transformer model for the WSI classification problem. Additionally, the Multi-Scale Feature Augmentation method is proposed to enhance permutation invariance on a large scale, which is important in the MIL assumption, while preserving local fine-grained contextual correlations. We extensively evaluated HiViT on several real-world datasets of diverse tasks, and our model outperformed the current state-of-the-art Transformer and MIL-based methods. Code will be available at https://github.com/BUAA-SMART-Med-CV/HiViT.