Trainable Nonlinear Reaction Diffusion: A Flexible Framework for Fast and Effective Image Restoration

Introduction

Image restoration is the process of estimating uncorrupted images from noisy or blurred ones. It is one of the most fundamental operations in image processing, video processing, and low-level computer vision. For several decades, image restoration remains an active research topic and hence new approaches are constantly emerging. There exists a huge amount of literature addressing the topic of image restoration problems, see for example for a survey.

In recent years, the predominant approaches for image restoration are non-local methods based on patch modeling, for example, image denoising with (i) Gaussian noise , (ii) multiplicative noise , or (iii) Poisson noise , image interpolation , image deconvolution , etc. Most state-of-the-art techniques mainly concentrate on achieving utmost image restoration quality, with little consideration on the computational efficiency, e.g., , despite the fact that it is a critical factor for real applications. However, there are a few exceptions. For example, there are two notable exceptions for the task of Gaussian denoising, BM3D and the recently proposed Cascade of Shrinkage Fields (CSF) model, which simultaneously offer high efficiency and high image restoration quality.

It is well-known that BM3D is a highly engineered Gaussian image denoising algorithm. It involves a block matching process, which is challenging for parallel computation on GPUs, alluding to the fact that it is not straightforward to accelerate BM3D algorithm on parallel architectures. In contrast, the recently proposed CSF model offers high levels of parallelism, making it well suited for GPU implementation, thus owning high computational efficiency.

In this paper, we propose a flexible learning framework to generate fast and effective models for a variety of image restoration problems. Our approach is based on learning optimal nonlinear reaction diffusion models. The learned models preserve the structural simplicity of these models and hence it is straightforward to implement the corresponding algorithms on massive parallel hardware such as GPUs.

Partial differential equation (PDEs) have become a standard approach for various problems in image processing. On the one hand they come along with a sound mathematical framework that allow to make clear statements about the existence and regularity of the solutions. On the other hand, efficient numerical algorithms have been developed, that allow to compute the solution of PDEs in very short time . Although recent PDE approaches have shown good performance for a number of image processing task, they still fail to produce state-of-the-art quality for classical image restoration tasks.

In the seminal work , Perona and Malik (PM) proposed a nonlinear diffusion model, which is given as the following PDE

where $\nabla$ is the gradient operator, $t$ denotes the time, $f$ is a initial image. The function $g$ is known as edge-stopping function or diffusivity function , and a typical $g$-function is given by $g(z)={1}/{(1+z^{2})}$. The proposed PM diffusion model (1) leads to a nonlinear anisotropic Anisotropic diffusion in this paper is understood in the sense that the smoothing induced by PDEs can be favored in some directions and prevented in others. The diffusivity is not necessary to be a tensor. It should be noted that this definition is different from Weickert’s terminology , where anisotropic diffusion always involves a diffusion tensor, and the PM model is regarded as an isotropic model. diffusion model which is able to preserve and enhance image edges. Hence, it is well suited for a number of image processing tasks such as image denoising and segmentation.

A first variant of the PM model is the so-called biased anisotropic diffusion (also known as reaction diffusion) proposed by Nordström , which introduces a bias term (forcing term) to free the user from the difficulty of specifying an appropriate stopping time for the PM diffusion process. This additional term reacts against the strict smoothing effect of the pure PM diffusion, therefore resulting in a nontrivial steady-state.

Subsequent works consider modifications of the diffusion or the reaction term for the reaction diffusion model , e.g., Acton et al. exploited a more complicated reaction term to enhance oriented textures; proposed to replace the ordinary diffusion term with a flow equation based on mean curvature. A notable work is the forward-backward diffusion model proposed by Gilboa et al. , which incorporates explicit inverse diffusion with negative diffusivity coefficient by carefully choosing the diffusivity function. The resulting diffusion processes can adaptively switch between forward and backward diffusion processes. In subsequent work , the theoretical framework for discrete forward-and-backward diffusion filtering has been investigated. Researchers also propose to exploit higher-order nonlinear diffusion filtering, which involves larger linear filters, e.g., fourth-order diffusion models . Meanwhile, theoretical properties about the stability and local feature enhancement of higher-order nonlinear diffusion filtering are established in .

It should be noted that the above mentioned diffusion models are handcrafted, including elaborate selections of the diffusivity functions, optimal stopping times and proper reaction forces. It is a generally difficult task to design a good-performing PDE for a specific image processing problem, as good insights into this problem and a deep understanding of the behavior of the PDEs are usually required. Therefore, an attempt to learn PDEs from training data via an optimal control approach was made in , where the PDEs to be trained have the form of

Coefficients $a(t)$ are free parameters (i.e., combination weights) to train. $\kappa(u)$ is related to the TV regularization and $\mathcal{O}(u)$ denotes a set of operators (invariants) over $u$, e.g., $\|\nabla u\|_{2}^{2}=u_{x}^{2}+u_{y}^{2}$.

1.2 Improvements from the side of image statistics/regularization

As shown in previous works, e.g., , the diffusion step (3) corresponds to a gradient descent step to minimize the energy functional given as

the functions $\rho$ (e.g., $\rho(z)=\text{log}(1+z^{2})$) is the so-called penalty function.Note that $\rho^{\prime}(z)=\phi(z)$. It is worthwhile to mention that the matrix-vector product $\nabla_{x}u$ can be interpreted as a 2D convolution of $u$ with the linear filter $k_{x}=$ ($\nabla_{y}$ corresponds to the linear filter $k_{y}=^{\top}$). The energy functional (4) can be also understood from the aspects of image statistics, image prior and image regularization. As a consequence, a lot of efforts listed below have been made to improve the capability of model (4).

