Stratify or Die: Rethinking Data Splits in Image Segmentation

We thank you for your valuable feedback. We answer the questions posed as follows.

Is the problem NP-Hard?

Creating the best split for segmentation tasks is formally NP-hard because it can be reduced from the multi-way partition [3] problem, a known NP-hard problem. In segmentation, each image contains pixel-level label distributions across multiple classes, and the goal is to partition the dataset into K folds such that the per-fold class distributions closely match the overall dataset, typically by minimizing a distributional distance like the Wasserstein distance. This is analogous to dividing a multiset of weighted elements into K subsets with balanced sums; except here, the "weights" are pixel counts per class and the balance is across multiple classes simultaneously. Computing the Wasserstein distance for a candidate split can be done in polynomial time, placing the problem in NP, but the underlying optimization is as hard as multi-dimensional multi-way partitioning, making it NP-hard. We will clarify this in the paper.

3. https://en.wikipedia.org/wiki/Multiway_number_partitioning

Exploration of ILP or Mixed-ILP

We thank the reviewer for this suggestion. In response, we experimented with formulating the stratification problem using mixed integer quadratic programming (MIQP) by minimizing the Linear-kernel Maximum Mean Discrepancy (L-MMD) between individual fold distributions and the global distribution. Unlike the Wasserstein distance, which requires sorting, L-MMD yields a solvable constrained objective function.

Despite significant effort, the Gurobi solver failed to converge to an optimal solution within the imposed time limit of 12 hours. We report the best solution found (for a single seed) below for comparison. This result was notably worse than those obtained using WDES or even IPS, both in terms of stratification quality and runtime feasibility.

Volume Frequency vs Pixel Frequency

We thank the reviewer for this insightful question. After carefully considering the example provided, we understand the concern to be about not just total pixel frequency per class, but also the distribution of object volumes across samples, specifically, ensuring that folds reflect not only the average volume but also the diversity of object sizes (e.g., x, 2x, 3x) for a given class.

Our current formulation focuses on pixel-level frequency distributions and does not explicitly control for the sample-wise volume profile of each class. While this could be extended, for example by stratifying on quantized volume bins per class, it would introduce an additional layer of complexity to the objective function, likely requiring multi-dimensional matching or joint constraints across volume profiles and class frequencies.

We agree this is an interesting direction and potentially important in settings where object scale is a key factor. However, we have not explored it in this work, and the effect on stratification quality and downstream model performance remains an open question.

Formal Description of WDES

We have included a formal description of the WDES genetic algorithm to the appendix in the manuscript and will be visible in the camera-ready version.

Thank you for your thoughtful comments.

Regarding the questions raised:

Class Co-occurence

The paper from Szymański et al [2] introduced a metric to track class co-occurrence called Label Pair Distribution (LPD) aimed at tracking 2nd order relationship between classes appearing in the same sample. We re-implemented this metric by considering pixel counts and show the results the following table under the column PLPD (Pixel LPD). A previously introduced method, CS-IPS, as a reponse to reviewer PqJg has also been included in the table.

[2] P. Szymański et al., "A network perspective on stratification of multi-label data," First International Workshop on Learning with Imbalanced Domains: Theory and Applications., pp. 22–35, 2017.

Open Source Code

Yes, we intend to release the full codebase as open source upon acceptance, under a permissive open-source license to support reproducibility and further research. The current restriction noted in the supplementary material applies only during the review period.

Using UNet + Resnet34 encoder

We thank the reviewer for this important observation. Our choice of a U-Net with a ResNet-34 encoder was guided by the goal of using a widely adopted, representative architecture to demonstrate the effect of stratification on evaluation stability during training convergence. While we did not conduct extensive experiments across architectures in this submission, we do not expect our main conclusions to be architecture-specific. Our primary focus was to show that stratification leads to more consistent and reliable evaluation.

Scalability and Feasibility of WDES to Large Datasets

We appreciate the reviewer’s comment. While WDES introduces more computational overhead than random splitting or simple heuristics, it is intended for small to medium-sized datasets where sampling variance and class imbalance can undermine evaluation reliability.

For larger datasets, random splits often suffice, as supported by Theorem A.1, which shows that class proportions per fold converge to the global distribution as dataset size increases.

We also point the reviewer to our response to Reviewer PqJg for a detailed discussion of WDES's time and space complexity. Finally, since stratification is typically done once during dataset preparation, we believe the added cost is often acceptable in practice.

Under-representative classes

We thank the reviewer for this thoughtful suggestion. To assess the impact of stratification on rare classes, we analyzed the segmentation performance of three underrepresented classes in PascalVOC: bicycle, boat, and potted plant.

Across all three classes, WDES consistently yields lower standard deviation in accuracy, F1 score, and IoU compared to random and IPS-based splits. This mirrors the overall trend observed in our main results, suggesting that WDES not only improves stability at the global level but also benefits evaluation consistency for rare classes in low-entropy datasets. We will emphasize these findings in the final version of the paper.

Correction 1

We thank the reviewer for pointing this out. Upon rechecking, we realized that it is actually a typographical error. The correct mean PLD for WDES on the Cityscapes dataset is 67, not 671 as originally shown. We have corrected this mistake in the manuscript.

We appreciate the reviewer’s attention to detail and apologize for the oversight.

Correction 2

We thank the reviewer for catching this. The label “PCIS” in Figure 4 is a typo and should correctly read “IPS”, referring to the Iterative Pixel Stratification baseline. We have corrected this mistake and updated the legend in the manuscript.

