GMM-based VAE model with Normalising Flow for effective stochastic segmentation
Authors: Conghui Li
Chern Hong Lim
Xin Wang
Venue: NeurIPS 2025 poster ()
Keywords: Stochastic segmentation
Variational Inference
Medical segmentation
Curvilinear structure segmentation
Primary area: probabilistic_methods

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.

GMM-based VAE model with Normalizing Flow for effective stochastic segmentation
Conghui Li
Chern Hong Lim"
School of IT
School of IT
Monash University Malaysia
Monash University Malaysia
Bandar Sunway, Selangor, 47500, Malaysia conghui li @monash edu
Bandar Sunway, Selangor, 47500, Malaysia
Lim. CwwrnHongomnnash. edu
Xin Wang"
School of Engineering
Monash University Malaysia
Bandar Sunway, Selangor, 47500, Malaysia Wang. Xinimonash...mmmmmmmoooo edu
Abstract
While deep neural networks possess the capability to perform semantic segmenta- tion, producing single deterministic output limits reliability in safety-critical ap- plications caused by uncertainty and annotation variability. To address this, stochas- tic segmentation models using Conditional Variational Autoencoders (CVAE), Bayesian networks, and diffusion have been explored. However, existing ap- proaches suffer from limited latent expressiveness and interpretability. Furthermore, our experiments showed that models like Probabilistic U-Net rely excessively on high latent variance, leading posterior collapse. This work propose a novel frame- work by integrating Gaussian Mixture Model (GMM) with Normalizing Flow (NF) in CVAE for stochastic segmentation. GMM structures the latent space into mean- ingful 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- curvilinear structure and medical image segmentation.
Introduction
In recent years, deep neural networks have made remarkable progress in semantic segmentation tasks. The main goal of most methods lies in generating a single segmentation result that is highly consistent with the content of the image. However, this deterministic segmentation has limitations in safcty-critical areas such as medical diagnosis, industrial testing, or autonomous driving. This due the images in such applications are often accompanied by inherent uncertainty, leading to the presence of multiple reasonable yet contradictory manual annotations. [41,88] 28). In this context, providing a single deterministic segmentation result may not be able to fully express the possible true semantic distribution of the image, thus limiting the reliability and practicability of the model in these critical tasks.
Currently, many studies have proposed probabilistic approches for sampling from output distributions, including models based Conditional Variational Autoencoders (CVAE) 25 24]. Normalizing
"Corresoondigg author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Flows [35]. Bayesian neural network [16] 14 8 and Diffusion B2 41]. Most of the study recognized the Gaussian distribution's 's limitation in capturing true feature distributions. Therefore, in recent years, researchers have been dedicated to enhancing the ability of variational autoencoders (VAE) in complex data modeling by replacing the traditional simple Gaussian distribution with a more complex distribution to improve the expressive power of the distribution and capture more complex semantic information [36] 33 37 (29]. An alternative to use the diffusion model to capture the data distribution more flexibly through a multi-step denoising process [32 41 However, the drawback of diffusion the unexplainability of distribution, which is crucial for security-sensitive tasks. Therefore, we advocate returning to the VAE framework and further enhancing the latent space modeling mechanism to make up for the deficiencies of the original Gaussian hypothesis.
We performed deeper analysis to prove the ineffective of Gaussian distribution by experimenting the existing Probabilistic segmentation model, Probabilistic UNet 255. The results show that the variance of its latent variables was significantly higher than the mean after training, reaching 500-fold gap in the curvilinear structure segmentation task. This suggests that the latent spaces fail to learn compact semantic representations and relying on high variance cover a wide range of possibilities, leading models to favor guesswork over deterministic predictions. To examine the effectiveness of Gaussian latent variables, we modified the Probabilistic U-Net by replacing the latent samples with purely random vectors. The results in Figure | show that the model's performance does not degrade significantly which indicates that the latent variables contribute limited semantic information.
To address the aforementioned limitations, we propose a novel framework that integrates Gaussian Mixture Model (GMM) and Normalizing Flow (NF) on a Conditional Varia- tional Autoencoder (CVAE). The GMM can define semantic clusters in the latent space more effectively, where each Gaussian component captures distinct mode of variation corresponding to meaningful visual concepts or feature pat- terns. In detail, each Gaussian component learns distinct features through a parameter competition mechanism, miti- gating the prediction ambiguity caused by single Gaussian prior distribution. Based on the latent space provided by GMM, NF learns the reasonable deformation of image fea- tures, and its Jacobian determinant directly quantifies the local uncertainty. Instead of only let the posterior distribu- tion fitting the data, we apply this combined model to both the prior and posterior distributions, and introduce constraints on the range of means and variances to stabilize NF's training gradients.
Imput
Qualitative comparison of using random number and gaussian distribution for latent variable in Prob. Unet [25]. 1-2 is the result from gaus- sian distribution, 3-4 is the result from random number
Figure 1:
Worth to mention that We bring this approach to curvilinear structures for the first time. Experimental results demonstrate superior performance in wide range of applications such as medical image segmentation, infrastructure crack detection, and city scene segmentation using LIDC-IDRI, Crack500 dataset, and cityscape dataset respectively. In infrastructure crack detection tasks (using Crack500 dataset), method achieves a better balance between recall and precision, improving the F1-score by 11.0% compared to the standard CVAE- stochastic segmentation. Furthermore, evaluations on the multi-label segmentation datasets, LIDC-IDRI and Cityscapes datasets show that our approach consistently outperformed the existing baseline models.
In summary, the contributions of this paper are as follows:
1. We propose a combination of GMM through a Multiple-Input Multiple-Output (MIMO) mechanism construct structured latent space in CVAE. GMM provides multiple potential semantic clusters to reduce the transformation complexity of NF and thereby improve performance.
2. We introduce NF into both the prior and posterior distributions to addresses the semantic ambiguity inherent in traditional VAEs caused by overly simplistic priors.
3. We are the first to apply a stochastic segmentation model to curvilinear structure segmenta- tion and achieve state-of-the-arr performance compared to the best existing baselines.

---