More filters of larger kernel size were considered in , instead of relatively small kernel size, such as usually used pair-wise kernels. The resulting regularization model leads to the so-called fields of experts (FoE) image prior, which works well for many image restoration problems.

Instead of hand-crafted ones with fixed shape, more flexible penalty functions were exploited, and they were learned from data . Especially, as shown in that those unusual penalties such as inverted penalties (i.e., $\rho(z)$ decreasing as a function of $|z|$) were found to be necessary.

In order accelerate the inference phase related to the model (4), in , it was proposed to truncate the gradient descent procedure to fixed iterations/stages, and then train this truncated optimization algorithm based on the FoE prior model.

In order to further increase the flexibility of multi-stage models, Schmidt and Roth considered varying parameters per stage.

2 Motivations and Contributions

In this paper we concentrate on nonlinear diffusion process due to its high efficiency. Taking into consideration the improvements mentioned in Sec. 1.1.2, we propose a trainable nonlinear diffusion model with (1) fixed iterations (also referred to as stages), (2) more filters of larger kernel size, (3) flexible penalties in arbitrary shapes, (4) varying parameters for each iteration, i.e., time varying linear filters and penalties. Then all the parameters (i.e., linear filters and penalties) in the proposed model are simultaneously trained from training data in a supervised way. The proposed approach results in a novel learning framework to train effective image diffusion models. It turns out that the trained diffusion processes leads to state-of-the-art performance, while preserve the property of high efficiency of diffusion based approaches. In summary, our proposed nonlinear diffusion process offers the following advantages:

It is conceptually simple as it is merely a standard nonlinear diffusion model with trained filters and influence functions;

It has broad applicability to a variety of image restoration problems. In principle, all the diffusion based models can be revisited with appropriate training;

It yields excellent results for several tasks in image restoration, including Gaussian image denoising,single image super resolution and JPEG deblocking;

It is highly computationally efficient, and well suited for parallel computation on GPUs.

A shorter paper has been presented as a conference version . In this paper, we incorporate additional contents listed as follows

We investigate more details of the training phase, such as the influence of (a) initialization, (b) the model capacity and (c) the number of training samples;

We consider more detailed analysis of trained models, such as how the trained models generate the patterns;

We exploit an additional application of single image super resolution to further illustrate the potential breadth of our proposed learning framework.

Proposed reaction diffusion process

In this section, we first describe our learning based reaction diffusion model for image restoration, and then we show the relations between the proposed model and existing image restoration models.

The fundamental idea of our proposed learning based reaction diffusion model is described in Sec. 1.2. In addition, we incorporate a reaction term in order to apply our model for different image processing problems, as shown later. As a consequence, our proposed nonlinear reaction diffusion model is formulated as

Note that our proposed diffusion process is truncated after a few stages, usually less than 10. Moreover, the linear filters and influence functions are adjustable and vary across stages. As our proposed approach is inspired by nonlinear reaction diffusion model but with trainable filters and influence functions, we coin our method Trainable Nonlinear Reaction Diffusion (TNRD).

In practice, we train the proposed nonlinear diffusion model (5) for specific image restoration problem by exploiting application specific reaction terms $\psi(u)$. For classical image restoration problems, such as Gaussian denoising, image deblurring, image super resolution and image inpainting, we can set the reaction term to be the gradient of a data term, i.e. $\psi(u)=\nabla_{u}\mathcal{D}(u)$.

For example, if we choose $\mathcal{D}^{t}(u,f)=\frac{\lambda^{t}}{2}\|Au-f\|_{2}^{2}$, we have $\psi^{t}(u)=\lambda^{t}A^{\top}(Au-f)$, where $f$ is the degraded input image, $A$ is the associated linear operator, and $\lambda^{t}$ is related to the strength of the reaction term. In the case of Gaussian denoising, $A$ is the identity matrix; for image super resolution, $A$ is related to the down sampling operation and for image deconvolution, $A$ corresponds to the linear blur kernel.

2 A more general formulation of the proposed diffusion model

Note that the data term $\mathcal{D}(u)$ related to (5) should be differentiable. In order to handle the problems involving a non-differentiable data term, e.g., the JPEG deblocking problem investigated in Section 6, we consider a more general form of our proposed diffusion model as follows

where $\text{Prox}_{\mathcal{G}^{t}}(\hat{u})$ is the proximal mapping operation related to the non-differentiable function $\mathcal{G}^{t}$, given as

3 Related image restoration models

As mentioned in Sec. 1.1.2, there exist natural connection between anisotropic diffusion and image regularization based energy functional. Therefore, Eq. (6) can be interpreted as a forward-backward step The forward step is a gradient descent step w.r.t the function $\mathcal{D}$ and the backward step is the proximal operation w.r.t the function $\mathcal{G}$. at $u_{t-1}$ for the energy functional given by

where $\mathcal{R}^{t}_{i}(u)=\sum\nolimits_{p=1}^{N}\rho_{i}^{t}((K_{i}^{t}u)_{p})$ are the regularizers. Since the parameters {$K_{i}^{t},\rho_{i}^{t}$} vary across the stages, (7) is a dynamic energy functional, which changes at each iteration.

If {$K_{i}^{t},\rho_{i}^{t}$} keep the same across stages, the functional (7) with $\mathcal{G}=0$ corresponds to the FoE prior regularized variational model for image restoration . In our work, we do not exactly solve this minimization problem anymore. In contrast, we run the gradient descent step for several iterations, and each gradient descent step is optimized by training. More importantly, we are allowed to investigate more generalized penalties.

To the best of our knowledge, Zhu and Mumford were the first to consider learning reaction diffusion models from data. The linear filters appearing in the prior were chosen (not fully trained) from a general filter bank by minimizing the entropy of probability distributions of natural images. The associated penalty functions were learned based on the maximum entropy principle. However, our proposed diffusion model is a multi-stage diffusion process with multiple image priors, where the filters and penalties are fully trained from clean/degraded pairs.

