Retaining Informative Latent Variables in Probabilistic Segmentation

Conditional latent-variable models can successfully quantify annotation variability in segmentation. Training such models involves tuning the dimensionality of the latent space to optimally capture the inherent data ambiguity. Nevertheless, we discover after careful tuning, that the latent space does not always reflect this. In fact, some latent dimensions are completely neglected. For such segmentation models the latent dimensionality is often poorly motivated or based on computational constraints. In this paper, we offer an information-theoretic approach to optimally leverage all latent dimensions. We adapt and improve the Probabilistic U-Net to maximize the mutual information between the latent and output variables, leading to improved latent space properties and higher segmentation performance.