Learning Non-local Image Diffusion for Image Denoising

TL;DR

This paper introduces a trainable non-local reaction diffusion (TNLRD) model that integrates non-local self-similarity priors into image denoising, outperforming local models and achieving state-of-the-art results.

Contribution

The paper proposes a novel TNLRD model that embeds non-local filters into the diffusion framework, enhancing denoising performance and reducing artifacts compared to local models.

Findings

TNLRD produces images with more textures and fewer artifacts.

The model achieves competitive PSNR and SSIM scores.

It outperforms local diffusion models in denoising quality.

Abstract

Image diffusion plays a fundamental role for the task of image denoising. Recently proposed trainable nonlinear reaction diffusion (TNRD) model defines a simple but very effective framework for image denoising. However, as the TNRD model is a local model, the diffusion behavior of which is purely controlled by information of local patches, it is prone to create artifacts in the homogenous regions and over-smooth highly textured regions, especially in the case of strong noise levels. Meanwhile, it is widely known that the non-local self-similarity (NSS) prior stands as an effective image prior for image denoising, which has been widely exploited in many non-local methods. In this work, we are highly motivated to embed the NSS prior into the TNRD model to tackle its weaknesses. In order to preserve the expected property that end-to-end training is available, we exploit the NSS prior by a set of non-local filters, and derive our proposed trainable non-local reaction diffusion (TNLRD) model for image denoising. Together with the local filters and influence functions, the non-local filters are learned by employing loss-specific training. The experimental results show that the trained TNLRD model produces visually plausible recovered images with more textures and less artifacts, compared to its local versions. Moreover, the trained TNLRD model can achieve strongly competitive performance to recent state-of-the-art image denoising methods in terms of peak signal-to-noise ratio (PSNR) and structural similarity index (SSIM).

Tables

Table 1. TABLE I: Block matching results

index	1	2	3	$\cdots$	$n$	$\cdots$	$p$
$Q_2$	$q_{12}$	$q_{22}$	$q_{32}$	$\cdots$	$q_{n2}$	$\cdots$	$q_{p2}$
$Q_3$	$q_{13}$	$q_{23}$	$q_{33}$	$\cdots$	$q_{n3}$	$\cdots$	$q_{p3}$
$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$
$Q_j$	$q_{1j}$	$q_{2j}$	$q_{3j}$	$\cdots$	$q_{nj}$	$\cdots$	$q_{pj}$
$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$	$\cdots$
$Q_L$	$q_{1L}$	$q_{2L}$	$q_{3L}$	$\cdots$	$q_{nL}$	$\cdots$	$q_{pL}$

Table 2. TABLE II: Influence of different loss function L2 and SSIM for image denoising with σ = 15 𝜎 15 \sigma=15 , 25 25 25 and 50 50 50 . Average PSNR (dB) and SSIM on 68 images for each σ 𝜎 \sigma .

Inference	Training
	$\sigma =15$		$\sigma =25$		$\sigma =50$
	L2	SSIM	L2	SSIM	L2	SSIM
PSNR	31.50	0.8852	29.01	0.8201	26.06	0.7094
SSIM	31.31	0.8864	28.83	0.8219	25.80	0.7113

Table 3. TABLE III: Average PSNR(dB) on 68 images from [ 25 ] for image denoising with σ = 15 𝜎 15 \sigma=15 , 25 25 25 and 50 50 50 .

$\sigma$	Method
	BM3D	WNNM	NLR-MRF	TNRD	TNLRD*	TNLRD**
	[21]	[19]	[20]	[25]
15	31.08	31.37	30.97	31.42	31.50	31.31
25	28.56	28.83	28.48	28.91	29.01	28.83
50	25.62	25.83	25.38	25.96	26.06	25.80

Table 4. TABLE IV: Average SSIM on 68 images from [ 25 ] for image denoising with σ = 15 𝜎 15 \sigma=15 , 25 25 25 and 50 50 50 .

$\sigma$	Method
	BM3D	WNNM	NLR-MRF	TNRD	TNLRD*	TNLRD**
	[21]	[19]	[20]	[25]
15	0.8717	0.8759	0.8699	0.8821	0.8852	0.8864
25	0.8013	0.8084	0.7972	0.8152	0.8201	0.8219
50	0.6864	0.6981	0.6665	0.7024	0.7094	0.7113

Equations

$$f=u+v,$$

$$\text{E}(u|f)=\frac{\lambda }{2}{\left|\left|u-f\right|\right|}_2^2+\stackrel{N_k}{\sum _{i=1}}{\rho }_i(k_i\ast u),$$

$$\frac{u_t-u_{t-1}}{\Delta t}=-\stackrel{N_k}{\sum _{i=1}}{\bar{k}}_i\ast {\varphi }_i(k_i\ast u_{t-1})-\lambda (u_{t-1}-f),$$

