Title: Stratify or Die: Rethinking Data Splits in Image Segmentation
Venue: NeurIPS 2025
Accepted: yes
Keywords: segmentation

stratification

genetic algorithm

cross-validation
Abstract: 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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDatasets used in machine are often randomly partitioned into disjoint sets that are used as folds for training and testing. This paper introduces methods that try to enforce that the proportion of pixels in each class for each fold approximates the number of pixels in associated classes for the whole dataset. The objective with this is to get more representable folds that yield more representative training results. To test the hypothesis, variance of the performance over the folds is recorded for the different fold-generating processes, and it is shown that the folds generated with the introduced method WDES produces models with lower variance in comparison to the introduced alternative method IPS as well as the classical random split method.\n\nSTRENGTHS AND WEAKNESSES:\n**Quality:**\nThe work is technically sound and backed by reasonable experiments. The questions that come to mind are related to methodology.\n1) It would be interesting to have some discussion on whether the problem Is NP-hard and if there could be a polynomial and non-random method for solving it. \n2) A genetic algorithm is introduced, when, I would assume, a integer or mixed integer problem approach equally well could have been formulated and passed to a standard solver depending on branch and bounds techniques. It would be interesting to see this as a baseline reference.\n3) Whereas the pixel frequency is of importance, I think it would be even more interesting if it handled the distribution of the volumes (number of pixels per sample). That is, say that the dataset contains three distinct volumes for a class: x, 2x and 3x and that they occur equally often (with volume frequency 1/3), which means that the pixel frequency for that class would be 2x. Then it would be nice if the proposed method would differentiate between a set having the volume frequencies (1/2,0,1/2) and a set having the volume frequencies (0,1,0) with respect to the volumes (x,2x,3x). In other words, so that one fold does not only get samples with average volume 2x whereas another fold gets 50% samples with volume x and 50% with volume 3x.\n\n**Clarity:**\nThe writing is excellent, and the ideas are very clearly formulated. Even though the genetic algorithm was explained in text, it would have been nice to include a more formal description of the algorithm because of how central it is to the paper. \n\n**Significance:**\nIf the variance of performance of the number of folds necessary for testing a method in image segmentation can be reduced, then the reliability of the performance is increased. Potentially also the number of necessary folds can be decreased. These types of results are of very high interest to the community.\n\n**Originality:**\nI haven’t seen this kind of approach taken in image segmentation before. I am however not familiar with the literature in other domains of application so I cannot tell how close it is to previous work in other domains of application.\n\nQUESTIONS:\n**Question 1:** The studied problem looks like it has a bin-packing kind of structure – do you think it is NP-hard? \n\n**Question 2:** Why was a genetic algorithm approach used instead of a integer/mixed integer problem approach?\n\n**Question 3:** Would it be easy to adapt the technique to represent different volumes with the right frequency instead of only looking at pixel-frequency? What effect would you suspect this to have?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces 2 new techniques, namely Iterative Pixel Stratification (IPS) and Wasserstein-Driven Evolutionary Stratification (WDES) to address  the shortcomings of existing sampling strategies, such as random splitting, for generation of representative folds from the original dataset for image segmentation. \n\nWhile IPS is a greedy approach where predefined class distribution is achieved based on the number of positive pixels in each fold and inspired by Iterative Stratification (IS) [23], WDES, a genetic algorithm inspired by Evosplit [28], optimizes the class distribution by minimizing the Label Wassertein Distance (LWD) and increasing the similarity between folds and the original dataset. \n\nThe paper evaluates the performances of IPS and WDES with 5 public segmentation datasets by using 3 metrics, and compares them against that of the random splitting. The results show that WDES generally produces better splits than the random splitting and IPS, especially when the evaluation data is a low-entropy segmentation dataset.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n-- Generating representative test sets by splitting the original dataset is crucial for unbiased evaluation. The paper aims to address this common challenge in the field of machine learning, which has not been fully explored yet, especially for image segmentation tasks.\n-- To the best of my knowledge, the proposed methodology is a novel idea. IPS nicely extends the existing IS and WDES looks sensible approach for finding the similar class distributions across folds.\n-- This paper is well-written with a strong structure and thorough content that provides a good level of details. \n-- The evalution is supported by several public datasets from various domains including medical, satellite and street images as well as general-purpose dataset. The paper provides results from 10-fold cross validation evaluated using multiple metrics, namely accuracy, F1 score and Intersection over Union (IoU), along with their means and standard deviations.\n-- In my opinion, the paper presents important insights that may stimulate further discussion in the field.\n\nWeaknesses:\n-- I think there are a couple of points that need to be addressed for better understanding of the proposed method.\n-- Firstly, each evaluation dataset has a large number of classes. However, the provided metrics are the means of the fold means, so the label imbalance is not evaluated in terms of co-occurence imbalance as some label combinations may appear only in one fold, which limits the ability to accurately gauge the true performance of the proposed method. Although the authors mentioned this issue as \"class co-occurence\" in the conclusion section, this important limitation should be thoroughly examined for the full picture.\n-- Secondly, the paper uses a UNet with a resnet34 encoder as its underlying architecture. This is a specific architectural choice by the authors so I wonder whether the results would hold if a different architecture or even a different encoder were used.\n-- Thirdly, the addition of precision and recall into the evaluation would be helpful.\n-- Although WDES might provide better stratification under certain conditions, its computational cost is high. It would be beneficial for the readers if the authors could elaborate on the scalability and the feasibility of applying this method to large datasets.\n-- Another improvement of better stratification might be for the performances of the machine learning models for underrepresentative classes, especially in low-entropy classes. Although this is not necessary, it would be great to examine if WDES can specifically improve the model performance for rare classes in PascalVOC or CamVid.\n-- Finally, I am surprised about the mean PLD of WDES for the Cityscapes dataset. The value is unexpectedly high and warrants further investigation.\n\nQUESTIONS:\nAn important question:\n-- **Will the authors release the code as open-source as the supplementary material says it is only for the academic review?**\n\nBased on the weaknesses above, \n-- Have the authors already evaluated the proposed method using another architecture? If not, could the authors provide some perspective on the effect of the architecture for the evaluation of the stratification method?\n-- Could the authors elaborate on the scalability and the feasibility of applying WDES to large datasets?\n-- What do the authors think about the performance improvements of the models for the underrepresented classes for low-entropy classes with WDES or IPS? Have they done any evaluation in this regard?\n-- Could the authors please comment on the unexpected high mean PLD for WDES for Cityscapes dataset?\n-- Can the authors check the Figure 4 in D Runtime Analysis and make sure the legend is correct (what does PCIS correspond to)?\n\nLIMITATIONS:\nAs mentioned in the weaknesses, I would encourage the authors to investigave the co-occurance imbalance as this is highly relevant for the segmentation tasks. Incorporating different architectures into the evaluation, along with additional metrics would offer a more compherensive understanding of the proposed method. Finally, I would suggest that the authors should add some comments about the application of WDES to large datasets in terms of the scalability and the feasibility, considering its high computation cost.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors introduce two stratification methods for image segmentation: iterative pixel stratification, a pixel-level adaptation of iterative stratification, and Wasserstein-driven evolutionary stratification, a genetic algorithm that minimizes Wasserstein distance between class distributions in dataset splits. These methods address the limitations of random splitting, which can cause unrepresentative folds and biased evaluation, especially in small or imbalanced datasets. Results show that the latter proposed method made more balanced splits and lower variance in accuracy for more reliable model evaluation.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The authors include a proof for how WDES approaches the optimal stratification as the number of generations and the population size increase.\n- The paper is clear, well-written, and provides good motivation for the proposed method.\n- The focus on segmentation is significant since it is a method that is generally used but has limited labelled data due to the cost of pixel-level labels.\n- The authors provide results on five different datasets of vastly varying domains to show the effectiveness of their proposed method.\n\nWeaknesses:\n- IPS performs quite poorly compared to both random splitting and WDES. It is unclear why this was used and why it was not tuned further to use a non-greedy strategy.\n- Random splitting seems to perform quite a bit better than WDES in high-entropy datasets, which restricts the effectiveness of the proposed method to low-entropy datasets.\n- Lower performance variance is the main metric used to see if the proposed method is effective. There seems to be no effect or a slightly negative effect on accuracy when using the method.\n- With larger datasets, the runtime and computational load may be significantly higher for WDES.\n\nQUESTIONS:\n1. Can the authors provide evidence for how WDES leads to better model selection or generalization after doing cross-validation with the folds that WDES found most optimal?\n2. Can the authors potentially improve the IPS method to be more competitive or explain why it was chosen over other potentially stronger iterative baselines?\n3. Is there a reason why Wasserstein difference was chosen as the metric for this problem as opposed to any other similarity metrics?\n4. How do the time and space complexity of WDES grow asymptotically?\n\nAddressing these questions effectively may help raise my score.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors argue that random splitting of image segmentation datasets into training, testing and validation data leads to unrepresentative test sets, resulting in biased evaluations and poor model generalization. Building on existing concepts from the stratification literature, which have addressed label distribution imbalance in classification tasks (in applications outside computer vision), the authors investigate stratification for image segmentation. This has specific challenges   due to the multi-label structure per sample (image) and the severe class imbalance that often occur in segmentation data. Essentially, the authors minimize criteria, which encourage the pixel-level class distributions in the data subsets (e.g., training, validation and testing subsets or folds) to match as closely as possible the class distribution of the overall dataset (Eqs. 2 and 3). The first criterion (Eq. 2) is the L1 norm between distributions and is optimized with the iterative stratification method in [23], which uses a greedy approach. The latter assigns samples to data subsets, to match a given number of pixels per class within each subset. The second criterion is the Wasserstein distance between distributions (Eq.3), optimized with an evolutionary algorithm like Evosplit [28], which evaluates a population of potential subset assignments (individuals), with the fittest individuals undergoing selective crossover and mutations for refinement. The authors evaluate the proposed stratification methods and random splitting across five image segmentation benchmarks, and used the standard deviation of accuracy, F1-score, and Intersection over Union across 10-fold cross-validation as measure of data splitting quality (with lower deviation indicating better performance).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- The paper is clear, technically sound and well written. \n\n- Semantic image segmentation and its evaluation remain important problems. \n\nWeaknesses\n\n-\tMy main concern with this paper is the following: The proposed methods are essentially removing (artificially) the class-distributions shifts between the testing and training data in segmentation, making the problem easier. I believe this does not reflect real and practical settings in image segmentation, where class proportions could change significantly (e.g. due to camera motion) and it is often the case that all possible class proportion are not covered in the training data. Therefore, artificially removing label-distribution shifts could lead to over-optimistic performance evaluations of segmentation algorithms (making the real-world problem much easier). In fact, there is a whole literature on unsupervised domain adaptation (UDA) for segmentation, which addresses specifically the challenges in distribution shifts (both in label distributions and image data distributions), which occur often in practice. In fact, there are even UDA benchmarks for segmentation that simulate such shifts and class imbalance to simulate real situations (and this is essentially the opposite of what this paper is doing). So, I am not sure that the main target of the paper (ensuring consistent class distributions across data subsets) is relevant in practice. This might explain why stratification has not been investigated before for image segmentation. I understand that stratification strategies are well used in other applications domains (other than computer vision).\n\n-\tThe main conclusion following the experiments is not surprising at all: Essentially, the proposed data splits methods yield lower performance variance across splits. Of course, such target data splits are artificially removing shifts between training and testing data, effectively making the problem easier. \n\n-\tI might be missing something but, in this specific case, sample-to-subset assignment by minimizing distances between distributions amounts to a discrete subset assignment problem under cardinality constraints. This is due to the fact that all the subsets have known cardinality (so, we just need to enforce constraints on class cardinality, not class distributions). Discrete subset assignment with cardinality constraints could be solved with a simple ranking, unless I am missing something. Essentially, for each class, you rank images in decreasing order using the number of class pixels within the image. Then you assign one or more images to each subset following this order and matching the required number of pixels for each subset.     \n\n-\tMinor: Theorem 3.1 (Convergence Rate of WDES to Empirical Optimum). I am not an expert in evolutionary algorithms, but the technical ingredients of the proof (Appendix 2) seem to follow standard textbook knowledge in the area, in which convergence results for evolutionary algorithms are established under some common assumptions (e.g., elitist model, ergodic dynamics) – see, for instance, [G. Rudolph, Convergence of evolutionary algorithms in general search spaces, Proceedings  of the IEEE International Conference on Evolutionary Computation, 1996]. So stating the result as a theorem in the main paper and claiming the derivation of convergence bound is, in my humble opinion, a bit of an over-claim. It gives readers the impression that Theorem 3.1 and its proof are new technical results. I would suggest presenting Theorem 3.1 as a Proposition and just state that the proof follows standard arguments in the literature (and cite proper references on the covergence of evolutionary algorithms).\n\nQUESTIONS:\n-\tFollowing on my main concern above, could the authors provide applications’ examples where it is relevant in image segmentation to ensure label distribution consistency across data splits (I do not see any). There are efforts in the segmentation literature that build benchmarks with distributions shifts (to simulate realistic settings), which is actually the opposite of what this paper is advocating. \n\n-\tFollowing on my comment above, did the author explore deterministic optimizers that enable subset assignments under cardinality constraints (I am not sure that genetic algorithms are the best option for optimizing the considered objective functions).\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-05"}