【LLIE专题】BEM：一对多映射，贝叶斯增强模型

---

本文将对《Bayesian Neural Networks for One-to-Many Mapping in Image Enhancement》这篇文章进行解读。该论文利用贝叶斯神经网络（BNN）来捕捉数据不确定性并产生多样化的输出，以表征实际退化图的一对多映射，为不同暗光增强任务或需求提供灵活的推理方案。
参考资料如下：
[1]. BEM论文地址
[2]. 源码地址

---

专题介绍

在低光照环境下，传统成像设备往往因画面昏暗、细节丢失而受限。LLIE（低照度暗光增强）技术应运而生，它通过提升图像亮度、对比度，减少噪点并恢复色彩细节，让暗夜变得清晰可见。

LLIE技术从传统方法如直方图均衡化、Retinex模型等起步，近年来借助深度学习，尤其是卷积神经网络（CNN），GAN模型，扩散模型实现了质的飞跃。这些算法能自动学习图像特征，精准处理低光照图像，效果显著优于传统技术。

本专题将聚焦LLIE技术的核心原理、应用案例及最新进展，让我们一起见证LLIE如何点亮暗夜，开启视觉新视界！欢迎一起探讨交流！

系列文章如下
【1】ZeroDCE
【2】HVI
【3】CLIP-LIT
【4】GLARE
【5】Retinexformer
【6】SG-LLIE
【7】GPP-LLIE
【8】FusionNet
【9】EnligtenGAN
【10】PTG

一、研究背景

在低光/水下图像增强任务中，同一退化图像可能对应多种合理增强结果（因光照、色彩等动态条件变化），而传统DNN只能生成单一输出。同时现有数据集（如LOL、UIEB）包含部分低质量标注，DNN会过拟合到劣质标签，导致增强效果不佳。

另一方面图像增强质量高度主观，需模型具备生成多样化输出的能力。确定性模型（如RetinexFormer等）在动态场景下泛化性差，而生成模型（如扩散模型）计算开销大，难以实时应用。

为解决这些问题，该论文提出了一种贝叶斯增强模型（BEM），关键点总结如下。

• 首次将BNN应用于图像增强的一对多映射问题
  • 利用贝叶斯神经网络（BNN）来捕捉数据不确定性并产生多样化的输出，即通过权重采样生成多样化增强结果。
• 动态动量先验（Momentum Prior）
  • 针对BNN在高维权重空间中易欠拟合的问题，提出动态更新的先验分布，通过指数移动平均（EMA）平滑参数更新，加速收敛并提升性能。
• 两阶段BNN+DNN框架
  • 为实现实时推理，引入了一种两阶段方法：第一阶段采用BNN在低维空间中建模一对多映射，基于策略选择合适解作为第二阶段输入的；第二阶段使用确定性神经网络（DNN）细化图像细节。
• 灵活推理策略
  • 蒙特卡洛（MC）预测：平均多个采样结果，鲁棒性强。
  • 基于排序的预测：利用无参考质量评估（如CLIP-IQA）选择最优候选，适应主观审美需求。

二、方法细节

整体框架如下：

图中有两个流程，虚线训练流程和实线推理流程。我们先来看看理论推导部分，如果不感兴趣的话，可以直接跳到2.1节

对于确定性DNN而言，当输入$x$不相同时，输出$y$肯定也不相同。作者认为，当一个输入能够对应多个似乎合理的输出时，这种确定性就不适用了。于是作者通过贝叶斯估计将不确定性引入网络权重$w$，从而获得权重$\mathbf{w}\sim P(\mathbf{w}|\mathbf{y},\mathbf{x})$的后验分布。在推理过程中，从该分布中采样权重。权重的后验分布表示为
$$P(\mathbf{w}|\mathbf{y},\mathbf{x})=\frac{P(\mathbf{y}|\mathbf{x},\mathbf{w})P(\mathbf{w})}{P(\mathbf{y}|\mathbf{x})}$$

为了便于求解，基于变分推理（variational inference，VI）思路，用更易于处理的分布$q(\mathbf{w}|\theta )$近似$P(\mathbf{w}|\mathbf{y},\mathbf{x})$。

• 真实后验分布：像一片复杂的山脉，高高低低很难描述。
• 变分后验分布：用一座光滑的小山丘（高斯分布）去近似这片山脉。虽然不完全一样，但能够尽量接近就行，且容易计算。

这样就可以从$q(\mathbf{w}|\theta )$分布中采样得到权重$w$。对于$\theta$，可以用KL散度求解
$$\begin{array}{rl}{\theta }^{\star }& =\arg \underset{\theta }{\min }\text{KL}\left[q(\mathbf{w}|\mathbf{\theta })\parallel P(\mathbf{w}|\mathbf{y},\mathbf{x})\right]\\ & =\arg \underset{\theta }{\min }\int q(\mathbf{w}|\mathbf{\theta })\log \frac{q(\mathbf{w}|\mathbf{\theta })}{P(\mathbf{w})P(\mathbf{y}|\mathbf{x},\mathbf{w})}\text{d}\mathbf{w}\\ & =\arg \underset{\theta }{\min }-{\mathbb{E}}_{q(\mathbf{w}|\mathbf{\theta })}\left[\log P(\mathbf{y}|\mathbf{x},\mathbf{w})\right]\\ & +\text{KL}\left[q(\mathbf{w}|\mathbf{\theta })\parallel P(\mathbf{w})\right]\end{array}$$

