A Residual Solver and Its Unfolding Neural Network for Total Variation Regularized Models

Among all regularization terms, Total Variation (TV) might be the most popular one because of its simplicity and effectiveness [1]. Assuming the unknown ground truth signal is $U(\vec{x})$ where $\vec{x}$ is the spatial coordinate, the traditional total variation regularization term is

where $\nabla$ is the gradient operator and $p=\{1,2\}$ indicates the standard $\ell_{p}$ norm. This term is isotropic when $p=2$ and becomes anisotropic when $p=1$. Choosing which $\ell_{p}$ norm is problem specific. In the rest of this paper, we use $p=1$ because it has been shown preserving edges better in many practical applications [10].

Although there are some variants of TV, such as spatial weighted TV and $\ell_{0.5}$ norm TV (Hyper Laplace), the standard TV with $\ell_{1}$ norm still is the most popular regularization term. The reasons come from three aspects. First, from optimization point of view, this term is convex. This means the global optimal solution is unique. And the convex optimization algorithms can be used for this term. Second, the $\ell_{1}$ norm is a simple computation operation. Thus, it can be efficiently computed and evaluated in practice. Third, the anisotropy property of this term is important in preserving edges in image processing. In general, preserving edges is always preferred in many image processing tasks.

When TV is used as the regularization term, many image processing problems, such as denoising, segmentation, inpainting and fusion, etc, can be formulated as the following minimization problem:

where $A$ is an imaging matrix, $U$ is the unknown signal to be estimated, $\vec{x}$ is a spatial coordinate, $f$ is the observed signal and $\lambda>0$ is a scalar that balances the imaging and the TV prior ($\lambda$ usually is related with noise level). The imaging matrix $A$ models a linear imaging process, such as blurring, down sampling, up sampling, etc. A linear matrix already covers a large range of image processing problems.

Many solvers have been developed for this model (Equation 2). The earliest iterative scheme was proposed in  [1]. This method is slow due to the stability constraints in the time step size. To accelerate the optimization procedure, a number of alternative methods have been proposed. One popular method is primal dual solver [11, 12]. Primal dual method simultaneously finds the optimal solution for this model and its dual, improving the numerical stability. Another technique is to split the imaging model and the TV regularization, by introducing a variable $W$:

This new model can be solved by alternating direction method of multipliers (ADMM) [13], split Bregman [14], etc. Equation 2 can also be solved by Iterative Shrinkage and Thresholding Algorithm (ISTA)  [15]. These methods are based on nonlinear partial differential equations and thus time consuming, compared with recently developed neural network solvers.

With the development of deep learning, several neural networks have been proposed to solve Eq. 2. One of the earliest work is the Learned ISTA (LISTA), which unfolds the ISTA as neural networks to perform the optimization algorithm [16]. One layer in the network corresponds to one iteration in ISTA. Different from ISTA, all the parameters in LISTA implicitly communicate with each other during the training process. Such communication can be considered as high order scheme from the partial differential equation point of view. Therefore, the network can solve the problem within a small number of layers, in contrast with the large number of iterations in iterative algorithm ISTA.

ISTA-net is another extension of ISTA [17]. It unfolds the ISTA but keeps the symmetric network structure, requiring the encoding and decoding are inverse functions for each other. Such constraint makes the network invertable, leading to better numerical stability than LISTA.

Another unfolding method is the primal dual network, which unfolds the primal dual iteration for Eq. 2 [18]. It has been shown converging much faster than the original primal dual iteration.

ADMM-net unfolds the classical ADMM iteration, leading to a neural network representation [19]. It first trains one layer. Then it fixes this layer, adds another layer and only trains the new added layer. By repeating this procedure, it shows good performance for MRI reconstruction problem. In this network, after 10 layers, adding more layers does not improve the result anymore.

Recently, a CNN based image reconstruction method for inverse problem is proposed in [20]. It shows that neural networks are very efficient in solving image reconstruction problems thanks to the parameter communication between different layers.

