# Learning a Sampling-Free Variational DNN Plugin from Tiny Training Sets to Refine OOD Segmentation With Uncertainty Estimation

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

## Key facts

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

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [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)
- [Segment Anything](https://mlsanity.com/paper/2304.02643/) (2023)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Squeeze-and-Excitation Networks](https://mlsanity.com/paper/1709.01507/) (2017)
- [Neural Discrete Representation Learning](https://mlsanity.com/paper/1711.00937/) (2017)
- [DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs](https://mlsanity.com/paper/1606.00915/) (2016)
- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://mlsanity.com/paper/1802.02611/) (2018)
- [Attention U-Net: Learning Where to Look for the Pancreas](https://mlsanity.com/paper/1804.03999/) (2018)
- [Generating Diverse High-Fidelity Images with VQ-VAE-2](https://mlsanity.com/paper/1906.00446/) (2019)
- [Rethinking Semantic Segmentation from a Sequence-to-Sequence Perspective with Transformers](https://mlsanity.com/paper/2012.15840/) (2020)
- [Segmenter: Transformer for Semantic Segmentation](https://mlsanity.com/paper/2105.05633/) (2021)
- [Albumentations: fast and flexible image augmentations](https://mlsanity.com/paper/1809.06839/) (2018)
- [Measuring Calibration in Deep Learning](https://mlsanity.com/paper/1904.01685/) (2019)
- [U-Mamba: Enhancing Long-range Dependency for Biomedical Image Segmentation](https://mlsanity.com/paper/2401.04722/) (2024)
- [Using Pre-Training Can Improve Model Robustness and Uncertainty](https://mlsanity.com/paper/1901.09960/) (2019)
- [A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts](https://mlsanity.com/paper/2303.15361/) (2023)
- [Road Extraction by Deep Residual U-Net](https://mlsanity.com/paper/1711.10684/) (2017)
- [FedDG: Federated Domain Generalization on Medical Image Segmentation via Episodic Learning in Continuous Frequency Space](https://mlsanity.com/paper/2103.06030/) (2021)
- [VM-UNet: Vision Mamba UNet for Medical Image Segmentation](https://mlsanity.com/paper/2402.02491/) (2024)
- [A Probabilistic U-Net for Segmentation of Ambiguous Images](https://mlsanity.com/paper/1806.05034/) (2018)
- [Monte Carlo Gradient Estimation in Machine Learning](https://mlsanity.com/paper/1906.10652/) (2019)
- [SegMamba: Long-range Sequential Modeling Mamba For 3D Medical Image Segmentation](https://mlsanity.com/paper/2401.13560/) (2024)
- [A survey of loss functions for semantic segmentation](https://mlsanity.com/paper/2006.14822/) (2020)
- [ResUNet++: An Advanced Architecture for Medical Image Segmentation](https://mlsanity.com/paper/1911.07067/) (2019)
- [Boundary loss for highly unbalanced segmentation](https://mlsanity.com/paper/1812.07032/) (2018)

## Abstract (arXiv)

Deep neural networks (DNNs) frequently fail to generalize to out-of-distribution (OOD) medical images because of variations in scanners and acquisition protocols. Retraining DNN models to address these distribution shifts is often impractical due to the high cost of acquiring and annotating new medical datasets. To address this, we introduce VarDeepPCA, a novel lightweight variational DNN framework designed to restore/refine degraded segmentation maps by leveraging intrinsic geometric priors. Unlike existing approaches that require target-domain data or extensive pre-training, our VarDeepPCA explicitly learns a distribution of valid anatomical geometries using only small in-distribution (ID) datasets. Theoretically, our novel variational learning framework leverages a reinterpretation of the softmax mapping to implicitly perform exact distribution modeling, thereby enabling computationally efficient, sampling-free learning and inference. This also enables VarDeepPCA to provide uncertainty estimates associated with its restored segmentation maps. We empirically validate our framework across 4 distinct clinical applications, using 14 publicly available datasets, involving segmentation of the myocardium, neuroretinal rim, prostate, and fetal head. Comparisons against 15 existing methods demonstrate that VarDeepPCA consistently restores segmentation maps produced by the existing methods on OOD data to (i) significantly improve anatomical plausibility of geometries and clinical utility of the segmentations, and (ii) significantly reduce errors, without needing any more training data than that used by existing methods.