Robust Representation Learning with Feedback for Single Image Deraining论文解读

Abstract

一个去雨网络可以被视为一个移除图片中雨水的条件生成器。大多数现有的去雨方法忽略了由于不确定性带来的模型误差，从而降低了embedding质量。和现有的直接将低质量特征嵌入到模型中不同，这篇文章将用latent高质量特征取代了低质量特征。同时文章使用了

Feedback based Representation Learning

3.1 Problem Formulation

在[15,4]中，作者介绍了两种类型的不确定性，认知不确定性（模型不确定性或系统不确定性），随机不确定性（统计不确定性，这种不确定性表示观测过程中固有的噪音）。变分推理可以用来衡量认知不确定性以计算方差，最大后验或者最大似然推断可用来衡量随机不确定性。为了精准描述我们的条件优化任务，我们对两个输出（conditional输出以及target输出）进行建模并且假设两个输出的概率分布均为高斯分布。我们的最优化目标 $P=−logp(y1,y2∣fW(x))\mathcal{P}=-\text{log}p(y_1,y_2|f^W(x))$ 具体表述如下：

$P=−log⁡p(y1,y\mathcal{P}=-\log p(y_1,y$

这里 $p(⋅)p(\cdot)$ 代表概率函数， $fWi(⋅)(i=1,2)f^{W_i}(\cdot)(i=1,2)$ 指代对应的神经网络， $x$ 指代输入雨水图片， $y_i(i=1,2)$ 是输出。我们将分布 $p(yi∣fWi(⋅))(i=1,2)p(y_i|f^{W_i}(\cdot))(i=1,2)$ 的均值表示为 $y_i$ 并将方差表示为 $σi2\sigma_i^2$ 。下标1的操作是指在利用embedding residual图的情况下生成residual map，最终的去雨图片就可以通过用雨水减去residual map得到。我们通过寻找basic embedding residual图以及latent mapping空间中的最优embedding residual图间的函数关系来提升去雨效果。

反馈机制

在模型训练后，一些error较大的embeddings通常会增加不确定性，使去雨效果退化[15]。为了获取high-quality特征，我们将closed-loop反馈机制引入了CNN中。如图2所示，在生成embedding残差图 $ϕ1(It)\phi_1(I_t)$ 后，我们使用error detector（ED），将embedding残差图 $ϕ1(It)\phi_1(I_t)$ 以及雨水图片 $I_t$ 作为输入来学习embedding残差图 $ϕ1(It)\phi_1(I_t)$ 和对应的残差图真值 $R_t$ 间的error。最终得到的error图用来补偿embedding残差图来进行更好的特征表达。

effect analysis

error detector的作用体现在训练过程以及结果中。对于训练过程，error compensation在embedding残差图上的作用使得R-to-R模型（connecting这个embedding residual map以及最终输出的residual map）能够被更加精准地学习。随着R-to-R模型在训练中收敛，那么最终输出的loss function对 O-to-R 模型（connecting输入雨水图片和embedding残差图）有着更小的影响，因此O-to-R模型主要基于embedding feature loss function进行学习并且acts as a feature guidance module。The decoupled learning is helpful to obtain appealing results[36]。To this end，由不确定性引起的model error可以被error compensation降低。 For the results, the rectified embedding residual map are always more conducive to generate better deraining results as shown in Fig. 5.

3.2 基于CNN的反馈过程

Error distribution complexity

通常来讲，由于error分布的复杂性，由训练过程中的不确定性所带来的variable errors很难被CNN remapped。通过将误差的倒数抽象为在0上下波动的点，图4展示了取倒数的绝对值能降低误差倒数分布的复杂度。此外，误差分布的上确界同样简化了误差倒数分布的复杂度。小的error（倒数很大）对于结果的影响很小，这些值被上确界截断，这些值同样被传导用于remap error map。

Error detector

特别地，我们使用误差绝对值的倒数乘以threshold参数 $θ1\theta_1$ 作为训练error detector的训练值，并且在error detector的末尾使用sigmoid函数。在这种情况下，error detector对上确界输出为1，并且error越大，error detector输出越接近0。对于embedding残差图以及error detector而言，对应的优化问题是：

$min⁡ϕ1(⋅)∥Rt×0.5−ϕ1(It)∥1\min\limits_{\phi_1(\cdot)}\Vert R_{t\times0.5}-\phi_1(I_t)\Vert_1$

$min⁡φ(⋅)∥\min\limits_{\varphi(\cdot)}\Vert$

这里 $I_t$ 指代雨水图片， $It×0.5指代I_{t\times0.5}指代$ 0.5 scale size的雨水图片 $I_t$ ， $R_t$ 表示residual map truth， $Rt×0.5R_{t\times 0.5}$ 表示0.5 scale size的 $R_t$ ， $ϕ(⋅)\phi(\cdot)$ 是error detector的function， $θ1\theta_1$ 表示threshold参数。由于雨水条纹的稀疏性，将residual图的大小降低为一半可以很好的代表原始的residual图，并且可以很大程度上减少计算量。注意 $φ(⋅)\varphi(\cdot)$ 是sigmoid函数，小的error会被上确界截断。此外，需要强调的一点是优化问题中的 $∣Rt×0.5−ϕ1(It)∣|R_{t\times0.5}-\phi_1(I_t)|$ 是固定的，并且取error倒数的绝对值对于减小分布复杂度是十分重要的。absolute error mao（公式4中的 $e r r$ 以及rectified embedding residual map通过如下的rectification过程计算得到：

$err=θ1φ(It×0.5,ϕ1(It))−θ1(4)err=\frac{\theta_1}{\varphi(I_{t\times0.5},\phi_1(I_t))}-\theta_1\quad\quad\quad\quad\quad\quad(4)$

$ϕ1′(It)=ϕ1(It)−err(1−2ϕ1(It))\phi_1'(I_t)=\phi_1(I_t)-err(1-2\phi_1(I_t))$

这里 $1$ 是全1矩阵， $ϕ1′(It)\phi_1'(I_t)$ 代表修正后的embedding residual map。通常来讲，absolute error（公式4）并不能很好地compensate 特征。出于这个原因，我们引入了这种方法来合理使用absolute error。

对于生成的去雨图片，应当没有雨水的区域总是存在雨水条纹，这是由于在训练过程中包含了大量的雨水区域。Specifically，（基本是介绍上面公式是怎么来的，懒得看了）

Feature compensator

类似地，我们对于feature compensator的目标是设计一个轻量的模型来获取high-quality embedding残差图以指导其他的embedding残差图。通过这种方式，需要被error detector修正的embedding 残差图会变得更加精确，从而降低了error remapping的难度。根据公式1，我们希望获取受到不确定性影响较小的embedding残差图。CNN是一个复杂的非线性函数，我们使用基于CNN的模型（图3（d））来transform残差图真值 $R_t$ 用于error compensation，这个过程表示如下：

$$$$

这里 $ζ=P(wi(Rt×i))=λ×∥wi(Rt×i)∥22\zeta=P(w_i(R_{t\times i}))=\lambda\times\Vert w_i(R_{t\times i})\Vert_2^2$ ， $P(⋅)P(\cdot)$ 表示regularizer， $λ\lambda$ 是setup parameter， $Rt×iR_{t\times i}$ 表示不同scale size的 $R_t$ （ $i$ =0.25或0.5）。 $ϕi(⋅)\phi_i(\cdot)$ is the function trained to map $I_t$ as the transformed residual map at $i$ scale size (as shown in Fig. 6(d)(e)), $wi(⋅)w_i(\cdot)$ represents the transformation function and $θ2\theta_2$ is the transformation parameter that determines the level of transformation. As $wi(⋅)w_i(\cdot)$ constrains the error by transforming features, the loss of details caused by uncertainty will be reduced.

3.3 网络结构

Feature Fusion Residual Block

我们使用group normalization[29] 以及 SE block[17]来更好地进行特征融合。如图3b所示，特征融合残差块（feature fusion residual block，FFRB）公式化表达如下：

$F F R B (x) = S E (G N (C o n v (R e s (x))))$

这里 $x$ 指代输入信号， $Res(⋅)Res(\cdot)$ 指代包含group normalization（GN）的残差块。由于当batch很小时group normalization的效果要比batch normalization以及instance normalization好，因此这里使用group normalization来减少内部协方差偏移。通过使用SE块[17]，有着更多语义信息的特征通道将会获得更大的权重。数学上来说，SE块可以表示如下：

$S(x)=Sigmoid(FC(Relu(FC(GAP(x)))))⊗xS(x)=Sigmoid(FC(Relu(FC(GAP(x)))))\otimes x$

这里 $GAP(⋅)GAP(\cdot)$ 是global average pooling操作， $FC(⋅)FC(\cdot)$ 指代全连接层。

Multi-stream Residual Architecture

我们提出的去雨网络基于编码解码器结构[24,31]。由编码解码器的U-FFRB模块带来的较大的感知域能够获取语义信息。由于雨水条纹通常有不同的形状、密度以及大小，仅使用one-stream CNN[23,24,31]可能会损失一些有用的信息。受到多尺度卷积核[33,28]的启发，我们使用multi-stream残差网络来提取多尺度特征。属于各自stream的FFRB的核大小分别被设定为3，5和7，我们利用这些核来寻找丢失的细节。如图3a所示，我们可以将multi-stream结构描述如下：

$M(x)=[U3×3(x),U5×5(x),U7×7(x)]M(x)=[U_{3\times3}(x),U_{5\times5}(x),U_{7\times7}(x)]$

这里 $Ui×iU_{i\times i}$ 代表核大小为 $i$ 的U-FFRB模块， $[⋅][\cdot]$ 代表concatenation操作。

（核心思想是使用不同大小的核来表示不同形状、密度、大小的雨水条纹）

Error detector 以及 feature compensator

对于特征补偿器，我们引入了一个额外的特征变换模块来自适应地对残差图真值进行变换（右端部分），并且如图3d中所示，我们构建了一个简明的编码解码branch来学习转换后的残差图（左端部分）。对于不同stream间的相关性，学习到的转换后残差图被拷贝了三份并且分别使用skip-connection嵌入到了三个stream内，如图3a所示（F1，F2）。对于error detector，两个包含FFRB的简明的branch被构建来学习残差图以及误差图。注意一个为了误差图生成所构建的branch包含编码解码器模块。误差图被用来补偿对应的embedding残差图（如图3-a所示）。

Refinement module

refinement模块的主要目的是调整聚合在一起的特征图。我们使用spatial pyramid pooling[9]来进一步获取多尺度特征。尺度因子被分别设定为4，8，16以及32。对于不同尺度的特征图，我们使用逐点卷积（1*1卷积）来降低它们的通道数并且使用最近邻插值的上采样来恢复它的原本大小。如图3a所示，down-up 结构可以进行如下公式化表达：

3.4 Loss Function

Mean Square Error（MSE）的平方惩罚会惩罚较大的error并且容忍较小的error，这会导致产生over-smoothed图片。因此，我们使用Mean Absolute Error（MAE）来提高鲁棒性。对于修正embedding residual map的error detector，我们最小化如下两个loss function：

$Le1=∥Rt×0.5−ϕ1(It)∥1L_{e1}=\Vert R_{t\times 0.5}-\phi_1(I_t)\Vert_1$

$Le2=∥θ1Rt×0.5−ϕ1(It)−φ(It×0.5,ϕ1(It))∥1(12)L_{e2}=\Vert\frac{\theta_1}{R_{t\times0.5}-\phi_1(I_t)}-\varphi(I_{t\times0.5},\phi_1(I_t))\Vert_1\quad\quad\quad\quad(12)$

注意公式12中的 $∣Rt×0.5−ϕ1(It)∣|R_{t\times0.5}-\phi_1(I_t)|$ 是固定的。对于feature compensator，loss function以及regularizer可以进行如下公式化表述：

$Lc=∑i∈{0.25,0.5}∥Rt×i+θ2wi(Rt×i)Rt×i−ϕi(It)∥1L_c=\sum\limits_{i\in\{0.25,0.5\}}\Vert R_{t\times i}+\theta_2w_i(R_{t\times i})R_{t\times i}-\phi_i(I_t)\Vert_1$

$Lp=∥wi(Rt×i)∥22L_p=\Vert w_i(R_{t\times i})\Vert_2^2$

在proposed模型中，最终输出的residual map通过如下的loss function来进行refined：

$Lf=∥Rt−f(It)∥1L_f=\Vert R_t-f(I_t)\Vert_1$

这里 $f(⋅)f(\cdot)$ 指代总体网络方程。此外，motivated by the goal of generating derained image to approximate its ground-truth image, we utilize the SSIM function [37] as the additional evaluation metric for the generated clean image. It is formulated as follows:

$L_{SSIM}=-SSIM(B_t,I_t-f(I_t))$

这里 $B_t$ 指代clean image truth。MAE损失函数以及SSIM损失函数被组合在一起可以保留per-pixel相似性并且同时保留全局结构。用于训练我们提出的RLNet所使用的overall loss function公式化表达如下：

$Lall=Lf+LSSIM+λLp+λ1Le1+λ2Le2+λ3LcL_{all}=L_f+L_{SSIM}+\lambda L_p+\lambda_1L_{e1}+\lambda_2L_{e2}+\lambda_3L_c$

这里 $λ,λ1,λ2\lambda,\lambda_1,\lambda_2$ 以及 $λ3\lambda_3$ 都是设定好的参数。