# Context-Aware Semantic Segmentation via Stage-Wise Attention

arXiv: https://arxiv.org/abs/2601.11310 · HTML: https://mlsanity.com/paper/2601.11310/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 29 of its 54 references that are in the map.
- First posted on arXiv 16 Jan 2026.

## Builds on (most cited first, 29 of 29)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Masked-attention Mask Transformer for Universal Image Segmentation](https://mlsanity.com/paper/2112.01527/) (2021)
- [Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions](https://mlsanity.com/paper/2102.12122/) (2021)
- [Unified Perceptual Parsing for Scene Understanding](https://mlsanity.com/paper/1807.10221/) (2018)
- [Swin Transformer V2: Scaling Up Capacity and Resolution](https://mlsanity.com/paper/2111.09883/) (2021)
- [SimMIM: A Simple Framework for Masked Image Modeling](https://mlsanity.com/paper/2111.09886/) (2021)
- [Multimodal Machine Learning: A Survey and Taxonomy](https://mlsanity.com/paper/1705.09406/) (2017)
- [CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification](https://mlsanity.com/paper/2103.14899/) (2021)
- [PVT v2: Improved Baselines with Pyramid Vision Transformer](https://mlsanity.com/paper/2106.13797/) (2021)
- [DeepGlobe 2018: A Challenge to Parse the Earth through Satellite Images](https://mlsanity.com/paper/1805.06561/) (2018)
- [Deep Deterministic Uncertainty: A Simple Baseline](https://mlsanity.com/paper/2102.11582/) (2021)
- [Boundary IoU: Improving Object-Centric Image Segmentation Evaluation](https://mlsanity.com/paper/2103.16562/) (2021)
- [Focal Modulation Networks](https://mlsanity.com/paper/2203.11926/) (2022)
- [Rethinking BiSeNet For Real-time Semantic Segmentation](https://mlsanity.com/paper/2104.13188/) (2021)
- [CascadePSP: Toward Class-Agnostic and Very High-Resolution Segmentation via Global and Local Refinement](https://mlsanity.com/paper/2005.02551/) (2020)
- [STU-Net: Scalable and Transferable Medical Image Segmentation Models Empowered by Large-Scale Supervised Pre-training](https://mlsanity.com/paper/2304.06716/) (2023)
- [MILD-Net: Minimal Information Loss Dilated Network for Gland Instance Segmentation in Colon Histology Images](https://mlsanity.com/paper/1806.01963/) (2018)
- [Collaborative Global-Local Networks for Memory-Efficient Segmentation of Ultra-High Resolution Images](https://mlsanity.com/paper/1905.06368/) (2019)
- [FLAIR: a Country-Scale Land Cover Semantic Segmentation Dataset From Multi-Source Optical Imagery](https://mlsanity.com/paper/2310.13336/) (2023)
- [Progressive Semantic Segmentation](https://mlsanity.com/paper/2104.03778/) (2021)
- [Ultra-High Resolution Segmentation with Ultra-Rich Context: A Novel Benchmark](https://mlsanity.com/paper/2305.10899/) (2023)
- [Guided Patch-Grouping Wavelet Transformer with Spatial Congruence for Ultra-High Resolution Segmentation](https://mlsanity.com/paper/2307.00711/) (2023)
- [Full Contextual Attention for Multi-resolution Transformers in Semantic Segmentation](https://mlsanity.com/paper/2212.07890/) (2022)
- [Ultra-high Resolution Image Segmentation via Locality-aware Context Fusion and Alternating Local Enhancement](https://mlsanity.com/paper/2109.02580/) (2021)
- [Memory-Constrained Semantic Segmentation for Ultra-High Resolution UAV Imagery](https://mlsanity.com/paper/2310.04721/) (2023)

## Built on by (most cited first, 1 of 1)

- [WeaveEarth: Structured Evidence Construction and Reasoning for Training-Free UHR Remote Sensing Understanding](https://mlsanity.com/paper/2607.10120/) (2026)

## Abstract (arXiv)

Semantic ultra-high-resolution (UHR) image segmentation is essential in remote sensing applications such as aerial mapping and environmental monitoring. Transformer-based models remain challenging in this setting because memory grows quadratically with the number of tokens, limiting either spatial resolution or contextual scope. We introduce CASWiT (Context-Aware Stage-Wise Transformer), a dual-branch Swin-based architecture that injects low-resolution contextual information into fine-grained high-resolution features through lightweight stage-wise cross-attention. To strengthen cross-scale learning, we also propose a SimMIM-style pretraining strategy based on masked reconstruction of the high-resolution image. Extensive experiments on the large-scale FLAIR-HUB aerial dataset demonstrate the effectiveness of CASWiT. Under our RGB-only UHR protocol, CASWiT reaches 66.37% mIoU with a SegFormer decoder, improving over strong RGB baselines while also improving boundary quality. On the URUR benchmark, CASWiT reaches 49.2% mIoU under the official evaluation protocol, and it also transfers effectively to medical UHR segmentation benchmarks. Code and pretrained models are available at https://huggingface.co/collections/heig-vd-geo/caswit