WINNet: Wavelet-inspired Invertible Network for Image Denoising

Blind image denoising can be achieved by learning from training samples with varying noise levels [21, 22, 23], however, the learned deep models still have poor generalization ability on images with noise levels beyond the training levels. Hence, one may wish to learn a universal denoising model with good performances on arbitrary noise levels.

Mohan et al. [30] found that the bias term in a denoising CNN will lead to a model bias to the training noise levels. Hence, they propose to remove all the bias terms in both ReLU activation functions and the batch normalization layers. The resultant bias-free CNN (BF-CNN) has a scale homogeneity property, i.e., $f_{\text{BF}}(\alpha\bm{y})=\alpha f_{\text{BF}}(\bm{y})$ where $f_{\text{BF}}(\cdot)$ denotes a bias-free CNN and $\alpha>0$, and shows consistently better generalization ability than its counterpart with bias terms. Helou and Süsstrunk [31] proposed a blind universal image fusion denoiser (BUIFD) whose structure is derived from a Bayesian framework and consists of a noise level CNN, a prior CNN and a fusion network. The information of noisy image and image prior are fused based on signal-to-noise ratio (SNR). By explicitly incorporating the noise level, BUIFD shows stronger generalization ability than that of conventional denoising CNNs.

The lifting scheme introduced in [28, 29] proposes to construct a wavelet transform by first splitting the signal into an even and an odd part. A predictor is used to predict the odd signal from the even part, therefore leading to a sparse residual signal. The update step is used to adjust the even signal based on the prediction error of the odd part to make it a better coarse version of the original signal. There can be multiple pairs of predictor and updater. The lifting scheme can represent a transform with perfect reconstruction condition and any intermediate representation can be used to infer the input signal and the final representation.

Inspired by this scheme, the invertible neural networks (INNs) [32, 33, 34, 35] was proposed for memory-efficient backpropagation model and used for constructing flow-based generative models. INNs are bijective function approximators and the intermediate feature representation of INNs can be used to perfectly reconstruct the input signal. Different invertible architectures have been proposed, including coupling layer [32], affine coupling layer [33], reversible residual network [34] and i-RevNet architecture [35]. INNs have been applied for low-level computer vision tasks, including image rescaling [36], image segmentation [37], image colorization [38, 39], image compression [40, 41, 42], and image denoising [43, 44].

The splitting operator is used to separate the input image into two parts, and the merging operator performs the inverse of the splitting operator. The splitting/merging operator in current methods [32, 33, 34, 35, 46, 38, 40, 36, 37] keep the same number of input and output coefficients.

In this paper, we propose to learn LINNs with redundant representations in analogy with the redundant wavelet transform which provides better performance compared to the decimated wavelet transform in image restoration tasks [47, 11]. Therefore, the splitting operator will lead to a redundant representation, and merging operator recovers the input image from the redundant representation.

For an input image $\bm{y}^{k}\in\mathbb{R}^{1\times N_{1}\times N_{2}}$ at the $k$-th scale, the splitting operator $\mathcal{S}$ is parameterized by a convolutional kernel $\bm{K}_{s}\in\mathbb{R}^{c\times 1\times p\times p}$ where $c$ denotes the number of channels and $p$ denotes the spatial filter size. A convolution is performed to obtain multi-channel features $\bm{f}^{k}\in\mathbb{R}^{c\times N_{1}\times N_{2}}$:

where $\otimes$ denotes the convolution operator.

The split operation then divides the first $h$ channels of $\bm{f}^{k}$ into coarse part $\bm{c}^{k}\in\mathbb{R}^{h\times N_{1}\times N_{2}}$, and the remaining $c-h$ channels into detail part $\bm{d}^{k}\in\mathbb{R}^{(c-h)\times N_{1}\times N_{2}}$. In this paper, we set the number of coarse channels to $h=1$, and we denote the splitting operator as:

The merging operator $\mathcal{M}$ represents the inverse of the splitting operator, and is parameterized by a convolutional kernel $\bm{K}_{m}\in\mathbb{R}^{1\times c\times p\times p}$. It first concatenates $\bm{c}^{k}$ and $\bm{d}^{k}$ and this results in $\hat{\bm{f}}^{k}$. A convolution is performed to recover $\bm{y}^{k}$:

The merging operator can then be denoted as:

To achieve invertibility, we can pick $\bm{K}_{s}$ so that it can be reshaped to a tight frame with $c\geq p^{2}$. The merging convolutional kernel then simply corresponds to the transpose of the splitting convolutional kernel. This will ensure the splitting and merging operators are invertible. We can use orthogonal transforms of size $c\times c$ such as orthogonal wavelet transforms, Discrete Cosine Transforms (DCT) and other analytical transforms as the splitting and merging operators. The tight frame is then obtained by picking $p^{2}$ columns of the original orthogonal matrix.

Besides the analytical transforms, it is also possible to learn an orthogonal matrix with proper parameterization methods, for example, the Cayley transform [48] of a skew-symmetric matrix $\bm{\Theta}-\bm{\Theta}^{T}$ is an orthogonal matrix:

where $\mathrm{I}$ is the identity matrix. The corresponding tight frame $\bm{K}_{s}$ can therefore be a sub-matrix (i.e., extracting $p^{2}$ columns) of the learned orthogonal matrix $\bm{K}$ parameterized by learnable parameters $\bm{\Theta}\in\mathbb{R}^{c\times c}$.

At each scale, the Predict/Update networks are shared among the forward and inverse transform of LINN, but are connected with different signs and direction.

Fig. 3(a) and Fig. 3(b) shows the schematics of the forward transform and the inverse transform of the $k$-th scale LINN, respectively. In the forward transform, $\left(\bm{c}^{k},\bm{d}^{k}\right)$ will be non-linearly transformed to a representation which is more suitable for denoising. After denoising, the denoised detail part and the original coarse part will be transformed back to the original domain using the inverse transform of LINN.

In the forward transform (shown in Fig. 3(a)), a Predict network conditioned on the coarse part aims to predict the detail part to make the resultant residual of the detail part sparse. The Update network conditioned on the detail part is used to adjust the coarse part to make it a smoother version of the input image. The Predict and Update networks are applied alternatively to update the coarse and detail parts. Let us denote $\left(\bm{c}^{k}_{0},\bm{d}^{k}_{0}\right)=\left(\bm{c}^{k},\bm{d}^{k}\right)$. The $m$-th pair ($m\in[1,M]$) of update and predict operation can be expressed as:

where $\bm{d}^{k}_{m}$ and $\bm{c}^{k}_{m}$ denotes the updated detail part and coarse part using the $m$-th Predict network $P_{m}^{k}(\cdot)$ and Update network $U_{m}^{k}(\cdot)$, respectively.

In the inverse transform of the $k$-th scale LINN (shown in Fig. 3(b)), $\left(\bm{d}^{k}_{m-1},\bm{c}^{k}_{m-1}\right)$ for $m\in[1,\cdots,M]$ can be estimated based on $\left(\bm{d}^{k}_{m},\bm{c}^{k}_{m}\right)$ and $P_{m}^{k}(\cdot)$ and $U_{m}^{k}(\cdot)$ as follows:

When no lossy operation is applied on $\left(\bm{c}^{k}_{M},\bm{d}^{k}_{M}\right)$, the inverse transform of LINN can, by construction, perfectly recover the inputs of the forward transform with any Predict and Update networks. In this paper, we enforce LINN to learn a sparsifying transformation through denoising.

The Predict and Update networks can be any functions and their properties will not affect the invertibility of LINN. In this paper, we use the same structure for each Predict/Update networks which we therefore denote with PUNets. To achieve a better complexity and performance trade-off, we construct PUNets with depth-wise separable convolution (SepConv) layers [49] and soft-thresholding non-linearity [8]. As shown in Fig. 4(c), a PUNet consists of an input convolutional layer $\text{Conv}_{i}$ with a soft-thresholding, $J$ residual blocks with SepConv layers, and an output convolutional layer $\text{Conv}_{o}$.