In general, neural network methods are more efficient than the iterative algorithms because of two reasons. First, the parameters in the neural networks can communicate with each other during the training process. This corresponds to the high order scheme from partial differential equation point of view. Second, the parameters in the neural networks are more adaptive to the training data. In contrast, the operations in the iterative algorithms are independent from the input image. As a result, the neural network usually only needs a small number of layers to solve Eq. 2 (computational efficiency).

Most of these neural network algorithms are supervised, requiring the ground truth to be explicitly given. Such requirement is not easily satisfied in most practical applications. This motivates us to develop an unsupervised neural network. Our contributions fall in following two aspects:

• •

  We develop a new iterative algorithm that solves the Eq. 2. Instead of directly finding the optimal solution, we find its related residual (details will be explained in later sections). Therefore, our solver is called Residual Solver (RS).

• •

  We unfold the proposed iterative residual solver into a neural network and our network is unsupervised. Therefore, it can be trained without knowing the optimal solution. Another advantage of our network is that it can be trained on low resolution images but applied on high resolution images. Our experiments confirm this property.

Since this procedure is iterative, we can modify this iteration by updating the dual variable $d$ first, leading to a more efficient algorithm

In this new iteration, the term $2d^{t+1}-d^{t}$ in Eq. 8 can be considered as extrapolation. And we also avoid the matrix multiplication with $A^{T}A$ in Eq. 4. In fact, Eq. 8 is simpler to compute, compared with Eq. 4. Thus, it is faster in practice.

In both iteration methods, Eq. 9 is the computation bottleneck. This equation is known as the standard ROF model [1]. Many algorithms have been developed to efficiently solve Eq. 9. Among these algorithms, following iteration was proposed in [22]

where $0<\beta<1/4$ is a scalar parameter, $in$ is an inner iteration index, $\vec{b}=(b_{x},b_{y})$ corresponds to the gradient field, $\vec{b}^{0}=0$, and $cut$ function is defined as

This iteration can theoretically guarantee the numerical convergence and has shown good performance in theory and practice [22]. We refer this algorithm as FastSolver (FS) in this paper.

This $cut$ function proposed in FS is the residual of the well-known $shrink$ operation. As shown in Fig. 1, $cut$ function plus $shrink$ function equals the identity function. Their difference forms a nonlinear transfer function that also appears in deep learning community. Moreover, the $cut$ function is similar with several well-known activation functions from neural network field, such as sigmoid, $atan$, soft sign, etc.

Until Section IV-D, we will focus on this bottleneck problem (Eq. 9). We will show a residual solver for this problem. Then we will unfold this iterative solver into a neural network, which is an enhanced version of the iterative solver. In Section IV-D, we will show how to embed this neural network into a larger network for the general TV regularized models.

To show that our RS is numerical stable, we need to find a unique vector field $\vec{b}$ for the unique residual $R$. It is well-known that Eq. 9 is convex [1]. Therefore, its optimal solution $U^{t+1}$ is unique, leading to the uniqueness of the residual $R=-\lambda\beta\nabla^{T}\vec{b}^{*}$. Although the vector field $\vec{b}^{*}$ is not unique, it belongs to a unique set in terms of divergence $\nabla^{T}$ equivalency

This set $\cal S$ is convex because $\nabla^{T}$ is a linear operation. Let $0\leq\beta\leq 1$ be a scalar. Then, for any $\vec{b}_{1}$ and $\vec{b}_{2}$ in this set, since $\nabla^{T}(\beta\vec{b}_{1}+(1-\beta)\vec{b}_{2})=\nabla^{T}\vec{b}^{*}$, we have

This result indicates that this unique set is a convex set.

Since the ground truth image $U$ has bounded intensity values (usually $U(\vec{x})\in[0,1])$ or $U(\vec{x})\in[0,255])$), $\nabla^{T}\vec{b}^{*}$ is also bounded ($<+\infty$). Therefore, we can require $\|\vec{b}\|_{2}<+\infty\,,\forall\vec{b}\in{\cal S}$.