于是，损失函数可以定义为
$$\mathcal{L}(\mathbf{x},\mathbf{y},\mathbf{\theta })=\underset{\text{data-dependent term}}{\underbrace{-{\mathbb{E}}_{q(\mathbf{w}|\mathbf{\theta })}\left[\log P(\mathbf{y}|\mathbf{x},\mathbf{w})\right]}}+\underset{\text{prior matching term}}{\underbrace{\text{KL}\left[q(\mathbf{w}|\mathbf{\theta })\parallel P(\mathbf{w})\right]}}$$

在训练BNN神经网络时，需要对参数随机采样（比如从高斯分布中抽一个随机数）。但直接随机采样来获取$w$会导致梯度无法回传（因为随机操作是“不可导的”）。于是要把随机性转移到一个固定的“噪声源”上： $w=\mu +\sigma \times ϵ$，其中 $ϵ\sim N(0,I)$。$\mu$ 和 $\sigma$ 是网络可学习的参数。$ϵ$ 是从标准正态分布采样的随机噪声（固定分布，与网络无关）。这样既保留了随机性，又能正常训练网络！

进一步，通过将先验参数更新为变分后验参数的指数移动平均（EMA）来提出动态先验。将$P(\mathbf{w})=\mathcal{N}(\mathbf{w};{\mathbf{\mu }}_t^{\text{EMA}},{{\mathbf{\sigma }}_t^{\text{EMA}}}^2\mathbf{I})$变为
$$\begin{array}{rl}{\mathbf{\mu }}_0^{\text{EMA}}& =0,{\mathbf{\sigma }}_0^{\text{EMA}}={\sigma }^{\text{o}}1,\\ {\mathbf{\mu }}_t^{\text{EMA}}& =\beta {\mathbf{\mu }}_{t-1}^{\text{EMA}}+(1-\beta ){\mathbf{\mu }}_t,t=\mathrm{1...}T,\\ {\mathbf{\sigma }}_t^{\text{EMA}}& =\beta {\mathbf{\sigma }}_{t-1}^{\text{EMA}}+(1-\beta ){\mathbf{\sigma }}_t,t=\mathrm{1...}T,\end{array}$$
也就是在每一轮迭代中，针对每个minibatch，都会按上式来做更新。其实是比较常见的策略，让参数更新更加稳定。其中$\beta$为EMA衰减率。

最终损失函数可以表征为下式
$$\begin{array}{rl}{\mathcal{L}}^{\text{mini}}(\mathbf{x},\mathbf{y},\mathbf{\theta })& =\underset{\text{data-dependent term}}{\underbrace{-{\mathbb{E}}_{q(\mathbf{w}|\mathbf{\theta })}\left[\log P(\mathbf{y}|\mathbf{x},\mathbf{w})\right]}}+\underset{\text{prior matching term}}{\underbrace{\frac{1}{M}\text{KL}\left[q(\mathbf{w}|\mathbf{\theta })\parallel P(\mathbf{w})\right]}},\\ & =\frac{1}{M}[\underset{\text{data-dependent term}}{\underbrace{\sum _i^M{\mathbb{E}}_{\mathbf{w}\sim q(\mathbf{w}|\mathbf{\theta })}{\Vert F({\mathbf{x}}^i;\mathbf{w})-{\mathbf{y}}^i\Vert }_2^2}}\\ & +\underset{\text{prior matching term}}{\underbrace{\log \frac{{\mathbf{\sigma }}_t^{\text{EMA}}}{\sigma }+\frac{{\sigma }^2+(\mathbf{\mu }-{\mathbf{\mu }}_t^{\text{EMA}}{)}^2}{2{{\mathbf{\sigma }}_t^{\text{EMA}}}^2}-\frac{1}{2}}}],\end{array}$$

BNN在每次前向过程中，可以从变分后验分布中采样到一组权重，从而可针对同一个输入获得多个不同的输出。

直接用BNN实现高分辨率的端到端增强会非常耗时，于是作者设计了双阶段方案，第一阶段利用BNN对低维粗信息中的一对多映射进行建模，第二阶段利用DNN对原始高维空间中的高频细节进行细化。

2.1 Stage1 BNN阶段

BNN低维度信息建模，实现1对多映射。

可以这样做的主要原因是亮度和颜色在低维度空间中也能有效表征。当然，太小肯定也不行。

