Deep K-SVD Denoising

一、预备知识

1.1 Iterative Shrinkage Thresholding Algorithm(ISTA)

这个可以看这篇文章：ISTA

1.2 K-SVD

对于模型：
$$\underset{D,X}{\min }\Vert Y-DX{\Vert }_F^2,\text{ s.t. }\forall i,{\Vert x_i\Vert }_0\le \epsilon$$
首先随机初始化一个字典$D$，利用OMP算法，计算得到稀疏编码矩阵$X$。模型可以简化为：
$$\underset{D}{\min }\Vert Y-DX{\Vert }_F^2$$
逐列更新字典。下面以更新字典的第$k$列为例，记$d_k$为字典$D$的第$k$列向量，记$x_T^k$为稀疏矩阵$X$的第$k$行向量，所以有：
$$\begin{array}{rl}\Vert Y-DX{\Vert }_F^2& ={\Vert Y-\sum _{j=1}^K{d}_j{x}_T^j\Vert }_F^2={\Vert \left(Y-\sum _{j\ne k}{d}_j{x}_T^j\right)-d_k{x}_T^k\Vert }_F^2={\Vert E_k-d_k{x}_T^k\Vert }_F^2\end{array}$$
上式中残差$E_k=Y-{\sum }_{j\ne k}{d}_j{x}_T^j$。

此时优化问题可描述为：
$$\underset{d_k,x_T^k}{\min }{\Vert E_k-d_k{x}_T^k\Vert }_F^2$$
因此我们需要求出最优的$d_k$, $x_T^k$，这里利用SVD的方式求解出两个优化变量。

但是，在这里需要注意的是，不能直接利用$E_k$进行求解，否则求得的新的$x_T^k$不稀疏。因此我们需要将$E_k$中对应的$x_T^k$不为0的位置提取出来，得到新的$E_k^{\prime }$。

上图很好的说明了如何求得$x_T^{\prime k}$和$E_k^{\prime }$。其实也很容易理解，这样求解出来的$x_T^{\prime k}$和$E_k^{\prime }$只会越来越稀疏，因为每次都是取非0的数进行计算的！

此时优化问题可描述为：
$$\underset{d_{k,}{x}_T^k}{\min }{\Vert E_k^{\prime }-d_k{x}_T^{\prime k}\Vert }_F^2$$
因此我们需要求出最优的$d_k$, $x_T^{\prime k}$，根据SVD分解，可得到：
$$E_k^{\prime }=U\Sigma V^T$$
取左奇异矩阵$U$的第1个列向量作为$d_k$，取右奇异矩阵的第1个行向量与第1个奇异值的乘积作为$x_T^{\prime k}$，然后将其对应地更新到$x_T^k$。

这样迭代一定数目之后，算法就可以训练出一个性能比较好的字典了。K-SVD算法的流程图如下：

1.3 Conv2d: Unfold and Fold

平时使用卷积操作时，既卷积核滑动窗口操作，调用nn.Conv2d就能完成对输入的卷积操作。但有时，可能要探究卷积核对应的某一通道的单个窗口的卷积操作，或显式地进行卷积操作。此时，就需要nn.Unfold和nn.Fold。

一般来说，Conv2d = Unfold + matmul + Fold。

• nn.Unfold按照官方的说法，就是从一个batch样本中，提取出滑动的局部区域块，也就是卷积操作中的提取Filter对应的滑动窗口。
• nn.Fold的操作与nn.Unfold相反，将提取出的滑动局部区域块还原成batch的张量形式。

下图很清晰的说明了nn.Unfold是如何工作的。

nn.Unfold的参数如下：

torch.nn.Unfold(kernel_size, dilation=1, padding=0, stride=1)

nn.Unfold的输入为$(N,C,H,W)$，那么nn.Unfold的输出为$(N,K,L)$，计算公式如下：
$$K=C\times \prod (kernel\_size)$$

$$L=\prod _d\lfloor \frac{\text{spatial\_size}[d]+2\cdot \text{padding}[d]-\text{dilation}[d]\cdot (\text{kernel\_size}[d]-1)-1}{\text{stride}[d]}+1\rfloor$$

其中，spatial_size表示输入张量的空间维度。

Example: input:(1,2,4,4)+kernel_size:(3,3)–>output:(1,18,4)

二、论文解读

这篇paper主要的贡献是提出了一个端到端的深度学习算法框架（LKSVD）。这个框架保留K-SVD算法原来的计算路径的同时，重新设计了一个基于监督学习的架构。这套架构需要学习的参数量少，并且保留了K-SVD算法的本质，其性能大大优于经典的K-SVD算法，并且非常接近最先进的基于深度学习的去噪方法。

