Wasserstein-Driven Evolutionary Stratification (WDES)

• The paper introduces WDES, a method that uses a genetic algorithm to minimize the 1-D Wasserstein distance between pixel-level label distributions, ensuring representative dataset splits.
• WDES is shown to achieve lower statistical heterogeneity and reduced fold-to-fold performance variance compared to random splitting and IPS, especially for small or imbalanced datasets.
• The method leverages cumulative Wasserstein distance and gene-level operations while enforcing fold-size constraints, making its optimization both efficient and scalable.

Wasserstein-Driven Evolutionary Stratification (WDES) is a split-assignment method for image segmentation datasets that frames data partitioning as a discrete optimization problem over pixel-label distributions. Introduced in "Stratify or Die: Rethinking Data Splits in Image Segmentation," WDES uses a genetic algorithm to minimize the average 1-D Wasserstein distance between the global class histogram of a dataset and the class histogram of each fold, with the goal of producing more representative train/test or cross-validation splits. In the reported study, WDES is evaluated against random splitting and Iterative Pixel Stratification (IPS), and is associated with lower statistical heterogeneity, lower fold-to-fold performance variance, and particular utility on small, imbalanced, and low-diversity segmentation datasets (Jami et al., 25 Sep 2025).

1. Problem setting and definition

The method is motivated by the observation that random splitting of datasets in image segmentation often leads to unrepresentative test sets, resulting in biased evaluations and poor model generalization. Although stratified sampling has proven effective for addressing label distribution imbalance in classification tasks, extending these ideas to segmentation remains challenging because segmentation labels are pixel-wise, multi-label at the image level, and often severely imbalanced. WDES addresses this by operating directly on pixel-level class counts rather than on image-level labels alone (Jami et al., 25 Sep 2025).

The formal setting uses a dataset

$D=\{(x_n,y_n)\}_{n=1}^{N},$

where each image has pixel-wise labels

$y_n \in \{1,\ldots,C\}^{P^n}.$

A split consists of a partition

$S=(S^1,\ldots,S^K),$

with $K$ folds and target proportions

$r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$

such that $|S^k|\approx r^kN$. Let $P_c$ denote the total number of pixels of class $c$ in the full dataset and $P_c^k$ the total number of pixels of class $c$ in fold $y_n \in \{1,\ldots,C\}^{P^n}.$0.

In this formulation, the central objective is not merely to balance the number of samples per fold, but to make each fold’s pixel-level label distribution resemble the global dataset distribution. This distinguishes WDES from procedures that satisfy sample-count constraints while leaving pixel-label composition substantially heterogeneous.

2. Wasserstein objective and label-distribution matching

WDES treats each fold’s pixel-level class histogram as a discrete distribution over $y_n \in \{1,\ldots,C\}^{P^n}.$1. The global distribution is defined as

$y_n \in \{1,\ldots,C\}^{P^n}.$2

where

$y_n \in \{1,\ldots,C\}^{P^n}.$3

and the fold-specific distribution is

$y_n \in \{1,\ldots,C\}^{P^n}.$4

Using cumulative sums

$y_n \in \{1,\ldots,C\}^{P^n}.$5

the 1-D Wasserstein distance of order 1 between $y_n \in \{1,\ldots,C\}^{P^n}.$6 and $y_n \in \{1,\ldots,C\}^{P^n}.$7 is written in closed form as

$y_n \in \{1,\ldots,C\}^{P^n}.$8

The paper describes this as proportional to the classic definition

$y_n \in \{1,\ldots,C\}^{P^n}.$9

specialized to the discrete 1-D setting (Jami et al., 25 Sep 2025).

The fitness used by WDES is the average over folds, termed the Label Wasserstein Distance (LWD):

$S=(S^1,\ldots,S^K),$0

The stratification problem is therefore

$S=(S^1,\ldots,S^K),$1

where $S=(S^1,\ldots,S^K),$2 is the set of all $S=(S^1,\ldots,S^K),$3-way splits of $S=(S^1,\ldots,S^K),$4 that respect $S=(S^1,\ldots,S^K),$5.