Some more works have also been dedicated to train the penalties in a diffusion model. In , the authors trained the penalty functions in the way that they first computed the image statistics appropriate to certain chosen filters (e.g., horizontal and vertical image derivatives), and then considered mixture models of fixed shape having a peak near zero and two heavy tails in both sides to fit to image statistics. Analogously, in , potential functions of fixed shape are chosen to resemble zero mean filter responses.

In very recent work , Schmidt and Roth exploited an additive form of half-quadratic optimization to solve problem (7). The resulting model shares similarities with classical wavelet shrinkage and hence it is termed “cascade of shrinkage fields” (CSF). The CSF model relies on the assumption that the data term in (7) is quadratic and the operator $A$ can be interpreted as a convolution operation, such that the corresponding subproblem can be solved in closed-form using the discrete Fourier transform (DFT). However, our proposed diffusion model does not have this restriction on the data term. In principle, any smooth data term is appropriate. Moreover, as shown in the following sections, we can even handle the case of non-smooth data terms.

As already mentioned, also proposed to train an optimized gradient descent algorithm for the energy functional similar to (7). However, their model is much more constrained, since they exploited the same filters for each gradient descent step and more importantly, they make use of a hand-selected influence function. This clearly restricts the model capability, as demonstrated in Sec. 4.

Comparing our model to the model of (cf. (2)), one can see that this approach learns only a linear model based on pre-defined image invariants. Therefore, this model can be interpreted as a simplified version of our model (5) with fixed linear filters and influence functions, and only the weight of each term is optimized.

The proposed diffusion model also bears an interesting link to convolutional networks (CNs) applied to image restoration problems in . One can see that each iteration (stage) of our proposed diffusion process involves convolution operations with a set of linear filters, and thus it can be treated as a convolutional network. The architecture of our proposed diffusion model is shown in Figure 1, where it is represented as a common feed-forward network. We refer to this network in the following as diffusion network.

However, we can introduce a feedback step to explicitly illustrate the special architecture of our diffusion network that we subtract “something” from the input image. Therefore, our diffusion model can be represented in a more compact way in Figure 2, where one can see that the structure of our CN model is different from conventional feed-forward networks. Due to this feedback step, it can be categorized into recurrent networks . It should be noted that the nonlinearity (i.e., influence functions in the context of nonlinear diffusion) in our proposed network are trainable. However, conventional CNs make use of fixed activation function, e.g., the ReLU function or sigmoid functions .

Learning framework

We train our diffusion networks in a supervised manner, namely we firstly prepare the input/output pairs for a certain image processing task, and then exploit a loss minimization scheme to learn the model parameters $\Theta_{t}$ for each stage $t$ of the diffusion process. The training dataset consists of $S$ training samples $\{f^{s},u_{gt}^{s}\}_{s=1}^{S}$, where $f^{s}$ is a noisy observation and $u_{gt}^{s}$ is the corresponding ground truth clean image. The model parameters $\Theta_{t}$ of each stage include the parameters of (1) the reaction force weight $\lambda$, (2) linear filters and (3) influence functions, i.e., $\Theta_{t}=\{\lambda^{t},\phi_{i}^{t},k_{i}^{t}\}$.

In the supervised manner, a training cost function is required to measure the difference between the output of the diffusion network and the ground-truth image. As our goal is to train a diffusion network with $T$ stages, the cost function is formulated as

where $u_{T}$ is the output of the final stage $T$. In our work we exploit the usual quadratic loss function This loss function is related to the PSNR quality measure. Note that as shown in , other quality measures, such as structural similarity (SSIM) and mean absolute error (MAE) can be chosen to define the loss function. At present we only consider the quadratic loss function due to its simplicity., defined as

As a consequence, the training task is generally formulated as the following optimization problem

where $\Theta=\{\Theta^{t}\}_{t=1}^{t=T}$ and $I_{0}$ is the initial status of the diffusion process. Note that the above loss function only depends on the output of the final stage $T$, i.e., the parameters in all stages are simultaneously trained such that the output of the diffusion process - $u_{T}$ is optimized. We call this training scheme joint training similar to . The joint training strategy is a minimization problem with respect to the parameters in all stages $\{\Theta_{1},\Theta_{2},\cdots,\Theta_{T}\}$.

One can see that our training model is also a deep model with many stages (layers). It is well-known that deep models are usually sensitive to initialization, and therefore training from scratch is prone to getting stuck at bad local minima. As a consequence, people usually consider a greedy layer-wise pre-training to provide a good initialization for the joint training (fine tune).

In our work, we also consider a greedy training scheme similar to , to pre-train our diffusion network stage-by-stage, where each stage is greedily trained such that the output of each stage is optimized, i.e., for stage $t$, we minimize the cost function

where $u_{t}^{s}$ is the output of stage $t$ of the diffusion process. Note that this is a minimization problem only with respect to the parameters $\Theta_{t}$ in stage $t$.

In this paper, we aim to investigate arbitrary influence functions. In order to conduct a fast and accurate training, an effective function parameterization method is required. Following the work of , we parameterize the influence function via standard radial basis functions (RBFs), i.e., each function $\phi$ is represented as a weighted linear combination of a family of RBFs as follows

where $\varphi$ represents different RBFs. In this paper, we exploit RBFs with equidistant centers $\mu_{j}$ and unified scaling $\gamma_{j}$. We investigate two typical RBFs : (1) Gaussian radial basis $\varphi_{g}$ and (2) triangular-shaped radial basis $\varphi_{t}$, given as

