Improved PDE Models for Image Restoration through Backpropagation

TL;DR

This paper introduces a novel approach to image restoration by optimizing PDE-based models using deep learning, specifically backpropagation, leading to more effective denoising with a solid mathematical basis.

Contribution

The paper presents a new method that combines PDE models with deep learning to optimize parameters for improved image filtering performance.

Findings

Enhanced denoising quality demonstrated through numerical results.

Stable learning process with optimized cross-diffusion coefficients.

Better image restoration compared to traditional PDE methods.

Abstract

In this paper we focus on learning optimized partial differential equation (PDE) models for image filtering. In this approach, the grey-scaled images are represented by a vector field of two real-valued functions and the image restoration problem is modelled by an evolutionary process such that the restored image at any time satisfies an initial-boundary-value problem of cross-diffusion with reaction type. The coupled evolution of the two components of the image is determined by a nondiagonal matrix that depends on those components. A critical question when designing a good-performing filter lies in the selection of the optimal coefficients and influence functions which define the cross-diffusion matrix. We propose the use of deep learning techniques in order to optimize the parameters of the model. In particular, we use a back propagation technique in order to minimize a cost function related to the quality of the denoising processe, while we ensure stability during the learning procedure. Consequently, we obtain improved image restoration models with solid mathematical foundations. The learning framework and resulting models are presented along with related numerical results and image comparisons.

Figures

Tables

Table 1. Table 1 : Time-step △ ​ t △ 𝑡 \triangle t and number of iterations M 𝑀 M fixed for the numerical experiments. Different choices lead to different stopping times.

	$△t$	M	T
$\sigma =10$	0.05	10	0.5
	0.05	15	0.75
	0.1	10	1
	0.125	10	1.25
	0.125	12	1.5

Equations

