Neural Architecture Distributions: A New Paradigm for Stochastic Segmentation
Reviewed by Pith T0 review T1 audit T2 compute T3 formal T4 kernel 2026-06-26 14:54 UTCgrok-4.3pith:Y6ULYUYRrecord.jsonopen to challenge →
The pith
Stochastic segmentation is achieved by sampling discrete architectures from a learned distribution over operator choices rather than continuous latent variables.
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
Core claim
Formulating stochastic segmentation as learning an architecture distribution and realizing output diversity through architecture sampling produces multiple plausible masks, each traceable to a specific configuration, trained by matching prediction sets to annotation sets with set-level supervision.
What carries the argument
Architecture distribution: learned probabilities over discrete operator choices at multiple searchable positions in a segmentation backbone, sampled at inference to execute distinct sub-networks that each yield one mask.
If this is right
- Each generated mask corresponds to an explicit architecture configuration, enabling provenance tracking.
- The candidate bank is built by evolutionary search so that the support of the distribution can be optimized independently of distribution learning.
- Set-level IoU-based energy-distance supervision prevents collapse toward averaged masks.
- The approach reports state-of-the-art distribution matching and hypothesis coverage on LIDC-IDRI.
- Performance holds on two additional extension tasks beyond the primary benchmark.
Where Pith is reading between the lines
- The discrete sampling mechanism could make uncertainty estimates easier to audit or constrain by hardware cost in deployed medical systems.
- Extending the same distribution-over-architectures idea to other dense-prediction tasks such as depth or surface-normal estimation might improve multi-hypothesis calibration without continuous noise.
- Hybrid models that combine architecture sampling with a small continuous latent variable could be tested to see whether they further increase coverage without sacrificing traceability.
Load-bearing premise
Sampling from the learned distribution over discrete operator choices will produce meaningfully diverse and well-calibrated masks that outperform continuous latent-variable or denoising baselines.
What would settle it
Direct measurement on LIDC-IDRI of whether the set of architecture-sampled masks yields lower energy-distance to the annotation set than masks from latent-variable or denoising baselines while preserving per-mask accuracy.
Figures
read the original 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The manuscript introduces Neural Architecture Distributions as a stochastic source for segmentation: a distribution is learned over discrete operator choices at searchable positions in a backbone; architectures are sampled at inference to produce diverse masks; the candidate bank is built via evolutionary search; and training uses set-level supervision that matches sampled predictions to annotation sets via an IoU-based energy-distance surrogate. The work claims state-of-the-art distribution matching and hypothesis coverage on LIDC-IDRI plus effectiveness on two extension tasks, positioning itself as the first formulation of stochastic segmentation via architecture distributions.
Significance. If the empirical claims and the independence of the evolutionary bank from training-set bias are substantiated, the contribution would be significant: it supplies a discrete, auditable, and provenance-bearing stochastic mechanism that can be searched before distribution learning, offering an alternative to continuous latent variables or denoising trajectories in safety-critical domains. The set-level energy-distance objective and the separation of bank construction from distribution learning are distinctive elements that, if shown to yield genuinely richer hypothesis coverage, would advance calibration and diversity modeling.
major comments (2)
- [Abstract and §3] Abstract and §3 (candidate-bank construction): the evolutionary search that populates the fixed candidate bank occurs prior to distribution learning and is driven by a fitness signal on the LIDC training annotations; without an explicit control (held-out search, post-search coverage diagnostic, or diversity regularizer), any reported gain in IoU-based energy distance or hypothesis coverage may reflect search-induced bias toward annotation-set means rather than the architecture-distribution paradigm itself. This is load-bearing for the SOTA claim.
- [Abstract] Abstract (set-level supervision): the IoU-based energy-distance surrogate is presented as the mechanism that prevents collapse to averaged masks, yet no derivation, equation, or ablation is visible showing that the term supplies independent grounding beyond matching the annotation set; if the surrogate reduces to a fitted quantity by construction, the circularity concern raised in the stress-test note applies directly to the central training procedure.
minor comments (1)
- [Abstract] The abstract states results on 'two extension tasks' without naming them or providing even summary metrics; adding these details would improve readability.
Simulated Author's Rebuttal
We thank the referee for the constructive comments on our work introducing architecture distributions for stochastic segmentation. We address each major comment below with point-by-point responses, providing the strongest honest defense of the manuscript while noting where revisions can strengthen the presentation.
read point-by-point responses
-
Referee: [Abstract and §3] Abstract and §3 (candidate-bank construction): the evolutionary search that populates the fixed candidate bank occurs prior to distribution learning and is driven by a fitness signal on the LIDC training annotations; without an explicit control (held-out search, post-search coverage diagnostic, or diversity regularizer), any reported gain in IoU-based energy distance or hypothesis coverage may reflect search-induced bias toward annotation-set means rather than the architecture-distribution paradigm itself. This is load-bearing for the SOTA claim.
Authors: The evolutionary search constructs an optimizable support for the architecture distribution by identifying discrete operator configurations that perform well on the task, which is a deliberate design choice to make the stochastic source searchable and auditable before learning the distribution over it. The subsequent distribution learning step then assigns probabilities via set-level supervision on sampled architectures, separating the bank construction from the probabilistic modeling. While we agree that explicit controls would further isolate the contribution, the reported gains arise from the combination of searchable discrete support and set-level matching rather than search bias alone, as the energy-distance objective operates on the sampled predictions independently of the search fitness. We will add a post-search coverage diagnostic on held-out data in the revision to address this concern directly. revision: partial
-
Referee: [Abstract] Abstract (set-level supervision): the IoU-based energy-distance surrogate is presented as the mechanism that prevents collapse to averaged masks, yet no derivation, equation, or ablation is visible showing that the term supplies independent grounding beyond matching the annotation set; if the surrogate reduces to a fitted quantity by construction, the circularity concern raised in the stress-test note applies directly to the central training procedure.
Authors: The IoU-based energy-distance surrogate is a proper metric between the empirical distribution of architecture-sampled predictions and the annotation set, chosen specifically because energy distance quantifies discrepancies in both location and spread without reducing to direct fitting of means. This provides grounding independent of the annotation set itself by penalizing mismatches in the full distribution of outputs. The set-level formulation ensures that multiple samples are matched collectively rather than individually, which is what prevents collapse. We acknowledge that an explicit derivation and ablation were not included in the abstract and will add both the mathematical formulation of the energy-distance term and a targeted ablation on its effect versus standard losses in the revised manuscript. revision: yes
Circularity Check
No significant circularity detected; derivation is self-contained
full rationale
The paper describes a method that first uses evolutionary search to construct a fixed candidate bank of architectures, then learns a distribution over operator choices within that bank, and trains via set-level supervision matching sampled predictions to annotation sets with an IoU-based energy-distance surrogate. No equations, self-citations, or derivations are provided that reduce any claimed result (e.g., distribution matching or hypothesis coverage) to its inputs by construction. The evolutionary search and energy-distance term are presented as standard components whose outputs are evaluated empirically on held-out data (LIDC-IDRI), with no load-bearing self-citation chains or renaming of known results. The central claims remain independent empirical assertions rather than tautological restatements of fitted quantities.
Axiom & Free-Parameter Ledger
free parameters (1)
- architecture distribution parameters
axioms (1)
- domain assumption Standard supervised neural network optimization converges to a useful distribution over architectures
Reference graph
Works this paper leans on
-
[1]
Simon A. A. Kohl, Bernardino Romera-Paredes, Clemens Meyer, Jeffrey De Fauw, Joseph R. Ledsam, Klaus H. Maier-Hein, S. M. Ali Eslami, Danilo Jimenez Rezende, and Olaf Ronneberger. A probabilistic u-net for segmentation of ambiguous images. InProceedings of the 32nd International Conference on Neural Information Processing Systems, NIPS’18, page 6965–6975,...
2018
-
[2]
Stochastic segmentation with conditional categorical diffusion models
Lukas Zbinden, Lars Doorenbos, Theodoros Pissas, Adrian Thomas Huber, Raphael Sznitman, and Pablo Márquez- Neila. Stochastic segmentation with conditional categorical diffusion models. In2023 IEEE/CVF International Conference on Computer Vision (ICCV), pages 1119–1129, 2023
2023
-
[3]
Stochastic segmentation networks: modelling spatially correlated aleatoric uncertainty
Miguel Monteiro, Loïc Le Folgoc, Daniel Coelho de Castro, Nick Pawlowski, Bernardo Marques, Konstantinos Kamnitsas, Mark van der Wilk, and Ben Glocker. Stochastic segmentation networks: modelling spatially correlated aleatoric uncertainty. InProceedings of the 34th International Conference on Neural Information Processing Systems, NIPS ’20, Red Hook, NY ,...
2020
-
[4]
U-net: Convolutional networks for biomedical image segmentation
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. InMedical image computing and computer-assisted intervention–MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234–241. Springer, 2015
2015
-
[5]
Unet++: A nested u-net architecture for medical image segmentation
Zongwei Zhou, Md Mahfuzur Rahman Siddiquee, Nima Tajbakhsh, and Jianming Liang. Unet++: A nested u-net architecture for medical image segmentation. InDeep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support, pages 3–11, Cham, 2018. Springer International Publishing
2018
-
[6]
Swin-unet: Unet-like pure transformer for medical image segmentation
Hu Cao, Yueyue Wang, Joy Chen, Dongsheng Jiang, Xiaopeng Zhang, Qi Tian, and Manning Wang. Swin-unet: Unet-like pure transformer for medical image segmentation. InProceedings of the European Conference on Computer Vision Workshops(ECCVW), 2022
2022
-
[7]
Aimon Rahman, Jeya Maria Jose Valanarasu, Ilker Hacihaliloglu, and Vishal M. Patel. Ambiguous medical image segmentation using diffusion models. InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11536–11546, June 2023
2023
-
[8]
GMM-based V AE model with normalising flow for effective stochastic segmentation
Conghui Li, Chern Hong Lim, and Xin Wang. GMM-based V AE model with normalising flow for effective stochastic segmentation. InThe Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025
2025
-
[9]
Normalizing flows for probabilistic modeling and inference.J
George Papamakarios, Eric Nalisnick, Danilo Jimenez Rezende, Shakir Mohamed, and Balaji Lakshminarayanan. Normalizing flows for probabilistic modeling and inference.J. Mach. Learn. Res., 22(1), January 2021
2021
-
[10]
M. M. Amaan Valiuddin, Christiaan G. A. Viviers, Ruud J. G. van Sloun, Peter H. N. de With, and Fons van der Sommen. Improving aleatoric uncertainty quantification in multi-annotated medical image segmentation with normalizing flows. InUncertainty for Safe Utilization of Machine Learning in Medical Imaging, and Perinatal Imaging, Placental and Preterm Ima...
2021
-
[11]
Denoising diffusion probabilistic models
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. InProceedings of the 34th International Conference on Neural Information Processing Systems, NIPS ’20, Red Hook, NY , USA, 2020. Curran Associates Inc
2020
-
[12]
Esteban Real, Alok Aggarwal, Yanping Huang, and Quoc V . Le. Regularized evolution for image classifier architecture search. InProceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence,...
2019
-
[13]
Dropout as a bayesian approximation: Representing model uncertainty in deep learning
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In Maria Florina Balcan and Kilian Q. Weinberger, editors,Proceedings of The 33rd International Conference on Machine Learning, volume 48 ofProceedings of Machine Learning Research, pages 1050–1059, New York, New York, USA, 20–22 Jun 2016...
2016
-
[14]
Simple and scalable predictive uncertainty estimation using deep ensembles
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. In I. Guyon, U. V on Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors,Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017
2017
-
[15]
Baumgartner, Kerem C
Christian F. Baumgartner, Kerem C. Tezcan, Krishna Chaitanya, Andreas M. Hötker, Urs J. Muehlematter, 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
2019
-
[16]
Simon A. A. Kohl, Bernardino Romera-Paredes, Klaus H. Maier-Hein, Danilo Jimenez Rezende, S. M. Ali Eslami, Pushmeet Kohli, Andrew Zisserman, and Olaf Ronneberger. A hierarchical probabilistic u-net for modeling multi-scale ambiguities.CoRR, abs/1905.13077, 2019
work page internal anchor Pith review Pith/arXiv arXiv 1905
-
[17]
Uncertainty quantification in medical image segmentation with normalizing flows
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
2020
-
[18]
Ishaan Bhat, Josien P. W. Pluim, and Hugo J. Kuijf. Generalized probabilistic u-net for medical image segementa- tion. In Carole H. Sudre, Christian F. Baumgartner, Adrian Dalca, Chen Qin, Ryutaro Tanno, Koen Van Leemput, and William M. Wells III, editors,Uncertainty for Safe Utilization of Machine Learning in Medical Imaging, pages 113–124, Cham, 2022. S...
2022
-
[19]
Flow stochastic segmentation networks
Fabio De Sousa Ribeiro, Omar Todd, Charles Jones, Avinash Kori, Raghav Mehta, and Ben Glocker. Flow stochastic segmentation networks. InProceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pages 14754–14765, October 2025
2025
-
[20]
Lalith Bharadwaj Baru, Kamalaker Dadi, Tapabrata Chakraborti, and Raju S. Bapi. Ambiguous medical im- age segmentation using diffusion schrödinger bridge. In James C. Gee, Daniel C. Alexander, Jaesung Hong, Juan Eugenio Iglesias, Carole H. Sudre, Archana Venkataraman, Polina Golland, Jong Hyo Kim, and Jinah Park, editors,Medical Image Computing and Comput...
2025
-
[21]
Springer Nature Switzerland
-
[22]
DARTS: Differentiable architecture search
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. InInternational Conference on Learning Representations, 2019
2019
-
[23]
SNAS: stochastic neural architecture search
Sirui Xie, Hehui Zheng, Chunxiao Liu, and Liang Lin. SNAS: stochastic neural architecture search. InInterna- tional Conference on Learning Representations, 2019
2019
-
[24]
Efficient neural architecture search via parameters sharing
Hieu Pham, Melody Guan, Barret Zoph, Quoc Le, and Jeff Dean. Efficient neural architecture search via parameters sharing. In Jennifer Dy and Andreas Krause, editors,Proceedings of the 35th International Conference on Machine Learning, volume 80 ofProceedings of Machine Learning Research, pages 4095–4104. PMLR, 10–15 Jul 2018
2018
-
[25]
Once for all: Train one network and specialize it for efficient deployment
Han Cai, Chuang Gan, Tianzhe Wang, Zhekai Zhang, and Song Han. Once for all: Train one network and specialize it for efficient deployment. InInternational Conference on Learning Representations, 2020
2020
-
[26]
Categorical reparameterization with gumbel-softmax
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. InInternational Conference on Learning Representations, 2017
2017
-
[27]
The cityscapes dataset for semantic urban scene understanding
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The cityscapes dataset for semantic urban scene understanding. InProc. of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2016
2016
-
[28]
Road crack detection using deep convolutional neural network
Lei Zhang, Fan Yang, Yimin Daniel Zhang, and Ying Julie Zhu. Road crack detection using deep convolutional neural network. In2016 IEEE International Conference on Image Processing (ICIP), pages 3708–3712, 2016
2016
-
[29]
Mind marginal non-crack regions: Clustering- inspired representation learning for crack segmentation
Zhuangzhuang Chen, Zhuonan Lai, Jie Chen, and Jianqiang Li. Mind marginal non-crack regions: Clustering- inspired representation learning for crack segmentation. InProceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 12698–12708, June 2024
2024
-
[30]
Calibrated adversarial refinement for stochastic semantic segmentation
Elias Kassapis, Georgi Dikov, Deepak K Gupta, and Cedric Nugteren. Calibrated adversarial refinement for stochastic semantic segmentation. InProceedings of the IEEE/CVF International Conference on Computer Vision, pages 7057–7067, 2021
2021
-
[31]
A probabilistic model for controlling diversity and accuracy of ambiguous medical image segmentation
Wei Zhang, Xiaohong Zhang, Sheng Huang, Yuting Lu, and Kun Wang. A probabilistic model for controlling diversity and accuracy of ambiguous medical image segmentation. InProceedings of the 30th ACM International Conference on Multimedia, MM ’22, page 4751–4759, New York, NY , USA, 2022. Association for Computing Machinery. 18 Neural Architecture Distributions
2022
-
[32]
Pixelseg: Pixel-by-pixel stochastic semantic segmentation for ambiguous medical images
Wei Zhang, Xiaohong Zhang, Sheng Huang, Yuting Lu, and Kun Wang. Pixelseg: Pixel-by-pixel stochastic semantic segmentation for ambiguous medical images. MM ’22, page 4742–4750, New York, NY , USA, 2022. Association for Computing Machinery
2022
-
[33]
Zhitong Gao, Yucong Chen, Chuyu Zhang, and Xuming He. Modeling multimodal aleatoric uncertainty in segmentation with mixture of stochastic experts.arXiv preprint arXiv:2212.07328, 2022
-
[34]
Diffusiondet: Diffusion model for object detection
Ting Chen, Ruixiang Zhang, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning.arXiv preprint arXiv:2208.04202, 2022
-
[35]
Rethinking Atrous Convolution for Semantic Image Segmentation
Liang-Chieh Chen, George Papandreou, Florian Schroff, and Hartwig Adam. Rethinking atrous convolution for semantic image segmentation.arXiv preprint arXiv:1706.05587, 2017
work page internal anchor Pith review Pith/arXiv arXiv 2017
-
[36]
Unified perceptual parsing for scene understanding
Tete Xiao, Yingcheng Liu, Bolei Zhou, Yuning Jiang, and Jian Sun. Unified perceptual parsing for scene understanding. InProceedings of the European conference on computer vision (ECCV), pages 418–434, 2018
2018
-
[37]
Deep high-resolution representation learning for visual recognition.IEEE transactions on pattern analysis and machine intelligence, 43(10):3349–3364, 2020
Jingdong Wang, Ke Sun, Tianheng Cheng, Borui Jiang, Chaorui Deng, Yang Zhao, Dong Liu, Yadong Mu, Mingkui Tan, Xinggang Wang, et al. Deep high-resolution representation learning for visual recognition.IEEE transactions on pattern analysis and machine intelligence, 43(10):3349–3364, 2020
2020
-
[38]
Swin transformer: Hierarchical vision transformer using shifted windows
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. InProceedings of the IEEE/CVF international conference on computer vision, pages 10012–10022, 2021
2021
-
[39]
Beyond the pixel-wise loss for topology-aware delineation
Agata Mosinska, Pablo Marquez-Neila, Mateusz Kozinski, and Pascal Fua. Beyond the pixel-wise loss for topology-aware delineation. In2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 3136–3145, 2018
2018
-
[40]
Topology-preserving deep image segmentation
Xiaoling Hu, Fuxin Li, Dimitris Samaras, and Chao Chen. Topology-preserving deep image segmentation. Advances in neural information processing systems, 32, 2019
2019
-
[41]
Recurrent u-net for resource- constrained segmentation
Wei Wang, Kaicheng Yu, Joachim Hugonot, Pascal Fua, and Mathieu Salzmann. Recurrent u-net for resource- constrained segmentation. In2019 IEEE/CVF International Conference on Computer Vision (ICCV), pages 2142–2151, 2019
2019
-
[42]
Crackformer: Transformer network for fine-grained crack detection
Huajun Liu, Xiangyu Miao, Christoph Mertz, Chengzhong Xu, and Hui Kong. Crackformer: Transformer network for fine-grained crack detection. In2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 3763–3772, 2021
2021
-
[43]
Joint topology-preserving and feature- refinement network for curvilinear structure segmentation
Mingfei Cheng, Kaili Zhao, Xuhong Guo, Yajing Xu, and Jun Guo. Joint topology-preserving and feature- refinement network for curvilinear structure segmentation. In2021 IEEE/CVF International Conference on Computer Vision (ICCV), pages 7127–7136, 2021. 19
2021
discussion (0)
Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.