respectively. The basis functions are shown in Figure 3, together with an example of the function approximation by using two different RBF methods. In our work, we have investigated both function approximation methods, and we find that they lead to similar results. We only present the results obtained by the Gaussian RBF in this paper.

In our work, the linear kernels $k_{i}^{t}$ related to the linear operators $K_{i}^{t}$ are defined as a linear combination of Discrete Cosine Transform (DCT) basis kernels $b_{r}$, i.e.,

where the kernels $k_{i}^{t}$ are normalized to get rid of an ambiguity appearing in the proposed diffusion model. More details can be found in the supplemental material. The kernels are formed in this way in order to keep the expected property of zero-mean.

The weights $\lambda^{t}$ in our model are constrained to be positive. To this end, we set $\lambda^{t}\leftarrow e^{\lambda^{t}}$ in the training phase for our implementation.

3 Computing gradients

For both greedy training and joint training, we make use of gradient-based algorithms (e.g., the L-BFGS algorithm ) for optimization. The key point is to compute the gradients of the loss function with respect to the training parameters. In greedy training, the gradient of the loss function at stage $t$ with respect to the model parameters $\Theta_{t}$ is computed using standard chain rule, given as

In the joint training, we compute the gradients of the loss function with respect to $\Theta_{t}$ by using the standard back-propagation technique widely used in the neural networks learning , namely, $\frac{\partial u_{t}}{\partial\Theta_{t}}$ is computed by using

Compared to the greedy training, we additionally need to calculate $\frac{\partial u_{t+1}}{\partial u_{t}}$. For the investigated image processing problems in this paper, we provide all necessary derivations in the supplemental material.

4 Experimental setup and implementation details

In our convolution based diffusion network, the image size stays the same when an image goes through the network, and we use the symmetric boundary condition for convolution calculation. In our original diffusion model (6), there is matrix transpose $K^{\top}v$, which exactly corresponds to the convolution operation $\bar{k}*v$ ($\bar{k}$ is obtained by rotating the kernel $k$ 180 degrees) in the cases of periodic and zero-padding boundary conditions. It should be noted that in the case of symmetric boundary condition used in this paper, this result holds only in the central image region. However, we still want to explicitly use the formulation $\bar{k}*v$ to replace $K^{\top}v$, because the former can significantly simplify the derivation of the gradients required for training.

We find that the direct replacement introduces some artifacts at the image boundary. In order to avoid these artifacts, we symmetrically pad the input image before it is sent to the diffusion network, and then we discard those padding pixels in the final output. More details are found in the supplemental material.

4.2 RBF kernels

Images exploited in this paper have the dynamic range in $$, and the filters have unit norm. In order to cover most of the filter response, we consider influence functions in the range$$. We use 63 Gaussian RBFs with equidistant centers at $[-310:10:310]$, and set the scaling parameter $\gamma=10$.

4.3 Experimental setup

Model capacity: In our work, we train the proposed diffusion network with at most 8 stages to observe its saturation behavior after certain stages. We first greedily train $T$ stages of our model with specific model capacity, then conduct a joint training to refine the parameters of the whole $T$ stages.

In this paper, we mainly consider four different diffusion networks with increasing capacity:

where $\text{TNRD}_{m\times m}^{T}$ denotes a nonlinear diffusion process of stage $T$ with filters of size $m\times m$. The filters number is $m^{2}-1$, if not specified. For example, $\text{TNRD}_{7\times 7}^{T}$ model contains $(48\times 48~{}(\text{filters})+48\times 63~{}(\text{penalties})+1~{}(\lambda))~{}\cdot T=5329\cdot T$ free parameters.

Training and test dataset: In order to make a fair comparison to previous works, we make use of the same training datasets used in previous works for our training, and then evaluate the trained models on commonly used test datasets. For image processing problems investigated in this paper, i.e., Gaussian denoising, single image super resolution and JPEG deblocking, we consider the following training and test datasets, respectively.

Gaussian denoising. Following , we use the same 400 training images, and crop a $180\times 180$ region from each image, resulting in a total of 400 training samples of size $180\times 180$, i.e., roughly 13 million pixels. We then evaluate the denoising performance of a trained model on a standard test dataset of 68 natural images, which is suggested by , and later widely used for Gaussian denoising testing. Note that the test images are strictly separate from the training datasets.

Single image super resolution. The publicly available framework of Timofte et al. provides a perfect base to compare single image super resolution algorithms. It includes 91 training images and two test datasets Set5 and Set14. Many recent state-of-the-art learning based image super resolution approaches accomplish their comparison based on this framework. Therefore, we also use the same 91 training images. We crop 4-5 sub-images of size $150\times 150$ from each training image, according to its size, and this finally gives us 421 training samples. We then evaluate the performance of the trained models on the Set5 and Set14 dataset.

JPEG deblocking. We train the diffusion models using the same training images as in the case of Gaussian denoising. In the test phase, we follow the test procedure in for performance evaluation. The test images are converted to gray-value, and scaled by a factor of 0.5, resulting 200 images of size $240\times 160$.

4.4 Approximate training time

Note that the calculation of the gradients of the loss function in (13) is very efficient even using a simple Matlab implementation, since it mainly involves 2D convolutions. The training time varies greatly for different configurations. Important factors include (1) model capacity, (2) number of training samples, (3) number of iterations taken by the L-BFGS algorithm, and (4) number of Gaussian RBF kernels used for function approximation. We report the most time consuming cases as follows.