Thank you for taking the time to provide valuable feedback. We’d like to address your questions and related weaknesses as follows.

How WDES leads to better model selection or generalization

We appreciate the reviewer's question. Our primary concern in this work is to improve the evaluation of models given a single train-test split, which is well-established in the image segmentation benchmarking literature. We used cross-validation only as a means to quantify the variance in evaluation metrics (e.g., accuracy, F1, IoU) across different proposed split that could be generated by each algorithm. Stratification makes performance estimation more reliable, which in turn strengthens downstream tasks such as model comparison and hyperparameter tuning. Without this, fluctuations caused by skewed class distributions can obscure true model performance. Finally, WDES is equally applicable to single train-test splits, where class imbalance in the test set can be even more problematic. In both settings, it provides a principled way to construct evaluation splits that better reflect the underlying dataset.

Potentially stronger iterative baseline

We thank the reviewer for the suggestion. We took the spirit of the suggestion and implemented an alternative iterative strategy that assigns samples to folds by prioritizing class rarity. Specifically, for all classes, samples are ranked based on the number of pixels of the class. We then start with the rarest class to assign all corresponding samples to folds. This process is repeated for the next rarest class until all samples are placed.

Preliminary results show that this variant performs slightly better than IPS on our stratification metrics, though it still falls short of WDES as it doesn't take class distributions inside folds into consideration. We call it CS-IPS in the following tables where we report the performance on the evaluation metrics. We are currently also evaluating its performance on the cross validation experiments on five datasets and hope to report these results before the end of the rebuttal period.

Why was Wasserstein Distance chosen?

We thank the reviewer for this question. We selected Wasserstein distance because it:

• Remains well-defined even when some classes are absent from a fold.
• Captures both how much and how far mass must be moved to align distributions.
• Naturally fits histogram-based distributions from pixel-level label counts.

To ensure our results are not tied to a specific similarity metric, we also evaluate stratification quality using Linear-kernel Maximum Mean Discrepancy (L-MMD) and Energy Distance (L-ED). As shown in the following table, WDES consistently outperforms IPS, random splitting, and the newly introduced CS-IPS on average across these alternative metrics. We note that IPS performs particularly well on the PascalVOC dataset, likely due to its high proportion of single-class samples. These results indicate that the benefits of WDES generalize beyond Wasserstein distance and remain consistent under other well-established distributional similarity measures.

Time and Space Complexity of WDES

We thank the reviewer for this question. The time and space complexity of WDES, as a genetic algorithm, depend on the population size M, number of generations G, number of samples N, number of classes C, and number of folds K.

Each individual encodes an assignment of all N samples into K folds, and the fitness evaluation involves computing the Wasserstein distance over C-dimensional class histograms. This results in a time complexity of O(M · G · K · C), with additional linear overhead in N from crossover and mutation. The space complexity is O(M · N), as we maintain a population of M individuals.

We also note that the actual runtime is sensitive to the choice of mutation and crossover probabilities, which influence how much of the population changes from one generation to the next.

Thank you for your feedback. We'd like to address your main questions as follows:

Unsupervised Domain Adaptation and Distribution Shifts

We agree with the sentiment that distribution shifts can never be fully ruled out in real-world conditions and remain underexplored. However, we see our work as tangential to benchmarks or algorithms aiming to improve performance in the presence of these shifts. Single train-test splits remain common in practice, and any given test split is more representative when applying our stratification method. It thus allows practitioners to evaluate the generalizability of their models as drops in performance occur due to an inability to adapt to similar images rather than to different images. In practice, both the in-distribution performance (capability to interpolate) and the out-of-distribution performance (capability to extrapolate) you mention are of relevance. By stratifying, these two effects can be disentangled as we get a better estimate for in-distribution performance. As a final note, we'd like to point out that even in the domain of computer vision, standard stratification remains an important tool to evaluate and eliminate bias in image classification, in particular in the medical domain [1]. We thus continue to believe that this topic is relevant and will further emphasize these points in the revised paper.

[1] C. Mosquera et al., "Class imbalance on medical image classification: towards better evaluation practices for discrimination and calibration performance," Imaging Informatics and Artificial Intelligence, vol. 34, pp. 7895–7903, 2024.

Class-sensitive ranking-based stratifier

Thank you for your comment. A simple ranking-based algorithm does not account for the fact that, once an image with a specific class is chosen, the other classes also present in that image may negatively influence the overall distribution. This interplay is exactly why we propose the genetic algorithm, and presumably why our proposed IPS, which follows a similar greedy line of thinking, generally does not perform as well.

To explore this further, we implemented an additional iterative stratifier, CS-IPS. It assigns samples to folds by prioritizing class rarity. Specifically, for all classes, samples are ranked based on the number of pixels of the class. We then start with the rarest class to assign all corresponding samples to folds. This process is repeated for the next rarest class until all samples are placed.

We show its results alongside those of random splitting, IPS, and WDES in the following table. While CS-IPS performs slightly better than IPS in some cases, WDES consistently outperforms all three methods on average across datasets. This supports our motivation for using a more flexible and globally optimized approach like the genetic algorithm.

Theorm -> Proposition

We thank the reviewer for this helpful observation. We agree that the convergence argument relies on standard results from the evolutionary algorithms literature, and we will revise the presentation accordingly. Specifically, we will rename Theorem 3.1 to a Proposition and clarify in the text that the result follows from well-established arguments under common assumptions. We will also update the appendix to include appropriate citations to properly acknowledge prior work.