LKSVD的框架如下：

在这个框架中，首先利用前面讲的Unfold的方式，将图片展开成很多重叠的小块，然后对每一小块做之后的处理。图中Patch Decomposition就是负责这一块的。再之后，在Patch Denoiser中，就是求解sparse coding的过程。也就是已知$y$和字典$D$求解稀疏向量$x$。这里采用的是ISTA算法。由于ISTA算法中的$\lambda$极大的影响算法的性能，这个框架中，专门使用一个神经网络来估计$\lambda$，其他的和ISTA算法一致。最后，采用平均的方式，用Fold将各个子块恢复为图片。整个架构的损失函数就是MSE。

下面分别说明每一部分。

2.1 Patch Denoiser：Sparse Coding

将${\ell }_0$范数换为${\ell }_1$范数，稀疏编码问题可以转换为如下：
$$\hat{\alpha }=\arg \underset{\alpha }{\min }\frac{1}{2}\Vert \mathbf{D}\alpha -\mathbf{y}{\Vert }_2^2+\lambda \Vert \alpha {\Vert }_1$$
其中，$\lambda >0$是正则化系数。利用ISTA算法来学习稀疏编码矩阵$\alpha$，其迭代式如下：
$${\hat{\alpha }}_{t+1}=S_{\lambda /c}\left(\widehat{{\alpha }_t}-\frac{1}{c}{\mathbf{D}}^T\left(\mathbf{D}{\hat{\alpha }}_t-\mathbf{y}\right)\right);{\hat{\alpha }}_0=0$$
其中， $c$ 是$D$的平方谱范数，$S_{\lambda /c}$是软阈值操作函数：
$${\left[S_{\theta }(\mathbf{v})\right]}_i=\mathrm{sign}\left(v_i\right){\left(\left|v_i\right|-\theta \right)}_{+}$$
通过近端梯度下降方法，将稀疏编码部分变成一个可学习的版本，其中$c$和$D$是可学习的参数。

2.2 Patch Denoiser: $\lambda$ Evaluation

为了让误差可控，对每个$y_k$都学习一个${\lambda }_k$。利用神经网络学习一个回归函数，也就是做如下映射：
$$\lambda =f_{\theta }(y)$$
其中，$\theta$为神经网络的参数。

具体网络架构如下：
$$\begin{array}{rl}MLP:y& \to [p\times 2p]\to ReLU\\ & \to [2p\times p]\to ReLU\\ & \to [p\times \frac{1}{2}p]\to ReLU\\ & \to [\frac{1}{2}p\times 1]\to \lambda \end{array}$$
其中，$p$为字典$D$的行数。在整个网络中，一共也只有$4p^2$个参数。

2.3 Patch Denoiser: Patch Reconstruction

$\hat{y}$是经过降噪后得到的图片矩阵，根据已经求得的字典$D$和稀疏向量$\hat{\alpha }$就可以得到降噪后的图片：
$$\hat{y}=D\hat{\alpha }$$
在神经网络里面，$D$是一组可学习的参数。

整个Patch Denoising和Patch Reconstruction过程与Convolutional Sparse Coding方法高度相关。

2.4 End-to-End Architecture

这个完整的端到端的算法架构可以描述如下：

• 将输入图像分割成完全重叠的patch；
• 通过上述Patch Denoising阶段对每个损坏的patch进行处理；
• 对这些patch的干净版本取平均来重构图像。

在最后一步中，舍弃原始的K-SVD，设计一个可学习的方法来重构图像。令$\mathbf{w}\in {\mathbb{R}}^{\sqrt{p}\times \sqrt{p}}$为每一块的权重系数，通过下式得到重构后的图像：
$$\hat{\mathbf{Y}}=\frac{{\sum }_k{\mathbf{R}}_k^T\left(\mathbf{w}\odot {\hat{y}}_k\right)}{{\sum }_k{\mathbf{R}}_k^T\mathbf{w}}$$
在整个算法中，可学习的参数有：$\theta$，$c$，$D$，$\mathbf{w}$。

三、参考文献

• Scetbon M, Elad M, Milanfar P. Deep k-svd denoising[J]. IEEE Transactions on Image Processing, 2021, 30: 5944-5955.

• Aharon M, Elad M, Bruckstein A. K-SVD: An algorithm for designing overcomplete dictionaries for sparse representation[J]. IEEE Transactions on signal processing, 2006, 54(11): 4311-4322.

• Beck A, Teboulle M. A fast iterative shrinkage-thresholding algorithm for linear inverse problems[J]. SIAM journal on imaging sciences, 2009, 2(1): 183-202.