In training, computing the gradients $\frac{\partial\mathcal{L}}{\partial\Theta}$ with respect to the parameters of one stage for 400 images of size $180\times 180$ takes about 35s ($\text{TNRD}_{5\times 5}$), 75s ($\text{TNRD}_{7\times 7}$) or 165s ($\text{TNRD}_{9\times 9}$) using Matlab implementation on a server with CPUs: Intel(R) Xeon E5-2680 @ 2.80GHz (eight parallel threads using parfor in Matlab, 63 Gaussian RBF kernels for the influence function parameterization). We typically run 200 L-BFGS iterations for optimization. Therefore, the total training time, e.g., for the $\text{TNRD}^{5}_{7\times 7}$ model is about $5\times(200\times 75)/3600=20.8h$. Code for learning and inference is available on the authors’ homepage www.GPU4Vision.org http://gpu4vision.icg.tugraz.at/binaries/nonlinear-diffusion.zip#pub94. For the training of the Gaussian denoising task, we have also accomplished a GPU implementation, which is about 4-5 times faster than our CPU implementation.

Training for Gaussian denoising

For the task of Gaussian denoising, we consider the following energy functional

By setting $\mathcal{D}(u)=\frac{\lambda}{2}\|u-f\|^{2}_{2}$ and $\mathcal{G}(u)=0$, we arrive at the following diffusion process with $u_{0}=f$

where we explicitly use a convolution kernel $\bar{k}_{i}$ (obtained by rotating the kernel $k_{i}$ 180 degrees) to replace the $K_{i}^{\top}$ for the sake of model simplicity, but we have to pad the input image. The gradients $\frac{\partial u_{t}}{\partial\Theta_{t}}$ and $\frac{\partial u_{t}}{\partial u_{t-1}}$ required in training are computed from this equation. Detailed derivations are presented in the supplemental material.

We started with the training for $\text{TNRD}_{5\times 5}^{T}$. We first considered the greedy training phase to train a diffusion process up to 8 stages (i.e., $T\leq 8$), in order to observe the asymptotic behavior of the diffusion network. In the greedy training, only parameters in one stage were trained at a time. We exploited a plain initialization to start the training, namely linear filters and influence functions were initialized from the modified DCT filters and the function $\phi(z)=2z/(1+z^{2})$, respectively.

After the greedy training was completed, we conducted joint training for a diffusion model of certain stages (e.g., $T=2,5,8$), to simultaneously tune the parameters in all stages. We initialized the joint training using parameters learned in greedy training, as this is guaranteed not to degrade the training performance.

We first trained our diffusion models for the Gaussian denoising problem with standard deviation $\sigma=25$. The noisy training images were generated by adding synthetic Gaussian noise with $\sigma=25$ to the clean images. Once we obtained a trained model, we evaluated its denoising performance on 68 natural images following the same test protocol as in . Figure 4 shows a denoising example for noise level $\sigma=25$ to illustrate the denoising process of our learned $\text{TNRD}_{5\times 5}^{5}$ diffusion network.

We present the final results of the joint training in Table I, together with a selection of recent state-of-the-art denoising algorithms, namely BM3D , LSSC , EPLL-GMM , opt-MRF , RTF model , the CSF model and WNNM , as well as two similar approaches ARF and opt-GD , which also train an optimized gradient descent inference. We downloaded these algorithms from the corresponding author’s homepage, and used them as is. Unfortunately, we are not able to present comparisons with , as their codes are not available.

From Table I, one can see that (1) the performance of the $\text{TNRD}_{5\times 5}^{T}$ model saturates after stage 5, i.e., in practice, 5 stages are typically enough; (2) our $\text{TNRD}_{5\times 5}^{5}$ model has achieved significant improvement (28.78 vs.28.60), compared to a similar model $\text{CSF}_{5\times 5}^{5}$, which has the same model capacity and (3) moreover, our $\text{TNRD}_{5\times 5}^{8}$ model is on par with so far the best-reported algorithm - WNNM. It turns out that our trained models perform surprisingly well for image denoising. Then, a natural question arises: what is the critical factor for the effectiveness of the trained diffusion models?

There are actually two main aspects in our training model: (1) the linear filters and (2) the influence functions. In order to have a better understanding of the trained models, we went through a series of experiments to investigate the impact of these two aspects.

Concentrating on the model capacity of 24 filters of size $5\times 5$, we considered the training of a diffusion process with 10 steps, i.e., $T=10$ for the Gaussian denoising of noise level $\sigma=25$. We exploited two main classes of configurations: (A) the parameters of every stage are the same and (B) every diffusion stage is different from each other. In both configurations, we consider two cases: (I) only train the linear filters with fixed influence function $\phi(z)=2z/(1+z^{2})$ and (II) simultaneously train the filters and influence functions.

Based on the same training dataset and test dataset, we obtained the following results: (A.I) every diffusion step is the same, and only the filters are optimized with fixed influence function. This is a similar configuration to previous works . The trained model achieves a test performance of 28.47dB. (A.II) with additional tuning of the influence functions, the resulting performance is boosted to 28.60dB. (B.I) every diffusion step can be different, but only the linear filters are trained with fixed influence functions. The corresponding model obtains a result of 28.56dB, which is equivalent to the variational model with the same model capacity. Finally (B.II) with additional optimization of the influence functions, the trained model leads to a significant improvement with the result of 28.86dB.

The analytical experiments demonstrate that without the training of the influence functions, there is no chance to achieve significant improvements over previous works, no matter how hard we tune the linear filters. Therefore, we believe that the most critical factor of our proposed training model lies in the adjustable influence functions. A closer look at the learned influence functions of the $\text{TNRD}_{5\times 5}^{5}$ model in Sec.4.2 strengthens our argument.

Comparing our proposed TNRD model to the CSF model , one can see that the degree of freedom is in principle the same, since in both models the filters and the non-linear functions can be learned. Therefore, one would expect a similar performance of both models in practice. However, it turns out the performance of the CSF model is inferior to our TNRD model in the case of Gaussian denoising task. The reason for the performance gap is still unclear and we plane to investigate it in future work.