$$\{\begin{array}{l}{u}_0=f,t=1,\cdots ,T\\ u_t=u_{t-1}-\left(\stackrel{N_k}{\sum _{i=1}}{\bar{k}}_i^t\ast {\varphi }_i^t(k_i^t\ast u_{t-1})+{\lambda }^t(u_{t-1}-f)\right).\end{array}$$

$$\{\begin{array}{l}{\Theta }^{\ast }={\text{argmin}}_{\Theta }\mathcal{L}(\Theta )=\sum _{s=1}^S\ell \left(u_T^s,u_{gt}^s\right)\\ \text{s.t.}\{\begin{array}{l}{u}_0^s=f^s\\ u_t^s=u_{t-1}^s-\left(\stackrel{N_k}{\sum _{i=1}}{\bar{k}}_i^t\ast {\varphi }_i^t(k_i^t\ast u_{t-1}^s)+{\lambda }^t(u_{t-1}^s-f^s)\right)\\ t=1\cdots T,\end{array}\end{array}$$

$$\frac{\partial \ell (u_T,u_{gt})}{\partial {\Theta }_t}=\frac{\partial u_t}{\partial {\Theta }_t}\cdot \frac{\partial u_{t+1}}{\partial u_t}\cdots \frac{\partial \ell (u_T,u_{gt})}{\partial u_T}.$$

$$v_{ka}(n)=\stackrel{L}{\sum _{j=1}}{a}_j\cdot v_k(q_{nj}).$$

$$v_{ka}=W_a\cdot v_k,$$

$$W_a=\stackrel{L}{\sum _{j=1}}{a}_j{V}_j.$$

$$\text{E}(u|f)=\frac{\lambda }{2}{\left|\left|u-f\right|\right|}^2+\stackrel{N_k}{\sum _{i=1}}{\rho }_i(W_i{K}_{i}u),$$

$$k_i\ast u\iff K_{i}u.$$

$$\frac{\partial \text{E}}{\partial u}=\stackrel{N_k}{\sum _{i=1}}{\bar{K}}_i{W_i}^{\top }{\varphi }_i(W_i{K}_{i}u)+\lambda (u-f),$$

$$\{\begin{array}{l}{u}_0=f,t=1,\cdots ,T\\ u_t=u_{t-1}-\left(\stackrel{N_k}{\sum _{i=1}}{\bar{K}}_i^t{W_i^t}^{\top }{\varphi }_i^t(W_i^t{K}_i^t{u}_{t-1})+{\lambda }^t(u_{t-1}-f)\right).\end{array}$$

$$k_i=\mathcal{B}\frac{c_i}{\Vert c_i{\Vert }_2},$$

$$a_i=\frac{b_i}{\Vert b_i{\Vert }_2},$$

$${\varphi }_i(z)=\stackrel{M}{\sum _{j=1}}{\alpha }_{ij}\phi (\frac{\left|z-{\mu }_j\right|}{\gamma }),$$

$$\phi \left(\frac{|z-\mu |}{\gamma }\right)=\text{exp}\left(-\frac{(z-\mu {)}^2}{2{\gamma }^2}\right)$$

$$\{\begin{array}{l}{\Theta }^{\ast }={\text{argmin}}_{\Theta }\mathcal{L}(\Theta )=\sum _{s=1}^S\ell \left(u_T^s,u_{gt}^s\right)\\ \text{s.t.}\{\begin{array}{l}{u}_0^s=f^s\\ u_t^s=u_{t-1}^s-\left(\stackrel{N_k}{\sum _{i=1}}{\bar{K}}_i^t{W_i^t}^{\top }{\varphi }_i^t(W_i^t{K}_i^t{u}_{t-1}^s)+{\lambda }^t(u_{t-1}^s-f^s)\right)\\ t=1\cdots T.\end{array}\end{array}$$

$$\frac{\partial \ell (u_T,u_{gt})}{\partial {\Theta }_t}=\frac{\partial u_t}{\partial {\Theta }_t}\cdot \frac{\partial u_{t+1}}{\partial u_t}\cdot \cdot \cdot \frac{\partial \ell (u_T,u_{gt})}{\partial u_T}.$$

$$\ell (u_T,u_{gt})=\frac{1}{2}{\left|\left|u_T-u_{gt}\right|\right|}_2^2,$$

$$\frac{\partial \ell (u_T,u_{gt})}{\partial u_T}=u_T-u_{gt}.$$

$$\frac{\partial u_{t+1}}{\partial u_t}=(1-{\lambda }^{t+1})I-\stackrel{N_k}{\sum _{i=1}}{K_i^{t+1}}^{\top }{W_i^{t+1}}^{\top }{\Lambda }_i{W}_i^{t+1}({\bar{K}}_i^{t+1}{)}^{\top },$$

$$\frac{\partial \ell (u_T,u_{gt})}{\partial u_t}=e.$$

$$\frac{\partial u_t}{\partial {\lambda }^t}=-(u_{t-1}-f{)}^{\top }.$$

$$\frac{\partial \ell }{\partial {\lambda }^t}=-(u_{t-1}-f{)}^{\top }e.$$

$$\frac{\partial u_t}{\partial k_i^t}=-\left({\Xi }_{inv}^{\top }{Y}^{\top }+U_{t-1}^{\top }{W_i^t}^{\top }{\Lambda }_i{W}_i^t({\bar{K}}_i^t{)}^{\top }\right),$$

$$k_i\ast u_{t-1}\iff U_{t-1}{k}_i.$$

$${\Xi }_{inv}^{\top }k⟺rot90(rot90(k)).$$

$$\frac{\partial \ell }{\partial k_i^t}=-\left({\Xi }_{inv}^{\top }{Y}^{\top }+U_{t-1}^{\top }{W_i^t}^{\top }{\Lambda }_i{W}_i^t({\bar{K}}_i^t{)}^{\top }\right)e.$$

$$\frac{\partial k_i^t}{\partial c_i^t}=\frac{1}{{\left|\left|c_i^t\right|\right|}_2}\left(I-\frac{c_i^t}{{\left|\left|c_i^t\right|\right|}_2}\cdot \frac{(c_i^t{)}^{\top }}{{\left|\left|c_i^t\right|\right|}_2}\right)\cdot {\mathcal{B}}^{\top }.$$

Full text

I Introduction

Image denoising is one of the most fundamental processing in image processing and low-level computer vision. While it has been extensively studied, image denoising is still an active topic in image processing and computer vision. The goal of image denoising is to recover the clean image $u$ from its noisy observation $f$, which is formulated as

[TABLE]

where $v$ is the noise. In this paper, we assume $v$ is the additive Gaussian noise with zero mean and standard derivation $\sigma$.

During the past decades, a large number of new image denoising methods are continuously emerging. It is a difficult task to precisely categorize existing image denoising approaches. Generally speaking, most image denoising approaches can be categorized as spatial domain and transform domain based methods. Transform domain based methods first represent an image with certain orthonormal transform, such as wavelets [1], curvelets [2], contourlets [3], or bandelets [4], and then attempt to separate noise from the clean image by manipulating the coefficients according to the statistical characteristics of the clean image and noise.

Spatial domain based approaches attempt to utilize the correlations between adjacent pixels in an image. Depending on the way how to select those adjacent pixels, spatial domain based methods can be categorized as local and non-local methods. In local methods, only those adjacent pixels in a spatial neighborhood (probably with fixed shape and size) of the test pixel are investigated. Pixels in this small spatial range are named as an image patch, and the clean pixel value is estimated from this local patch. A large number of local algorithms have been proposed, including filtering based methods [5, 6, 7, 8], anisotropic diffusion based methods [9, 10, 11], variational methods with various image regularizers [12, 13, 14], and patch-based models via sparse representation [15, 16].

Local methods concentrate on the modeling of the patch itself. Nowadays, it is widely-known that another type of image prior is very effective for image denoising - nonlocal self-similarity (NSS) prior; that is, in natural images there are often many similar patches (i.e., nonlocal neighbors) to a given patch, which may be spatially far from it. Inspired by the seminal work of nonlocal means [17], the NSS prior has been widely exploited for image denoising in various framework, such as K-SVD algorithm with nonlocal modeling [18], nuclear norm minimization with nonlocal modeling [19], and Markov Random Fields with nonlocal modeling [20]. Usually, NSS prior based models can significantly improve their corresponding local versions. As a consequence, many state-of-the-art image denoising algorithms are built on the NSS prior, such as BM3D [21], LSSC [18], NCSR [22], and WNNM [19].

Usually, local methods cannot perform very well when the noise level is high, because the correlations between neighboring pixels have been corrupted by the severe noise. Therefore, it is generally believed that local models are not expected to compete with those nonlocal models, especially those state-of-the-art ones, in terms of restoration quality. However, with the help of techniques from machine learning, a few local models, such as opt-MRF [23], Cascade Shrinkage Fields (CSF) [24], and recently proposed Trainable Non-linear Reaction Diffusion (TNRD) [25], succeed achieving state-of-the-art denoising performance via appropriate modeling and supervised learning. It is noticeable that the TNRD model has demonstrated strongly competitive, even better performance against the best-reported nonlocal algorithm - WNNM, meanwhile with much higher computational efficiency.

As mentioned earlier, incorporating the NSS prior has succeeded to boost many image denoising algorithms. Therefore, we are highly motivated to introduce the NSS prior to the best-performing diffusion framework - TNRD to investigate whether it can also boost the TNRD model as usual.

I-A Our contributions

The goal of this paper is to embed the NSS prior into the TNRD model for the task of image denoising. To this end, we propose trainable non-local reaction diffusion (TNLRD) models. The contributions of this study are four-fold:

• a)

We propose a compact matrix form to exploit the NSS prior, which can facilitate the subsequent formulations and derivations associated with the nonlocal modeling. In this work, the NSS prior is defined by a set of non-local filters. In a TNLRD model, the filter responses of $L$ similar patches generated by a local spatial filter are further filtered by its corresponding non-local filter.

• b)

We construct the nonlocal diffusion process with fixed $T$ iterations, which is parameterized by iteration-varying local spatial filters, non-local filters and nonlinear influence functions. Deriving the gradients of the training loss function w.r.t those learning parameters is not trivial, due to the involved nonlocal structure. We provide detailed derivations, which greatly differ from the original TNRD model.

• c)

The training phase is accomplished in a loss-specific manner, where a loss function measuring the difference between clean image $u_{gt}$ and denoised image $u_{T}$ is utilized to optimize the model parameters. In this study, we investigate two different loss functions, namely PSNR-oriented quadratic loss and SSIM related loss.

• d)

We conduct comprehensive experiments to demonstrate the denoising performance of the proposed TNLRD models. As illustrated in Section IV, the proposed TNLRD models outperform recent state-of-the-art methods in terms of PSNR and SSIM.

The following section are organized as follows. In Section 2, we give a brief review of the related works. In Section 3, we introduce the proposed TNLRD models and the training issue. In Section 4, we discuss the influence of the parameters in the proposed TNLRD models, then show the denoising comparison with the previous state-of-the-arts. Finally in Section 5, we draw the conclusion.

II Background and Related work

In this section, we first give a brief review of the TNRD model for image denoising, then introduce the NSS scheme.

II-A Trainable non-linear reaction diffusion model

Chen et. al [25] proposed a simple but effective framework for image restoration - TNRD, which is derived from the following energy functional

[TABLE]

where the regularization term $\overset{N_{k}}{\underset{i=1}{\sum}}\rho_{i}(k_{i}\ast u)$ is a high-order MRFs - Fields of Experts (FoE) [26], defined by a set of linear filters $k_{i}$ and the penalty function $\rho_{i}$. $N_{k}$ is the number of filters, $*$ denotes the 2D convolution operator. $\lambda$ is the strength of data term.

With appropriate modeling of the regularization term, minimizing the energy functional (2) can lead to a denoised image. The steepest-descent procedure for minimizing the energy (2) reads as

[TABLE]

where convolution kernel $\bar{k}_{i}$ is obtained by rotating the kernel $k_{i}$ 180 degrees, $\phi_{i}(\cdot)=\rho_{i}{{}^{\prime}}(\cdot)$ is the influence function [27] or flux function [10], $\Delta t$ donates the time step.

The TNRD model truncates the gradient descent procedure (3) to $T$ iterations, and it then naturally leads to a multi-layer diffusion network with $T$ layers. This modification introduces additional flexibility to the diffusion process, as it becomes easier to train the influence function in this framework. Moreover, as it can be considered as a multi-layer network, we can exploit layer-varying parameters. Therefore, the TNRD model is given as the following diffusion network with $T$ layers.

[TABLE]

Note that the parameters $\{k_{i}^{t},\phi_{i}^{t},\lambda^{t}\}$ vary across the layers. $f$ is the input of the diffusion network. It is clear that each layer of (4) only involves a few image convolution operations, and therefore, it bears an interesting link to the convolutional networks (CN) employed for image restoration problems, such as [28].

The parameters of TNRD models in (4) is trained in a supervised manner. Given the pairs of noisy image $f$ and its ground-truth $u_{gt}$, the parameters $\Theta^{t}=\{k_{i}^{t},\phi_{i}^{t},\lambda^{t}\}$ are optimized by minimizing certain loss function $\ell(u_{T},u_{gt})$, where $u_{T}$ is given by the inference procedure (4). The training procedure is formulated as

[TABLE]

where $\Theta=\{\Theta^{t}\}_{t=0}^{t=T-1}$. The training problem can be solved via gradient based algorithms, e.g., commonly used L-BFGS algorithm [29]. The gradients of the loss function with respect to $\Theta_{t}$ are computed using the standard back-propagation technique widely used in the neural networks learning [30]. There are two training strategies to learn the diffusion processes: 1) the greedy training strategy to learn the diffusion process stage-by-stage; and 2) the joint training strategy to joint train all the stages simultaneously. Generally speaking, the joint training strategy performs better [25], and the greedy training strategy is often employed to provide a good initialization for the joint training. For simplicity, we just consider the joint training scheme to train a diffusion process by simultaneously tuning the parameters in all stages. The associated gradient $\frac{\partial\ell(u_{T},u_{gt})}{\partial\Theta_{t}}$ is presented as follows,