Thank to the bounded norm and convexity, we can obtain a unique solution if we add an additional constraint, such as minimizing $||\vec{b}||_{1}$. More specifically, we have

Thanks to the convexity of $\cal S$, $\vec{b}_{1}^{*}$ is unique.

In practice, we always initialize $\vec{b}^{0}=0$. This initialization leads to the unique $\vec{b}_{1}^{*}$ since it is the “minimum” vector field in this convex set. We observed that the same input always find the same vector field $\vec{b}_{1}^{*}$ in practice.

The convexity of the solution set and uniqueness of vector field $\vec{b}_{1}^{*}$ provide the theoretical guarantees for our RS. They also reveal the numerical stability of our RS, which is confirmed in following experiment on BSDS500 data set.

Our RS can reach the optimal solution. We run RS and FS with 200 iterations on 500 natural images from BSDS500 data set. Then we compute their final loss function value for each image, respectively. For the same image, we compute the loss function value difference. The distribution of such differences for all 500 images is shown in Fig. 2 (a). The average difference is about $-2\times 10^{-3}$. The negative values indicate that our residual solver always reaches lower energy level for these images. The results confirm that RS can reach the optimal solution.

We also compute the SSIM between the original and processed images. For FastSolver and RS, the SSIM difference is small ($\approx 10^{-8}$) and its distribution is shown in Fig. 2 (b).

In addition, our RS runs faster. The average running time of RS and FastSolver for the 500 color images on the same hardware with the same iteration number is 0.89 seconds and 1.07 seconds, respectively. This indicates that our RS can reduce about $17\%$ of the running time.

This new residual solver has three important properties:

• •

  implicit: Our RS does not require to compute $U^{t}$ during the iteration. In contrast, $U^{t}$ has to be explicitly computed in Eq. 11 for each iteration in FastSolver.

• •

  gradient domain: Our RS solves the problem in the gradient domain, which is different from previous methods that work in the intensity domain.

• •

  fast: In Eq. 14, $\nabla V^{t+1}$ is fixed during the inner iterations. Therefore, $\nabla V^{t+1}$ can be computed before the loop, which reduces about $17\%$ running time in practice.

First, the gradient operator for $V$ is relaxed to a convolution kernel in our neural network. And this convolution kernel is shared among each unfolding block. Second, the linear operator $I-\lambda\beta\nabla\nabla^{T}$ ($I$ is the identity matrix) for $\vec{b}^{t}$ in Eq. 14 is relaxed to a convolution kernel. But it is not shared for the network blocks in order to improve the neural network’s capacity. Finally, the $-\lambda\beta\nabla^{T}$ operator in Eq. 15 is relaxed to a convolution kernel. The network architecture is shown in Fig. 3. One block corresponds to one iteration in RS.

In this network, $V$ is the input for our neural network. The convolution of $V$ is shared for all blocks, which is shown in shaded blue. The step symbol indicates the $cut$ activation function, which is implemented by clip operation. The $\otimes$ indicates a convolution kernel that is not shared between blocks. The initial $\vec{b}^{0}=0$ is the same as the iterative algorithm. Different from previous related neural networks, our network does not adopt any Batch Normalization. Therefore, it can be trained on low resolution images but applied on high resolution images.

The proposed RSnet has three important features:

• •

  Our network is unsupervised. It does not require the ground truth images. It is suitable for practical applications where the ground truth is unknown, such as biomedical image reconstruction.

• •

  Our network can be trained on low resolution images but applied on high resolution images. This property is important for training process on hardware with limited memory. Our designed loss function is independent from image resolution and our network does not contain any batch normalization nor drop out layer.

• •

  The result from our network is comparable with the counterpart from classical iterative algorithms. As shown in Section IV-B, we can set the parameter in our loss function for this network such that results from our method are comparable with the counterpart from classical iterative algorithms.

