Title: GMM-based VAE model with Normalising Flow for effective stochastic segmentation
Venue: NeurIPS 2025
Accepted: yes
Keywords: Stochastic segmentation

Variational Inference

Medical segmentation

Curvilinear structure segmentation
Abstract: While deep neural networks possess the capability to perform semantic segmentation, producing a single deterministic output limits reliability in safety-critical applications, caused by uncertainty and annotation variability. To address this, stochastic segmentation models using Conditional Variational Autoencoders (CVAE), Bayesian networks, and diffusion have been explored. However, existing approaches suffer from limited latent expressiveness and interpretability. Furthermore, our experiments showed that models like Probabilistic U-Net rely excessively on high latent variance, leading to posterior collapse. This work propose a novel framework by integrating Gaussian Mixture Model (GMM) with Normalizing Flow (NF) in CVAE for stochastic segmentation. GMM structures the latent space into meaningful semantic clusters, while NF captures feature deformations with quantified uncertainty. Our method stabilizes latent distributions through constrained variance and mean ranges. Experiments on LIDC, Crack500, and Cityscapes datasets show that our approach outperformed state-of-the-art in curvilinear structure and medical image segmentation.


Peer reviews:
comments: SUMMARY:
The presented paper introduces a way to modify probabilistic VAEs (specifically tested with Probabilistic UNet, Hierarchical Probabilistic UNet, and PHiSeg) to create a more "expressive" latent space by encoding the latent variables into a Gaussian Mixture Model, followed by using Normalizing Flow with a spline model to transform the variables into a more informative distribution.

STRENGTHS AND WEAKNESSES:
## Quality:
- Strength: the authors' approach improves on state-of-the-field baselines for several disparate tasks, including newer architectures based on diffusion

## Clarity
- Weakness: In figure 3, putting the standard deviation in the middle position leads to confusion - a better structure would be something like "input, stdev, real x4, predicted x4)" , and having labels instead of just a - i

## Significance
- Strength: the improved baselines in many tasks, plus the ability to apply this approach to several other models, mean it is likely to have impact
- Strength: the incorporation of the experiment where the Prob UNet latent variable was randomized was a clever way to look at the issues with this class of models

## Originality
- I am unaware of other work proposing a similar approach but this is the least confident part of my review

QUESTIONS:
- Are there any other sorts of architectures is this approach suited for integration with besides CVAE?
- How much training time is added by the incorporation of these components over training of the networks alone (as in the LIDC data)?
- Do PHiSeg and Hierarchial Prob UNet, when trained on LIDC, exhibit the same response to randomization of the latent variable as the ProbUNet does?