首先对输入的做预处理，包括低通滤波（LP）+ 下采样（r=1/16）至低维空间，避免空间混叠，减少计算量。其中的LP是通过FFT实现的。然后将低维度数据输入至BNN，这里使用了Instance Normalization，以更好地捕捉高对比度局部光照。其他网络结构并不太重要，关键在于其权重服从变分后验分布 $q(w\mid \theta )$就行。经过模型后的输出再做一次上采样，得到粗映射结果$z$。采样方式为双线性，整个过程可用下式表征，
$$\mathbf{z}=\text{Up}(F(\text{Down}(\text{LP}(\mathbf{x}),r);\mathbf{w})),\mathbf{w}\sim q(\mathbf{w}|\mathbf{\theta })$$

第一阶段的$z$相当于增强后的光照条件，进一步将输入$x$与粗预测$z$相结合，得到目标y的第一阶段预测，记为$\tilde{\mathbf{y}}$
$$\tilde{\mathbf{y}}=(\mathbf{x}+\alpha \mathbf{z})\ast \mathbf{z}$$
其中$\alpha$是一个比较小的系数。直接简单相加（$x+z$）会丢失原始图像的高频细节；直接相乘（$x\ast z$）会放大输入图像中的噪声。最终$\tilde{\mathbf{y}}$和$y$进行${\mathcal{L}}^{\text{mini}}$损失计算。

2.2 Stage2 DNN阶段

这个阶段就比较简单了，直接选个好一点的模型作为Backbone，论文分别采用了transformer和mamba作为Backbone来测试。其中需要关注的点就是训练和推理两个阶段的输入有所不同。

• 推理阶段：拼接原始图像 $x$与上采样的粗预测 $Up(z\ast )$，作为第二阶段模型输入
• 训练阶段：使用真实粗标签$\text{LP}\left(\frac{\sqrt{{\mathbf{x}}^2+4\alpha \mathbf{y}-\mathbf{x}}}{2\alpha }\right)$代替第一阶段的结果，避免DNN将多输出回归为单一结果。（这公式就是解一元二次方程得到的）

其实这个阶段就专注于微调，保留多样性，若用不同的结果去训练，相当于前面的多样性就没价值了。

2.3 推理策略

真正推理时，在第一阶段和第二阶段中间还有一步，那就是从第一阶段的多个输出中选择一个。选择的策略可以根据任务应用去制定，论文里测试了两种方式，

• 蒙特卡洛（MC）：平均 $K$ 个 $z_k$ 得 $z\ast$，适合鲁棒性场景。
• 基于排序：通过IQA（如CLIP-IQA、NIQE）评分选择最优 $z\ast$，支持文本引导的个性化增强（论文中采用的默认文本为“明亮、自然、质量好的照片”）。

完整的双阶段推理伪码如下，还是比较直观的。

通过这种设计，消除了产生多个冗余输出的需要，显著加快了推理速度。且此策略可以用于大多数LLIE网络方案。

三、实验论证

针对暗光增强和水下增强分别做了实验。

• 暗光增强（LLIE）：LOL-v1/v2（配对）、LIME/NPE/MEF（非配对）。
• 水下增强（UIE）：UIEB-R90（配对）、U45/UCCS（非配对）。

定量分析
BEM在PSNR/SSIM/LPIPS上全面超越SOTA（如RetinexFormer、GLARE）。且Mamba骨干优于Transformer。灰色部分是通过100次采样选择最优结果，展示了该方法的潜力。

无参考指标的表现也不错。

另外，初始化的σ对结果也是有影响，太大太小都会导致效果下降，图中0.05时指标最佳。

定性分析
下图是从输入到输出的一对多映射的可视化。这些预测是根据三个指标进行排序的：PSNR， CLIP-IQA和NIQE，它们反映了图像质量的不同方面。

不确定性可视化：阴影区域不确定性低，光照区域高，符合一对多映射的物理意义。

不同策略的可视化：当数据集有噪声干扰时，虽然由于平均效果，MC输出避免了最差的结果，但它仍然可能无法满足审美期望。相比之下，CLIP-IQA在不需要数据清洗的情况下获得了美观的结果，这使得它对带有低质量数据的数据集兼容度更高。

其他方案视觉效果对比

论文末尾还有一些消融实验的分析，证明了所用策略的有效性。这里不在赘述，感兴趣的可以看一下原文附录部分。

作者还提到了很关键的一个点，要想PSNR高，就让亮度分布与GT尽量接近即可。但GT如果质量也差，那么增强结果就无法带来感知质量的提升。因此，没必要一味追求PSNR。

四、总结和思考

该论文提出了首个基于BNN的图像增强框架BEM，解决一对多映射问题。为图像增强的不确定性建模提供新思路，尤其适用于动态场景和标签噪声数据。更进一步，可将该策略扩展至其他视觉任务（如超分或降噪）

或许会有人觉得这种跑多次选最佳的方式有点“赖皮”，这就仁者见仁智者见智了。

或许还会有人觉得现实中不仅有一对多映射，还有多对一映射（同一高质量场景受到不同的暗光退化），这种情况如何解决？关注本专栏，后续将介绍多对一的新思路。❤️

---

感谢阅读，欢迎留言或私信，一起探讨和交流。
如果对你有帮助的话，也希望可以给博主点一个关注，感谢。