Here, we propose a loss function that is independent from the image resolution. Although Eq. 2 is well-known from the literature, it depends on the resolution of input image because the $\lambda$ is for the whole image rather than for each pixel. To relax this constraint, we propose to use following loss function (normalized version) for our neural network

This loss function is independent from the image resolution. Therefore, this loss function is comparable for different resolution images. In contrast, Eq. 2 is NOT comparable for images with different resolution. When we compare the energy levels between iterative algorithms and our neural network method, we simply scale them by the total number of pixels ($\int_{\vec{x}}1\mathrm{d}\vec{x}$). With this loss function and our network architecture, our neural network can be trained on low resolution images but applied on high resolution images. This property is different from traditional neural networks that involve batch normalization layers or drop out layers. They can only deal with the same resolution images (training and inference images must have the same resolution).

We randomly crop 10,000 image patches with $128\times 128$ resolution from 500 natural images in the BSDS dataset [23]. Each patch is turned into gray scale. Eighty percent of these patches are used as training set and the rest is the validation set. During training, we fix ${\lambda}=10$, all kernel size is set to $3\times 3$, learning rate is 0.0003, batch size is 64, and the epoch number is 3000. We increase the number of block from 3 to 20 and set the number of channels as 8, 16 and 32, respectively.

After training, the final loss function values are shown in Fig. 4. When increasing the number of blocks or the number of channels, the loss is reduced after training. After some value, increasing the them does not improve the result anymore. As comparison, the loss function value for the optimal solutions from FastSolver is 126. This result shows that our neural network can reach the optimal results.

Keeping in mind the complexity of network and the effectiveness, we fix the number of block as 10 and the number of channel as 16 in the rest of our paper. This configuration leads to 21K learnable parameters. Such small number of parameters can be easily trained with a modern GPU hardware.

The loss function values are compared with the counterpart from FastSolver. The distribution of their difference is shown in Fig. 5 (a). We further compare their SSIM difference and the distribution is shown in Fig. 5 (b). RSnet keeps more structural information (Fig. 5 (b)).

Some learned convolution kernels are visualized in Fig. 6. These kernels already contain the classical gradient and divergence operators. Therefore, this network is an enhanced version of our iterative algorithm that only has gradient and divergence operators. As a result, this network has better performance in some applications (Section V-B).

RSnet takes only 0.064 second for each color image with 481$\times$321 resolution from BSDS500 data set on a GTX 1080 Ti GPU card. In contrast, the average running time of RS and FastSolver is 0.89 seconds and 1.07 seconds, respectively.

We have shown the residual solver and its unfolding RSnet that can efficiently solve the bottleneck problem in TV regularized models.

In this section, we show how to embed RSnet into a larger network that can solve the general total variation regularized models (Eq. 2). More specifically, we unfold the dual first iteration algorithm (Eq. 7 to Eq. 9) into a neural network, which contains RSnet as a sub network for Eq. 9. The network structure is shown in Fig. 7.

In image smoothing, we show that RS and RSnet can reach the optimal solution as FastSolver [22]. This experiment also confirms that RSnet can be trained on low resolution images but applied on other resolution images.

In image smoothing, the imaging matrix $A$ is simply the identity matrix. We apply FastSolver, RS or RSnet on 500 color images from BSDS dataset. And the algorithms perform on each color channel separately. The iteration number of FastSolver and RS is set to 200. The RSnet is trained on 10,000 image patches with resolution $128\times 128$, but applied on $321\times 481$ or $481\times 321$ resolution images. Two examples from the data set are shown in Fig. 8. As can be told from Fig. 8, the smoothing results are comparable. To evaluate the quality of these results, structural similarity index measure (SSIM) is used to calculate the similarity between the input and the output images [24].

For these two examples, we also calculate the energy for both methods with first 20 iterations. The results are shown in Fig. 9. RS converges faster than FastSolver at the first few iterations. We observed that this fact is also true for other images from the BSDS500 data set. This is because the operations in RS are independent from the image content.

