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: Accurate segmentation is essential in error-critical domains such as medical imaging, where outputs support clinical decisions. Probabilistic models like the Stochastic Segmentation Network (SSN) enable uncertainty quantification, but existing methods typically use low-rank plus diagonal covariance structures that struggle to capture both global and local spatial correlations, limiting performance gains over deterministic models. We revisit low-rank formulations and introduce two approaches - Single-Basis and Mixture-of-Bases decompositions - that project predicted noise onto learned covariance bases, either globally or within partitioned volume blocks. This yields richer, more flexible uncertainty modeling with minimal parameter overhead. On the most challenging organs in the 3D TotalSegmentator CT dataset, our methods significantly improve Dice scores over deterministic and baseline stochastic models while preserving strong calibration, with the Mixture-of-Bases performing best. These findings show that basis-driven covariance modeling can enhance segmentation accuracy and uncertainty estimation in 3D medical imaging.

Git: https://github.com/andreasaspe/mob-seg3d

Serve As Reviewer: ~Peter_J.T._Kampen1

Submission Number: 52