[TABLE]

II-B Non-local self-similarity scheme

Based on the observation that one can always find a few similar patches to a reference patch in the same image, which might be significantly apart from the reference patch, an image prior named non-local self-similarity (NSS) was introduced in [17]. As described in Fig. 2, similar patches to a reference patch can be found in a significantly larger spatial range than the patch size. The non-local similar patches can be collected by using a k-nearest neighbor (k-NN) algorithm [20] or using a kernel function to map the patch distance to coefficients [31].

The NSS prior has proven highly effective for many image restoration problems, and it becomes greatly popular nowadays. A lot of state-of-the-art image restoration algorithms exploit this type of image prior, such as image denoising algorithms BM3D [21] and WNNM [19], image interpolation approaches NARM [31] and ANSM [32]. As a consequence, many local models also attempt to incorporate the NSS prior to boost the performance of the local versions, such as the LSSC method [18], which is a nonlocal extension of the K-SVD algorithm.

We want to especially emphasize the a NSS prior induced method - the NLR-MRF model proposed in [20], which extends the spatial range of the original FoE model [26], as it is highly related to our work. As described in [20], in the NLR-MRF model, several similar patches are firstly collected for each reference patch, and then the responses of these similar patches to a local filter are filtered by a cross-patch filter, generating more sparse responses compared with the local filter responses. With the extend spatial range, NLR-MRF models surpassed the original FoE models in both quality and quantity performance.