This objective makes explicit that WDES is optimizing similarity of label distributions across dataset splits. A plausible implication is that the method is best understood as a constrained distribution-matching procedure over folds, rather than as a generic evolutionary heuristic for balanced partitioning.

3. Genetic algorithm design

Each individual in the population is represented by a vector

$S=(S^1,\ldots,S^K),$6

where $S=(S^1,\ldots,S^K),$7 indicates which fold sample $S=(S^1,\ldots,S^K),$8 is assigned to. The implementation also stores each sample’s per-class pixel counts $S=(S^1,\ldots,S^K),$9 so that fold totals $K$0 can be computed by summation. Initial individuals are drawn by assigning each sample randomly to fold $K$1 with probability $K$2, followed by random reassignment until each fold’s size is either $K$3 or $K$4. The population size in the paper is $K$5 (Jami et al., 25 Sep 2025).

For each candidate assignment $K$6, the method computes

$K$7

and then evaluates

$K$8

Lower $K$9 is better, and the paper states that no further normalization is used.

Selection uses tournament selection of size $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$0, with $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$1 in practice: to choose one parent, the algorithm samples $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$2 individuals uniformly with replacement and selects the one with lowest fitness. Crossover is gene-level one-point or two-point crossover. In the two-point case, two crossover points $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$3 are chosen at random, offspring are formed by exchanging the middle segment, and a correction step then enforces fold-size constraints by reassigning excess samples from over-populated folds to under-populated folds in proportion to target deficits.

Mutation operates by swap:

$r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$4

applied with probability $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$5, where two distinct sample indices are selected at random. Because mutation swaps fold assignments, it preserves fold-size exactly. In practice, the paper uses $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$6. The algorithm runs for $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$7 generations, with $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$8 in the reported configuration, and employs elitism by always carrying forward the single best individual found so far. After $r=(r^1,\ldots,r^K)\in[0,1]^K,\qquad \sum_k r^k=1,$9 generations, it returns the best $|S^k|\approx r^kN$0.

4. Convergence, optimality, and computational profile

The theoretical analysis is stated under standard assumptions for a non-elitist but ergodic genetic algorithm, with Rudolph 1994 and 1996 cited for convergence in finite search spaces. The relevant assumptions are that the mutation operator has a positive probability of reaching any feasible stratification in one step and that elitist replacement is employed so that the best individual never disappears. Under these conditions, if $|S^k|\approx r^kN$1 is the random set of stratifications examined across $|S^k|\approx r^kN$2 individuals over $|S^k|\approx r^kN$3 generations, and

$|S^k|\approx r^kN$4

then the expected suboptimality of the returned solution satisfies

$|S^k|\approx r^kN$5

Moreover, because mutation is ergodic, as $|S^k|\approx r^kN$6 the probability of visiting every $|S^k|\approx r^kN$7 tends to $|S^k|\approx r^kN$8, so WDES converges in probability to the global minimizer of $|S^k|\approx r^kN$9 over all assignments (Jami et al., 25 Sep 2025).

The reported complexity separates fitness evaluation from genetic operators. Fitness evaluation is $P_c$0 per individual once per generation, yielding $P_c$1. Crossover plus correction is $P_c$2 per new individual, for $P_c$3 total. The overall complexity is therefore

$P_c$4

For typical $P_c$5–$P_c$6 K images, $P_c$7–$P_c$8 classes, $P_c$9, and $c$0, the method is described as feasible on a modern CPU/GPU, and fitness evaluations can be parallelized over $c$1.

The implementation notes also emphasize that the 1-D Wasserstein distance, computed as a cumulative distance, is $c$2 and does not require a full $c$3 optimal transport solver. By caching each sample’s $c$4, fitness can be updated incrementally in $c$5 per sample move, and all $c$6 fitness computations per generation can be vectorized or parallelized, for example through Python multiprocessing or PyTorch. This suggests that the practical bottleneck is not the Wasserstein computation itself but repeated evaluation across a population under fold-size constraints.

