# Neural Architecture Distributions: A New Paradigm for Stochastic Segmentation

arXiv: https://arxiv.org/abs/2606.21061 · HTML: https://mlsanity.com/paper/2606.21061/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 35 of its 42 references that are in the map.
- First posted on arXiv 19 Jun 2026.

## Builds on (most cited first, 30 of 35)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Rethinking Atrous Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1706.05587/) (2017)
- [DARTS: Differentiable Architecture Search](https://mlsanity.com/paper/1806.09055/) (2018)
- [UNet++: A Nested U-Net Architecture for Medical Image Segmentation](https://mlsanity.com/paper/1807.10165/) (2018)
- [Deep High-Resolution Representation Learning for Visual Recognition](https://mlsanity.com/paper/1908.07919/) (2019)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [Unified Perceptual Parsing for Scene Understanding](https://mlsanity.com/paper/1807.10221/) (2018)
- [Swin-Unet: Unet-like Pure Transformer for Medical Image Segmentation](https://mlsanity.com/paper/2105.05537/) (2021)
- [Regularized Evolution for Image Classifier Architecture Search](https://mlsanity.com/paper/1802.01548/) (2018)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (2018)
- [Once-for-All: Train One Network and Specialize it for Efficient Deployment](https://mlsanity.com/paper/1908.09791/) (2019)
- [Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning](https://mlsanity.com/paper/2208.04202/) (2022)
- [SNAS: Stochastic Neural Architecture Search](https://mlsanity.com/paper/1812.09926/) (2018)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Topology-Preserving Deep Image Segmentation](https://mlsanity.com/paper/1906.05404/) (2019)
- [PHiSeg: Capturing Uncertainty in Medical Image Segmentation](https://mlsanity.com/paper/1906.04045/) (2019)
- [Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty](https://mlsanity.com/paper/2006.06015/) (2020)
- [Ambiguous Medical Image Segmentation using Diffusion Models](https://mlsanity.com/paper/2304.04745/) (2023)
- [Beyond the Pixel-Wise Loss for Topology-Aware Delineation](https://mlsanity.com/paper/1712.02190/) (2017)
- [A Hierarchical Probabilistic U-Net for Modeling Multi-Scale Ambiguities](https://mlsanity.com/paper/1905.13077/) (2019)
- [Stochastic Segmentation with Conditional Categorical Diffusion Models](https://mlsanity.com/paper/2303.08888/) (2023)
- [Recurrent U-Net for Resource-Constrained Segmentation](https://mlsanity.com/paper/1906.04913/) (2019)
- [Uncertainty quantification in medical image segmentation with normalizing flows](https://mlsanity.com/paper/2006.02683/) (2020)
- [Calibrated Adversarial Refinement for Stochastic Semantic Segmentation](https://mlsanity.com/paper/2006.13144/) (2020)

## Abstract (arXiv)

Stochastic segmentation seeks to represent multiple plausible masks for a single image, which is essential in safety- and quality-critical applications such as medical imaging or building defect inspection. Most existing methods introduce stochasticity by injecting continuous latent variables or by iterative denoising trajectories, whose stochastic sources are difficult to search or audit directly. We propose architecture distributions as a new stochastic source for segmentation: instead of sampling a latent variable or noise, we sample a discrete architecture from a learned distribution over operator choices at multiple searchable positions in a segmentation backbone. Each sampled architecture yields one mask through the selected active path, so inference depends on the executed subnet rather than the complete candidate bank. This approach also supports architectural provenance, since each output corresponds to a specific architecture configuration. To reduce collapse toward averaged masks, we train with set-level supervision by matching a set of architecture-sampled predictions to the annotation set using an IoU-based energy-distance surrogate. We further construct the candidate bank with evolutionary search, making the support of the stochastic source optimizable before distribution learning. The proposed method achieves state-of-the-art distribution matching and hypothesis coverage on LIDC-IDRI, and remains effective on two extension tasks. To the best of our knowledge, this is the first work to formulate stochastic segmentation as learning an architecture distribution and realizing output diversity through architecture sampling.