Our NSS prior extended TNRD model is also derived from a FoE prior based model. Compared with the TNLRD model to exploit in this paper, the NLR-MRF model is much more constrained in two aspects:

• a)

It employs unchanged parameters for each iteration. However, our NLRD model makes use of iteration-varying parameters.

• b)

Although the penalty functions in the NLR-MRF model are adjustable, they are functions of fixed shape (heavy-tailed functions with a single minimum at the point zero), such as Gaussian Scale Mixtures (GSM) or Student-t distribution. In the TNLRD model, the influence functions are parameterized via radial basis functions, which is able to generate functions of arbitrary shapes. As demonstrated in [25], those seemingly unconventional influence functions found by the training phase play a key role for the success of the TNRD model.

III Trainable non-local reaction diffusion models for image denoising

In this section, we first describe the non-local filter, then introduce the trainable non-local reaction diffusion for image denoising, coined as TNLRD. Finally we give the gradient derivation in the training issue.

III-A Compact matrix form to model the NSS prior

In this work, we make use of k-NN algorithm to collect a fixed number of similar patches. Similar patches are collected by block matching with mean squared error as patch similarity metrics in a large searching window. For the sake of computational efficiency, the size of searching window is set to be several times larger than that of the local spatial filters, as that in [21, 20, 19]. For each possible patch in an image $u$ of size $N\times M$ ($p=N\times M$, and $u$ is represented as a column vector $u\in\mathcal{R}^{p}$), we collect $L$ similar patches (including the reference patch itself) via block matching. Therefore, after running block matching, we can obtain results summarized in Table I.

In Table I, the numbers in each column indicate the indexes of the found similar patches to the corresponding reference patch. For example, in the column $n$, the numbers $\{q_{n2},q_{n3},\cdots,q_{nL}\}$ indicate the indexes of $L-1$ similar patches to the reference patch $n$, and the similar patches are sorted according to the distance to the reference patch, i.e., $d(P_{q_{n2}},P_{n})\leq d(P_{q_{n3}},P_{n})\leq\cdots\leq d(P_{q_{nL}},P_{n})$, where $P_{q_{n2}}$ denotes an image patch centered at the point $q_{n2}$, and function $d$ is a distance measurement of two image patches.

Based on the results in Table I, we construct $L$ highly sparse matrices of size $p\times p$, namely, $\{V_{1},V_{2},\cdots,V_{j},\cdots,V_{L}\}\in\mathcal{R}^{p\times p}$. $V_{j}$ only involves the information from the $j^{th}$ row of Table I. Each row of $V_{j}$ contains merely a non-zero number (exactly one), and its position is given by one of the indexes $Q_{j}$. For example, in the $n^{th}$ row of matrix $V_{j}$, only the element at position $q_{nj}$ is one, and the remaining elements are all zeros. It is easy to see that matrix $V_{1}$ is the identity matrix, i.e., $V_{1}=I\in\mathcal{R}^{p\times p}$. As shown later, the NSS prior can be easily embedded into the TNRD framework with the help of matrices $V_{j}$.

In our work, we introduce a set of non-local filters to embed NSS priors into the TNRD model. A non-local filter is represented as a vector with $L$ elements, e.g., $a=\{a_{1},a_{2},\cdots,a_{L}\}$, whose $j^{th}$ value $a_{j}$ is assigned to the $j^{th}$ similar patch. In the TNLRD model, the filter response map $v_{k}$ generated by a spatial filter $k$ (i.e., $v_{k}=u*k$) is further filtered by a non-local filter $a$, resulting a response map $v_{ka}$, then for the reference patch $n$, its non-local filter response is given as

