Uncertainty-Aware Regularization for Image-to-Image Translation

TL;DR

The paper addresses the challenge of reliable uncertainty estimation in medical image-to-image translation by introducing Uncertainty-Aware Regularization (UAR), a lightweight, model-agnostic term that regularizes the spatial distribution of aleatoric uncertainty. By modeling residuals with spatially varying parameters $\alpha$ and $\beta$ under a Generalized Normal Distribution and enforcing a total-variation penalty on $\hat{\beta}$, the method yields more coherent and edge-preserving uncertainty maps while improving reconstruction quality. The approach is validated on two medical-imaging datasets, including a new paired WCE-to-FICE capsule endoscopy dataset, showing improved PSNR, SSIM, LPIPS, and RRMSE, and more accurate uncertainty localization, particularly under noise and artifacts. Overall, UAR enables robust, interpretable uncertainty guidance within a single end-to-end model, with potential for adoption across more advanced I2I architectures and clinical validation.

Abstract

The importance of quantifying uncertainty in deep networks has become paramount for reliable real-world applications. In this paper, we propose a method to improve uncertainty estimation in medical Image-to-Image (I2I) translation. Our model integrates aleatoric uncertainty and employs Uncertainty-Aware Regularization (UAR) inspired by simple priors to refine uncertainty estimates and enhance reconstruction quality. We show that by leveraging simple priors on parameters, our approach captures more robust uncertainty maps, effectively refining them to indicate precisely where the network encounters difficulties, while being less affected by noise. Our experiments demonstrate that UAR not only improves translation performance, but also provides better uncertainty estimations, particularly in the presence of noise and artifacts. We validate our approach using two medical imaging datasets, showcasing its effectiveness in maintaining high confidence in familiar regions while accurately identifying areas of uncertainty in novel/ambiguous scenarios.

Figures (7)

• Figure 1: Non-probabilistic image translation (green) optimizes point-estimates for the residuals between the predicted and the target images. On the contrary, the probabilistic approach (orange) models the residuals with a distribution, allowing the variance of errors to change spatially. Our method takes this a step further (purple) by regularizing the predicted variances (or distribution parameters) to achieve more precise uncertainty estimation (The discriminator is omitted from the diagram for simplicity).
• Figure 2: The figure shows the shape $\beta$ parameter and predicted aleatoric uncertainty of an image at different epochs during the training. Without regularization (first row), the variances in the predictions (uncertainty) remain relatively the same throughout training. In contrast, with regularization (second row), the predicted uncertainty gets progressively less noisy and more semantically refined over the course of the training.
• Figure 3: Uncertainty estimate is sensitive to small changes in input and network parameters, resulting in potentially noisy maps.
• Figure 4: Impact of noise on uncertainty prediction and image reconstruction. The figure illustrates the impact of varying levels of Gaussian noise on image reconstruction and predicted aleatoric uncertainty. Row (a) depicts the non-regularized version, whereas row (b) shows the regularized variant. In the presence of noise, the non-regularized version is more significantly affected, with uncertainty maps rapidly diverging from the original regions of uncertainty. In contrast, the regularized version demonstrates greater robustness to noise.
• Figure 5: Qualitative Comparison. As can be seen from columns 3 and 6, while the regions of residual errors are consistent between the regularized and non-regularized variant, UAR shows consistently low residual errors. Correspondingly, the uncertainty maps are less-noisy and more structurally coherent.