Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Masked Attention as a Mechanism for Improving Interpretability of Vision Transformers

Clément Grisi, Geert Litjens, Jeroen van der Laak

TL;DR

This paper tackles the interpretability and robustness challenges of Vision Transformers on whole-slide images by introducing a tissue-aware masked attention mechanism. It embeds a three-stage Hierarchical Vision Transformer to process WSIs at macro, meso, and slide scales, and uses a tissue-per-patch mask to zero out background contributions in self-attention via $pct$-based masking. On the PANDA prostate dataset, masked self-attention achieves comparable ISUP grading performance to plain self-attention while delivering significantly more coherent, background-free attention heatmaps. The approach enhances the reliability and clinical relevance of ViT-based models in digital pathology, with potential to reduce artefacts in heatmaps and improve diagnostic contextualization.

Abstract

Vision Transformers are at the heart of the current surge of interest in foundation models for histopathology. They process images by breaking them into smaller patches following a regular grid, regardless of their content. Yet, not all parts of an image are equally relevant for its understanding. This is particularly true in computational pathology where background is completely non-informative and may introduce artefacts that could mislead predictions. To address this issue, we propose a novel method that explicitly masks background in Vision Transformers' attention mechanism. This ensures tokens corresponding to background patches do not contribute to the final image representation, thereby improving model robustness and interpretability. We validate our approach using prostate cancer grading from whole-slide images as a case study. Our results demonstrate that it achieves comparable performance with plain self-attention while providing more accurate and clinically meaningful attention heatmaps.

Masked Attention as a Mechanism for Improving Interpretability of Vision Transformers

TL;DR

This paper tackles the interpretability and robustness challenges of Vision Transformers on whole-slide images by introducing a tissue-aware masked attention mechanism. It embeds a three-stage Hierarchical Vision Transformer to process WSIs at macro, meso, and slide scales, and uses a tissue-per-patch mask to zero out background contributions in self-attention via -based masking. On the PANDA prostate dataset, masked self-attention achieves comparable ISUP grading performance to plain self-attention while delivering significantly more coherent, background-free attention heatmaps. The approach enhances the reliability and clinical relevance of ViT-based models in digital pathology, with potential to reduce artefacts in heatmaps and improve diagnostic contextualization.

Abstract

Vision Transformers are at the heart of the current surge of interest in foundation models for histopathology. They process images by breaking them into smaller patches following a regular grid, regardless of their content. Yet, not all parts of an image are equally relevant for its understanding. This is particularly true in computational pathology where background is completely non-informative and may introduce artefacts that could mislead predictions. To address this issue, we propose a novel method that explicitly masks background in Vision Transformers' attention mechanism. This ensures tokens corresponding to background patches do not contribute to the final image representation, thereby improving model robustness and interpretability. We validate our approach using prostate cancer grading from whole-slide images as a case study. Our results demonstrate that it achieves comparable performance with plain self-attention while providing more accurate and clinically meaningful attention heatmaps.
Paper Structure (15 sections, 6 figures, 3 tables)

This paper contains 15 sections, 6 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Proposed Method
  3. Hierarchical Vision Transformer.
  4. Masked Attention.
  5. Experimental Results
  6. Dataset.
  7. Data Preprocessing & Evaluation Metric.
  8. Prostate Cancer Grading.
  9. Model Interpretability.
  10. Conclusion
  11. Architecture Overview
  12. Masked Attention Pseudo Code
  13. Data Preprocessing
  14. PANDA Dataset Details
  15. Stitched Attention Heatmaps

Figures (6)

  • Figure 1: Region-level attention maps.
  • Figure 2: Overview of our Hierarchical Vision Transformer for whole-slide image analysis. This figure illustrates the multi-scale processing of whole-slide images.
  • Figure 3: Example result of data preprocessing pipeline
  • Figure 4: Unrolling a $2048\times2048$ region into non-overlapping $256\times256$ patches
  • Figure 5: PANDA label distribution
  • ...and 1 more figures