[TABLE]

It turns out that the above formulation can be given in a more compact way, which reads as

[TABLE]

where the matrix $W_{a}$ is defined by $V_{j}$ and the non-local filter $a$, given as

[TABLE]

In the following subsections, we will see that formulating the NSS prior in the way of (7) can significantly simplify the corresponding formulations, thus easier to understand and to follow, when compared to the formulations in [20]. In addition, the non-local filter $W_{a}$ in matrix form is also highly sparse, as each row of $W_{a}$ only has $L$ non-zero elements. As a result, $W_{a}$ can be efficiently stored via sparse matrix.

III-B Trainable non-local reaction diffusion

Following the formulation in the previous subsection to exploit the NSS prior, it is easy embed the NSS prior into the TNRD framework, and then we arrive at our proposed trainable non-local reaction diffusion.

In order to explain our proposed TNLRD model more clearly, we start from the following energy functional, which incorporates the NSS prior in a natural way

[TABLE]

where $u,f\in\mathcal{R}^{p}$ is the latent image and the noisy image respectively, and $p$ is the number of pixels in image. The local convolution kernel $k_{i}$ in (2) is represented as its corresponding matrix form $K_{i}\in\mathcal{R}^{p\times p}$ that is a highly sparse matrix, such that

[TABLE]

$W_{i}$ is highly sparse matrix defined as in (7) to model the NSS prior, which is related to a non-local filter.

We follow the basic idea of TNRD that unfolds the gradient descent process as a multi-layer network model with layer-varying parameters, to derive the proposed TNLRD model. It is easy to firstly check that the gradient of the energy functional (8) with respect to $u$ is given as

[TABLE]

where function $\phi_{i}$ is given as $\phi_{i}=\rho_{i}^{\prime}$, matrix $\bar{K}_{i}$ is the matrix form related to the linear kernel $\bar{k}_{i}$, which is obtained by rotating kernel $k_{i}$ 180 degrees 111 It should be noticed that the exact formulation for the first matrix in (9) should be $K_{i}^{\top}$. We make use of $\bar{K}_{i}$ to simplify the model complexity. More details can be found in [25].

Therefore, our proposed non-local diffusion model is given as the following multi-layer network with layer-varying parameters

[TABLE]

Note that the parameters in layer $t$ include local filters $k_{i}$, non-local filters $a_{i}$ (i.e., matrix $W_{i}$), nonlinear functions $\phi_{i}$ and the trade-off parameter $\lambda$. The parameter set in layer $t$ is given as $\Theta_{t}=\{\lambda^{t},\phi_{i}^{t},K_{i}^{t},W_{i}^{t}\}$, where $i=1,...,N_{k}$ and $t=1,...,T$.

According to the diffusion process (10), one can see that the TNRD model can be treated as a special case of the TNLRD models with $L=1$, as the corresponding non-local diffusion model clearly degenerates to the local version given in (4), if we set $W=a_{1}V_{1}=I\in\mathcal{R}^{p\times p}$.

In this work, we parameterize the local filters, non-local filters, nonlinear functions in the following way. Concerning the local filters, we follow the TNRD model, and exploit zero-mean filters of unit norm. This is accomplished by constructing the filter $k_{i}$ in the way of

[TABLE]

where $\|\cdot\|_{2}$ denotes the $\ell_{2}$-norm, and $\mathcal{B}$ is a filter bank. Therefore, it is clear that the filter $k$ is a linear combination of the basis filters in the filter bank. In order to achieve the property of zero-mean, the filter bank $\mathcal{B}$ in this work is chosen as a modified DCT basis, which is obtained by removing the filter with constant entries from the complete DCT filters.

The non-local filters in the TNLRD models are vectors with unit length constraint. Therefore, we construct the non-local filter $a_{i}$ as

[TABLE]

where $b_{i}$ is completely free of any constraint.

Following the work of TNRD, the nonlinear functions $\phi_{i}$ are parameterized via radial basis function (RBFs), i.e., function $\phi_{i}$ is represented as a weighted linear combination of a set of RBFs as follows,

[TABLE]

where $\varphi(\cdot)$ here is Gaussian RBFs with equidistant centers $\mu_{j}$ and unified scaling $\gamma$. The Gaussian radial basis is defined as

[TABLE]

As described above, the proposed TNLRD model contains plenty of free parameters, which can be learned from training samples. In this work, we train the TNLRD model parameters in a loss-based learning manner. Given the pairs of degraded image $f$ and the ground-truth original image $u_{gt}$, the parameters are optimized by minimizing certain loss function $\ell(u_{T},u_{gt})$, which is defined to measure the difference between the output $u_{T}$, given by the inference procedure (10) and the desired output, i.e., the ground-truth $u_{gt}$.

In summary, the training procedure is formulated as

[TABLE]

where the parameters $\Theta=\{\Theta_{t}\}_{t=1}^{T}$. Note that we do not specify the form of the loss function in the training phase at present. The basic requirement for the loss function is that it should be differentiable. In our study, we consider two different loss functions for training, see Section IV-E.

III-C Gradients in the training phase for the TNLRD model

Usually, gradient-based algorithms are exploited to solve the corresponding optimization problem (14) in the training phase. Therefore, it is important to compute the gradients of cost function with respect to model parameters $\Theta$ for TNLRD.

The gradient of loss function $\ell(u_{T},u_{gt})$ with respect to parameters in the layer $t$, i.e., $\Theta_{t}$, is computed using back-propagation technique widely used in neural networks learning [30],

[TABLE]

In the case of quadratic loss function, i.e.,

[TABLE]

$\frac{\partial\ell(u_{T},u_{gt})}{\partial u_{T}}$ is directly derived from (16),

[TABLE]

$\frac{\partial u_{t+1}}{\partial u_{t}}$ is computed from (10),

[TABLE]

