Bayesian Neural Networks for One-to-Many Mapping in Image Enhancement

TL;DR

This work tackles the inherent one-to-many ambiguity in low-light and underwater image enhancement by modeling uncertainty with Bayesian Neural Networks. It introduces the Bayesian Enhancement Model (BEM), which uses a two-stage BNN-DNN pipeline to capture coarse variability in a low-dimensional latent space and then refine high-frequency details, enabling fast, high-quality outputs. An Adaptive Prior stabilizes Bayesian training, and two inference modes—ranking-based selection and Monte Carlo sampling—offer flexible, uncertainty-aware predictions. Extensive experiments on paired and unpaired LLIE/UIE datasets show that BEM outperforms deterministic baselines and competitive probabilistic methods in both fidelity (PSNR/SSIM/LPIPS) and perceptual/no-reference metrics, with significantly reduced latency. The approach is backbone-agnostic and scalable, offering practical impact for real-time enhancement applications while providing principled uncertainty estimates.

Abstract

In image enhancement tasks, such as low-light and underwater image enhancement, a degraded image can correspond to multiple plausible target images due to dynamic photography conditions. This naturally results in a one-to-many mapping problem. To address this, we propose a Bayesian Enhancement Model (BEM) that incorporates Bayesian Neural Networks (BNNs) to capture data uncertainty and produce diverse outputs. To enable fast inference, we introduce a BNN-DNN framework: a BNN is first employed to model the one-to-many mapping in a low-dimensional space, followed by a Deterministic Neural Network (DNN) that refines fine-grained image details. Extensive experiments on multiple low-light and underwater image enhancement benchmarks demonstrate the effectiveness of our method.

Figures (10)

• Figure 1: One-to-Many Mapping where an image crop $\mathbf{x}$ associated with multiple targets $\{\mathbf{y}^1, \ldots, \mathbf{y}^6\}$. A DNN (left) can only predict one of the targets. In contrast, a BNN (right) can produce many predictions according to a learned probability distribution.
• Figure 2: The two-stage pipeline. In Stage I, the BNN with weights $\mathbf{w} \sim q(\mathbf{w}|\bm{\theta})$ is trained by minimizing the minibatch loss $\mathcal{L}^\text{mini}$ in \ref{['eq:final_loss']}. In Stage II, the DNN with weights $\mathbf{w}^\text{G}$ is trained by minimizing the L1 loss, $L1(\mathbf{y}, \hat{\mathbf{y}})$. The inference process is denoted by $\rightarrow$, while the training process for each stage is indicated by $\rightarrow$.
• Figure 3: Visual comparisons of the DNN baseline, BEM$\text{MC}$, and BEM$\text{Rank}$ with CLIP-IQA. The rightmost patches highlight the diverse unselected predictions reflecting BEM’s one-to-many modeling capability.
• Figure 4: One-to-many mapping from input to outputs. The predictions are sorted by CLIP-IQA and NIQE.
• Figure 5: Inference speed on an Nvidia RTX 4090.