5. Statistical heterogeneity metrics and empirical findings

To quantify how balanced a splitting method is, the paper introduces three statistical heterogeneity metrics. The first is Sample Distribution (SD),

$c$7

The second is Pixel Label Distribution (PLD),

$c$8

The third is Label Wasserstein Distance (LWD),

$c$9

Lower SD, PLD, and LWD all indicate more representative splits. In the experiments, WDES achieves the lowest PLD and LWD across five datasets and matches the ideal SD of random splitting, which already enforces $P_c^k$0 by construction (Jami et al., 25 Sep 2025).

The five reported benchmarks span street scenes, medical endoscopy, satellite imagery, and general-purpose semantic segmentation:

Dataset	Size and classes	Reported statistics
PascalVOC 2012	2 913 images, 21 classes	CC=2.48, CU=344, AIR=79, entropy=1.82
CamVid	701 images, 32 classes	CC=17.44, CU=370, AIR=7.56 × 10⁶, entropy=1.94
EndoVis 2018	2 235 images, 12 classes	CC=6.00, CU=1 117, AIR=1 492, entropy=2.38
LoveDA	2 522 images, 8 classes	CC=4.97, CU=1 568, AIR=4.5, entropy=2.61
Cityscapes	2 974 images, 35 classes	CC=16.6, CU=1 412, AIR=467, entropy=3.17

For downstream evaluation, the study trains a U-Net with a ResNet34 backbone using 10-fold cross-validation under each splitting strategy, and measures fold-to-fold standard deviation in accuracy, F1, and IoU. WDES yields the lowest variance in all metrics on low-entropy datasets, specifically PascalVOC and EndoVis, indicating more stable test splits. On high-entropy datasets such as Cityscapes and CamVid, random splits become competitive, because class diversity alone already evens out random folds.

The paper further states that small or low-diversity datasets see the largest PLD/LWD gains, reported as 50–70% relative, together with variance reductions of up to $P_c^k$1 lower standard deviation. As class cardinality CC grows, all methods converge; WDES retains a small edge, but with diminishing returns.

6. Relation to IPS, prior stratification work, and interpretation

WDES is introduced alongside Iterative Pixel Stratification (IPS), described as a straightforward, label-aware sampling method tailored for segmentation tasks. In the heterogeneity results, IPS often overshoots SD because it does not enforce exact sample counts, and its PLD and LWD are worse than WDES. This comparison is important because it separates two notions of balance: exact sample-count control and pixel-label representativeness. The reported experiments indicate that a method may satisfy or nearly satisfy one of these properties without optimizing the other (Jami et al., 25 Sep 2025).

The source paper positions WDES as building on existing stratification concepts. Its references explicitly include Sechidis et al. 2011 on Iterative Stratification for multi-label data and Florez-Revuelta 2021 on Evosplit, a genetic algorithm for multi-label splitting. In this sense, WDES can be situated at the intersection of multi-label stratification and evolutionary search, but specialized for segmentation through pixel-level label accounting and a Wasserstein-based fitness.

A common misconception in this area is that segmentation splits can be treated as a direct extension of classification splits. The framing of WDES argues against that simplification: segmentation combines multi-label image composition with strong class imbalance at pixel resolution, so image-level randomization can yield folds whose label histograms are statistically misaligned even when fold sizes are nominally correct. Conversely, the results also show that WDES is not uniformly dominant in every regime. On high-entropy datasets, random splits become competitive, which suggests that naturally diverse datasets can partially mitigate split bias without explicit optimization.

Within the paper’s own claims, the strongest case for WDES is therefore regime-dependent: its value is most pronounced when datasets are small, imbalanced, or low-diversity, and less pronounced as class diversity and class cardinality increase. That interpretation is consistent with the reported pattern that WDES consistently produces more representative splits, while the magnitude of its empirical advantage narrows when random folds are already stabilized by dataset heterogeneity.