Separable Convolution: A depth-wise separable convolution layer [49] consists of a depth-wise convolution layer and a $1\times 1$ convolution layer (shown in Fig. 4(a)). It can reduce the number of parameters, especially when the channel number and the spatial filter size are large. Compared to the standard convolution of size $c_{out}\times c_{in}\times q\times q$, the SepConv layer only requires $c_{in}\times q\times q+c_{out}\times c_{in}$ parameters.

Soft-thresholding: The soft-thresholding operator is used as the non-linearity in PUNets. The soft-thresholding operator has a tight connection with sparse coding. It can also be considered as a two-sided ReLU non-linearity and can be more effective than ReLU for some image processing applications. The soft-thresholding operator can be expressed as:

To ensure $\bm{\lambda}$ are non-negative, a Softplus function¹¹1$\text{Softplus}(x)=\frac{1}{\beta}\cdot\log(1+\exp{(\beta\cdot x)})$. is applied on each learned thresholds.

Multi-scale Property: Multi-resolution signal decomposition is an essential property of the wavelet transform. For an input image, the wavelet transform provides a multi-scale analysis which captures the information at different scales. In order to mimic the wavelet transform, we iterate the LINN-based decomposition on the coarse part. Moreover, we apply à trous convolution with dilation rate $2^{k-1}$ for the $k$-th scale LINN even though the redundancy factor in our setting is not two. As a result, the larger scale LINN will have a bigger receptive field. Fig. 5 shows the exemplars of dilated $3\times 3$ filter at two different levels.

Any existing image denoiser can be applied here, however, as our objective is to achieve an image denoising algorithm with high interpretability, we propose to use a sparsity-driven denoising architecture.

Let us assume that the learned forward transform by LINN is a sparsifying transform and the small number of large coefficients correspond to the image contents. (The forward transformed images by LINN will be shown in Section VI Simulation Results to validate this assumption). The denoising process can then be modeled as a basis pursuit denoising problem. To enrich the representation ability of the sparsity-driven denoising network, we over-parameterize with a convolutional dictionary as:

where $\bm{g}=[\bm{g}_{1},\cdots,\bm{g}_{N}]$ are the sparse features, $\bm{D}=[\bm{D}_{1},\cdots,\bm{D}_{N}]$ is the over-parameterized convolutional dictionary, and $\bm{\lambda}=[\lambda_{1},\cdots,\lambda_{N}]$ are the regularization parameters.

The above $l_{1}$-norm minimization problem can be solved using Iterative Shrinkage-Thresholding Algorithm (ISTA) [50]:

where $\mu$ is the step size.

The LISTA [51] algorithm parameterizes the unknown dictionaries in ISTA as learnable parameters. We apply a convolutional LISTA (CLISTA) as our sparsity-driven denoising network. The schematic of the CLISTA denoising network is shown in Fig. 6. There are $T$ layers of soft-thresholding operations in CLISTA. The convolutional dictionaries $\bm{W}_{a}$ and $\bm{W}_{s}$ are shared across layers. For the $t$-th layer, there is a soft-thresholding operator with learnable soft-thresholds:

where $\bm{W}_{a}\in\mathbb{R}^{N\times(c-h)\times r\times r}$ is the analysis convolutional dictionary, $\bm{W}_{s}\in\mathbb{R}^{(c-h)\times N\times r\times r}$ is the synthesis convolutional dictionary, and $\bm{\lambda}_{t}\in\mathbb{R}^{N}$ is the soft-threshold vector.

For the first layer, $\bm{g}_{0}$ is initialized as $\bm{W}_{a}\otimes\bm{d}^{k}_{M}$. At the last layer, the synthesis convolution layer $\bm{W}_{s}$ is used to reconstruct the denoised detail channels from the estimated sparse feature $\bm{g}_{T}$, i.e., $\hat{\bm{d}}^{k}_{M}=\bm{W}_{s}\otimes\bm{g}_{T}$.

To achieve robust image denoising, we propose to adapt the soft-thresholds in PUNets and CLISTA to the noise level $\sigma$ in the image. For images with larger noise levels, the soft-thresholds will be larger leading to more effective denoising. This is similar to the model-based image denoising methods in [8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18] where the soft-thresholds scale with the noise levels. For noisy image super-resolution problem, DeepAM [52] also shows that better generalization can be achieved by adjusting soft-thresholds with respect to noise level.