where matrix $\Lambda_{i}$ is a diagonal matrix given as $\Lambda_{i}=\text{diag}({\phi_{i}^{t+1}}^{\prime}(z_{1}),\cdots,{\phi_{i}^{t+1}}^{\prime}(z_{p}))$, $z=W_{i}^{t+1}K_{i}^{t+1}u_{t}$.

Combining the result of (18) and (15), we arrive at the result of $\frac{\partial\ell(u_{T},u_{gt})}{\partial u_{t}}$, denoted by

[TABLE]

Now, we focus on the computation of the gradients $\frac{\partial u_{t}}{\partial\Theta_{t}}$, which are derived from the diffusion procedure (10).

III-C1 Computing $\frac{\partial\ell}{\partial\lambda^{t}}$

$\frac{\partial u_{t}}{\partial\lambda^{t}}$ is computed as

[TABLE]

Therefore, $\frac{\partial\ell}{\partial\lambda_{t}}$ is given as

[TABLE]

III-C2 Computing $\frac{\partial\ell}{\partial c_{i}^{t}}$

Firstly, $\frac{\partial u_{t}}{\partial k_{i}^{t}}$ is computed as

[TABLE]

where matrix $\Lambda_{i}$ is a diagonal matrix given as $\Lambda_{i}=\text{diag}({\phi_{i}^{t}}^{\prime}(z_{1}),\cdots,{\phi_{i}^{t}}^{\prime}(z_{p}))$, $z=W_{i}^{t}K_{i}^{t}u_{t-1}$. Matrix $U_{t-1}$ and $Y$ are constructed from the images $u_{t-1}$ and $y$ of 2D form, respectively. For example, $U_{t-1}$ is constructed in the way that its rows are vectorized local patch extracted from image $u_{t-1}$ for each pixel, such that

[TABLE]

The matrix $Y$ is defined in the same way, and the image $y$ is given as $y={W_{i}^{t}}^{\top}\phi_{i}^{t}(W_{i}^{t}K_{i}^{t}u_{t-1})$. Matrix $\Xi_{inv}^{\top}$ is a linear operator which inverts the vectorized kernel $k$. In the case of a square kernel $k$, it is equivalent to the Matlab command

[TABLE]

As a consequence, $\frac{\partial\ell}{\partial k_{i}^{t}}$ is given as

[TABLE]

As the filter $k_{i}$ is parameterized by coefficients of $c_{i}$, we need to additionally calculate $\frac{\partial k_{i}^{t}}{\partial c_{i}^{t}}$, which is computed from (11),

[TABLE]

Combining the results of (23) and (22), we can obtain the required gradients $\frac{\partial\ell}{\partial c_{i}^{t}}$.

III-C3 Computing $\frac{\partial\ell}{\partial\alpha_{i}^{t}}$

$\frac{\partial u_{t}}{\partial\alpha_{i}^{t}}$ is computed from (10) and (13),

[TABLE]

where $G_{rc}=\varphi(\frac{\left|z_{r}-\mu_{c}\right|}{\gamma})$, $G\in\mathcal{R}^{p\times M}$, $r=1,\cdots,p$ and $c=1,\cdots,M$, $z=W_{i}^{t}K_{i}^{t}u_{t-1}$.

Therefore, $\frac{\partial\ell}{\partial\alpha_{i}^{t}}$ is given as

[TABLE]

III-C4 Computing $\frac{\partial\ell}{\partial b_{i}^{t}}$

Firstly, $w_{i}^{t}$ and $\bar{w}_{i}^{t}$ are defined as the vectorized form of matrix $W_{i}^{t}$ and ${W_{i}^{t}}^{\top}$ respectively, holding that $w_{i}^{t}=\mathbf{vec}(W_{i}^{t})\in\mathcal{R}^{p^{2}}$ and $\bar{w}_{i}^{t}=\mathbf{vec}({W_{i}^{t}}^{\top})\in\mathcal{R}^{p^{2}}$. The relation between $w_{i}^{t}$ and $\bar{w}_{i}^{t}$ reads as

[TABLE]

where matrix $P$ is a rearrange matrix.

$\frac{\partial u_{t}}{\partial w_{i}^{t}}$ is computed from (10),

[TABLE]

where $y={W_{i}^{t}}^{\top}h$, $h=\phi_{i}^{t}(z)$, $z=W_{i}^{t}\hat{u}$, and $\hat{u}=K_{i}^{t}u_{t-1}$. In the computation of $\frac{\partial y}{\partial w_{i}^{t}}$, the following relations are useful, namely

[TABLE]

and

[TABLE]

where matrix $H$ and $\hat{U}_{t-1}$ are highly sparse, given as

[TABLE]

and

[TABLE]

respectively.

Given that $y={W_{i}^{t}}^{\top}h=H\bar{w}_{i}^{t}$, $\frac{\partial y}{\partial w_{i}^{t}}$ is computed as

[TABLE]

where $\frac{\partial h}{\partial w_{i}^{t}}$ is given as

[TABLE]

where $\Lambda_{i}=\text{diag}({\phi_{i}^{t}}^{\prime}(z_{1}),\cdots,{\phi_{i}^{t}}^{\prime}(z_{p}))$.

Combining these derivation, $\frac{\partial u_{t}}{\partial w_{i}^{t}}$ is computed as

[TABLE]

As $W_{i}^{t}$ is computed from (7), $\frac{\partial w_{i}^{t}}{\partial a_{i}^{t}}$ is given as,

[TABLE]

where $v_{j}=\mathbf{vec}(V_{j})\in\mathcal{R}^{p^{2}},j=1,2,\cdots,L$. Then, we can obtain $\frac{\partial\ell}{\partial a_{i}^{t}}$ from (26) and (27), given as

[TABLE]

As the non-local filter $a_{i}$ is parameterized by the coefficients $b_{i}$ as shown in (12), we need to additionally compute $\frac{\partial a_{i}^{t}}{\partial b_{i}^{t}}$

[TABLE]

Combining the gradients in (29) and (28), we can obtain the required gradient $\frac{\partial\ell}{\partial b_{i}^{t}}$.