In this experiment, we show the transferring ability of RSnet. We train our network on the noise free image patches with resolution $128\times 128$, but apply the trained network on noisy images with higher resolution. This indicates that the learned RSnet can be transferred for other tasks. We add noises with Gaussian distribution ($N(0,15^{2})$) on the images from BSDS dataset. We apply FastSolver, RS and RSnet on these noisy images, respectively.

Two examples are shown in Fig. 10. As can be seen in Fig. 10, the results of FastSolver and RS are comparable but the result from RSnet is better. PSNR is calculated to show this quantitatively. There are two reasons for this. First, RSnet is an enhanced version of RS because RSnet has multi channels (16 in our case) while RS only has two channels (gradient operator). Second, RSnet was trained on the noise free images, which make the learned parameters produce similar results.

We evaluate the denoising performance on 500 images of the dataset. In Fig. 11(a) and (b), we show the difference of the converged energy between RS, RSnet and FastSolver, respectively. We show the difference of their PSNR in Fig. 11 (c) and (d), respectively . As can be seen, RS can converge to a smaller energy than FastSolver. Although RSnet gets higher energy, it can reach better PSNR. In this application, RSnet has the best performance due to the two reasons mentioned above.

In previous experiments, we have shown the performance of RS and RSnet on the problems where pixels are only affected by their local neighbors. In this section, we use our RS and RSnet to reconstruct ultrasonic images where each pixel has an influence on all other pixels (global behavior).

The imaging principle of ultrasound is shown in Fig. 12 (a) and the imaging domain is discretized, shown in Fig. 12 (b). The average speed of sound in human tissue $S$ is 1540 meters per second. We denote the contrast speed of sound image as $C(\vec{x})$ and then the speed of sound is $S+C(\vec{x})$. Letting $P(\vec{x})$ denote the sound wave traveling path, we can calculate its traveling time by

Letting $U=\frac{1}{S+C(\vec{x})}-\frac{1}{S}$ be the time delay compared with the average speed $S=1540$, we can get the total time delay

In practice, $\mathrm{d}P$ can be discretized on a mesh grid to form an imaging matrix $A$ (see Fig. 12 (b)). And the time delay $f$ can be measured by an ultrasound device.

To recover the $U$ (and $C$), we minimize following model

where $0<w<1$ is a weight parameter since the uncertainty of each pixel along $x$ and $y$ axes are different [25]. We set $w=0.9$ according to this work.

We use FastSolver, RS and RSnet to solve this problem, respectively. We generate 500 simulated images as ground truth. To generate one image, we first generate a random region number $n$ and then generate $n$ random seeds in 2D. We use Voronoi diagram to obtain $n$ regions from these seeds. Then, we give each region a random number in $[0,255]$. For such image, according to the discretization in Fig. 12 (b), we obtain the corresponding observations.

Our RSnet is trained on the same 10,000 natural image patches with $128\times 128$ resolution (the same data set as described in previous sections), but the loss function is changed to Eq. 23. The learning rate, batch size, epoch, etc. are the same as described in previous sections. This also shows the transfer property of RSnet since the training data set is natural images but the testing data set is piecewise constant biomedical images.

Two examples are shown in Fig. 13, including the ground truth, observation and the reconstruction results from FastSolver, RS and RSnet, respectively. The reconstructed results from FastSolver and RS are almost the same since they reach the same global optimal solution. To evaluate the quality of the results, PSNR and SSIM are calculated and the values are shown in Fig. 13.

The reconstruction results have some issue at left and right boundaries, which is indicated by red arrows. The reason is that the ultrasound device does not have enough samples in the left and right boarder regions. This issue is not caused by solving algorithms but is the limitation from hardware. In practice, such boundary regions should be removed from the reconstructed images and the central region is the main concern.

In this application, RSnet can achieve real time performance, much faster than the other two iterative algorithms because of the small number of blocks. Therefore, RSnet can be used in real-time imaging in our future work.