2 Learned influence functions

A close inspection of the learned 120 penalty functionsThe penalty function $\rho(z)$ is integrated from the influence function $\phi(z)$ according to the relation $\phi(z)=\rho^{\prime}(z)$ $\rho$ in the $\text{TNRD}_{5\times 5}^{5}$ model demonstrated that most of the penalties resemble four representative shapes shown in Figure 5.

Truncated convex penalty functions with low values around zero to promote smoothness.

Negative Mexican hat functions, which have a local minimum at zero and two symmetric local maxima.

Truncated concave functions with smaller values at the two tails.

Double-well functions, which have a local maximum (not a minimum any more) at zero and two symmetric local minima.

At first glance, the learned penalty functions (except (a)) differ significantly from the usually adopted penalty functions used in PDE and energy minimization methods. However, it turns out that they have a clear meaning for image regularization.

Regarding the penalty function (b), there are two critical points (indicated by red triangles). When the magnitude of the filter response is relatively small (i.e., less than the critical points), probably it is stimulated by the noise and therefore the penalty function encourages smoothing operation as it has a local minimum at zero. However, once the magnitude of the filter response is large enough (i.e., across the critical points), the corresponding local patch probably contains a real image edge or certain structure. In this case, the penalty function encourages to increase the magnitude of the filter response, alluding to an image sharpening operation. Therefore, the diffusion process controlled by the influence function (b), can adaptively switch between image smoothing (forward diffusion) and sharpening (backward diffusion). We find that the learned influence function (b) is closely similar to an elaborately designed function in a previous work , which leads to an adaptive forward-and-backward diffusion process.

Similar forms of the learned penalty function in (c) with a concave shape are also observed in previous work on image prior learning . This penalty function also encourages to sharpen the image edges. Concerning the learned penalty function (d), as it has local minima at two specific points, it prefers specific image structure, implying that it helps to form certain image structure. We also find that this penalty function is exactly the type of bimodal expert functions for texture synthesis employed in .

Therefore, our learned penalty functions confirmed existing robust penalties based prior models and many priors exploiting some unusual penalties, which can produce patterns and enhance preferred features. As a consequence, the diffusion process involving the learned influence functions does not perform pure image smoothing any more for image processing. In contrast, it leads to a diffusion process for adaptive image smoothing and sharpening, distinguishing itself from previous commonly used image regularization techniques.

3 Pattern formation using the learned influence functions

In the previous work on Gibbs reaction diffusionThe terminology of “reaction diffusion” in is a bit different from ours. In our formulation, “reaction term” is related to the data term, while in , it means the diffusion term controlled by those downright penalty functions. , it is shown that those unconventional penalty functions such as Figure 5(c) have significant meaning in visual computation, as they can produce patterns. We also find that those unconventional penalty functions learned in our models can produce some interesting image patterns.

We consider the following diffusion process involving our learned linear filters and the corresponding influence functions

where the filters $k_{i}$ and influence functions $\phi_{i}$ are chosen from a certain stage of the learned models. Note that we do not incorporate a reaction term in this diffusion model. We run (16) from starting points $u_{0}$ (uniform noise images in the range $$), and it converges to a local minimumThe corresponding diffusion processes are unstable, and therefore we have to restrict the image dynamic range to$$.. Some synthesized patterns are shown in Figure 7. One can see that the diffusion model with our learned influence functions and filters can produce edge-like image structure and repeated patterns from fully random images. This kind of diffusion model is known as Gibbs reaction diffusion in . We provide another example in Figure 13 to demonstrate how our learned diffusion models can generate meaningful patterns for image super resolution.

4 Important aspects of the training framework

Our training model is also a deep model with many stages (layers), but we find that it is not very sensitive to initialization. Based on the training for Gaussian denoising, we conducted experiments with fully random initializations and some plain settings.

We firstly investigated the case of greedy training where the model was trained stage by stage and the parameters of one stage were trained at a time. We initialized the parameters using fully random numbers in the range $[-0.5,0.5]$. It turns out that the resulting models with different initializations lead to a deviation within 0.01dB in the test phase. That is to say, the greedy training strategy is not sensitive to initialization.

Then, we considered the case of joint training, where all the parameters in all stages were trained at a time. We also initialized the training with fully random numbers in the range $[-0.5,0.5]$. In this case, it turns out that the resulting models lead to inferior results, e.g., in the case of $\text{TNRD}^{5}_{5\times 5}$ (28.61 vs.28.78). However, plain initializations can generate equivalent results. For example, we considered a plain initialization (all stages were initialized from the modified DCT filters and an unified influence function $\phi(z)=2z/(1+z^{2})$), the resulting models performed almost the same as those models trained from some good initializations such as parameters obtained from greedy training, e.g.,$\text{TNRD}^{5}_{5\times 5}$, (28.75 vs.28.78) and $\text{TNRD}^{5}_{7\times 7}$, (28.91 vs.28.92).

We believe that this appealing property of our training framework is attributed to the well-distributed gradients across stages. We show in Figure 6 an example to illustrate the gradients of the training loss function with respect to the parameter of all stages. One can see that the well-known phenomenon of “vanishing gradient” in the back-propagation phase of a usual deep model does not appear in our training model. We believe that the reason for the well-distributed gradients is that our training model is more constrained. For example, in a more general sense, the rotated kernel $\bar{k}_{i}$ in our formulation is not necessary to be the rotated version of the kernel $k_{i}$, and it can be an arbitrary kernel. However, we stick to this form, as it has a clear meaning derived from energy minimization.

4.2 Influence of the number of training samples

In our training, we do not consider any regularization for the training parameters, and we finally reach good-performing models. A probable reason is that we have exploited sufficient training samples (400 samples of size $180\times 180$). Thus an interesting question arises: how many samples are sufficient for our training?