The direct computation of $\frac{\partial\ell}{\partial a_{i}^{t}}$ is quite time-consuming and memory-inefficient. Benefit from the sparse matrix structure, matrix $H$, $\hat{U}$ and $\frac{\partial w_{i}^{t}}{\partial a_{i}^{t}}$ are not constructed explicitly. Therefore, the computation of $\frac{\partial\ell}{\partial a_{i}^{t}}$ is quite efficient. As mentioned above, matrix $\{V_{j}\}_{j=1}^{T}$ are highly sparse, each row of $V_{j}$ has precisely one non-zero value, and others are zeros. Therefore, the computation of $\frac{\partial\ell}{\partial a_{i}^{t}}$ can be interpreted as picking up values from $\frac{\partial\ell}{\partial w_{i}^{t}}$ indexed by $\frac{\partial w_{i}^{t}}{\partial a_{i}^{t}}$. The computation of $\frac{\partial\ell}{\partial w_{i}^{t}}$ can be further simplified as

[TABLE]

where $\hat{e}=(\bar{K}_{i}^{t})^{\top}e\in\mathcal{R}^{p}$, $h^{\prime}$ is the derivative of $h$ w.r.t. $z$. Considering the sparse structure of $V_{j}$, the indexing of $\frac{\partial\ell}{\partial w_{i}^{t}}$ is actually the indexing of $\hat{u}$ and $\hat{e}$ using $\frac{\partial w_{i}^{t}}{\partial a_{i}^{t}}$ in forms that described in Table I. The computational complexity of $\frac{\partial\ell}{\partial a_{i}^{t}}$ can be greatly reduced.

Implementation will be made publicly available after acceptance.

IV Experimental results

IV-A Training of TNLRD models

Concerning the model complexity 222While TNLRD models with more stages provide better denoising performance, they cost more time in both training and inference phase., the stages of inference $T$ is set to 5; the local filter size $m\times m$ is set to $5\times 5$ and $7\times 7$; the number of non-local similar patches $L$ is set to 3, 5, 7, 9. The size of searching window is $31\times 31$. The size of block matching is $7\times 7$.

We trained the TNLRD models for Gaussian denoising problem with different standard deviation $\sigma$. We minimize (14) to learn the parameters of the TNLRD models with commonly used gradient-based L-BFGS [33]. The gradient of loss function with respect to parameters can be derived from (15) - (29). The training dataset of original and noisy image pairs is constructed over 400 images as [25] [24]. We cropped a $180\times 180$ region from each image, resulting in a total of 400 training images of size $180\times 180$. In the training phase, computing the gradients of one stage for 400 images of size $180\times 180$ takes about 480$s$ on a server with CPUs: Intel(R) Xeon E5-2650 @ 2.00GHz (eight cores). We run 200 L-BFGS iterations for optimization. Therefore, the total training time for ${\text{TNLRD}}_{7\times 7\times 3}^{5}$ model is $5\times(200\times 480)/3600=133.3h$. Implementation will be made publicly available after acceptance.

In order to perform a fair comparison to previous works, i.e., BM3D [21], WNNM [19], NLR-MRF [20] and TNRD [25], we used the 68 test images in [25], which are original introduced by [26] and are widely used in image denoising. We evaluated the denoising performance using PSNR as [25] and SSIM as [34]. SSIM provides a perceptually more plausible image error measure, which has been verified in psychophysical experiments. SSIM values range between 0 and 1, where 1 is a perfect restoration. We also test TNRD and our TNLRD models on a 9 image set which are collecting from web, as shown in Fig. 3. The codes of the comparison methods were downloaded from the authors’s homepage.

IV-B Influence of parameters initialization

The TNLRD models with different parameters configuration are denoted as, ${\text{TNLRD-}method}_{m\times m\times L}^{T}$. The $method$ denotes the parameters initialization method, $tnrd$ for initializing from the TNRD models, and $plain$ for initializing from plain settings. In [20], the author trained NLR-MRF models starting from MRF models with local spatial clique, i.e., FoE models, using NSS setting. We followed the same training scheme as that in [20] for training NLR-MRF. We started from the local TNRD models by setting $a_{i1}^{t}=1$ and $a_{ij}^{t}=0$, $j=2,...,L$, and conducted a joint training for parameters of the $T$ steps inference (10), denoted as ${\text{TNLRD-}tnrd}_{m\times m\times L}^{T}$.

We also trained the parameters of TNLRD models via the greedy training from plain initialization, then jointly trained the $T$ steps inference (10), denoted as ${\text{TNLRD-}plain}_{m\times m\times L}^{T}$. Greedy training means a strategy that greedily trains a multi-layer diffusion network layer by layer. In the plain initialization training, we observed that TNLRD models with joint training surpass models obtained in greedy training by 0.55dB in average. Therefore, it is recommended that joint training should be conducted after greedy training.

We trained TNLRD models using both parameters initialization method, and got two models, namely ${\text{TNLRD-}tnrd}_{7\times 7\times 5}^{5}$ and ${\text{TNLRD-}plain}_{7\times 7\times 5}^{5}$. We evaluated their denoising performance on the 68 test images. Models trained by $tnrd$ and $plain$ initialization achieve almost the same denoising performance, i.e., 29.01dB in average. This conclusion holds for our models with other model capacities. For the sake of training efficiency 333The plain initialization with greedy and joint training is more time consuming than the TNRD initialization which only conducts joint training., in the following experiments, we mainly discuss the models trained via TNRD initialization, which is coined as ${\text{TNLRD}}_{m\times m\times L}^{T}$ omitting the $method$ in ${\text{TNLRD-}method}_{m\times m\times L}^{T}$.

IV-C Influence of number of non-local similar patches

In this subsection, we investigate the influence of different number of non-local similar patches $L$ for both ${\text{TNLRD}}_{5\times 5\times L}^{5}$ and ${\text{TNLRD}}_{7\times 7\times L}^{5}$.