$$\{\begin{array}{l}{u}_t=\nabla \cdot (d_1(\mathbf{w})\nabla u+d_2(\mathbf{w})\nabla v)-\lambda (u-u^0)\text{ in }\Omega \times {\mathbb{R}}^{+},\\ v_t=\nabla \cdot (d_3(\mathbf{w})\nabla u+d_4(\mathbf{w})\nabla v)\text{ in }\Omega \times {\mathbb{R}}^{+},\\ u(\mathbf{x},0)=u^0(\mathbf{x}),v(\mathbf{x},0)=v^0(\mathbf{x})\text{ in }\Omega ,\\ u_{\eta }=0,v_{\eta }=0\text{ on }\Gamma \times {\mathbb{R}}^{+},\end{array}$$

$$D(\mathbf{w})=\left(\begin{array}{cc}{d}_1(\mathbf{w})& d_2(\mathbf{w})\\ d_3(\mathbf{w})& d_4(\mathbf{w})\end{array}\right).$$

$$x_{j_1}=a_1+h_1{j}_1,x_{j_2}=a_2+h_2{j}_2,j_k=0,1,\dots ,N_k,h_k=\frac{b_k-a_k}{N_k},k=1,2,$$

$${\delta }_k{Z}_{\mathbf{j}}=\frac{Z_{\mathbf{j}+(1/2){\mathbf{e}}_k}-Z_{\mathbf{j}-(1/2){\mathbf{e}}_k}}{h_k},{\delta }_k{Z}_{\mathbf{j}+(1/2){\mathbf{e}}_k}=\frac{Z_{\mathbf{j}+{\mathbf{e}}_k}-Z_{\mathbf{j}}}{h_k}.$$

$$\frac{U_{\mathbf{j}}^{m+1}-U_{\mathbf{j}}^m}{\Delta t}$$

$$\frac{V_{\mathbf{j}}^{m+1}-V_{\mathbf{j}}^m}{\Delta t}$$

$$D({\mathbf{W}}^m{)}_{\mathbf{j}+(1/2){\mathbf{e}}_k}=\frac{D({\mathbf{W}}_{\mathbf{j}}^m)+D({\mathbf{W}}_{\mathbf{j}+{\mathbf{e}}_k}^m)}{2}$$

$$(U,V{)}_h$$

$$(U,V{)}_{h_1^{\ast }}$$

$$(U,V{)}_{h_2^{\ast }}$$

$$\displaystyle\left(\frac{{U}^{m+1}-{U}^{m}}{\Delta t},{U}^{m+\theta}\right)_{h}+\left(\frac{{V}^{m+1}-{V}^{m}}{\Delta t},{V}^{m+\theta}\right)_{h}+\sum_{k=1}^{2}\Big{(}||(d_{1})^{\frac{1}{2}}\delta_{k}U^{m+\theta}||_{h_{k}^{*}}^{2}$$

$$\displaystyle+||(d_{4})^{\frac{1}{2}}\delta_{k}V^{m+\theta}||_{h_{k}^{*}}^{2}+(d_{2}\delta_{k}V^{m+\theta},\delta_{k}U^{m+\theta})_{h_{k}^{*}}+(d_{3}\delta_{k}U^{m+\theta},\delta_{k}V^{m+\theta})_{h_{k}^{*}}\Big{)}$$

$$=-{\lambda }^{m+\theta }(U^{m+\theta }-U^0,U^{m+\theta }{)}_h.$$

$${\mathbf{W}}^{m+\theta }=\frac{{\mathbf{W}}^{m+1}+{\mathbf{W}}^m}{2}+(\theta -\frac{1}{2})\Delta t\frac{{\mathbf{W}}^{m+1}-{\mathbf{W}}^m}{\Delta t}$$

$$\displaystyle\frac{||{\bf W}^{m+1}||_{h}^{2}-||{\bf W}^{m}||_{h}^{2}}{2\Delta t}+(\theta-\frac{1}{2})\Delta t\bigg{|}\bigg{|}\frac{{\bf W}^{m+1}-{\bf W}^{m}}{\Delta t}\bigg{|}\bigg{|}_{h}^{2}$$

$$\displaystyle+\sum_{k=1}^{2}\Big{(}||(d_{1})^{\frac{1}{2}}\delta_{k}U^{m+\theta}||_{h_{k}^{*}}^{2}+||(d_{4})^{\frac{1}{2}}\delta_{k}V^{m+\theta}||_{h_{k}^{*}}^{2}+(d_{2}\delta_{k}V^{m+\theta},\delta_{k}U^{m+\theta})_{h_{k}^{*}}$$

$$\displaystyle+(d_{3}\delta_{k}U^{m+\theta},\delta_{k}V^{m+\theta})_{h_{k}^{*}}\Big{)}\leq\lambda^{m+\theta}\|U^{m+\theta}-U^{0}\|_{h}\|U^{m+\theta}\|_{h}.$$

$${\lambda }^{m+\theta }\Vert U^{m+\theta }-U^0{\Vert }_h\Vert U^{m+\theta }{\Vert }_h$$

$$\le {\lambda }^{m+\theta }\Vert U^{m+\theta }{\Vert }_h^2+\frac{({\lambda }^{m+\theta }{)}^2}{4ϵ}\Vert U^0{\Vert }_h^2+ϵ\Vert U^{m+\theta }{\Vert }_h^2.$$

$$\displaystyle\sum_{k=1}^{2}\Bigl{(}||(d_{1})^{\frac{1}{2}}\delta_{k}U^{m+1}||_{h_{k}^{*}}^{2}+||(d_{4})^{\frac{1}{2}}\delta_{k}V^{m+1}||_{h_{k}^{*}}^{2}$$

$$\displaystyle+(d_{2}\delta_{k}V^{m+1},\delta_{k}U^{m+1})_{h_{k}^{*}}+(d_{3}\delta_{k}U^{m+1},\delta_{k}V^{m+1})_{h_{k}^{*}}\Bigr{)}>0,$$

$$(1-2\Delta t{\lambda }^{m+1}-2\Delta tϵ)||{\mathbf{W}}^{m+1}|{|}_h^2\le ||{\mathbf{W}}^m|{|}_h^2+2\Delta t\frac{({\lambda }^{m+1}{)}^2}{4ϵ}\Vert U^0{\Vert }_h^2.$$

$${\lambda }^{m+1}\le {\lambda }_{\text{max}},0<\zeta <1-2\Delta t{\lambda }^{m+1}-2\Delta tϵ,m=1,2,\dots$$

$$||{\mathbf{W}}^{m+1}|{|}_h^2\le (1+2\Delta t(\lambda +ϵ){\zeta }^{-1})||{\mathbf{W}}^m|{|}_h^2+\Delta t\frac{{\lambda }_{\text{max}}^2}{2ϵ\zeta }\Vert U^0{\Vert }_h^2.$$

$$\displaystyle||{\bf W}^{m+1}||_{h}^{2}\leq e^{(1+2(\lambda_{\text{max}}+\epsilon)\zeta^{-1})t^{m+1}}\Bigl{(}1+t^{m+1}\frac{\lambda_{\text{max}}^{2}}{2\epsilon\zeta}\Bigr{)}||{\bf W}^{0}||_{h}^{2}$$

$$\displaystyle\bigg{|}\bigg{|}\frac{U^{m+1}-U^{m}}{\Delta t}\bigg{|}\bigg{|}_{h}^{2}\leq$$

$$\displaystyle+2(d_{1}\delta_{k}U^{m},d_{2}\delta_{k}V^{m})_{h_{k}^{*}}\Big{)}$$

$$+2(1+{\eta }^{-1})({\lambda }^m{)}^2(\Vert U^m{\Vert }_h^2+\Vert U^0{\Vert }_h^2),$$