In order to answer this question, we re-train the $\text{TNRD}^{5}_{5\times 5}$ model using different size of training dataset, and then evaluate the test performance of trained models. We summarize the results in Figure 8. One can see that (1) too few training samples (e.g., 40 images) will clearly lead to over-fitting, thus inferior test performance, and (2) 200 images are typically enough to prevent over-fitting.

4.3 Influence of the filter size

In our model, the size of involved filters is a free parameter. In principle, we can exploit filters of any size, but in practice, we need to consider the trade-off between run time and accuracy.

In order to investigate the influence of the filter size, we increase the filter size to $7\times 7$ and $9\times 9$. We find that increasing the filter size from $5\times 5$ to $7\times 7$ brings a significant improvement of 0.14dB ( $\text{TNRD}^{5}_{7\times 7}$ vs.$\text{TNRD}^{5}_{5\times 5}$) as show in Table I. However, when we further increase the filter size to $9\times 9$, the resulting $\text{TNRD}^{5}_{9\times 9}$ only leads to a performance of 28.96dB (a slight improvement of 0.05dB relative to the $\text{TNRD}^{5}_{7\times 7}$ model). We can conjecture that further increasing the filter size to $11\times 11$ might bring negligible improvements.

Note that the above conclusion is drawn from a relatively small training data set of 400 images of size $180\times 180$. It should be mentioned that when the size of the model increase, the size of training data set should also increase to avoid over-fitting. However, our current CPU implementation for training prevents us from training with larger model and large-scale data sets (millions). A faster implementation on GPUs together with the stochastic gradient descent optimization strategy is left to future work.

We also consider a model with smaller filters, $3\times 3$. We summarize the results of different model capacities in Figure 9. In practice, we prefer the $\text{TNRD}^{5}_{7\times 7}$ model as it provides the best trade-off between performance and computation time. Therefore, in later applications, we only consider $\text{TNRD}^{T}_{7\times 7}$ models.

Fig. 10 shows the trained filters of the $\text{TNRD}^{5}_{7\times 7}$ model in the first and last stage for the task of Gaussian denoising. One can find many edge and image structure detection filters along different directions and in different scales.

5 Training for different noise levels and comparison to recent state-of-the-arts

The above training experiments are based on Gaussian noise of level $\sigma=25$. We also trained diffusion models for the noise levels $\sigma=15$ and $\sigma=50$. The test performance is summarized in Table I, together with comparison to very recent state-of-the-art denoising algorithms. In experiments, we observed that joint training can always gain an improvement of about 0.1dB over the greedy training for the cases of $T\geq 5$.

From Table I, one can see that for all noise levels, the resulting $\text{TNRD}_{7\times 7}$ model achieves the highest average PSNR. The $\text{TNRD}^{5}_{7\times 7}$ model outperforms the benchmark - BM3D method by 0.35dB in average. This is a notable improvement as few methods can surpass BM3D more than 0.3dB in average . Moreover, the $\text{TNRD}^{5}_{7\times 7}$ model also surpasses the best-reported algorithm - WNNM method, which is quite slow as shown in Table II.

6 Run time

The algorithm structure of our TNRD model is similar to the CSF model, which is well-suited for parallel computation on GPUs. We implemented our trained models on GPU using CUDA programming to speed up the inference procedure, and finally it leads to significantly improved performance, see Table II. We make a run time comparison to other denoising algorithms based on strictly enforced single-threaded CPU computation ( e.g., start Matlab with -singleCompThread) for a fair comparison, see Table II. We only present the results of some selective algorithms, which either have the best denoising result or run time performance. We refer to for a comprehensive run time comparison of various algorithms LSSC, EPLL, opt-MRF and $\text{RTF}^{5}$ methods are much slower than BM3D on the CPU, cf. . .

We see that our TNRD model is generally faster than the CSF model with the same model capacity. It is reasonable, because in each stage the CSF model involves additional DFT and inverse DTF operations, i.e., our model only requires a portion of the computation of the CSF model. Even though the BM3D is a non-local model, it still possesses high computational efficiency. In contrast, another non-local model - WNNM achieves compelling denoising results at the expense of huge computation time. Moreover, the WNNM algorithm is hardly applicable for high resolution images (e.g., 10 mega-pixels) due to its huge memory requirements. Note that our model can be also easily implemented with multi-threaded CPU computation.

In summary, our $\text{TNRD}^{5}_{7\times 7}$ model outperforms these recent state-of-the-arts, meanwhile it is the fastest method even with a CPU implementation. We present an illustrative denoising example in Figure 11 on an image from the test dataset. More denoising examples can be found in the supplemental material based on images from the test dataset and a megapixel-size natural image of size $1050\times 1680$.

Single image super resolution (SISR)

As demonstrated in the last section that our trained diffusion model can lead to explicit backward diffusion process, which sharpens image structures like edges. This is the very property demanded for the task of image super resolution. Therefore, we are motivated to investigate the SISR problem with our proposed approach.

We start with the following energy functional

where the linear operator $A$ is a bicubic interpolation which links the high resolution (HR) image $h$ to the low resolution (LR) image $f$ via $f=Ah$. Casting $\mathcal{D}(u)=\frac{\lambda}{2}\|Au-f\|^{2}_{2}$ and $\mathcal{G}(u)=0$, the energy functional (17) suggests the following diffusion process

where the starting point $u_{0}$ is given by the direct bicubic interpolation of the LR image $f$. Computing the gradients $\frac{\partial u_{t}}{\partial\Theta_{t}}$ and $\frac{\partial u_{t}}{\partial u_{t-1}}$ with respect to (18) can be done with little modifications to the derivations for image denoising. Detailed derivations are presented in the supplemental material.