LIMITATIONS:
Limitations are not discussed and should be. Potential negative societal impact is unlikely.
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-06-30

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the problem of stochastic semantic segmentation by extending the Conditional Variational Autoencoder (CVAE) framework. The authors integrate a Gaussian Mixture Model (GMM) with Normalizing Flows (NF) to model both prior and posterior distributions. A Multiple-Input, Multiple-Output (MIMO) GMM structure is used to induce a structured, multimodal latent space. NF transformations are applied with constraints on means and variances, and a β-annealed KL divergence term is introduced to regulate learning.  \nThe method is evaluated on three tasks: medical nodule segmentation (LIDC-IDRI), curvilinear crack detection (Crack500), and urban scene parsing (Cityscapes). It consistently outperforms probabilistic U-Net and diffusion-based baselines in Generalized Energy Distance (GED), Hungarian-matched IoU, F1-score, and mIoU.\n\nSTRENGTHS AND WEAKNESSES:\n*Strengths:*  \n- Methodological innovation: The method combines a structured GMM latent space with normalizing flows and β-annealed KL. This is a new approach to modeling multimodal uncertainty.  \n- Broad empirical evaluation: demonstrated gains across medical, structural-defect, and urban-scene datasets, which indicates the generality of the framework.\n\n*Weaknesses:*  \n- Computational cost is not reported: the additional runtime and memory demands of the GMM + NF pipeline remain unclear, making it hard to assess the practicality of the proposed method.  \n- Statistical robustness missing: results are presented for a single run per setting, without seed-based variance or significance testing.  \n- Limited discussion of failure modes: potential risks such as overconfident or spurious segmentations in safety-critical contexts, are not analyzed.\n\nQUESTIONS:\n1. *Statistical robustness across multiple random seeds.*  \n   Consider reporting mean ± standard deviation over at least three independent runs (using different random seeds for weight initialization and data shuffling), and include paired significance tests (e.g., Wilcoxon signed-rank or permutation tests) when comparing to the probabilistic U-Net on the primary metrics in Tables 1–3.  \n   *Score impact:* Consistent, significant gains will boost confidence; high variability or non-significance will weaken it.\n\n2. *Quantify computational overhead and memory footprint.*  \n   Please provide three measures: 1) training time per epoch; 2) inference latency per 3D volume / 2D image, and 3) peak GPU memory usage for the baseline probabilistic U-Net versus your GMM + NF variant on LIDC-IDRI and Cityscapes. Also, please clarify how the number of NF layers and GMM components affects these costs.  \n   *Score impact:* A clearer description of the method’s efficiency and scalability could positively influence the score, while substantial overhead without justification may raise concerns.\n\nLIMITATIONS:\nThe manuscript lacks an explicit discussion of the computational overhead, the robustness to hyperparameters, failure-case examples, and broader risks such as misinterpretation of stochastic masks in clinical or safety-critical settings. \nI recommend adding a dedicated “Limitations & Societal Impact” section that quantifies resource usage, analyzes error modes, and reflects on potential negative impacts and dataset biases.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a novel stochastic segmentation framework by integrating a Gaussian Mixture Model (GMM) with Normalizing Flow (NF) in a Conditional Variational Autoencoder (CVAE), aiming to overcome the limitations of standard Gaussian latent spaces in modeling semantic diversity. GMM helps structure the latent space into distinct semantic clusters, while NF adds flexibility to both prior and posterior distributions, addressing issues like posterior collapse and poor expressiveness. Extensive experiments across LIDC-IDRI, Crack500, and Cityscapes datasets demonstrate state-of-the-art performance in uncertain and fine-grained segmentation tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. The combination of GMM and Normalizing Flow can enhance interpretability and expressiveness of the latent space.\n\n2. The model produces competative evaluation results on three datasets.\n\nWeaknesses:\n\n1. The method seems like a simple combination of GMM, VAE, and Normalizing Flow. The authors should provide more insights and theoretical justifications of why we need to integrate all the three components.\n\n2. The motivation and application of this paper is not well-aligned. The paper is motivated to address the uncertainty and annotation variability issues, but only one study case of non-deterministic segmentation task is evaluated. The overall logic is not very convincing: the proposed approach aims to address a stochastic segmentation problem but why does it also outperform deterministic segmentation baselines?\n\n3. Important ablation studies are missing. (1) how do the core components, GMM and NF affect performance? I only see a very simple ablation on Crack500 in the Appendix (which is not mandatory to read for reviewers). (2) How do the components affect stochastic segmentation performance? I believe this comparison is more important than the table the authors show in the Appendix. (3) What happens if we remove GMM?\n\n4. The comparison of computational cost is missing. The new approach introduces additional computations in GMM, which might dominate the whole model. Details of throughput/memory are needd to guarantee fair comparisons.\n\n5. There is another concern regarding the GMM process. If the category count becomes higher, does the number of GMM clusters need to get higher accordingly?\n\nQUESTIONS:\nPlease refer to the Weaknesses.\n\nLIMITATIONS:\nThe current version of the paper needs to be improved due to the insufficient evaluation and the misalignment to the motivation.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a probabilistic model for stochastic image segmentation based on a Conditional Variational Autoencoder and U-Net framework that integrates Gaussian Mixture Model (GMM) and Normalizing Flow (NF) to improve latent variable representation.\nGMM helps create better structure of latent space while NF model image feature transforms to gauge local uncertainty.\nThe author also introduce some constraints on the range of means and variances to stabilize training of NF.\nThe proposed approach achieves state-of-the-art results on three different datasets.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n1. The authors were able to combine two existing components GMM and NF to make them work well together under Probabilistic U-Net framework and improve the overall segmentation results.\n2. Experimental results are performed on three datasets for both stochastic and deterministic segmentation and compared with various group of segmentation approaches.\n3. Having ablation study to help better understand the role of both GMM and NF contribute to the final results.\n\n# Weaknesses\n1. Missing some details on network structures, e.g. posterior and prior networks which were introduced in Probabilistic U-Net. Thus, it's hard to follow without reading prior works.\n2. Missing details about the latent space dimension, the number of distributions of the Gaussian mixture model being used in each dataset.   \nI'm curious to see if there is a correlation between number of semantic classes and the number of distributions.\n3. The idea of combining GMM and NF has been done in some prior works, e.g. FlowGMM.\n\nQUESTIONS:\n1. How is the single-sample inference conducted? More details should be provided.\n2. What are the missing details mentioned in the weaknesses section?\n\nLIMITATIONS:\nNo. The authors did not provide limitations of their work.\nThe authors should address some difficulties in training the model with NF or limitation in the dimension of the latent variables.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}