Stochastic semantic segmentation with Stochastic Mixture of Bottleneck Experts

Modelling aleatoric uncertainty in semantic segmentation is crucial for safety–critical applications, such as autonomous driving and medical imaging. However, capturing diverse yet consistent predictions under multimodal distributions remains a significant challenge. This paper introduces the Stochastic Mixture of Bottleneck Experts (SMoBE), a novel framework that addresses aleatoric uncertainty by integrating a mixture of experts at the model’s bottleneck. SMoBE models the complex multimodal predictive distribution by resolving it into a set of unimodal distributions, each handled by a distinct expert network within the latent space. Furthermore, SMoBE models the stochastic feature map directly from the learned distributions, rather than relying on latent code injection. This approach makes randomness an inherent part of the feature representation, which in turn helps in creating diverse, spatially coherent, and plausible segmentation hypotheses. We validate SMoBE across two distinct semantic segmentation scenarios, demonstrating that it achieves performance comparable to state-of-the-art methods.