Adapting Vision Transformers to Ultra-High Resolution Semantic Segmentation with Relay Tokens

TL;DR

This work tackles ultra-high-resolution semantic segmentation by enabling ViT-based models to reason across both fine local details and broad global context. It introduces Relay Tokens, a lightweight, plug-in mechanism that creates a shared local/global processing pathway, allowing cross-scale information exchange with minimal parameter overhead and without rearchitecting pretrained backbones. Training employs local, global, and consistency losses to ensure accuracy at both scales and coherent cross-resolution predictions. Across diverse datasets and backbones, Relay Tokens deliver consistent mIoU gains (up to 15%) and substantial memory savings, proving to be an efficient and practical alternative to more complex multi-scale designs.

Abstract

Current approaches for segmenting ultra high resolution images either slide a window, thereby discarding global context, or downsample and lose fine detail. We propose a simple yet effective method that brings explicit multi scale reasoning to vision transformers, simultaneously preserving local details and global awareness. Concretely, we process each image in parallel at a local scale (high resolution, small crops) and a global scale (low resolution, large crops), and aggregate and propagate features between the two branches with a small set of learnable relay tokens. The design plugs directly into standard transformer backbones (eg ViT and Swin) and adds fewer than 2 % parameters. Extensive experiments on three ultra high resolution segmentation benchmarks, Archaeoscape, URUR, and Gleason, and on the conventional Cityscapes dataset show consistent gains, with up to 15 % relative mIoU improvement. Code and pretrained models are available at https://archaeoscape.ai/work/relay-tokens/ .

Figures (14)

• Figure 1: Context and Details. Examples of scenes where both the high-resolution details (top-left) and the low-resolution context (bottom) are necessary to correctly predict the true label maps (top right).
• Figure 2: ViT with Cross-Resolution Relay Tokens. We simultaneously process a small high-resolution window $x^\text{local}$ and a larger low-resolution window $x^\text{global}$with a shared network augmented with relay tokens . In each of the $B$ consecutive transformer blocks, relay tokens are first added to the patch tokens of the global window and processed by the transformer. The updated relay tokens are then processed with the tokens of the local window and passed to the next block. The model applies supervision at each scale using losses $\mathcal{L}^\text{global}$ and $\mathcal{L}^\text{local}$, while enforcing consistency between resolutions via a cross-resolution loss $\mathcal{L}^\text{cons}$. We denote shared weights with and independent parameters with .
• Figure 3: Vision Transformer With Relay Tokens. PyTorch code to add relay tokens to a ViT, here '...' stands for standard ViT code. See Appendix for more details. D -- embedding dimensionR -- number of relays
• Figure 4: Predictions at Different Scales.
• Figure 5: Qualitative Results. We visualise a $2048 \times 1024$ region from each dataset, comparing our multi-scale approach to the baseline sliding window. We highlight the extents of the global and local images and mark points of interest with .