Structured Covariance Modeling Using Learned Mixture-of-Bases for Uncertainty in 3D Segmentation

Peter J.T. Kampen, Andreas With Aspe, Kristine Aavild Juhl, Anders Nymark Christensen, Morten Rieger Hannemose, Anders Dahl, Rasmus Reinhold Paulsen, Josefine Vilsbøll Sundgaard

Keywords: Segmentation, Uncertainty, Low-Rank Approximation

TL;DR: We seek to improve 3D segmentation performance using various low-rank approximations to the logit distribution.

Abstract:

Uncertainty quantification often plays a key role when deploying deep learning models in segmentation tasks, such as in medical imaging, where the results are used directly for clinical decision support. Existing stochastic segmentation methods, such as Stochastic Segmentation Networks (SSNs), typically rely on low-rank plus diagonal covariance structures to model predictive uncertainty. While computationally efficient, this parameterization often fails to capture both global and local spatial correlations, leading to limited improvements over deterministic models. In this work, we revisit low-rank formulations and introduce two new approaches: a Single-Basis and a Mixture-of-Bases decomposition. By projecting predicted noise structures onto learned covariance bases — either globally or, for the Mixture-of-Bases, within blocks obtained by partitioning the volume — we achieve richer and more flexible uncertainty modeling with negligible increases in the number of parameters. Evaluated on the 3D segmentation task of challenging anatomies from the TotalSegmentator CT dataset. Our approaches achieve significant Dice score improvements over deterministic and baseline stochastic models while maintaining competitive calibration, with the Mixture-of-Bases yielding the greatest improvement. These results demonstrate that basis-driven covariance modeling is a simple yet powerful way to improve both segmentation accuracy and uncertainty estimation in 3D medical image analysis.

All code and experiments will be made public upon acceptance.

Serve As Reviewer: ~Peter_J.T._Kampen1

Submission Number: 52