2
Related work
The goal of probabilisticc semantic segmentation is to capture both aleatoric uncertainty and the inherent ambiguity present in segmentation labels. A variety of methods have been proposed to address this. Early approaches employed Bayesian neural networks, in which model parameters were treated as probability distributions and sampled during inference [ 42 30
I [6] While this
enabled uncertainty modeling, it significantly increased inference overhead and model complexity due to repeated sampling and the need to maintain distributions for each parameter.
To address these limitations, researchers later adopted Conditional Variational Autoencoder (CVAB) frameworks for semantic segmentation. Initial CVAE models were typically built on U-Net architec- tures, where a latent variable was concatenated with the last feature layer to introduce randomness into the segmentation process [25]. Subsequent work extended this by incorporating hicrarchical latent structures, injecting randomness multiple levels of the encoder-decoder pipeline of U-Net to enhance expressiveness 24
However, a major limitation persisted: the use of simple Gaussian distributions in both the prior and posterior, which restricted the model's ability to capture complex, multi-modal semantic representa- tions. To overcome this, researchers engaged more expressive probabilistic structures [36] 23 33
For example, cFlow applies normalizing flow (NF) to the posterior to enhance flexibility and mit- igate posterior collapse. Despite these improvements, some studies argue that such models still underperform in accuracy-sensitive applications due to their limited prior modeling capabilities (35).
To address this, Zhang et al. proposed the Joint Probabilistic U-Net (JProb. Unet), which employs a Reversible multi-layer perceptron (MLP) to simultaneously transform both the prior and posterior distributions into more complex forms 4 However, the transformation capacity of Reversible MLPs limited compared to conventional flow-based methods, often requiring deep stacking to approximate complex distributions. Additionally, its grouped variable structure lacks the ability to model global dependencies effectively.
In this work, we propose stochastic segmentation framework that combines Gaussian Mixture Model (GMM) with Normalizing Flow (NF) to construct flexible and structured latent space. The GMM enables explicit modeling of multi-modal semantic priors, while the NF applies nonlinear transformations to each Gaussian component, significantly enhancing the capacity of both the prior and posterior distributions. Compared with traditional CVAE approaches that assume single Gaussian, our method can effectively mitigates issues such as posterior collapse and semantic ambiguity. It also demonstrates superior modeling capability for fine-grained structures such as cracks. Experimental results show that our approach achieves notable improvements on tasks invol curvilinear structure segmentation, medical image uncertainty modeling, and multi-label scene understanding highhihnnnnn its strong potential for nigh-precision and interpretable stochastic segmentation applications.
3 Methodology
3.1 Gaussian mixture distribution
In most existing VAE based image segmentation models, the latent variable is usually modeled as a unimodal Gaussian distribution N (u,,), However, our experiments in Table show that this simple distribution hypothesis is difficult to effectively capture complex and structured semantic information in images. To address these limitations, we introduce the Gaussian Mixture Model (GMM) as the base distribution for modeling latent variables to improve the initialization (Figure
GMM is multimodal probabilistic model defined as weighted sum of multiple Gaussian compo- nents, enabling the representation of diverse semantic clusters. Compared to a single Gaussian, GMM offers greater representational capacity, allowing the latent space to explicitly encode distinct semantic regions in the image. Through a parameter competition mechanism, each Gaussian component is encouraged to specialize in specific semantic attributes, thereby enhancing both the interpretability of the latent space and the clarity of the resulting segmentations.
In this work, we adopt MIMO mechanism to generate the parameters of Gaussian Mixture Model. Each input header receives the same image input but is designed extract different semantic perspectives or preferences from that input. A shared backbone network encodes the global contextual
3

---

Forward process
Ps.
Posterior
My.
qo (aloi s)
Inverse process
dieergence
Pv(y(z,
space
My
Prior
My
po(Sola)
Flow stepclorward)
x: Input image
Latent veriable
[" Flow stepcinverse)
y: Ground truth
Inference step
Figure 2: Overview of the GMM-based VAE with Normalizing Flow. During training, the input image and label y concatenated and duplicated inputs to the posterior network, which outputs the mean and covariance for each GMM component to construct the latent space qu(2))a x. y). multi-layer NF transforms this space into more expressive latent distribution (=2 y). KL divergence is computed using the inverse NF in the prior and the forward NF in the posterior to align their distributions. During inference, the prior network generates the latent space p(=[<r via GMM and NF, from which a latent variable is sampled and fused into the feature maps of Decoder pe ( 2) introduce stochastic segmentation output ÿ.
information the image, serving as a common feature foundation for all branches. Each output head then generates the parameters of the Gaussian components, producing means and covariances set:
(Un,, E1), (122, Ez)...., (Jo: En)}
Where represent the number of output in MIMO framework and the number of GMM component. Additionally, separate head predicts the component weights n, forming the complete mixture model. The resulting Gaussian Mixture distribution is defined as:
Po(2) Σ To N( Mos Eo)
(2)
Among them, single Gaussian distribution term is defined as:
(2m) E. 1/2 exp (-1(e 1o) Io Ho)
N( Mo: Eno)
3.2 Normalizing flow
NF enhances probabilistic modeling by transforming a simple base distribution into a more expressive and flexible through sequence of learnable, invertible transformations.
Taking the prior distribution in CVAE as an example, assume initial latent variable sampled from base distribution: 2g (x))). A normalizing flow function is typically constructed as composition of multiple transformations ft. where the subscript denotes the number of flow layers. The final transformed latent variable is given by:

---

2 == = fK-1 Q... fi(2))
In stochastic segmentation task, aligning the learned latent distribution with the true data distribution enhances the expressiveness of the posterior and alleviates posterior collapse, phenomenon where the approximate posterior q=a|xr y) degenerates the prior p(=[æ), causing the K Lddivergence O. This lead samples generated by VAE lack diversity, and the semantic features extracted by the latent variable Z not obvious. Secondly, incorporating Normalizing Flow (NF) into the prior distribution enables the sampling of latent variables that encode richer semantic information. By learning a more flexible and data-adaptive prior, NF encourages the latent variables align with meaningful structures in the imput space instead providing more noise for randomness, thus improving both the interpretability and effectiveness of the generative process.
In conventional NF frameworks, the distribution to be transformed referred to as the base distribu- tion, which is typically for example, standard normal distribution. However, the parameters of base distribution evolve during training in CVAE framework. This feature introduces challenges for NF application in CVAE. Specifically, in the early stages of training, before the flow is fully trained, the base distribution tends to exhibit high variance to cover the uncertain region due to its limited expressive capacity. As a result, the sampled latent variables are highly dispersed, making it difficult for the NF to learn effective and stable transformations. We applied segmented 8 annealing impose constraints
algorithm mitigate this issue, but it still happen by chance. Therefore, on the mean and variance of each Gaussian component in the GMM base distribution, enforcing the following range 20 < μ < 20. This constraint ensures the latent variables remain within reasonable range, promoting smoother and more stable NF transformations throughout training. In implementation, we utilize three-layer Neural Spline Flows (NSF)
1K to transform the GMM- based base distribution, leveraging the NSF's capability to model complex, non-linear mappings in the latent space. Different with the simple transformation, NSF use the monotonic rational-quadratt splines as the elementwise transformation, which offers much greater flexibility than affine transforms. Specific explanation of NSF can be found at appendix A.
To recap. we choose GMM + NF over Gaussian NF due to the limitations of a single Gaussian, which is unimodal and overly simplistic. While NFs are theoretically capable of mapping such base distribution into a complex, multimodal target, doing so requires deep and intricate transformations that can hinder convergence, especially during early training. In contrast, GMM naturally introduces a multimodal prior, allowing the NF to focus on localized refinements such as fine-tuning, stretching, or twisting individual modes rather than learning global semantic structure from scratch.
3.3 Evidence lower bound
The Evidence Lower Bound (ELBO) is the objective function used to train Variational Autoencoder (VAE). It provides a lower bound the true data log- likelihooo and balances Reconstruction loss and Regularization. In this task, Reconstruction loss is referred to Dice, and Regularization is referred as KI. divergence
Log probabilities: this method, we used normalizing flow. The probability density function of the posterior is:
K II
0 Thg-1
<<s|,, y) = 20(20/2 y)
det
A=1
For the flow log probability, we usually need to add Jacobian determinant the basis of the log probability:
K Σ
0fk Thg
logg(=/a, y) log (0(2 20 y)
log det
Where is the posterior probability density function and go represents the base distribution before the flow operation, which in this study the probability density function of GMM. det(-) stands for Jacobian determinant.
5

---

The significance of the Jacobian determinant in NF lies in ensuring density consistency between the base distribution and the transformed distribution. The term log 90 [2, y) corresponds the log-probon see-prbbabiltty of the sample under the base distribution, while the second term accounts for the log-volume change introduced by the flow transformation. This correction ensures that probability mass is preserved and that the final distribution accurately reflects the transformation applied to the base
In the posterior, since the flow is constructed as a sequence of transformations, we define it as a chain of invertible functions. Let:
fi(zo), 22 f2(21),
KK fx(=
This forward transformation is used in the posterior flow. In the prior, fK represents the inverse function and z' represent the inversed latent variable, the process is defined in the same way. The total log-volume change of the transformation is calculated by summing the log-determinants of the Jacobians at each layer. Since each transformation f& invertible, we can compute the probability density of the transformed latent variable 2 using the change-of ofannge-of-vriiaaee variables formula:
K
O5
p(=|>) = po(=)1>)
II
det
(8)
02k
Taking the logarithm yields the log bog--obbbbilty::
K logp(z|z) log po log det
0f 0zk_
KL divergence: Where is the prior probability density function and represents the base distribution. The Jacobian determinant is positive because we have performed the inverse operation and need to add the logarithmic volume change caused by the transformation to the base log probability.
The conventional KL formula is:
(g( v))p(((()))) = ()())))))))) [log ((22, y) log [(2=))
(10)
We insert the log probability of our calculated prior and posterior into the KI. formula.
K2
a5(9) (4)
(T(((=/x,))) p(=[=)) (x
log loggo(=o|=,.0)
log
det
Øzk-1
(11)
Kp
@0% 10p-1(p)
log Po (=6lar))
log
det
0z
1==1
Finally, the ELBO of proposed method is as follows:
LELBO (x) = ==< [loga =)] 33 KL. (g(2/2, y) p(+[)))
(12)
Where the Egsg(( =.y) [log p(y) x 2)] is refereed as reconstruction term which we applied the MSE. The weight 3 is used to control the influence of the KL term.
KL vanishing: In order to further mitigate the influence of posterior collapse, inspired by various VAE annealing algorithms [18 15 4 20], we set 3 to zero at the initial stage of training, allowing the model to only optimize the reconstruction terms the initial stage and learn to encode valid information at first. We call this step the Warm-up operation of the model. In the subsequent training, the weight of / gradually increased linearly until the set maximum value of is reached to achieve stability. The annealing algorithm formula of VAE is follows:

---

0,
if < Ti (Warm up),
t Ti
Bt
Bmax:
if S/S T2 (Increase),
(13)
T2
Bmax:
ift T1 + 72 (Stable).
Where represents the current training epoch, Ti is the number of warm ups, T2 is the number of epochs from warm up the stable stage, and mmaxis the set maximum value of B3.
4
Experiment
We evaluated our method on three tasks with varying structural and annotation characteristics: (1) cancer detection with multiple expert annotations capturing boundary uncertainty; (2) crack detection with single annotations focused on fine curvilinear structures; and (3) urban scene segmentation as multi-label task with single annotations. For fair comparison on the last two deterministic datasets, we evaluated the probabilistic model using single-sample inference to match deterministic baselines.
Setting:
We evaluate our proposed method by applying it to three representative VAE-based stochastic segmentation models: Probabilistic U-Net, Hierarchical Probabilistic U-Net, and PHISeg, All models trained using the Adam optimizer for 500 epochs with a batch size of 32 and all experiments programmed by Pytorch 2.4.1 and conducted using NVIDIA A100 Tensor Core GPU.
We evaluate performance using combination of conventional and stochastic segmentation metrics. For general segmentation quality, we report the mean Intersection over Union (mloU), precision, recall, and F1-score, which provide insight into overall accuracy and class-wise balance. To assess the consistency and diversity of probabilistic predictions, we include the Generalised Energy Distance (GED), proper scoring rule that evaluates sample quality relative ground truth distributions, and Hungarian-Mateen Matched Intersection over Union (HM-IoU), which measures alignment between sampled predictions and ground truth annotations through optimal bipartite matching.
Metrics:
4.1 LIDC-IDRI
Inout
Figure 3: Qualitative results on four LIDC images with proposed method. (a)-(d) represent ground truth label from four different experts. (e) represent the standard deviation of prediction. (f-(i) represent the prediction from proposed method.
The Lung Image Database Consortium image collection (LIDC-IDRI) consists of thoracic computed tomography (CT) scans for diagnostic and lung cancer screening, with annotated lesions provided by multiple radiologists 22 Following the experimental setting of CCDM, we extracted total of 15,096 slices of size 128 128, and divided the dataset into training, validation, and testing sets with ratio of 60:20:20.

---

Table 1: Quantitative results on LIDC, with methods listed in chronological order. Except for Prob. U-Net + random, all baseline results are taken from [4]]. Our proposed method is integrated into three baseline models, each showing substantial performance improvements. Note that lower values of GED indicate better performance, while higher values of HM-IoU are preferred.
Method
GED16
GEDs2
GED
GED1n 100
HM-IoU16
HM-IoU32
Prob. Unet [25
0.310
0.303
0.252
0.552
0.548
Prob. Unet + random
0.314
0.31 10.003
0.545
0.548
HProb. Unet
24
0.270
0.530
PhiSeg 3
0.262+0000
0.247+0.00
0.224 004
0.586
0.595
SSN 28
0.259
0.243
0.225 4-0.002 0000
0.558
0.555
cFlow 35
0.225 #0001
0.584
CAR 21
0. 009
JProb. Unet
0.20610.01
0.647.10.01
PixelSeg 44
0.24
0.614
MoSE 7
0.208
0,195 1955...0m
0.195
0.189 20.002
0.624
AB 6
0.213
0.196
0.193
0.61440.01
0.619
CIMD 32
0.234
0.218 +
0.210
0.587
0.592
CCDM 41
0,212
0.194 001
0.187
0.183 002
0.623
0,631
Prob. Unet + Our
0.196.10.000 000001
0.189 1:0. 001
0.185 3:0.0000000 002
0.184 000
0.641
0.644
HProh. Unet + Our
0.224.100002
0,219 1:00 001
0.2161.0000
0.215
0.562.
0.569
PhiSeg Our
0.25040.000
0.243 4:0 000
0.241
0.237+0.00
0.623
0.621 20.001
The model performance of baselines is derived from [1]] We used the same experimental data enhancement method as in the paper to rotate the training image by degrees, 90 degrees, 180 degrees and 270 degrees. The models involved in the experiment can be divided into three categories, VAB-based, diffusion-based. and some other different methods such as mixture of expert. Based on the results in the table [I our best model achieved the best results in 4 of the 6 metries. Among them, compared with the most popular diffusion- based models, the performance of the proposed method is improved all metires. Compared with cFlow using only posterior NF, GMM+NF also achieved better performance on both metrics in probabilistic U-net architectures. JProb. Unet outperforms all other models M IoUs2, but all three of our models outperformed it. The qualitative analysis results based on probabilistic U-net can be seen in Fig 3 (b)-(e) annotations created for the four experts, (g)-() show the sample distribution of the model, and f is the average value after multiple samples.
In addition to the analysis of the proposed method, we also designed an interesting experiment by replacing the latent variable with random number. The table [ shows the performance differences between Prob. Unet and Prob.Unet + random. This proves our hypothesis where the conventional mul- tivariate Gaussian distribution is too simple in distribution structure to represent useful information, and its latent variable can only provides some randomness the model.
4.2
Deterministic segmentation
4.2.1
Crack500
The Crack500 dataset is designed for pixel-wise pavement crack segmentation and consists of 500 high-resolution images, resulting in 3,368 cropped images of size 360x640 [43. The dataset is split into 1,897 training, 347 validation, and 1,124 testing images. Building on the setting (9) we employ horizontal flipping, random cropping, and random rotations of 90", 180", and 270' as our data augmentation strategies. Additionally, all training samples are cropped to 256 x 256 during training.
For this task, the results of five traditional deep learning baselines are taken from 9 As shown in the ( the performance of the Probabilistic U-Net is significantly lower than that of naive U-net with notably high recall. This supports our hypothesis: due to the large variance the latent distribution, the sampled latent variables are overly random, causing the model to produce over- inclusive segmentations. This issue is particularly pronounced in tasks with highly skewed semantics, such as cracks. Our experimental results demonstrate that our method effectively improves prior expressiveness. On the Probabilistic U-Net architecture, our approach is greatly increase in metrics compared with naive Probabilistic U-Net, and achieves state-of-the performance. Additionally,
8

---

(a) Comparison from Crack500 dataset. All methods trained and tested by 256 256 resolution. It can he known from the image that segmentation is more compact and close the boundary
(b) Comparison from Cityscape dataset. All methods trained and tested by 256 512 resolution. The hoxes highlights our advantage
Figure 4: Qualitative comparison of deterministic segmentation. All stochastic models (including the proposed method and Prob. U-Net) are evaluated using single sample. Despite this, our model produces more fine- ane-ggrined segmentation results compared mainstream deterministic methods across both tasks.
qualitative results are presented in the Fig 4a comparing the outputs of baseline models with the improved predictions generated using our proposed method.
Table 2: Quantitative results Cityscapes. All
Table 3: Quantitative results Crack500. All
the results of baselines are from 41. Param
the results of baselines are from 9
represent the parameter number of model; represent million
Method
Precision
Recall
F1
UNet [34]
62.22
68.85
61.83
Method
Backbone
Param
IoU
VGG-UNet 31
58. 18.18
60.26
51.79
DeepLabv3
5
ResNet50
39m
58.6
TopoNet 19
66.81
62.68
60.06
DeepLabv3 5
ResNet101
58m
59.2
DRU [39
61.94
71.43
62.82
UPerNet 40
ResNet101
83m
60.7
Crackformer 226
69.13
66.24
64.75
HRNet 33
w48v2
70m
63.3
JTFN [IO]
68.81
69.06
65.76
UPerNet 27
Swin-Tiny
58m
65.5
JTFN + CIRL 9
70.32
69.93
67.62
CCDM 41
30m
60.3
Prob. Unet 25
56.80
69.39
60.80
CCDM 41
Dino VIT-S
50m
65.8
Prob. Unet+Our 9
72.97
71.72
71.80
Prob. Unet
25
33m
63.2
Prob. Unet+Our
38m
73.0
4.2.2 Cityscape
Cityscapes is standard benchmark dataset for multi-class semantic segmentation 11 It contains 2,975 training images and 500 validation images, each with a resolution of 512 x 1024, annotated across 19 semantic classes.
In this task, the baseline model is taken from CCDM 41 Following proposed setting, we employ resolution 256 512 in training and testing set. Our proposed method achieves the best performance within the Probabilistic U-Net architecture. As shown in the table 2 even using a single sample, model outperforms the latest diffusion- based probabilistic models. Additionally, provide qualitative results the Fig 4b) including comparisons with CAVE, diffusion-based models, and quantitative visualization of method's segmentation distribution.
4.3 Ablation study
In the ablation study, we conducted three different experiments, comparing the performance of conventional Gauss and GMM in the case of NF; Performance comparison with or without NF; The performance different NF methods is compared. Among the three different tasks, the baseline model has the largest performance variance on the curvilinear task, so we use Crack500 to test the contribution of different modules. score is used to compare the performance of different models in ablation experiments.

---

Table Ablation study on Crack500 dataset (Fl score). Comparison among different NF types and GMM component numbers.
Experiment No.
c
d
NF Type
NSF
NSF
NSF
RealNVP
Glow 3
No. of GMM
6
FI Score
66.04
71.23
71.24
61.39
69.15
69.09
Param
37.9
38.0
38.1
34.0
38.0
37.9
Gaussian vs GMM: We set the number of distributions of the Gaussian mixture model to 1, and 6 respectively, where Gaussian mixture equivalent to multivariate Gaussian distribution when the number is 1. According to the table we can see that increasing the number of distributions from 1 to the Gaussian mixture model can bring great performance improvement, while increasing the number from 3 to 6 is relatively less important. This proves that because the Gaussian distribution too simple, more complex transformations are needed to build a multimodal structure. However, the Gaussian mixture distribution simulate the data form relatively better, reducing the variation required for NE
GMM vs NE: Next, we test the gain that NF brings to the model. According to the table the performance the model is significantly improved after NF used. This proves that the distribution structure close to the data can significantly solve the problems of posterior collapse and insufficient prior expression, and greatly improve the model performance.
Different NF: Finally, we experiment the effect of different NF methods on the model performance. These two methods are RealNVP 12 and Glow (22). Specific explanation of these NF methods
can be found at appendix A.2 and A.3 The results show that the method can better fit the model
performance and improve the expression ability of the model.
Conclusion
In this paper, we propose novel GMM-based framework to enhance the expressive power of latent space in CVAE-based stochastic segmentation. To validate our hypothesis, we replaced the latent variables with random noise and observed similar results to those using standard Gaussian prior, indicating that the Gaussian mainly injects randomness rather than meaningful structure. In the processing pipeline, a more informative prior compared to existing methods. To stabilize the training, we constrain the ranges of mean and covariance prevent extreme variance and apply KL annealing to mitigate KL. vanishing. The model demonstrates strong performance across three distinct tasks. Notably, we the first to apply stochastic segmentation to curvilinear structures, addressing the limitations of conventional stochastic models in such scenarios.
model incorporates NF in both the prior and posterior networks, yiclding
However, This study has two main limitations: First, compared to conventional CVAE models, our model introduces a significant increase in computational cost. Although inference speed is not a major concern in target application scenarios of stochastic segmentation, this is still an aspect to improve in our future work. Second, the training process is less stable. Specifically, the normalizing flow requires computing the log-determinant for gradient updates, which makes it more prone to KL vanishing compared to CVAE. As a result, careful tuning hyperparameters is necessary, along with additional strategies such annealing schedules and constraining the variance range of GMM components to ensure stable training.
10

---

References
[1] Fahim Irian Alam and Iker Gondra. A bayesian network-basel tunable image segmentation algorithm for object recognition. In 2011 IEEE International Symposium on Signal Processing and Information Technology (ISSPIT), pages 011-016, 2011.
[2] Samuel G Armato III. Geoffrey McLennan, Lue Bidaut, Michael McNitt-Gras Charles Meyer, Anthony P Reeves, Binsheng Zhao, Denise R Aberle, Claudia Henschke, Eric A Hoffman, et al. The lung image database consortium (lidc) and image database resource initiative (idri): a completed reference datahase of lung nodules on ct scans. Medical physics, 38(2)915-933) 2011.
[3] Christian F. Baumgartner, Kerem C. Tezcan, Krishna Chaitanya, Andreas M. Hötker, Urs J. Muchlematter, Khoschy Schawkat, Anton S. Becker, Olivio Donati, and Ender Konukoglu. Phiseg: Capturing uncertainty in medical image segmentation. page 119-127, Berlin, Heidelberg, 2019. Springer-Verlag.
[4] Dajie Chen, Qiyu Chen, Zhesi Cui, Ruyi Wang, and Gang Liu. Sa-vac: novel approach for reservoir characterization based on variational auto-encoder and seleetive attention mechanism. Earth Science Informatics, 16(4):328333301, 2023.
[5] Liang-Chieh Chen, George Papandreou, Florian Schroff, and Hartwig Adam. Rethinking atrous convolution for semantic image segmentation. arXiv preprint 1706.:/706.0555 05587, 2017.
[6] Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint arXiv:2208. 04202, 2022.
[7] Xi Chen, Diederik Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, 16/ 1.02731, 2016.
and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint 161
[8] Chen, Yue Zhao, and Chuancai Liu. Medical image segmentation using scalable functional variational bayesian neural networks with gaussian processes. Neurocomputing, 500:58-722 2022.
[9] Zhuangzhuang Chen, Zhuonan Lai, Jie Chen, and Jianqiang Li. Mind marginal non-crack regions: cesterinn--nspired representation learning for crack segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12698-12708, June 2024.
[10] Mingfei Cheng. Kaili Zhao, Xuhong Guo, Yajing Xu, and Jun Guo. Joint topology-preservingg and feature- refinement network for curvilincar structure segmentation. In 202/ IEEE/CVF International Conference Computer Vision (ICCV), pages 7127-7136, 2021.
[11] Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urhan scene understanding. In Proc. of the IEFF Conference on Computer Vision and Pattern Recognition (CVPR), 2016
[12] Laurent Dinh, Jascha Sohl-Dickee Dicksteii, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605 08803, 2016.
[13] Conor Durkan, Artur Bekasov, lain Murray, and George Papamakarios. Neural spline flows. Advances in neural information processing systems, 32, 2019.
[14] Martin Ferianc, Divyansh Manocha, Hongxiang Fan, and Miguel Rodrigues. Combinet: Compact convolutional bayesian neural network for image segmentation. In Igor Farkaš, Paolo Masulli, Sebastian Otte, and Stefan Wermter, editors, Artificial Neural Networks and Machine Learning - ICANN 2021, pages 483-494, Cham, 2021. Springer International Publishing.
[15] Hao Fu, Chunyuan Li, Xiaodong Liu, Jianfeng Gao, Asli Celikyilmaz, and Lawrence Carin. Cyclical anncaling schedule: A simple approach to mitigating kl vanishing. arXiv preprint /rXiv: 10145, 2019.
[16] Shangqi Gao, Hangqi Zhou, Yibo Gao, and Xiahai Zhuang. Bayeseg: Bayesian modeling for medical image segmentation with interpretable generalizability. Medical Image Analysis, 102889,, 2023.
[17] Zhitong Gao, Yucong Chen, Chuyu Zhang, and Xuming He. Modeling multimodal aleatoric uncertainty in segmentation with mixture of stochastic experts. arXiv preprint arXiv:2212.075: 07328, 2022.
[18] Irina Higgins, Loic Matthey Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick. Shakir Mohamed, and Alexander Lerchner. beta-vac: Learning basic visual concepts with constrained variational framework. In International conference on learning representations, 2017.
[19] Xiaoling Hu, Fuxin Li, Dimitris Samaras, and Chao Chen. dopology-preserving deep image segmentation. Advances in neural information processing systems, 32, 2019.

---

[20] Chin-Wei Huang. Shawn Tan, Alexandre Lacoste, and Aaron C Courville. Improving explorability in variational inference with annealed variational objectives. Advances in neural information processing systemms, 31, 2018.
[21] Elias Kassapis, Georgi Dikov, Deepak K Gupta, and Cedric Nugteren. Calibrated adversarial refinement for stochastic semantic segmentation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 7057-7067, 2021.
[22] Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1xl convolutions, Advances in neural information processing systems, 31, 2018.
[23] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow, Advances in neural information processing systems, 29, 2016.
[24] Simon A. A. Kohl, Bernardino Camera-Paredest Klaus H. Maier-Hein, Danilo Jimenez Rezende, S. M. Ali Eslami, Pushmeet Kohli, Andrew Zisserman, and Olaf Ronneberger. hierarchical probabilistic u-net for modeling multi-scale ambiguities. CoRR, 19077 2019.
[25] Simon A. A. Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph Ledsam, Klaus H. Maier-Hein S. M. Ali Eslami, Danilo Jimenez. Rezende, and Olaf Ronneherger. A probabilistic u-net for segmentation of ambiguous images. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS' 18, page 6965-6975. Curran Associates Inc., 2018.
[26] Huajun Liu, Xiangyu Miao, Christoph Mertz, Chengzhong Xu, and Hui Kong. Crackformer: Transformer network for fine- grained crack detection. In 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 3763 3772, 2021.
[27] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hicrarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pages 10012-10022, 2021.
[28] Miguel Monteiro, Loic Le Folgoe, Daniel Coelho de Castro, Nick Pawlowski. Bernardo Marques, Kon- stantinos Kamnitsas, Mark der Wilk, and Ben Glocker. Stochastic segmentation networks: modelling spatially correlated aleatoric uncertainty. Proceedings of the 34th International Conference on Neural Information Processing Systems, NIPS '20. Curran Associates Inc., 2020.
[29] Rogan Morrow and Wei-Chen Chiu. Variational autoencoders with normalizing flow decoders. arXiv preprint arx::22004. 05617, 2020.
[30] E.N. Mortensen and Jin Jia. Real-time semi-automatic segmentation using bayesian network. In 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'00) 06), volume 1, pages 1007-1014, 2006.
[31] Agata Mosinska, Pablo Marquez-Neila, Mateusz Kozinski, and Pascal Fua. Beyond the pixel-wise loss for topology-awan delincation. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3136-3145, 2018
[32] Aimon Rahman, Jeya Maria Jose Valanarasu, llker Hacihaliloglu, and Vishal M. Patel. Ambiguous medical image segmentation using diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11536-11546, June 2023.
[33] Danilo Rezende and Shakir Mohamed, Variational inference with normalizing flows. In International conference on machine learning, pages 1530-1538. PMIRR, 2015.
[34] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assisted intervention-MICCH. 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages Springer, 2015.
[35] Raghavendra Selvan, Frederik Faye, Jon Middleton, and Akshay Pai. Uncertainty quantification in medical image segmentation with normalizing flows. In Mingxia Liu, Pingkun Yan, Chunfeng Lian, and Xiaohuan Cao, editors, Machine Learning in Medical Imaging, pages 80-90, Cham, 2020.
[36] Jakub Tomczak and Max Welling. Vae with vampprior. In International conference artificial intelligence and statisticx, pages 1214-1223, PMLR, 2018
[37] Jakub M Tomczak and Max Welling. Improving variational auto- auto-encoders using householder flow, arXiv preprint arXiv:1611.0966 2016.
12

---

[38] Jingdong Wang, Ke Sun, Tianheng Cheng, Borui Jiang, Chaorui Deng, Yang Zhao, Dong Liu, Yadong Mu, Mingkui Tan, Xinggang Wang. al. Deep high-resolution representation learning for visual recognition. IEEE transactions on pattern analysis and machine intelligence, 43(10))3349-3366 2020.
[39] Wei Wang, Kaicheng Yu, Joachim Hugonot, Pascal Fua, and Mathieu Salzmann. Recurrent u-net for resource-constrane segmentation. 2019 IEEE/CVF International Conference on Computer Vision
(ICCV), pages 2142-2151, 2019.
[40] Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun. Unified perceptual parsing scene understanding. In Proceedings of the European conference on computer vision (ECCV), pages 418 434, 2018.
[41] Lukas Zbinden, Lars Doorenbos, Theodoros Pissas, Adrian Thomas Huber, Raphael Sznitman, and Pablo Márquez-Neie Ncila. Stochastic segmentation with conditional categorical diffusion models. In 2023 IEEE/CVF International Conference Computer Vision (ICCV), pages 1119-1129) 2023.
[42] Lei Zhang and Qiang Ji. A bayesian network model for automatic and interactive image segmentation. IEEE Transactions on Image Processing. 20(9):2582-255 2011.
[43] Lei Zhang, Fan Yang. Yimin Daniel Zhang, and Ying Julie Zhu. Road crack detection using deep convolutional neural network. In 2016 IEEE International Conference Image Processing (ICIP), pages 3708 3712, 2016.
[44] Wei Zhang, Xiaohong Zhang, Sheng Huang, Yuting Lu, and Kun Wang. Pixelseg: Pixel-by-pixce stochastic semantic segmentation for ambiguous medical images. MM '22, page 4742 4750, New York, NY, USA, 2022. Association for Computing Machinery.
[45] Wei Zhang, Xiaohong Zhang. Sheng Huang. Yuting Lu, Kun Wang. probabilistic model for controlling diversity and accuracy of ambiguous medical image segmentation. In Proceedings of the 30th ACM International Conference on Multimedia, MM '22, page 4751-4759, New York, NY, USA, 2022. Association for Computing Machinery.
13

---

A Appendix supplemental material
A.1 Neural Spline Flows (NSF)
Neural Spline Flows 13
propose replace the simple element-wise affine transformation in coupling (or autoregressive) layers with more expressive monotonic functions defined by splines. Specifically, they use monotonic rational-quadratic splines.
A.1.1 Monotonic Rational-Quadrate Spline Transformation (for single element x:;)
This transformation maps an input x; an output yi = g(zz). The function g is defined piecewise by K rational-quadratic segments over K bins.
1. Domain and Knots: The transformation is typically defined over a bounded interval
[- B, Bj. This interval is partitioned by K + knot points {(x((),
Boundary knots:
(-B, - B) and (x:(K), (e(K)y(()) (B.
Internal knots: K - internal knot coordinates (x((), y (*) for k I...., K are outputs of neural network (conditioned on the other part of the input, Ta. in a
coupling layer). These are constrained such that (0)) x1) <.. < x(() and ( <
y( < ... < (K) to ensure monotonicity. This is often achieved by parameterizing bin widths and heights via a softmax, then taking cumulative sums.
2. Derivatives at Knots: For the spline to be C1 continuous (continuous with continuous first derivative), derivatives 5(4) (dg (==))) are needed at each knot.
Boundary derivatives g(0) and ((K) are often fixed (e.g., to 1, implying linear "tails" there).
outside B, B) if g(x)
K 1 internal derivatives (k() > 0 for k = 1, K are also outputs of the neural network (passed through a softplus function to ensure positivity).
3. Rational-Quadratic Function in Bin k: Consider the k-th bin, from (z (()) to (x(k+1), , (y Let E(x)
(kp(()) So, 0 S g S 1 as x goes from ((k to (8+1))
1 Let (b)) 2 aI be the slope of the line segment connecting the knots of bin k. The
transformation gu (E) within this bin is given by (Gregory and Delbourgo, 1982):
Nx(E) (x(E) = (x) - (b++) - -y()) Dz(E)
where:
Nx(E) = s(0)g2 + 8(0) -E)
(E( (k)) [8(k+1) ((k) 25 (2(())E((1 - E)
This specific form ensures that (0) (()), (1( y(k+1), (dg/de) 8(4) (z(k+1) - x()),, and (dg/dE)l<=1 8(++1) positions and positive derivatives ensure (vr is monotonically increasing.
(1-1) (( (k)). The constraints on knot
A.1.2 Invertibility
gi (v) within bin): To find the inverse x = = (or =
1. First, identify the correct bin k such that (A) S (k+1), This can be done with binary search since the (A) values are sorted.
2. Then, solve y gk (E) for g. This equation is quadratic in E. Let Ay =y (x()))) and
Aybin (k++) (k). The equation becomes:
Ay (s(M) + [0(*+1) ((h) 23 20 - E)) -Aybnn (o10g2 + g(4)g(1 - 5))
14

---

Rearranging terms yields a quadratic equation of the form of + bf + 0:
Aybin S(A)
Av(o!k+1) + 5(h) )
= Ayhin ((M) Ay ((M) (x) Ays
C= = -yys(k) (Note: Error in previous formula, this term arises from & coefficient if not careful)
Corrected : The quadratic derived from = gx(E) typically written directly as:
(Aybins (R)
Ay( + 5(k) 28 (k) ))²²
+ (Aybin Ay(d(A+1) + 8(k) (k) Ays(K))
- -yss(
(This form seems more consistent with solving for & directly)
Since the function is monotonie within the bin, there will be exactly one root [0, 1]. The
solution is typically given as g -b-vb²-4o for numerical stability when 4ac is small.
2n
3. Once & is found, x = (*) + (xx(++1) - (*))).
The coefficients for the quadratic in to solve are (from Durkan et al., 2019, Appendix A, re-arranged for at² +  + 0):
a' (y(k+1) (s()) o(h)) + (v (0(k)() ((k) k2)
8' = (g(k+1) - (((h)) - (y - (8( + 5(()) - (A)
d' = gg(1) (y 11 (k))
Then {
20'
-V- (10)2-aad Aa'd"
A.1.3 Jacobian Determinant (Element-wise Derivative)
The derivative gx(x)= dg/ Ag/d for in bin le is:
dg
dgs
dE
dx
dc
dx
We have de dx 2(++1)-s0 -(())- The derivative
dg. dg
is the derivative of the rational expression.
The paper gives the final form of g' (xx as:
(((x)= (2()))2 1 + 28 (x ()) ((x) - E
(E)
where D&(E) (k) + [8(++1] + g(()) [s())][( - 5). Since g'(x) > (due to monotonicity
constraints), log [z)] = log(g' (x)). When used in coupling layer that transforms x6 based on Za.
the total oeg-dtterminanti is 2, (o),)( where of is the spline derivative for the i-th dimension of wa, parameterized by Ta.
A.1.4 Tails
Outside the interval |-B. B). the transformation is often set be the identity function (g(x) = 2) or a linear transformation matching the slope and value of the outermost spline segments at x BB to allow for unbounded inputs.
NSF significantly increases the flexibility of the clement-wise transformations compared to simple affine ones, allowing for more expressive models with fewer layers or simpler conditioning networks, while maintaining analytical invertibility and Jacobian computation.
A.2 RealNVP (Real Volume Preserving)
RealNVP 12 introduces affine coupling layers designed to have Jacobian matrix that is triangular, making its determinant easy to compute.
15

---

A.2.1 Affine Coupling Layer Mechanism
Consider an input vector E IRP.
Partitioning: The input a split into two disjoint parts:
To =
21:d (the first d dimensions)
Tb = 1:4 (the remaining D dimensions)
2. Transformation: The output E R (with parts Ya and yb) is defined as:
ya Ja (Identity transformation for the first part)
J. exp(s(ra)) + t(xa) (Affine transformation for the second part)
Here, @ denotes element-wise multiplication. The functions & IRd + RR-dd (scale) and 1: Rd -> R- (translation) are typically implemented neural networks. They take as input and produce the parameters for transforming Ib-
A.2.2 Invertibility
The inverse transformation =f===((o (y) ) is analytically computable:
1. From = -o, we directly get = Va-
2. Substitute Za into the second equation to solve for Ts:
Y/ 26 3 exp(s(Ma)) + ((Va)
(yb ((aa) exp(-s(ya))
A.2.3 Jacobian Determinant
The Jacobian matrix of the transformation f (i.e., 0y Oz) has the following structure:
(Vya, Uh)
Oya/ /rr 0yn/ /dxa 0
J = axra Ib.)
Analyzing the blocks:
dya/Or, Ia (Identity matrix of size d x d)
dya/Ozs 0 (Matrix of zeros of size (D-d)
/dr. This block involves derivatives of s(x.)) and (xra)) with respect F.- Its exact form is complex but not needed for the determinant.
0y,/0r, = ((xe(((xxs((x))) (A diagonal matrix of size (D- (D (D-dl d) whose diagonal elements exp(s; (xr)) for component j)
So, the Jacobian matrix J is:
diag diag(exp(((xx))))
Dea
This is a lower triangular block matrix. The determinant of a triangular matrix is the product of its diagonal elements.
D-d
D-d
det(J) dea( det((iigg(ep((((a))) =1 II exp(s;(2a) exp
Σ sy((r))
The log-determinant term required for Eq. A1 log det del (01- Using the property log (z)(0(= Therefore, for
det(AA 1/ det(A). we have log (v) density evaluation with f-1, the determeerminant is Σ sy (za). assuming s(x.) directly outputs the log-scale factors. (If s(x..) outputs scale factors, then it's Σ log (sy(x))). The crucial part is that s(x,) is conditioned on Za. the part of that is not transformed by this specific affine operation.
16

---

A.2.4
Composition and Alternation
Multiple coupling layers are stacked. To ensure all variables can influence each other, the roles of which part is passed through (identity) and which part is transformed are alternated in successive layers, or a permutation of dimensions is applied between layers.
A.3
Glow
Glow [22] enhances the RealNVP framework by introducing two new operations within each "step" of the flow, in addition to the affine coupling layer. A single step in Glow comprises:
Actnorm
2. Invertible 1x1 Convolution
3. Affine Coupling Layer
A.3.1 Actnorm (Activation Normalization)
This layer performs an element-wise affine transformation per channel:
Ychw = Ote Icchw + Be
where &ehw the activation at channel c, height h, width w. aa (scale) and B. (bias) are learnable parameters, one per channel
Initialization: a, and B. are initialized such that, for the first minibatch data, the output activations Yehar for each channel e have zero mean and unit variance across spatial and batch dimensions. After this initialization, Or, and 3r become regular trainable parameters independent of minibatch statistics.
Inverse: Fehw = (Nebw B,)/
Log- Determinant: The Jacobian is diagonal matrix. For an input tensor H x W spatial dimensions and C' channels, the total log-determinant is:
log (et = H W.
log loel
A.3.2 Invertible 1x1 Convolution
This operation generalizes fixed permutations. A convolution with input channels and C' output channels applies linear transformation to the channel vector at each spatial location independently. oh/h E R is the vector of channels spatial location (h, w). and W E RCCC is the convolution kernel weight matrix, then:
sha
Who
Invertibility: The matrix W must be invertible. The inverse operation is x//..- = -W V/aw Log- Determmnnt:n::: For an input tensor of H x W spatial dimensions, the total log- determinant is:
log (et((Ixtcom))) ) H W log det(WW)
Efficient Computation of det(W)) and W- To avoid the O(C*) cost of directly comput- ing (W)) and W-1, W parameterized its LU decomposition:
W = PIU
where:
P permutation matrix (fixed, not learned).
a lower triangular matrix with ones on its diagonal.
U is an upper triangular matrix whose diagonal elements Us are learned (and must be non-zero).
17

---

The determinant is then det W) = det(P) det(L) det(U) = (+1) 1- [ Us. So. log det(w) Σイ log U3: The parameters Lij for j and U.j for i S j are learned. - p-1 also efficient due the triangular nature of L and U.
Inversion W-1=
A.3.3 Affine Coupling Layer
This is the same as in RealNVP.
A.3.4 Multi-Scale Architecture
Glow also employs multi-scale architecture. After few steps of flow at given resolution:
1. Squeeze: Spatial dimensions are reduced and channel dimensions are increased. For 8X8X imput, it reshapes into (s/2) (s/2) 4c.
2. Factor Out: Half of the channels are split off and assumed to follow Gaussian distribution. These factored-out variables contribute directly to the latent u. The other half continue through subsequent flow steps at the new, coarser resolution.
The total log-determinant for Glow model is the sum of the log-determinants from all actnorm layers, all 1xl convolution layers, and all affine coupling layers, plus the bog-probability of the factored-out variables under their assumed Gaussian prior.
18

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how be expected to generalize other settings.
much the results
is fine to include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
19

---

Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
20

---

Answer: [No]
Because of the confidential consideration, repository will be access after published.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm ec/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https SonedeSubmissiooPP..
//nips.cc/pp wwwppbblccgguiicuceccccccoeooopp
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
21

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics httts::/neerrpppt http://aaarrg..cccbbbbbll/ccccle
Question:
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
22

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
23

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [TODO]
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
24


Reviews:
[{"review_id": "QddSSq4pvG", "reviewer": "Reviewer_gS3U", "summary": "The 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.", "questions": "1. How is the single-sample inference conducted? More details should be provided.\n2. What are the missing details mentioned in the weaknesses section?", "limitations": "No. 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.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# 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.", "was_revised": true, "final_justification": "The authors addressed all my concerns so I keep my initial positive rating.", "rebuttal": "### Q: Missing details about the latent space dimension, the number of distributions of the Gaussian mixture model being used in each dataset.\n\nThank you very much for your reminder. In our experiments, the number of NSF layers is set to 3 across all three datasets, and the dimensionality of the latent variable is 8. The number of GMM components $K$ is set to 3 for both the LIDC and Crack500 datasets, and 5 for the Cityscapes dataset.\n\n### Q: I'm curious to see if there is a correlation between number of semantic classes and the number of distributions.\n\nThank you for your question; it is indeed a very valuable point. The main purpose of using GMM in our model is to help NF learn better and improve performance. We do not model each semantic label separately; instead, we perform modeling directly on the input image. This implies that the GMM clusters mainly operate in the latent space and are not directly equivalent to the label space. Therefore, the number of GMM components is not inherently related to the number of segmentation labels.\n\nTo test this idea, we did a simple experiment on the Cityscapes dataset. We fixed the number of semantic labels to 19 and increased the number of GMM components $K$ step by step. If there was a strong link between $K$ and the number of labels, the model should perform better when $K$ gets close to 19. Due to huge data and experiments, limited time and resources, we only trained each one for 200 epochs to save time.\n\n**Table: Effect of GMM component number $K$ (Cityscapes, 19 labels).**\n\n| GMM Components ($K$) | mIOU    | Param |\n|----------------------|---------|--------|\n| 1                    | 68.31   | 38.0   |\n| 3                    | 71.72   | 38.0   |\n| 5                    | **71.94** | 38.0 |\n| 6                    | 71.73   | 38.1   |\n| 7                    | 71.55   | 38.1   |\n| 9                    | 71.62   | 38.2   |\n| 11                   | 71.68   | 38.2   |\n| 13                   | 71.19   | 38.3   |\n| 15                   | 70.26   | 38.3   |\n| 19                   | 69.13   | 38.4   |\n\nHowever, as shown in the results table, the model achieved the best performance when $K = 5$. After that, the performance fluctuated slightly and then continued to drop when $K > 11$, likely due to too many ineffective clusters. This result indirectly proves our hypothesis. If there were a positive or negative correlation between $K$ and the number of labels, we would expect the performance to steadily increase or decrease as $K$ increases. But instead, the performance first increased, then stayed stable, and finally decreased, which supports our assumption.\n\nMoreover, the ablation study in the appendix also confirms this. On the Crack500 dataset, the performance difference between $K=3$ and $K=6$ is very small.\n\nDue to time and computing limitations, we were only able to conduct a simple experiment and brief analysis on this question. However, we believe the current results provide initial evidence to support our hypothesis. In future work, we plan to perform more detailed experiments to further validate our theory.\n\n### Q: How is the single-sample inference conducted? More details should be provided.\n\nOur implementation supports two single-sample inference modes. The first is the maximum-confidence mode. In this mode, we first identify the GMM component with the highest posterior mixing coefficient for a given input. We then select the mean of this component $\\mu_k$, as the deterministic latent variable for inference. This represents the most likely semantic cluster for the input. The second is the random-sampling mode, where we first sample a component $k$ based on the mixing coefficients, and then sample a latent variable $z$ from that component's Gaussian distribution $\\mathcal{N}(\\mu_k, \\Sigma_k)$. Based on our experiments, there was not a significant difference in performance, so we selected the random-sampling mode in our main experiments to maintain the randomness of our model.\n\n### Q: The idea of combining GMM and NF has been done in some prior works, e.g. FlowGMM.\n\nThank you very much for your reminder. Our literature review only explored stochastic segmentation. To the best of our knowledge, we applied GMM+NF for the first time in stochastic segmentation. Papers with similar ideas in other fields have been updated in the literature review.\n\n### Q: What are the missing details mentioned in the weaknesses section?\n\nThank you very much for your reminder. This study has two main limitations: First, compared to conventional CVAE models, our model introduces a significant increase in computational cost. Although inference speed is not a major concern in target application scenarios of stochastic segmentation, this is still an aspect to improve in our future work. Second, the training process is less stable. Specifically, the normalizing flow requires computing the log-determinant for gradient updates, which makes it more prone to KL vanishing compared to CVAE. As a result, careful tuning of hyperparameters is necessary, along with additional strategies such as annealing schedules and constraining the variance range of GMM components to ensure stable training."}, {"review_id": "Lpg32Zs2bh", "reviewer": "Reviewer_VZXu", "summary": "The 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.", "questions": "Please refer to the Weaknesses.", "limitations": "The current version of the paper needs to be improved due to the insufficient evaluation and the misalignment to the motivation.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths:\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?", "was_revised": true, "final_justification": "Thanks for the authors' rebuttal. My major concerns are addressed and I'd like to raise the score to 4.", "rebuttal": "### Q: 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\nWe appreciate the comment from the reviewer. There may be a misunderstanding between the background and the motivation of our work. Specifically, stochastic segmentation mainly aims to address uncertainty and annotation variability. Many solutions have been proposed for this task, which can be broadly categorized into diffusion-based and CVAE-based approaches.\n\nThe main drawback of diffusion-based methods is the lack of interpretability of the distribution. This is critical in security-sensitive tasks such as medical image segmentation, where explainability is essential. On the other hand, CVAE-based methods often suffer from limited data modeling capacity. For instance, in Probabilistic U-Net, replacing the Gaussian distribution with pure random noise still yields comparable performance. This suggests that simple distributions struggle to capture the complexity of real-world data, and thus, the latent variables contribute limited semantic information. In practice, this limitation often manifests as large variances in latent space. There are some attempts to solve these problems, but they have all kinds of drawbacks. Please refer to the related work section for details\n\nOur key innovation lies in increasing the complexity of both the prior and posterior distributions by combining GMM with NF for the first time. This enables better modeling of complex real-world data distributions. The core component is the NF, which transforms a simple base distribution into a more complex one. The GMM is used to provide a more flexible and structured base distribution, thereby helping the NF to learn more effectively and achieve better performance. We adopt the GMM for the following reasons: Compared to a single Gaussian, GMM offers greater representational capacity, allowing the latent space to explicitly encode distinct semantic regions in the image. Through a parameter competition mechanism, each Gaussian component is encouraged to specialize in specific semantic attributes, thereby enhancing both the interpretability of the latent space and the clarity of the resulting segmentations.\n\n---\n\n### Q: The motivation and application of this paper is not well-aligned.\n\nThank you for your comment. The three selected datasets represent three different application scenarios, which allows for a more comprehensive evaluation of our model.\n\nMulti-annotator datasets are relatively scarce. To the best of our knowledge, only LIDC and QUBIQ are publicly available. Among them, LIDC is widely used and validated in the stochastic segmentation community. It provides a fair and consistent benchmark for evaluating model performance and enables direct comparison with previous works which have passed peer reviews. Therefore, we chose LIDC as the multi-annotator dataset for our experiments.\n\n---\n\n### Q: The overall logic is not very convincing[...]?\n\nAll baseline results in our paper are taken from state-of-the-art publications in their respective areas. Specifically, Table 1 and Table 2 in the main paper are cited from CCDM, a diffusion-based stochastic segmentation model, while Table 3 is cited from JTFN+CIRL.\n\nIn fact, many stochastic segmentation papers have evaluated their methods on the Cityscapes dataset and achieved competitive results compared to deterministic models (can check all the baseline methods we compared in table 1 of main paper). It is important to note that if the decoder of a stochastic model does not incorporate latent variables, it essentially becomes a deterministic model. When the latent space is well-modeled, the variance across different outputs is structured and meaningful, rather than being purely random. This indicates that the model’s stochastic behavior is controlled and interpretable.\n\n---\n\n### Q: Important ablation studies are missing.  \n(1) how do the core components, GMM and NF affect performance?  \n(2) How do the components affect stochastic segmentation performance?  \n(3) What happens if we remove GMM?\n\nThank you for your suggestions regarding the ablation experiment section. In fact, these questions have already been addressed in the ablation study on the Crack500 dataset presented in the appendix:\n\n1. In **Experiment (a)**, we set $K=1$, making the GMM equivalent to a single Gaussian. This allows us to observe the effect of removing GMM from the model.  \n2. In **Experiment (d)**, we remove the NF module and use only the GMM, which helps evaluate the contribution of NF alone.  \n3. **Experiments (b)** and **(c)** show how model performance changes as $K$ increases.  \n4. **Experiments (e)** and **(f)** demonstrate the impact of using different types of NF on the final performance.  \n\nIn addition, we also added additional ablation experiments on the LIDC and Cityscapes datasets in the appendix, including:  \n- The performance of the model with and without the GMM module (when $K=1$).  \n- The impact of removing the NSF module.  \n- The effect of increasing the number of GMM components $K$.  \n- The effect of increasing the number of NF layers.  \n\nIn addition to comparing model performance, we have also included training time per epoch, peak GPU memory usage, inference latency, and total parameter size to provide a more comprehensive analysis.\n\n**Table: Ablation study on Cityscapes dataset**\n| Experiment No. | a| b| c| d| e| f|\n|----------------|-------|-------|-------|-------|-------|-------|\n| No. of NSF layer | 3   | 3| 3| -| 4| 6|\n| No. of GMM| 1| 3| 6| 3| 3| 3|\n| mIOU| 68.04 | 71.23 | 71.97 | 65.39 | 69.15 | 69.09 |\n| Train time (s)| 323| 355| 370| 197| 569| 861|\n| Peak GPU (MB)| 2239.19 | 2239.36 | 2239.87 | 1988.27 | 2240.44 | 2242.48 |\n| Infer latency (s)| 0.4030 | 0.4129 | 0.4179 | 0.1705 | 0.5894 | 0.7171 |\n| Param size (M)| 38.0|38.0|38.1| 33.7|38.0| 38.0|\n\n**Table: Ablation study on LIDC dataset**\n\n| Experiment No. | a| b| c| d| e| f|\n|----------------|-------|-------|-------|-------|-------|-------|\n| No. of NSF layer | 3| 3| 3| -| 4| 6|\n| No. of GMM| 1   | 3| 6| 3| 3| 3|\n| HM-IOU$_{16}$| 0.621 | 0.640 | 0.637 | 0.574 | 0.607 | 0.582 |\n| Train time (s)| 217| 222| 245| 114| 466| 739|\n| Peak GPU (MB)| 738.55 | 738.74 | 738.99 | 687.78 | 739.77 | 741.82 |\n| Infer latency (s)| 0.3222 | 0.3294 | 0.3311 | 0.0899 | 0.4259 | 0.6114 |\n| Param size (M)| 37.9  | 38.0  | 38.1  | 33.6  | 38.0  | 38.0  |\n\nAs shown in the ablation study in Cityscapes and LIDC, incorporating the normalizing flow module significantly increases inference latency and training time. However, we believe this trade-off is acceptable for the following reasons:  Although NF adds computational overhead, it provides a much more useful latent space, more accurate uncertainty modeling, and more accurate segmentation outputs. These advantages are particularly valuable in high-reliability tasks such as medical image segmentation, where performance and interpretability are crucial. In typical stochastic segmentation scenarios, the focus is often on uncertainty estimation, diversity of predictions, and output reliability rather than on real-time inference speed. Lastly, the current state-of-the-art in stochastic segmentation is diffusion-based models, which require 50 to 1000 denoising steps per sample. For example, CCDM reports an inference latency of 15.67 seconds, significantly higher than ours.\n\n---\n\n### Q: 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\nThank you for your question; it is indeed a very valuable point. The main purpose of using GMM in our model is to help NF learn better and improve performance. We do not model each semantic label separately; instead, we perform modeling directly on the input image. This implies that the GMM clusters mainly operate in the latent space and are not directly equivalent to the label space. Therefore, the number of GMM components is not inherently related to the number of segmentation labels.\n\nTo test this idea, we did a simple experiment on the Cityscapes dataset. We fixed the number of semantic labels to 19 and increased the number of GMM components $K$ step by step. Because we needed to train many models, we only trained each one for 200 epochs to save time.\n\n**Table: Effect of GMM component number $K$ (Cityscapes, 19 labels)**\n| GMM Components ($K$) | mIOU    | Param |\n|----------------------|---------|--------|\n| 1| 68.31| 38.0|\n| 3| 71.72| 38.0|\n| 5| **71.94** | 38.0 |\n| 6| 71.73| 38.1|\n| 7| 71.55| 38.1|\n| 9| 71.62| 38.2|\n| 11| 71.68| 38.2|\n| 13| 71.19| 38.3|\n| 15| 70.26| 38.3|\n| 19| 69.13| 38.4|\n\nHowever, as shown in the results table, the model achieved the best performance when $K = 5$. After that, the performance fluctuated slightly and then continued to drop when $K > 11$, likely due to too many ineffective clusters. This result indirectly proves our hypothesis. If there were a positive or negative correlation between $K$ and the number of labels, we would expect the performance to steadily increase or decrease as $K$ increases. But instead, the performance first increased, then stayed stable, and finally decreased, which supports our assumption.\n\nMoreover, the ablation study in the appendix also confirms this. On the Crack500 dataset, the performance difference between $K=3$ and $K=6$ is very small.\n\nDue to time and computing limitations, we were only able to conduct a simple experiment and brief analysis on this question. However, we believe the current results provide initial evidence to support our hypothesis. In future work, we plan to perform more detailed experiments to further validate our theory."}, {"review_id": "DQRL0DQuNx", "reviewer": "Reviewer_4wT6", "summary": "This 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.", "questions": "1. *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.", "limitations": "The 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.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "*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.", "was_revised": true, "final_justification": "I reviewed the author response, the new multi-seed results, and the added compute/ablation tables. The paper introduces a structured CVAE with a GMM latent and constrained normalizing flows, and shows consistent gains over probabilistic U-Net and diffusion-based baselines on LIDC, Crack500, and Cityscapes. The authors addressed my key concerns by providing (i) three-seed results with mean ± std and (ii) clear reports of training time, inference latency, and peak memory. These additions increase my confidence in the empirical claims and do not change my positive assessment of the technical contribution.", "rebuttal": "### Q: Statistical robustness across multiple random seeds.\n\nThank you for the valuable suggestion. We re-evaluated our model using three different random seeds: 194, 683, and 2596. The results are as follows:\n\n1. On the LIDC dataset, the model achieved $GED_{16}$ to $GED_{100}$ scores of $0.194 \\pm 0.002$, $0.191 \\pm 0.003$, $0.181 \\pm 0.002$, and $0.186 \\pm 0.002$, respectively. The $HM\\text{-}IoU_{16}$ and $HM\\text{-}IoU_{32}$ scores were $0.640 \\pm 0.005$ and $0.642 \\pm 0.004$, respectively.\n2. On the Cityscapes dataset, the model achieved an IoU of $72.4 \\pm 0.2$.\n3. On the Crack500 dataset, the precision was $70.15 \\pm 0.41$, recall was $71.56 \\pm 0.28$, and F1-score was $70.42 \\pm 0.27$.\n\nWe also evaluated the Probabilistic U-Net under the same settings and conducted a Wilcoxon signed-rank test to compare it with our model. Due to the small sample size ($n=3$), the test yielded a p-value of 0.25, which is the lowest possible resolution at this scale.\n\n---\n\n### Q: Quantify computational overhead and memory footprint.\n\nThank you for your suggestion regarding the ablation study. In response, we have added additional ablation experiments on the LIDC and Cityscapes datasets in the appendix, including:\n\n1. The performance of the model with and without the GMM module (when $K=1$).  \n2. The impact of removing the NSF module.  \n3. The effect of increasing the number of GMM components $K$.  \n4. The effect of increasing the number of NF layers.\n\nIn addition to comparing model performance, we have also included training time per epoch, peak GPU memory usage, inference latency, and total parameter size to provide a more comprehensive analysis.\n\n**Table: Ablation study on Cityscapes dataset**\n\n| Experiment No.         | a     | b     | c     | d     | e     | f     |\n|------------------------|-------|-------|-------|-------|-------|-------|\n| No. of NSF layer       | 3     | 3     | 3     | -     | 4     | 6     |\n| No. of GMM             | 1     | 3     | 6     | 3     | 3     | 3     |\n| mIOU                   | 68.04 | 71.23 | 71.97 | 65.39 | 69.15 | 69.09 |\n| Train time (s)         | 323   | 355   | 370   | 197   | 569   | 861   |\n| Peak GPU memory (MB)   | 2239.19 | 2239.36 | 2239.87 | 1988.27 | 2240.44 | 2242.48 |\n| Infer latency (s)      | 0.4030 | 0.4129 | 0.4179 | 0.1705 | 0.5894 | 0.7171 |\n| Param size (M)         | 38.0  | 38.0  | 38.1  | 33.7  | 38.0  | 38.0  |\n\n**Table: Ablation study on LIDC dataset**\n\n| Experiment No.         | a     | b     | c     | d     | e     | f     |\n|------------------------|-------|-------|-------|-------|-------|-------|\n| No. of NSF layer       | 3     | 3     | 3     | -     | 4     | 6     |\n| No. of GMM             | 1     | 3     | 6     | 3     | 3     | 3     |\n| HM-IoU$_{16}$          | 0.621 | 0.640 | 0.637 | 0.574 | 0.607 | 0.582 |\n| Train time (s)         | 217   | 222   | 245   | 114   | 466   | 739   |\n| Peak GPU memory (MB)   | 738.55 | 738.74 | 738.99 | 687.78 | 739.77 | 741.82 |\n| Infer latency (s)      | 0.3222 | 0.3294 | 0.3311 | 0.0899 | 0.4259 | 0.6114 |\n| Param size (M)         | 37.9  | 38.0  | 38.1  | 33.6  | 38.0  | 38.0  |\n\nAs a reference, the peak GPU memory usage of the baseline Probabilistic U-Net is 669.61 MB on LIDC and 1624.30 MB on Cityscapes.\n\nAs shown in the ablation study in Cityscapes and LIDC, incorporating the normalizing flow module significantly increases inference latency and training time. However, we believe this trade-off is acceptable for the following reasons:  Although NF adds computational overhead, it provides a much more useful latent space, more accurate uncertainty modeling, and more accurate segmentation outputs. These advantages are particularly valuable in high-reliability tasks such as medical image segmentation, where performance and interpretability are crucial. In typical stochastic segmentation scenarios, the focus is often on uncertainty estimation, diversity of predictions, and output reliability rather than on real-time inference speed. Lastly, the current state-of-the-art in stochastic segmentation is diffusion-based models, which require 50 to 1000 denoising steps per sample. For example, CCDM reports an inference latency of 15.67 seconds, significantly higher than ours.\n\n\n---\n\n### Q: Limitations and Social impact\n\nThank you very much for your reminder. This study has two main limitations:\n\n1. Compared to conventional CVAE models, our model introduces a significant increase in computational cost. Although inference speed is not a major concern in target application scenarios of stochastic segmentation, this is still an aspect to improve in our future work.\n2. The training process is less stable. Specifically, the normalizing flow requires computing the log-determinant for gradient updates, which makes it more prone to KL vanishing compared to CVAE. As a result, careful tuning of hyperparameters is necessary, along with additional strategies such as annealing schedules and constraining the variance range of GMM components to ensure stable training.\n\nFailure cases are inevitable in any model. Unlike deterministic segmentation, due to its randomness, stochastic segmentation can determine which regions have high confidence and which have low confidence by outputting multiple possible results and calculating the standard deviation of each pixel point, thereby ensuring that the occasional failed segmentation samples do not cause serious social impact. This is the reason why it can be widely applied to highly sensitive tasks."}, {"review_id": "XwpUabfYhp", "reviewer": "Reviewer_kRCB", "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.", "questions": "- Are there any other sorts of architectures is this approach suited for integration with besides CVAE?\n- How much training time is added by the incorporation of these components over training of the networks alone (as in the LIDC data)?\n- 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.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "## Quality:\n- Strength: the authors' approach improves on state-of-the-field baselines for several disparate tasks, including newer architectures based on diffusion\n\n## Clarity\n- 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\n\n## Significance\n- 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\n- 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\n\n## Originality\n- I am unaware of other work proposing a similar approach but this is the least confident part of my review", "was_revised": true, "final_justification": "I thank the authors for their comments and am happy to keep my score at an accept.", "rebuttal": "### Q: Are there any other sorts of architectures this approach is suited for integration with besides CVAE?\n\nThe first thing that comes to our mind is that this combination can be applied in Bayesian Neural Networks (BNN), by replacing the posterior estimation of simple isotropic Gaussian with GMM+NF. Furthermore, according to the comment from reviewer gS3U, this design can also benefit generative semi-supervised learning architectures, such as FlowGMM.\n\n---\n\n### Q: Do PHiSeg and Hierarchical Probabilistic U-Net, when trained on LIDC, exhibit the same response to randomization of the latent variable as the ProbUNet does?\n\nAccording to our experiments, randomizing the latent variables in PHiSeg and Hierarchical Probabilistic U-Net results in a more significant performance drop compared to Probabilistic U-Net. Under the same experimental settings, after applying randomization:  \n- For PHiSeg, the $GED_{16}$ score increased from 0.284 to 0.311, and the $HM$-$IoU_{16}$ decreased from 0.570 to 0.524.  \n- For Hierarchical Probabilistic U-Net, the $GED_{16}$ score increased from 0.265 to 0.337, and the $HM$-$IoU_{16}$ decreased from 0.555 to 0.518.\n\nWe believe this is because both PHiSeg and Hierarchical Probabilistic U-Net adopt a hierarchical modeling strategy. Rather than using a single latent space, they model multiple levels of latent variables. This hierarchical structure distributes the diversity modeling capability across different layers of the network, which offsets the negative effect of sampling from a near-random Gaussian latent space slightly. As a result, we did not observe the same conclusion as reported for Probabilistic U-Net.\n\n**Note:** Since the baseline results in Table 1 of the main paper are directly taken from the CCDM paper, our reproduced results may slightly differ. However, the presented experiment addresses the concern about the effect of latent variable randomization across different models.\n\n---\n\n### Q: How much training time is added by the incorporation of these components over training of the networks alone (as in the LIDC data)?\n\nThank you for your suggestion. In response, we have added additional ablation experiments on the LIDC and Cityscapes datasets in the appendix, including:\n\n1. The performance of the model with and without the GMM module (when $K=1$)  \n2. The impact of removing the NSF module  \n3. The effect of increasing the number of GMM components $K$  \n4. The effect of increasing the number of NF layers\n\nWe also include training time per epoch, peak GPU memory usage, inference latency, and total parameter size for a more comprehensive comparison.\n\n**Ablation Study on Cityscapes**\n\n| Experiment No.          | a     | b     | c     | d     | e     | f     |\n|-------------------------|-------|-------|-------|-------|-------|-------|\n| No. of NSF layer        | 3     | 3     | 3     | -     | 4     | 6     |\n| No. of GMM              | 1     | 3     | 6     | 3     | 3     | 3     |\n| mIOU                    | 68.04 | 71.23 | 71.97 | 65.39 | 69.15 | 69.09 |\n| Train time (s)          | 323   | 355   | 370   | 197   | 569   | 861   |\n| Peak GPU memory (MB)    | 2239.19 | 2239.36 | 2239.87 | 1988.27 | 2240.44 | 2242.48 |\n| Inference latency (s)   | 0.4030 | 0.4129 | 0.4179 | 0.1705 | 0.5894 | 0.7171 |\n| Param size (M)          | 38.0  | 38.0  | 38.1  | 33.7  | 38.0  | 38.0  |\n\n**Ablation Study on LIDC**\n\n| Experiment No.          | a     | b     | c     | d     | e     | f     |\n|-------------------------|-------|-------|-------|-------|-------|-------|\n| No. of NSF layer        | 3     | 3     | 3     | -     | 4     | 6     |\n| No. of GMM              | 1     | 3     | 6     | 3     | 3     | 3     |\n| HM-IoU$_{16}$           | 0.621 | 0.640 | 0.637 | 0.574 | 0.607 | 0.582 |\n| Train time (s)          | 217   | 222   | 245   | 114   | 466   | 739   |\n| Peak GPU memory (MB)    | 738.55 | 738.74 | 738.99 | 687.78 | 739.77 | 741.82 |\n| Inference latency (s)   | 0.3222 | 0.3294 | 0.3311 | 0.0899 | 0.4259 | 0.6114 |\n| Param size (M)          | 37.9  | 38.0  | 38.1  | 33.6  | 38.0  | 38.0  |\n\nAs shown in the ablation study in Cityscapes and LIDC, incorporating the normalizing flow module significantly increases inference latency and training time. However, we believe this trade-off is acceptable for the following reasons:  Although NF adds computational overhead, it provides a much more useful latent space, more accurate uncertainty modeling, and more accurate segmentation outputs. These advantages are particularly valuable in high-reliability tasks such as medical image segmentation, where performance and interpretability are crucial. In typical stochastic segmentation scenarios, the focus is often on uncertainty estimation, diversity of predictions, and output reliability rather than on real-time inference speed. Lastly, the current state-of-the-art in stochastic segmentation is diffusion-based models, which require 50 to 1000 denoising steps per sample. For example, CCDM reports an inference latency of 15.67 seconds, significantly higher than ours.\n\n---\n\n### Q: Limitations are not discussed and should be\n\nThank you very much for your reminder. This study has two main limitations:\n\n1. Compared to conventional CVAE models, our model introduces a significant increase in computational cost. Although inference speed is not a major concern in target application scenarios, it remains a direction for future optimization.\n2. The training process is less stable due to the use of normalizing flow, which involves computing the log-determinant in gradients. This makes the model more susceptible to KL vanishing than CVAE-based models. Thus, careful hyperparameter tuning and additional stabilizing techniques, such as annealing schedules and variance constraints for the GMM, are necessary."}]

Author rebuttal:

Meta-review: a) The authors improve existing probabilistic segmentation models by combining Normalizing Flow (NF) and Gaussian Mixture Models (GMM) in a Conditional Variational Auto Encoder (CVAE) with the aim of improving the latent space expressiveness and capturing the feature uncertainty. 

b) 
- New and original framework (NF and GMM in CVAE)
- Extensively tested on multiple datasets and SOTA results
- Ablation study to highlight the role of GMM vs NF

c) 
- Missing details or references (FlowGMM)
- Missing numerical exp

d) All reviewers are positive about the paper and their reviews mostly overlap. The work is valuable yet it remains incremental as a combination of existing ideas therefore I do not recommend spotlight of oral.

e) The weaknesses cited above were resolved during the rebuttal. All reviewers are satisfy by the authors responses.
Decision: Accept (poster)