Different from BF-CNN which does not have bias terms and has restricted expressive power, our PUNets and CLISTA denoising network have the dynamically adjusted soft-thresholds with respect to the noise level. This gives us good generalization ability, flexibility as well as improved expressive power.

Spectral Norm Loss: When learning from noisy-clean image pairs of a single noise level $\sigma_{N}$, the parameters of PUNets are learned to optimize the performance of this training noise level. When the test noise level $\sigma_{T}<\sigma_{N}$, the rescaled soft-thresholds could make the PUNets change from a contractive function to an expansive function. We therefore use a spectral norm loss to regularize the property of the PUNets:

where $\|\cdot\|_{S}$ denotes the spectral norm of a matrix, and $P_{m,j}^{k}$ and $U_{m,j}^{k}$ are the effective linear filter of the $j$-th residual block on the $m$-th lifting step at the $k$-th scale when soft-thresholds are set to zeros.

The linear loss $\mathcal{L}_{{s}}$ is used to regularize the Lipschitz constant of the PUNets when $\sigma_{T}\rightarrow 0$, i.e., assuming that soft-thresholding operators in PUNets are identity functions. It can improve the generalization ability of WINNets to unseen small noise levels and can also stabilize the training of deep/multi-scale models by constraining the expansiveness of the PUNets.

Orthogonal Loss: For CLISTA denoising network, the soft-thresholds will also be adjusted with respect to noise level. We would like to enforce $\hat{\bm{d}}^{k}_{M}=\bm{d}_{M}^{k}$ when $\sigma_{T}=0$. That is, the CLISTA denoising network should be close to an identity function when there is no noise. This can be achieve when $\bm{W}_{s}$ and $\bm{W}_{a}$ are orthogonal. Therefore, an orthogonal loss is imposed between the synthesis and the analysis convolutional dictionary:

where $\bm{\delta}\in\mathbb{R}^{(c-h)\times(c-h)\times(2r-1)\times(2r-1)}$ is the Kronecker delta function which takes value 1 at coordinate $(i,i,r,r)$ $\forall i\in[1,\cdots,c-h]$ otherwise 0.

Since the PUNets are conditioned on the image noise level, we propose a simple model-inspired noise estimation network in order to achieve blind image denoising. The proposed noise estimation network is inspired by [53] here, authors use an iterative algorithm to select low-rank patches from the input noisy image and estimate the noise level as the smallest eigenvalue of the covariance matrix of the selected low-rank patches. Compared to CNNs used to infer noise levels, e.g., [23, 31], this model-based method can provide more interpretable results and robust generalization ability.

We propose a learnable noise estimation network (NENet) which is based on the principles of [53] and uses backpropagation through singular value decomposition (SVD) in PyTorch. The network is shown in Fig. 7. The input image is divided into $\bm{P}\in\mathbb{R}^{s^{2}\times N_{p}}$ patches where each column represents a vectorized $s\times s$ patch and $N_{p}$ is the number of patches. Instead of iteratively selecting low-rank patches, we propose to use a selection network (SENet) $\mathcal{S}:\mathbb{R}^{s^{2}}\rightarrow\mathbb{R}^{1}$ to estimate a scalar weight $w_{i}\in[0,1]$ for each patch. The SENet performs a soft selection and will learn to assign large weights to low-rank patches and vice versa. Given the estimated weights and the patches, the noise variance is estimated as the minimum singular value of the weighted covariance matrix of the patch matrix:

where ${\sigma}_{\text{min}}$ denotes the smallest singular value of the weighted covariance matrix $\bm{P}\text{diag}(\bm{w})\bm{P}^{T}$ with $\bm{w}=[w_{1},\cdots,w_{N_{p}}]^{T}$.

The SENet is composed of blocks with 1-D convolution layer and ReLU non-linearity without bias term, followed by a global averaging pooling and a Sigmoid function to map the weight in the range of $[0,1]$. The parameters of SENet can be learned using backpropagation. For the noise estimation network, the MSE noise loss is used:

