Enhanced Diagnostic Performance via Large-Resolution Inference Optimization for Pathology Foundation Models

TL;DR

This work tackles the mismatch between WSIs and fixed-input pathology foundation models by introducing a space- and time-efficient inference framework that sparsifies self-attention with local neighborhood blocks and uses global attention-derived token scores to prune non-informative patches. By combining local windowed attention with a handful of global tokens and a dynamic token-pruning strategy, the approach reduces memory from $O(N^2)$ and accelerates computation, enabling higher-resolution inference under the same GPU budget. Empirical results on PANDA and SegPath demonstrate substantial memory/time savings and improved ROI classification (up to 7.67% in some settings) while maintaining competitive segmentation performance, with ablations confirming robustness to pruning ratio $p$ and window size $w$. The findings highlight a practical path toward deploying high-resolution pathology foundation models in real-world workflows and motivate further integration with vision-language modeling for pathology.

Abstract

Despite their prominent performance on tasks such as ROI classification and segmentation, many pathology foundation models remain constrained by a specific input size e.g. 224 x 224, creating substantial inefficiencies when applied to whole-slide images (WSIs), which span thousands of resolutions. A naive strategy is to either enlarge inputs or downsample the WSIs. However, enlarging inputs results in prohibitive GPU memory consumption, while downsampling alters the microns-per-pixel resolution and obscures critical morphological details. To overcome these limitations, we propose an space- and time- efficient inference strategy that sparsifies attention using spatially aware neighboring blocks and filters out non-informative tokens through global attention scores. This design substantially reduces GPU memory and runtime during high-resolution WSI inference while preserving and even improving the downstream performance, enabling inference at higher resolutions under the same GPU budget. The experimental results show that our method can achieves up to an 7.67% improvement in the ROI classification and compatible results in segmentation.

Figures (4)

• Figure 1: The left figure shows the interplay between the resolution and the linear probing accuracy on the PANDA dataset; The right figure shows the interplay between the resolution and the GPU memory consumption.
• Figure 2: An overview of our inference optimization strategies. Sparse attention masks attention scores from distant tokens, while token pruning removes tokens with the lowest global attention scores.
• Figure 3: Qualitative examples of the sparse attention matrix and pruned image tokens.
• Figure 4: Qualitative examples of the UNI-2's performance on the segmentation task when the input pixel varies.