We considered the model capacity of $\text{TNRD}_{7\times 7}^{5}$, and trained diffusion models for three upscaling factors $\times 2$, $\times 3$ and $\times 4$, using exactly the same 91 training images as in previous works . The trained models are evaluated on two different test data sets: Set5 and Set14. Following previous works , the trained models are only applied to the luminance component of an image, and a regular bicubic upscaling method is applied to the color components.

The test results are summarized in Table III and Table IV. One can see that in terms of average PSNR, our trained diffusion model $\text{TNRD}_{7\times 7}^{5}$ leads to significant improvements over very recent state-of-the-arts in all cases, meanwhile it is still among the fast algorithmsNote that our approach is a Matlab implementation, while some of other algorithms are based on C++ implementations, such as SR-CNN.. A SISR example is shown in Figure 12 to illustrate its effectiveness. One can see that our approach can obtain more accurate image edges, as shown in the zoom-in parts. More SISR examples can be found in the supplemental material.

We apply the learned diffusion parameters to the diffusion equation (16). It turns out that the diffusion process can also generate some interesting patterns from random images, as shown in Figure 7. We believe that this ability to generate image patterns from weak evidence is the main reason for the superiority of our trained model for the SISR task. In order to further validate our argument, we carry out a toy SISR experiment based on a synthesized image with repeated hexagons. The results are shown in Figure 13, where one can see that our trained model can better reconstruct those repeated image structures.

JPEG deblocking experiments

In order to further demonstrate the applicability of our proposed framework for those problems with a non-smooth data term, we investigate the JPEG deblocking problem - suppressing the block artifacts in the JPEG compressed images, which is formulated as a non-smooth optimization problem. Motivated by , we consider the following variational model based on the FoE image prior

By setting $\mathcal{D}(u)=0$ and $\mathcal{G}(u)=\mathcal{I}_{Q}(Du)$, we obtain the following diffusion process

where $\text{proj}_{Q}(\cdot)$ denotes the orthogonal projection onto $Q$. More details can be found in the supplemental material.

We also trained diffusion models for the JPEG deblocking problem. We followed the test procedure in for performance evaluation. We distorted the images by JPEG blocking artifacts. We considered three compression quality settings $q=10,20~{}\text{and}~{}30$ for the JPEG encoder.

We trained three nonlinear diffusion $\text{TNRD}_{7\times 7}$ models for different compression parameter $q$. We found that for JPEG deblocking, 4 stages are already enough. Results of the trained models are shown in Table 6, compared with several representative deblocking approaches. We see that our trained $\text{TNRD}^{4}_{7\times 7}$ outperforms all the competing approaches in terms of PSNR. Concerning the run time, our model takes about 11.2s to handle an image of size $1024\times 1024$ with CPU computation, while the strongest competitor (in terms of run time) - SADCT consumes about 56.5s RTF is slower than SADCT, as it depends on the output of SADCT.. Furthermore, our model is extremely fast on GPUs. For the same image size the GPU implementation takes about 0.095s. See the supplemental material for JPEG deblocking examples.

We have proposed a trainable nonlinear reaction diffusion framework for effective image restoration. Its critical point lies in the additional training of the influence functions. We have trained our models for the problem of Gaussian denoising, single image super resolution and JPEG deblocking. Based on standard test datasets, the trained models result in the best-reported results. We believe that the effectiveness of the trained diffusion models is attributed to the following desired properties of the models

Anisotropy. In the trained filters, we can find rotated derivative filters in different directions, cf. Fig 10, which will make the diffusion happen in some special directions.

Higher order. The learned filters contain first, second and higher-order derivative filters, cf. Fig 10.

Adaptive forward/backward diffusion through the learned nonlinear functions. Nonlinear functions corresponding to explicit backward diffusion appear in the learned nonlinearity, cf. Fig 5.

Meanwhile, the structure of trained models is very simple and well-suited for parallel computation on GPUs. As a consequence, the resulting algorithms are significantly faster than all competing algorithms and hence are also applicable to the restoration of high resolution images.

2 Discussion

One possible limitation of the proposed TNRD approach is that one has to define the ground truth - the expected output of the diffusion network during training. For image restoration applications in this paper, this is not a problem as we have a clear choice for the ground truth. However, for those applications with ambiguous ground truth, e.g., image structure extraction , we will have to make efforts to define the ground truth.

Furthermore, the trained diffusion networks will only perform well in the way they are trained. For example, the trained model based on noise level $\sigma=25$ will break for an input image with noise $\sigma=50$, and the trained model for upscaling factor $\times 3$ will also lead to inferior performance when it is applied to the SISR problem of upscaling factor $\times 2$. It is generally hard to train a universal diffusion model to handle all the noise levels or all upscaling factors.

Our approach is to optimize a time-discrete PDE, which is inspired by FoE based model, but we do not aim to minimize a series of FoE based energies. Our model directly learns an optimal trajectory for a certain possibly unknown energy functional, the minimizer of which provides a good estimate of the demanded solution. Probably, such a functional can not be modeled by a single FoE energy, while our learned gradient descent/forward-backward steps provide good approximation to the local gradients of this unknown functional.

3 Future work

From an application point of view, we think that it will be interesting to consider learned nonlinear reaction diffusion based models also for other image processing tasks such as image inpainting, blind image deconvolution, optical flow. Moreover, since learning the influence functions turned out to be crucial, we believe that learning optimal nonlinearities (i.e., activation functions) in standard CNs could lead to a similar performance increase. There are actually two recent works to investigate a parameterized ReLU function in standard deep convolutional networks, which indeed brings improvements even with little freedom to tune the activation functions. Finally, it will also be interesting to investigate the unconventional penalty functions learned by our approach in usual energy minimization approaches.