A Study on Context Length and Efficient Transformers for Biomedical Image Analysis

TL;DR

This work investigates how context length affects biomedical image analysis using transformers, addressing the quadratic scaling of self-attention with context length. By systematically varying patch size (ViT) and attention window (Swin) across 2D/3D segmentation, denoising, and classification tasks, the authors reveal that preserving high-resolution information via smaller patches yields substantial performance gains, especially for pixel-level tasks. They further evaluate efficient long-context alternatives Hyena and MambaVision, demonstrating notable speedups (often >80%) and the ability to operate with longer contexts while maintaining or improving accuracy, highlighting their practical value for biomedical imaging pipelines. The findings guide architectural choices for efficient biomedical transformers and identify gaps requiring future exploration, such as broader operator testing, larger datasets, and extended context lengths in real-world clinical settings.

Abstract

Biomedical imaging modalities often produce high-resolution, multi-dimensional images that pose computational challenges for deep neural networks. These computational challenges are compounded when training transformers due to the self-attention operator, which scales quadratically with context length. Recent developments in long-context models have potential to alleviate these difficulties and enable more efficient application of transformers to large biomedical images, although a systematic evaluation on this topic is lacking. In this study, we investigate the impact of context length on biomedical image analysis and we evaluate the performance of recently proposed long-context models. We first curate a suite of biomedical imaging datasets, including 2D and 3D data for segmentation, denoising, and classification tasks. We then analyze the impact of context length on network performance using the Vision Transformer and Swin Transformer by varying patch size and attention window size. Our findings reveal a strong relationship between context length and performance, particularly for pixel-level prediction tasks. Finally, we show that recent long-context models demonstrate significant improvements in efficiency while maintaining comparable performance, though we highlight where gaps remain. This work underscores the potential and challenges of using long-context models in biomedical imaging.

Figures (9)

• Figure 1: Visualization of how context length changes with patch size and attention window size. When using ViT, we use smaller patches to tokenize the input image, resulting in longer context lengths. When using Swin, we use larger windows of attention, resulting in longer context lengths.
• Figure 2: Attention and alternative operators. Left, we show a standard transformer block. Right, we show the operators we evaluate in the transformer blocks: self-attention, Hyena, and MambaVision.
• Figure 3: Task visualization. We visualize a network input and ground truth output for each task. Starting from the upper left and moving clockwise: retinal vessel segmentation, microscopy denoising, pneumothorax classification, pulmonary embolism classification, CMR denoising, and abdominal CT organ segmentation.
• Figure 4: ViT performance. We visualize performance for each task, operator, and patch size with 95% confidence intervals. An X on the x-axis indicates the patch size exceeded our hardware capacity.
• Figure 5: Swin performance. We visualize performance for each task, operator, and patch size with 95% confidence intervals. An X on the x-axis indicates that the window size exceeded our hardware capacity.