# Stratify or Die: Rethinking Data Splits in Image Segmentation

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 40 references that are in the map.
- First posted on arXiv 25 Sep 2025.

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

- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Deep High-Resolution Representation Learning for Visual Recognition](https://mlsanity.com/paper/1908.07919/) (2019)
- [LoveDA: A Remote Sensing Land-Cover Dataset for Domain Adaptive Semantic Segmentation](https://mlsanity.com/paper/2110.08733/) (2021)
- [Hierarchical Multi-Scale Attention for Semantic Segmentation](https://mlsanity.com/paper/2005.10821/) (2020)
- [UNetFormer: A UNet-like Transformer for Efficient Semantic Segmentation of Remote Sensing Urban Scene Imagery](https://mlsanity.com/paper/2109.08937/) (2021)
- [DoubleU-Net: A Deep Convolutional Neural Network for Medical Image Segmentation](https://mlsanity.com/paper/2006.04868/) (2020)
- [A Novel Transformer Based Semantic Segmentation Scheme for Fine-Resolution Remote Sensing Images](https://mlsanity.com/paper/2104.12137/) (2021)
- [SeMask: Semantically Masked Transformers for Semantic Segmentation](https://mlsanity.com/paper/2112.12782/) (2021)
- [A Network Perspective on Stratification of Multi-Label Data](https://mlsanity.com/paper/1704.08756/) (2017)
- [Using DUCK-Net for Polyp Image Segmentation](https://mlsanity.com/paper/2311.02239/) (2023)
- [AerialFormer: Multi-resolution Transformer for Aerial Image Segmentation](https://mlsanity.com/paper/2306.06842/) (2023)
- [Stratified Sampling for Extreme Multi-Label Data](https://mlsanity.com/paper/2103.03494/) (2021)
- [EffiSegNet: Gastrointestinal Polyp Segmentation through a Pre-Trained EfficientNet-based Network with a Simplified Decoder](https://mlsanity.com/paper/2407.16298/) (2024)
- [EvoSplit: An evolutionary approach to split a multi-label data set into disjoint subsets](https://mlsanity.com/paper/2102.06154/) (2021)
- [Novel split quality measures for stratified multilabel Cross Validation with application to large and sparse gene ontology datasets](https://mlsanity.com/paper/2109.01425/) (2021)

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

- [Not All Nodes Are Created Equal: Homophily-Aware Stratification for Stable GNN Evaluation](https://mlsanity.com/paper/2609.19210/) (2026)

## Abstract (arXiv)

Random splitting of datasets in image segmentation often leads to unrepresentative test sets, resulting in biased evaluations and poor model generalization. While stratified sampling has proven effective for addressing label distribution imbalance in classification tasks, extending these ideas to segmentation remains challenging due to the multi-label structure and class imbalance typically present in such data. Building on existing stratification concepts, we introduce Iterative Pixel Stratification (IPS), a straightforward, label-aware sampling method tailored for segmentation tasks. Additionally, we present Wasserstein-Driven Evolutionary Stratification (WDES), a novel genetic algorithm designed to minimize the Wasserstein distance, thereby optimizing the similarity of label distributions across dataset splits. We prove that WDES is globally optimal given enough generations. Using newly proposed statistical heterogeneity metrics, we evaluate both methods against random sampling and find that WDES consistently produces more representative splits. Applying WDES across diverse segmentation tasks, including street scenes, medical imaging, and satellite imagery, leads to lower performance variance and improved model evaluation. Our results also highlight the particular value of WDES in handling small, imbalanced, and low-diversity datasets, where conventional splitting strategies are most prone to bias.