$$\bigg{|}\bigg{|}\frac{V^{m+1}-V^{m}}{\Delta t}\bigg{|}\bigg{|}_{h}^{2}\leq\sum_{k=1}^{2}\frac{8}{h_{k}^{2}}\Big{(}||d_{3}\delta_{k}U^{m}||_{h_{k}^{*}}^{2}+||d_{4}\delta_{k}V^{m}||_{h_{k}^{*}}^{2}+2(d_{3}\delta_{k}U^{m},d_{4}\delta_{k}V^{m})_{h_{k}^{*}}\Big{)}.$$

$$\displaystyle\frac{||{\bf W}^{m+1}||_{h}^{2}-||{\bf W}^{m}||_{h}^{2}}{2\Delta t}+\sum_{k=1}^{2}\Bigl{(}||(d_{1})^{\frac{1}{2}}\delta_{k}U^{m}||_{h_{k}^{*}}^{2}+||(d_{4})^{\frac{1}{2}}\delta_{k}V^{m}||_{h_{k}^{*}}^{2}$$

$$+(d_2{\delta }_k{V}^m,{\delta }_k{U}^m{)}_{h_k^{\ast }}+(d_3{\delta }_k{U}^m,{\delta }_k{V}^m{)}_{h_k^{\ast }}$$

$$\displaystyle-\frac{4\Delta t}{h_{k}^{2}}\big{(}(1+\eta_{1})(||d_{1}\delta_{k}U^{m}||_{h_{k}^{*}}^{2}+||d_{2}\delta_{k}V^{m}||_{h_{k}^{*}}^{2}+2(d_{1}\delta_{k}U^{m},d_{2}\delta_{k}V^{m})_{h_{k}^{*}})$$

$$\displaystyle+||d_{3}\delta_{k}U^{m}||_{h_{k}^{*}}^{2}+||d_{4}\delta_{k}V^{m}||_{h_{k}^{*}}^{2}+2(d_{3}\delta_{k}U^{m},d_{4}\delta_{k}V^{m})_{h_{k}^{*}}\big{)}\Bigr{)}$$

$$\displaystyle\leq\Big{(}\lambda^{m}+\epsilon+\Delta t(1+\eta^{-1})(\lambda^{m})^{2}\Big{)}\|U^{m}\|_{h}^{2}$$

$$\displaystyle+\Big{(}\frac{(\lambda^{m})^{2}}{4\epsilon}+\Delta t(1+\eta^{-1})(\lambda^{m})^{2}\Big{)}\|U^{0}\|_{h}^{2}.$$