where $\sigma_{i}$ and $\hat{\sigma}_{i}$ denotes the ground-truth noise level and the estimated noise level, respectively.

WINNet can learn from noisy-clean image pairs of a single noise level or a range of noise levels as in [21, 22, 23]. The average mean squared error (MSE) between the restored image and the clean image is used as the reconstruction loss:

where $\bm{x}_{i}$ and $\bm{\hat{x}}_{i}$ denotes the input clean image and the denoised image, respectively.

The overall training objective of our WINNet is:

where $\lambda_{1}$, and $\lambda_{2}$ are regularization parameters.

We follow the training and evaluation settings in [21] to perform experiments. Before training and testing, the clean images are normalized to $[0,255]$. The noisy images are obtained by adding additive white Gaussian noise to the clean image with respect to Eq. (1) with variance $\sigma^{2}$.

In this section, we will visualize the output produced by components of a WINNet with $K=2$ levels which is trained on data with noise level $\sigma_{N}=25$.

We compare the proposed WINNet with several state-of-the-art image denoising algorithms including the model-based methods: BM3D [16], WNNM [17], EPLL [57], and the learning-based methods: DnCNN [21], FFDNet [22], BUIFD [31], BF-CNN [30]. The model-based methods can handle images with arbitrary noise levels. DnCNN [21] and BUIFD [31] can perform blind image denoising when the test noise level is within the range of the training noise levels, BF-CNN [30] can handle unseen noise levels, while the FFDNet [22] requires as input a noise map of the test noise level.

With the plug-and-play technique, image denoisers can be applied to solve general image restoration problems [19, 20, 21, 22, 23], for example, image deblurring. In this case, the goal is to recover a sharp image $\bm{x}$ from the blurred and noisy observation $\bm{y}=\bm{k}\otimes\bm{x}+\bm{n}$ where $\bm{k}$ is the blurring kernel and $\bm{n}\sim\mathcal{N}(0,\sigma^{2})$ represents the measurement noise with variance $\sigma^{2}$. The image deblurring task can be formulated as the following optimization problem:

where $\Phi(\cdot)$ is a prior term and $\lambda$ is the regularization parameter. With half-quadratic splitting [6], the image deblurring problem can be solved by iteratively optimizing two sub-problems:

where $\beta=\sqrt{\lambda/\mu}$ is a hyper-parameter and can be interpreted as the noise level if the $\bm{z}$ sub-problem is treated as Gaussian denoising on $\bm{x}_{k}$.

In [6], a CNN-based Gaussian denoiser is used to solve the $\bm{z}$ sub-problem. The hyper-parameter $\lambda$ is set to be fixed during iterations, while $\mu$ is set to exponentially decay from a large value to the given noise level $\sigma$ with a fixed iteration number. Since the proposed NENet is an effective noise level estimator and WINNet can denoise images with noise beyond training noise levels, we propose to use WINNet for image deblurring. The $\bm{x}$ sub-problem can be solved with closed-form solution with the estimated noise level $\beta_{k}$ by NENet. The $\bm{z}$ sub-problem can be solved using the proposed blind WINNet with noise level $2\beta_{k}$ (perform denoising with a stronger strength to ensure convergence). With the proposed robust NENet and WINNet, we can achieve image deblurring without accessing the noise levels and using the pre-defined regularization parameters; $\lambda$ is the only free parameter and is set to 0.23 as in [6]. Algorithm 1 illustrates the plug-and-play image deblurring algorithm with the proposed WINNet.

Table III shows the average PSNR (dB) of the EPLL [57], IRCNN [6] and the proposed method evaluated on Set12 with 8 different kernels from [58] and noise level 2.55. We can see that the proposed method is able to achieve highly competitive performance. Fig. 13 shows an exemplar image deblurring process of the proposed method on image Cameraman which is blurred using the first kernel from [58]. We can see that the proposed method converges after 8 iterations and the PSNR of $\bm{x}_{k}$ and $\bm{z}_{k}$ consistently improves and finally reaches a similar result.