Neural Architecture Distributions: A New Paradigm for Stochastic Segmentation

Stochastic segmentation seeks to represent multiple plausible masks for a single image, a requirement in safety- and quality-critical applications such as medical imaging and building defect inspection. Most existing methods introduce stochasticity by injecting continuous latent variables or by following iterative denoising trajectories, making their stochastic sources difficult to search or audit. We study architecture distributions as a 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 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 architecture-sampled predictions to the annotation set using an IoU-based energy-distance surrogate. We construct the candidate bank with evolutionary search, making the support of the stochastic source optimizable before distribution learning. NADS improves distribution matching and hypothesis coverage on LIDC-IDRI, and generalizes to additional stochastic benchmarks (Hippocampus, QUBIQ Prostate) as well as deterministic segmentation tasks (Cityscapes, Crack500). The resulting formulation learns an input-conditioned architecture distribution and produces multi-annotator output diversity through architecture sampling with set-level supervision.