$$\displaystyle\sum_{k=1}^{2}\Bigl{(}||(d_{1})^{\frac{1}{2}}\delta_{k}U^{m}||_{h_{k}^{*}}^{2}+||(d_{4})^{\frac{1}{2}}\delta_{k}V^{m}||_{h_{k}^{*}}^{2}$$

$$+(d_2{\delta }_k{V}^m,{\delta }_k{U}^m{)}_{h_k^{\ast }}+(d_3{\delta }_k{U}^m,{\delta }_k{V}^m{)}_{h_k^{\ast }}$$

$$\displaystyle-\frac{4\Delta t}{h_{k}^{2}}\big{(}(1+\epsilon)(||d_{1}\delta_{k}U^{m}||_{h_{k}^{*}}^{2}+||d_{2}\delta_{k}V^{m}||_{h_{k}^{*}}^{2}+2(d_{1}\delta_{k}U^{m},d_{2}\delta_{k}V^{m})_{h_{k}^{*}})$$

Full text

1 Introduction

Nonlinear diffusion processes are well-known and widely used for image noise removal. Roughly speaking, the idea is to combine an effective noise reduction by diffusion with the preservation of the edges and other important image features. Amongst the better known approaches are those where the diffusion coefficient depends on the gradient with an inverse proportion, [15, 18].

The use of nonlinear complex diffusion filters (NCDF) where the image is represented by a complex function and the process of filtering is governed by a diffusion equation with a complex-valued diffusion coefficient, was investigated in [12]. Those filters, which bring the advantage of using the imaginary part of the solution as an edge detector avoiding the computation of the gradient to control the diffusion coefficient, can be successfully applied for denoising in particular for medical imaging despeckling [6, 16].

A complex diffusion equation can be written as a cross-diffusion system for the real and imaginary parts. The use of general cross-diffusion systems for image processing, which encompasses the complex-diffusion equations as particular cases, was investigated in [2]. The theoretical studies of the correspondent initial valued boundary problems were presented in [1] and [2] for the linear and nonlinear cases, respectively, where well-posedness and some scale-space representation properties were derived under hypothesis on the cross-diffusion coefficient matrix. However, besides these conditions, there is not much insight on the form that these coefficients should take in practical applications. In [3], the same authors applied the cross-diffusion model to several examples related to the problem of reducing the speckle noise in optical coherence tomography (OCT) images. The conclusion is that the performance of the model is greatly influenced by the choice of the cross-diffusion coefficient matrix.

All those methods based on diffusion filtering are highly dependent on some crucial parameters (e.g. [2, 17]). The parameters that lead to the most effective methods vary depending on the acquisition methods, the nature of the images and the associated noise profile. Furthermore, diffusion filters require the specification of a stopping time. To circumvent this issue, a reaction term can be introduced which keeps the steady-state solution close to the original image [18].

The challenge lies in the formulation of robust models with optimal parameters corresponding to each architecture. The use of neural networks for image processing tasks has set the pace of current research in this area. However, the general intrinsic “black box” nature of such algorithms is a drawback, particularly in the context of medical applications. In [10] the authors use a learning framework inspired in nonlinear reaction diffusion processes where the filters and the influence functions are simultaneously learned from training data. The results are auspicious. There are, however, some important questions remaining: the numerical solution is not related with the original reaction diffusion problem, and more importantly, the qualitative properties of the computed solution are not derived. Moreover, the method does not cover promising models as the one based on cross-diffusion.

We propose the use of deep learning techniques such as backpropagation [13] and Adam algorithm [14] to optimize the parameters of a system of PDEs for a specific task. In this work, we focus on nonlinear cross-diffusion with reaction filters in order to optimize the reaction parameter and the influence functions of the cross-diffusion matrix for the removal of gaussian noise.

The optimization process we suggest includes the stability of the inherent numerical method, which is crucial for obtaining successful results. We drive our attention first to synthetic images in order to develop our learning strategy. The methodology consists in the parametrization of the functions that govern the cross-diffusion system, the computation of the gradient of a cost function with respect to these parameters and the application of the backpropagation technique, widely used in neural network supervised learning, to learn the optimal parameters. Experiments in real noisy images show that the training process improves significantly the performance of the filters.

Our contributions can be summarized as follows:

• •

we derive the stability conditions for a broad class of cross-diffusion models;

• •

we build a recurrent neural network equivalent to a cross-diffusion system where its neurons are the points of the spatial discretization and the weights are the non-linear diffusion coefficients;

• •

we obtain improved stable cross-diffusion models for image denoising through backpropagation.

The article is organized as follows. In Section 2 the fully discrete cross-diffusion with reaction model is presented. Next, the stability conditions are derived. Section 4 is devoted to the learning model. We present the numerical experiments in Section 5. We end the paper with a section dedicated to conclusions.

2 Cross-diffusion reaction model

In this section we present a fully discrete cross-diffusion reaction model for image restoration. The image is represented by a two-component vector field, ${\bf w}=(u,v)^{\top}$, and the restoration process is governed by the nonlinear cross-diffusion reaction system

[TABLE]

where $\Omega=(a_{1},b_{1})\times(a_{2},b_{2})\subset\mathbb{R}^{2}$ is the domain of interest, $u^{0}$ and $v^{0}$ are the given initial conditions for $u$ and $v$ and $\eta$ denotes the outward normal vector to the boundary $\Gamma=\partial\Omega$. The cross-diffusion matrix of the model is given by

[TABLE]

$\lambda$ is a time dependent non-negative parameter.

The domain $\overline{\Omega}=\Omega\cup\Gamma$ is discretized by the points ${\bf x}_{{\bf j}}=(x_{j_{1}},x_{j_{2}})$, where

[TABLE]

for two given integers $N_{1},N_{2}\geq 1$, ${\bf j}=(j_{1},j_{2})$ and ${\bf h}=(h_{1},h_{2})$ . This spatial mesh on $\overline{\Omega}$ is denoted by $\overline{\Omega}_{\bf h}$ and $\Gamma_{\bf h}=\Gamma\cap\overline{\Omega}_{\bf h}$. Points halfway between two adjacent grid points are denoted by ${\bf x}_{{\bf j}\pm(1/2){\bf e}_{k}}={\bf x}_{\bf j}\pm\frac{h_{k}}{2}{\bf e}_{k}$, $k=1,2$, where $\{{\bf e}_{1},{\bf e}_{2}\}$ is the canonical basis, that is, ${\bf e}_{k}$ is the standard basis unit vector in the $k$th direction.

For the discretization in time, we consider a mesh with time step $\Delta t$, $0=t^{0}<t^{1}<t^{2}<\ldots$, where $t^{m+1}-t^{m}=\Delta t$.

We denote by $Z_{\bf j}^{m}$ the value of a mesh function $Z$ at the point $({\bf x}_{{\bf j}},t^{m})$. For the formulation of the finite difference approximations, we use the centered finite difference quotients in the $k$th spatial direction, for $k=1,2$,

[TABLE]

In order to formulate the discrete cross-diffusion restoration problem, let $u^{0}:\Omega_{\bf h}\to\mathbb{R}$ be a discrete real-valued function standing for the grey level values on $\Omega_{\bf h}$ of the noisy image to be restored. From $u^{0}$, an initial distribution ${\bf W}^{0}=(U^{0},V^{0})$, for the cross-diffusion, is required. This is given by two real-valued functions $U^{0},V^{0}:\Omega_{\bf h}\rightarrow\mathbb{R}$, that can be selected following different criteria. A simple choice, that will be considered in the experiments we present later in this paper, consists of taking $U^{0}({\bf x}_{{\bf j}})=u^{0}({\bf x}_{{\bf j}})$ and $V^{0}({\bf x}_{{\bf j}})=0$, ${\bf x}_{{\bf j}}\in\overline{\Omega}_{\bf h}$. A more detailed discussion on the initial data can be seen in [1, 2].

Let ${\bf W}^{m}_{\bf j}=(U_{{\bf j}}^{m},V_{{\bf j}}^{m})^{\top}$, such that ${\bf x}_{\bf j}\in\overline{\Omega}_{\bf h}$. Given the initial solution ${\bf W}^{0}_{\bf j}=(U^{0}_{\bf j},V^{0}_{\bf j})$, the numerical solution of (1) at the time $t^{m+1}$ is obtained considering the following finite difference scheme

[TABLE]

where

[TABLE]

and $\theta\in\{0,1\}$. If $\theta=0$ then (3) is an explicit method. If $\theta=1$ then method (3) is semi-implicit and its solution is obtained by solving a system of linear equations. In the next section we will derive the stability properties of both, explicit and semi-implicit, methods. Later, in sections 4.2 and 4.3, we will discuss their role in our learning strategy: to optimize the cross-diffusion model we will use the explicit scheme ; to denoise the images we will use the semi-implicit discretization of the optimized model.

Equations (3) and (4), for ${\bf j}$ such that ${\bf x}_{{\bf j}}\in\Gamma_{{\bf h}}$, are defined using points of the form ${\bf x}_{{\bf j}}+{\bf e}_{k}$ and ${\bf x}_{{\bf j}}-{\bf e}_{k}$, which don’t belong to $\Omega_{\bf h}$ and are not yet defined. For those points ${\bf x}_{{\bf j}}+{\bf e}_{k}$ and ${\bf x}_{{\bf j}}-{\bf e}_{k}$, we consider the solution in ${\bf x}_{{\bf j}}-{\bf e}_{k}$ and ${\bf x}_{{\bf j}}+{\bf e}_{k}$, respectively. This corresponds to the usual discretization of the homogeneous Neumann boundary conditions on $\Gamma_{{\bf h}}$ with central finite differences.

3 Stability of the numerical scheme

We will now investigate the stability of the finite difference scheme (3)-(4). The approach generalizes the strategies in [4] used to study the stability of the particular case of complex-diffusion equations written as a cross-diffusion system.

For each $x_{\bf j}=(x_{j_{1}},x_{j_{2}})\in\bar{\Omega}_{\bf h}$, we define the rectangle $\square_{\bf j}=(x_{j_{1}},x_{j_{1}+1})\times(x_{j_{2}},x_{j_{2}+1})$ and denote by $|\square_{\bf j}|$ the measure of $\square_{\bf j}$. We consider the discrete $L^{2}$ inner products

[TABLE]

[TABLE]

and

[TABLE]

Their correspondent norms are denoted by $\|.\|_{h}$, $\|.\|_{h_{1}^{*}}$ and $\|.\|_{h_{2}^{*}}$, respectively. For ${\bf W}=({U},{V})^{\top}$ we define $\|{\bf W}\|_{h}^{2}=\|{U}\|_{h}^{2}+\|{V}\|_{h}^{2}$.

To simplify the notation and where it is clear from the context, we write in what follows $d_{\ell}$ instead of $d_{\ell}({\bf W}^{m}),$ or $d_{\ell}({\bf W}^{m})_{\bf j}$, for $\ell=1,2,3,4.$

Multiplying both members of equations (3) and (4) by $U^{m+\theta}$ and $V^{m+\theta}$, respectively, according to the discrete inner product $(\cdot,\cdot)_{h}$, and using summation by parts we get

[TABLE]

We can write

[TABLE]

and then

[TABLE]

For any $\epsilon>0$ we have that

[TABLE]

3.1 Stability of the semi-implicit scheme

We start by considering the case where $\theta=1$. With the assumption that

[TABLE]

which corresponds to the natural assumption of the diffusion matrix $D$ to be uniformly positive definite, from (5) and (6) we get

[TABLE]

Assuming that

[TABLE]

for some $\lambda_{\text{max}},\zeta\in{\rm I\kern-1.69998ptR}^{+}$ where $\epsilon$ is a constant arbitrarily chosen, we get

[TABLE]

If (8) holds, using the Duhamel’s principle ([9], Lemma 4.1 in Appendix B) we get

[TABLE]

and we conclude that the method is stable.

3.2 Stability of the explicit scheme

We now consider the case where $\theta=0$. Applying the $||\cdot||_{h}$ norm on both sides of equations (3) and (4), and using the inequalities $(a\pm b)^{2}\leq 2a^{2}+2b^{2}$ and $(a\pm b)^{2}\leq(1+\eta)a^{2}+(1+\eta^{-1})b^{2}$ for $\eta>0$, we obtain, respectively,

[TABLE]

for any $\eta>0$ and

[TABLE]

Combining the inequalities (6), (9) and (10) with (5) leads to

[TABLE]

Let $\eta=\epsilon$. In order to obtain a stable scheme we impose that

[TABLE]

for some $\epsilon>0$.

If (11) holds, then

[TABLE]

We now take a_{\epsilon}=2\Big{(}\lambda^{m}+\epsilon+\Delta t(1+\epsilon^{-1})(\lambda^{m})^{2}\Big{)} and b_{\epsilon}=2\Big{(}\frac{(\lambda^{m})^{2}}{4\epsilon}+\Delta t(1+\epsilon^{-1})(\lambda^{m})^{2}\Big{)}. From (12) and using the Duhamel’s principle we get

[TABLE]

and we conclude that the method is stable.

Note that we can rewrite (11) as

[TABLE]

where $M$ is a square matrix of dimension $2N_{1}N_{2}\times 2N_{1}N_{2}$. In order for (13) to hold we require $M$ to be semi-positive definite, that is, the eigenvalues of $M$ are all non-negative. Using the Gershgorin Theorem, a sufficient condition for the method to be stable is to impose that the influence functions satisfy, for some $\epsilon>0$,

[TABLE]

4 Learning model

Our goal is to optimize the parameters of our model using deep learning techniques. In order to adapt the cross-diffusion scheme (3)-(4) to a neural network architecture, we concatenate $U$ and $V$ into a single column vector $w$ with $2N_{1}N_{2}$ entries:

[TABLE]

The $m$-th iteration of the scheme (3)-(4) can now be written in the vector formulation

[TABLE]

with $\lambda^{m+\theta}>0$, and where the first difference matrices ($K_{r*}^{*}$) are block matrices:

[TABLE]

$k_{r}^{x}$ and $k_{r}^{y}$ denote the backward difference operators with respect to $x$ and $y$, respectively. The matrices $D_{L}^{x,m},D_{L}^{y,m},D_{R}^{x,m}$ and $D_{R}^{y,m}$ are $2N_{1}N_{2}\times 2N_{1}N_{2}$ diagonal matrices. Each entry $(j,j)$ of $D_{L}^{x,m}$ and $D_{L}^{y,m}$ depend on the function $d_{1}$ for $j=1,\ldots,N_{1}N_{2}$ and on the function $d_{4}$ for $j=N_{1}N_{2}+1,\ldots,2N_{1}N_{2}$, at the time $t^{m}$. The entries of $D^{x,m}_{R}$ and $D^{y,m}_{R}$ follow the same pattern, replacing $d_{1}$ and $d_{4}$ by $d_{2}$ e $d_{3}$, respectively. It remains to define the difference matrices $K_{l}^{x}$ and $K_{l}^{y}$:

[TABLE]

where $k_{l}^{x}$ and $k_{l}^{y}$ denote the forward difference operators with respect to $x$ and $y$, respectively.

We reiterate that the scheme (15) is equivalent to the cross-diffusion scheme (3)-(4).

4.1 Parameterization of the influence functions

An important aspect of the methodology of our learning strategy consists in the parametrization of arbitrary influence functions. We parameterize the influence functions of the cross-diffusion matrix (2) through gaussian radial basis functions (RBFs) [8]. Here we consider that they only depend on the second component of the vector $w=(u,v)^{\top}$, i.e., they depend on the component which plays the role of edge detector. Each of these functions has the expression

[TABLE]

where $\mu_{i}$, for $i=1,...,P$, are equidistant points in the set $A$ of edge-detector range of values, $\nu$ is a positive scalar (the so called scale of the RBF), and $\delta_{\ell,i}$ are, for each $\ell=1,\dots,4$, the $P$ interpolants od $d_{\ell}$. A good approximation requires a careful balance between $\nu$ and $P$.

In the learning model, we consider in (15) $\lambda^{m+\theta}=\lambda$, where $\lambda$ is a positive scalar.

4.2 Constrained optimization

In order to optimize the influence functions of the cross diffusion matrix $d_{1},d_{2},d_{3}$ and $d_{4}$ and also the scalar $\lambda$ which characterizes the reaction term, a training cost function and a training dataset are required. We will use a set of $B$ gray-scale images which will serve as basis for our training set. We fix the time-step $\Delta t$ and a number of steps $M$, which will result in a cross-diffusion stopping time of $T=M\Delta t$. In the learning viewpoint, we will have a $M$-layered convolutional neural network. We aim to minimize the loss function $L$,

[TABLE]

with $w_{i}^{m}=(u_{i}^{m},v_{i}^{m})^{\top}$, $w_{i}^{nl}=(u_{i}^{nl},v_{i}^{nl})^{\top}$, where $w_{i}^{m}$ and $w_{i}^{nl}$ are, respectively, the $m$- th iteration of (15) of the $i$-th corrupted image and the non-corrupted image, respectively, $B$ is the learning batch size, and $\Theta$ is the set of parameters we want to optimize. As such, the set of parameters $\Theta$ is

[TABLE]

and we must now obtain $\frac{\partial L}{\partial\Theta}(\Theta,w_{1}^{M},\dots,w_{B}^{M},w_{1}^{nl},\dots,w_{B}^{nl})$ to solve the minimization problem.

We need to guarantee that the learning procedure ends up with a stable scheme. For that, we will make use of the stability conditions derived in Section 3. Although the use of the explicit method (15) with $\theta=0$ can bring some advantages in terms of computational effort, the stability conditions are much more restrictive when compared to the semi-implicit method (15) with $\theta=1$. Moreover, since we need to impose the stability conditions over the values of $\Theta$, the nonlinearities in the stability condition (14) corresponding to the explicit scheme carry out some issues that we want to avoid here. For that reason our goal is to learn a semi-implicit scheme for image denoising.

To obtain a stable semi-implicit scheme, instead of minimizing the loss function (17) we seek the solution of the constrained optimization problem

[TABLE]

This constraints ensure that (7) hold. We first note that (18b)-(18c) are not in the standard formulation over the values of $\Theta$. As the radial basis functions (16) are strictly positive, we replace (18b)-(18c) with

[TABLE]

using the fact that (19a)-(19d) implies (18b)-(18c). To solve the inequality constrained problem we define the augmented Lagrangian $\mathcal{L}$ as

[TABLE]

which is the usual Powell-Hestenes-Rockafellar function applied to problem (18a) with constraints (19a)-(19d). The vector $\mu$ is the set of Lagrange multipliers associated with constraints (19a)-(19d) and $\rho$ is the penalty parameter.

We follow the main model algorithm in [7] with some minor modifications. The learning procedure will iterate through $\mathcal{L}(\Theta_{k},\mu_{k},\rho_{k})$. At the iteration $k$ of the minimization algorithm, the Lagrange multipliers are updated through the formula

[TABLE]

where $\bar{\mu}>0$ is a multiplier cap, and the penalty parameter is increased or decreased if the infeasibility increases or decreases, respectively. The measure of infeasibility at iteration $k$ is given by the quantity

[TABLE]

The penalty parameter is updated in the following way: for some predefined $\tau\in(0,1)$ and $\gamma>1$, if $I_{k}\leq\tau I_{k-1}$ then $\rho_{k+1}=\rho_{k}/\gamma$, otherwise $\rho_{k+1}=\gamma\rho_{k}$.

Remark 1

The stability condition (8) is not very limitative and for that reason we don’t include it in the formulation of the optimization problem (18). Even so, we verify if the condition is satisfied by the optimal solution. In our numerical experiments, presented later in the paper, that was always the case.

4.3 Deriving the gradients

A central step to solve the optimization problem is to compute the gradients of the loss function with respect to the training parameters. We note that the direction of maximum growth of $L$ can be obtained via back-propagation:

[TABLE]

Although the objective is optimize a semi-implicit cross-diffusion with reaction scheme, we use the explicit scheme (15) with $\theta=0$ for the learning steps.

We can readily determine

[TABLE]

We will obtain $\frac{\partial w^{m}}{\partial\Theta}$ and $\frac{\partial w^{m+1}}{\partial w^{m}}$ using (15). We start by noticing that

[TABLE]

We observe that

[TABLE]

where $g_{L}^{x,m}$ is the vector with the entries of $D_{L}^{x,m}$. From (16), we can write $g^{x,m}_{L}=G^{x,m}\Lambda_{1,4}$, where $\Lambda_{1,4}$ is a $2P$ sized vector with the parameters $\delta_{1,i}$, $i=1,\ldots,P$ and $\delta_{4,i}$, $i=1,\ldots,P$ of $d_{1}$ and $d_{4}$, respectively, and

[TABLE]

where $S^{x}=|k^{x}_{r}|$ and with $\Phi^{x,m}$ being a $N_{1}N_{2}\times P$ matrix

[TABLE]

Analogously, we write $g^{y,m}_{L}=G^{y,m}\Lambda_{1,4}$, $g^{x,m}_{R}=G^{x,m}\Lambda_{2,3}$, $g^{y,m}_{R}=G^{y,m}\Lambda_{2,3}$, where $S^{y}=|k^{y}_{r}|$, and obtain

[TABLE]

[TABLE]

Consequently,

[TABLE]

To derive the expression of $\displaystyle\frac{\partial w^{m+1}}{\partial w^{m}}$, we consider (15). We have

[TABLE]

with

[TABLE]

[TABLE]

[TABLE]

[TABLE]

where $d_{\ell}^{\prime}(v^{m})$ is a vector of dimension $N_{1}N_{2}$ with entries $[d_{\ell}^{\prime}(v^{m})]_{j}=d_{\ell}^{\prime}(v_{j}^{m})$, being $v_{j}^{m}$ is the second component of the vector $w_{j}^{m}=(u_{j}^{m},v_{j}^{m})^{\top}$. We obtain the derivatives of the influence functions numerically with a centered difference scheme.

4.4 Back-propagation algorithm

We now describe how the unknown parameters are learned via back-propagation. We store $B$ gray-scale images to form the basis for our training set. In each iteration, we crop a random $N_{1}\times N_{2}$ section of each image and add the desired synthetic random noise to each section (for example noise of Gaussian type with zero mean and standard deviation $\sigma$). This set of freshly cut sections with noise will make up the batch for the iteration.

There is also the need to fix the step time $\Delta t$ and the number of layers $M$ before solving the minimization problem. This is a crucial choice, as $T=M\Delta t$ becomes the diffusion time for which the procedure will be optimized. The functions $d_{1},d_{2},d_{3}$ and $d_{4}$ are initialized as in the nonlinear complex diffusion case [2], while $\lambda$ is initialized with the value that provides the best average peak to signal noise ratio (PSNR) for the nonlinear complex diffusion process over the training set.

Given these particular considerations, we present the learning algorithm model (Algorithm 1), which is the Adam algorithm [14] applied to our particular problem.

5 Experimental results

Having previously established a learning framework for the optimization of a cross-diffusion process, we now test it with some examples. We set $B=50$ and store 50 gray-scaled images (Figure 1) that will serve as training set. The crop size for the batch extraction was fixed for all tests as $N_{1}=N_{2}=100$ and the influence functions were initialized as in the non-linear complex diffusion case [4, 5, 12],

[TABLE]

with

[TABLE]

We experimented with different values of $\triangle t$ and $T$ to see the effect of changing stopping times in this learning framework. The values chosen are summarized in (Table 1). For the RBF interpolation (16), we chose $A=[-20,20]$, $P=151$ and $\nu=0.2$. For the augmented Lagrangian parameters, we fixed for all tests $\bar{\mu}=2$, $\rho_{1}=6\times 10^{5}$, $\tau=0.5$ and $\gamma=2$. Finally, we set $K_{\text{max}}=2000$ for the number of iterations of the minimization problem.

In Figure 3 we average the values of (17) over groups of $10$ algorithm updates for all the training procedures carried out with $\sigma=10$. The results show a similar pattern for all combination of $\triangle t$ and $M$. The algorithm reaches a significant lower loss value relatively fast (in $50$ to $100$ updates) and proceeds to improve (although at a slower pace) throughout the rest of the optimization procedure.

We describe the learned influence functions in figures 4 and 7 for $\sigma=10$ and $\sigma=20$, respectively. We notice that the learning translates into the determination of the appropriate scale of the original function in the cases of influence functions $d_{1}$, $d_{2}$ and $d_{4}$, and into the radical reshaping of $d_{3}$. The results are consistent across the different stopping times and both levels of noise.

In figures 5 and 8 we test the learned parameters against the non-linear complex diffusion case and with the value of $\lambda$ that provides the best PSNR in the training set. The comparisons are performed in the test set (figure 2), which is made by images that were not integrated in the training set. We see in both cases a significant increase in PSNR and a slight increase in blur ([11]). The combinations of $\triangle t$ and $M$ that achieve better performances both for $\sigma=10$ and $\sigma=20$ are those that combine into a greater value of $T$, that is, a later stopping time. The reason for this is the introduction of the reaction term, which shifts the steady-state solution into a non-trivial one (when compared to diffusion filters without such reaction), and consequently the denoising process benefits from a longer and more controlled diffusion.

Finally, we compare the action of the nonlinear complex filter against the best learned parameters (that is, the ones trained with a larger $T$) on the widely used Lena image. The results are in figures 6 and 9 for $\sigma=10$ and $\sigma=20$, respectively. In both cases, we achieved significantly improved PSNR and blur values.

6 Conclusions

We have successfully adapted a cross-diffusion model for image restoration into a learning framework, obtaining a way to automatically parametrize the cross-diffusion matrix for different images and levels of noise. The numerical experiments show a significant improvement of the results, measured with the PSNR and Blur metrics, obtained with the proposed learning algorithm when compared with the initialization. By making the parallelism between neural networks and the parametrization of PDEs, we believe that this work can be transferred to a broad scope of related problems.

Acknowledgements

This work was partially supported by the Centre for Mathematics of the University of Coimbra – UID/MAT/00324/2019, funded by the Portuguese Government through FCT/MEC and co-funded by the European Regional Development Fund through the Partnership Agreement PT2020. The second author was supported by the FCT grant PD/BD/142956/2018.