As described above, the TNRD model can be treated as a special case of the TNLRD model with $L=1$. Therefore, in the training phase, the TNLRD model can be initialized from its local version. The denoising performance of the trained TNLRD models with different configurations are illustrated in Fig. 4. One can see that the performance of the trained models is improved when $L$ increases, and is degraded when $L$ continues to increase. A performance peak exists. for ${\text{TNLRD}}_{5\times 5\times L}^{5}$, it is $L=5$; for ${\text{TNLRD}}_{7\times 7\times L}^{5}$, it is $L=7$. While a peak exists, the performance gap is within 0.05dB. For the sake of computational efficiency 444Larger $L$ will take more time for both training phase and test phase., in the rest of this section, we set $L=5$. ${\text{TNLRD}}_{5\times 5\times 5}^{5}$ surpasses ${\text{TNRD}}_{5\times 5}^{5}$ about 0.14dB. ${\text{TNLRD}}_{7\times 7\times 5}^{5}$ surpasses ${\text{TNRD}}_{7\times 7}^{5}$ about 0.10dB.

Fig. 5 shows the trained local and non-local filters of ${\text{TNLRD}}_{7\times 7\times 5}^{5}$ in the first and last inference stage, in the training of Gaussian denoising with $\sigma=25$. In most of the non-local filters, the first element is near 1, while the rest are near zero, for example $[0.999,0.001,-0.018,-0.010,-0.027]^{\top}$; while in some of the non-local filters, the first element is of the same scale with the rest, for example $[0.915,0.185,0.194,0.178,0.243]^{\top}$. The former non-local filters are related with simple local filters, for example the directional derivatives. The later non-local filters are related with complex local filters, hence all the local filter response of the similar patches are useful.

IV-D Influence of filter size

We also investigate the influence of filter size, as shown in Fig. 4. The increasing of the filter size from $5\times 5$ to $7\times 7$ brings an average 0.11dB improvement. In the evaluating of denoising performance, we prefer ${\text{TNLRD}}_{7\times 7\times 5}^{5}$ model as it provides better trade-off between performance and run time.

IV-E Influence of loss function

In [35, 36], the loss function for discriminative training is SSIM instead of L2 for image inpainting and denoising respectively. The trained models with SSIM loss function may provide visually more plausible results. Inspired by these works, we trained our TNLRD models using SSIM loss function as [34]. In the case of $\sigma=25$, the trained TNLRD models via SSIM loss achieves SSIM result of 0.8219, while the corresponding average PSNR is 28.83dB, as shown in Table II. The TNLRD models with the same capacity trained via the L2 loss, achieves a result of SSIM = 0.8201 and PSNR = 29.01dB. As shown in Fig. 6, the TNLRD models trained via SSIM loss offer sharper image than that trained via L2 loss. SSIM loss function benefits the TNLRD models to produce more visually plausible denoising results. From Table III, we note that our TNLRD models trained via SSIM loss achieve competitive performance with WNNM in terms of PSNR, and provide better recovered images in terms of SSIM. We also note that, compared with models trained via SSIM loss, models trained via L2 loss achieve competitive performance in terms of SSIM, and superior performance in terms of PSNR. Bearing these in mind, in the following comparison with other image denoising methods, we prefer the models trained with L2 loss.

IV-F Denoising

The above training experiments are conducted on Gaussian noise level $\sigma=25$. We also trained the proposed TNLRD models for the noise level $\sigma=15$ and $\sigma=50$. After training the models, we evaluated them on the 68 test images used in [25]. We also tested the TNRD models and our TNLRD models on the 9 test image set.

The denoising performance on the 68 test images is summarized in Table III and IV, compared with some recent state-of-the-art denoising algorithms. As illustrated in Table III and Fig. 12, the proposed TNLRD models outperform the TNRD models by almost 0.1dB, BM3D by 0.45dB, WNNM by 0.18dB and NLR-MRF by 0.53dB. In Fig. 6 (i-p), we can see that our TNLRD models recover more clear stems in the sea anemone than the TNRD models. While BM3D and WNNM tend to over-smooth texture regions, our TNLRD models produce sharper recovered image. In Fig. 7 (i-p), we can also see that clear and straight steel structures are recovered by our TNLRD models, while the TNRD models tends to offer the over-smooth results in the texture regions. The same phenomenon can be also found in the recovered image produced by BM3D and WNNM. Taking a close look at the recovered images produced by BM3D, WNNM and TNRD, one can see some artifacts in the plain regions.

We also compared our TNLRD model with these methods for cases of $\sigma=15$ and $\sigma=50$, as shown in Table III. When the image is heavily degraded by the noise, i.e., $\sigma$ is getting larger, the local methods, e.g., the TNRD model, can not collect enough information for inference, and may create artifacts and remove textures. On the contrary, the non-local methods collect more information, and tackle the artifacts and preserve textures. We show some denoising examples with $\sigma=50$ in Fig. 8 and 9.

We also compare our TNLRD and TNRD on the 9 test images collected from web. In Fig. 10 (e-h), we can see that our TNLRD models recover the vertical lines more clear than the TNRD models. In Fig. 11 (e-h), we can see that our TNLRD models recover the window structures more precisely than the TNRD models. In Fig. 13, we can conclude that our TNLRD models surpass the TNRD models for each test image. The average PSNR produced by our TNLRD and TNRD models are 32.46dB and 32.24dB respectively.

From the detailed comparison with some state-of-the-art denoising methods, especially the newly proposed TNRD, we can conclude that our TNLRD models offer better quality and quantity performance in Gaussian denoising.

V Conclusion

In this paper, we propose trainable non-local reaction diffusion models for image denoising. We introduce the NSS prior as non-local filters to the TNRD models. We train the models parameters, i.e., local linear filters, non-local filters and non-linear influence functions, in a loss-based learning scheme. From the comparison with the state-of-the-art image denoising methods, we concluded that our TNLRD models achieve superior image denoising performance in terms of both PSNR and SSIM. Our TNLRD models also provide visually plausible denoised image with less artifacts and more textures.