Neural Architecture Distributions: A New Paradigm for Stochastic Segmentation

2026-06-18

AI Summary

This work addresses the challenge of effectively modeling uncertainty arising from multiple plausible masks in high-stakes medical image segmentation. It formulates stochastic segmentation as a neural architecture distribution learning problem, where diverse and traceable deterministic masks are generated by sampling discrete subnetworks from a searchable operator space. To prevent prediction collapse, the method introduces set-level supervision and employs an IoU-based energy distance proxy loss for set matching during training. An evolutionary search strategy is further integrated to optimize the candidate operator pool. Evaluated on the LIDC-IDRI dataset, the proposed approach achieves state-of-the-art performance in distribution alignment and hypothesis coverage, while maintaining strong results across two extended tasks.

Abstract

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.

Problem

• stochastic segmentation
• neural architecture distribution
• multiple plausible masks
• output diversity
• uncertainty modeling

Innovation

• Neural Architecture Distribution
• Stochastic Segmentation
• Architecture Sampling
• Set-level Supervision
• Evolutionary Search