Discriminative Transfer Learning for General Image Restoration

TL;DR

This paper introduces a discriminative transfer learning approach for general image restoration that enables a single trained model to handle multiple tasks and conditions efficiently, with easy transferability to new problems.

Contribution

The proposed method combines proximal optimization with discriminative learning, allowing one-pass training and transferability across various image restoration tasks and conditions.

Findings

Requires only a single training pass

Can be transferred to untrained tasks

Achieves efficiency comparable to existing methods

Abstract

Recently, several discriminative learning approaches have been proposed for effective image restoration, achieving convincing trade-off between image quality and computational efficiency. However, these methods require separate training for each restoration task (e.g., denoising, deblurring, demosaicing) and problem condition (e.g., noise level of input images). This makes it time-consuming and difficult to encompass all tasks and conditions during training. In this paper, we propose a discriminative transfer learning method that incorporates formal proximal optimization and discriminative learning for general image restoration. The method requires a single-pass training and allows for reuse across various problems and conditions while achieving an efficiency comparable to previous discriminative approaches. Furthermore, after being trained, our model can be easily transferred to new likelihood terms to solve untrained tasks, or be combined with existing priors to further improve image restoration quality.

Tables

Table 1. Table 1: Analysis of state-of-the-art methods. In the table, “Transferable” means the model can be used for different restoration tasks and problem conditions; “Modular” means the method can be combined with other existing priors at test time.

	FoE	EPLL	BM3D	TRD	ours
Runtime efficiency			✓	✓	✓
Easy to parallelize				✓	✓
Transferable	✓	✓	✓		✓
Modular	✓	✓	✓		✓

Table 2. Table 2: Average PSNR(dB) on 68 images from  [ 26 ] for denoising.

\rowcolor[HTML]EFEFEF KSVD	FoE	BM3D	LSSC	WNNM	EPLL
30.87	30.99	31.08	31.27	31.37	31.19
\rowcolor[HTML]EFEFEF opt-MRF	ARF	CSF5	TRD5	DTL${}_3{}^{}{}_{}{}^3$	DTL${}_3{}^{}{}_{}{}^5$
31.18	30.70	31.14	31.30	30.91	31.00

Table 3. Table 3: Runtime (seconds) comparison for image denoising.

Image size	${256}^2$	${512}^2$	${1024}^2$	${2048}^2$	${4096}^2$
WNNM	157.73	657.75	2759.79	-	-
EPLL	29.21	111.52	463.71	-	-
BM3D	0.78	3.45	15.24	62.81	275.39
CSF5	1.23	2.22	7.35	27.08	93.66
TRD5	0.39	0.71	2.01	7.57	29.09
DTL${}_3{}^{}{}_{}{}^3$	0.60	1.19	3.45	12.97	56.19
DTL${}_3{}^{}{}_{}{}^3$ (Halide)	0.11	0.26	1.60	5.61	20.85

Table 4. Table 4: Average PSNR (dB) on 32 images from  [ 20 ] for non-blind deconvolution.

\rowcolor[HTML]EFEFEF Input	Levin [19]	Schmidt [30]	CSF${}_3{}^{}{}_{}{}^{\text{pw}}$	DTL${}_3{}^{}{}_{}{}^3$
22.86	32.73	33.97	33.48	34.34

Table 5. Table 5: Test with different HQS iterations ( T 𝑇 T ) and model stages ( K 𝐾 K ) for image denoising. Average PSNR (dB) results on 68 images from  [ 20 ] with noise σ = 15 𝜎 15 \sigma=15 and 25 25 25 are reported (before and after “/” in each cell respectively).

		# HQS iterations
		1	3	5
# stages	1	29.80 / 26.81	30.89 / 28.12	30.96 / 28.28
	3	30.54 / 27.82	30.91 / 28.19	31.00 / 28.42
	5	30.54 / 27.83	30.92 / 28.18	-

Equations

$$\begin{array}{r}\frac{\lambda }{2}||\mathbf{b}-\mathbf{A}\mathbf{x}|{|}_2^2+\sum _{i=1}^N{\varphi }_i({\mathbf{F}}_i\mathbf{x})\end{array}$$

$$\begin{array}{r}\frac{\lambda }{2}||\mathbf{b}-\mathbf{A}\mathbf{x}|{|}_2^2+\frac{\rho }{2}||\mathbf{z}-\mathbf{x}|{|}_2^2+\sum _{i=1}^N{\varphi }_i({\mathbf{F}}_i\mathbf{z}),\end{array}$$

$${\mathbf{z}}^t=\underset{\mathbf{z}}{\mathrm{argmin}}\left(\frac{{\rho }^t}{2}||\mathbf{z}-{\mathbf{x}}^{t-1}|{|}_2^2+\sum _{i=1}^N{\varphi }_i({\mathbf{F}}_i\mathbf{z})\right),$$

$${\mathbf{x}}^t=\underset{\mathbf{x}}{\mathrm{argmin}}\left(\lambda ||\mathbf{b}-\mathbf{A}\mathbf{x}|{|}_2^2+{\rho }^t||{\mathbf{z}}^t-\mathbf{x}|{|}_2^2\right),$$

$$\begin{array}{r}{\mathbf{z}}^t:={\mathbf{prox}}_{\Theta }({\mathbf{x}}^{t-1},{\rho }^t),\end{array}$$

$$\begin{array}{rl}& {\mathbf{z}}_k^t={\mathbf{z}}_{k-1}^t-\sum _{i=1}^N{{\mathbf{F}}_i^k}^{\mathsf{T}}{\psi }_i^k({\mathbf{F}}_i^k{\mathbf{z}}_{k-1}^t),\\ & s.t.{\mathbf{z}}_0^t={\mathbf{x}}^{t-1},k=1,2,...,K,\end{array}$$

$$\begin{array}{r}{\mathbf{z}}^t:={\mathbf{prox}}_{\Theta }({\mathbf{x}}^{t-1})\end{array}$$

$$\begin{array}{r}{\mathbf{x}}^t=(\lambda \mathbf{b}+{\rho }^t{\mathbf{z}}^t)/(\lambda +{\rho }^t)\end{array}$$

$$\begin{array}{r}{\mathbf{x}}^t={\mathcal{F}}^{-1}\left(\frac{\mathcal{F}(\lambda {\mathbf{A}}^{\mathsf{T}}\mathbf{b}+{\rho }^t{\mathbf{z}}^t)}{\mathcal{F}(\lambda {\mathbf{A}}^{\mathsf{T}}\mathbf{A}+{\rho }^t)}\right),\end{array}$$

$$\begin{array}{r}\frac{\partial \ell }{\partial \Omega }=\sum _{t=1}^T\left(\frac{\partial {\mathbf{x}}^t}{\partial \lambda }\frac{\partial \ell }{\partial {\mathbf{x}}^t}+\frac{\partial {\mathbf{z}}^t}{\partial \Theta }\frac{\partial {\mathbf{x}}^t}{\partial {\mathbf{z}}^t}\frac{\partial \ell }{\partial {\mathbf{x}}^t}\right).\end{array}$$

$$\begin{array}{r}\frac{\lambda }{2}||\mathbf{b}-\mathbf{A}\mathbf{x}|{|}_2^2+\sum _{i=1}^N{\varphi }_i({\mathbf{F}}_i\mathbf{x})+\tau \mathcal{S}(\mathbf{x})\end{array}$$

$$\begin{array}{rl}\frac{\lambda }{2}||\mathbf{b}-\mathbf{A}\mathbf{x}|{|}_2^2& +\frac{\rho }{2}||\mathbf{z}-\mathbf{x}|{|}_2^2+\sum _{i=1}^N{\varphi }_i({\mathbf{F}}_i\mathbf{z})\\ & +\frac{{\rho }_s}{2}||\mathbf{v}-\mathbf{x}|{|}_2^2+\tau \mathcal{S}(\mathbf{v})\end{array}$$

$${\mathbf{z}}^t={\mathbf{prox}}_{\Theta }({\mathbf{x}}^{t-1})$$

$${\mathbf{v}}^t=\underset{\mathbf{v}}{\mathrm{argmin}}\frac{{\rho }_s^t}{2}||\mathbf{v}-{\mathbf{x}}^{t-1}|{|}_2^2+\tau \mathcal{S}(\mathbf{v})\approx \text{BM3D}({\mathbf{x}}^{t-1},\frac{\tau }{{\rho }_s^t})$$

$${\mathbf{x}}^t=\underset{\mathbf{x}}{\mathrm{argmin}}\lambda ||\mathbf{b}-\mathbf{A}\mathbf{x}|{|}_2^2+{\rho }^t||{\mathbf{z}}^t-\mathbf{x}|{|}_2^2+{\rho }_s^t||{\mathbf{v}}^t-\mathbf{x}|{|}_2^2,$$

$$\begin{array}{rl}& {\mathbf{z}}^t={\mathbf{prox}}_{\Theta }({\mathbf{x}}^{t-1}),\\ & {\mathbf{x}}^t=(\lambda {\mathbf{A}}^{\mathsf{T}}\mathbf{b}+{\rho }^t{\mathbf{z}}^t)/(\lambda \mathbf{a}+{\rho }^t)\end{array}$$

Full text

1 Introduction

Low-level vision problems, such as denoising, deconvolution and demosaicing, have to be addressed as part of most imaging and vision systems. Although a large body of work covers these classical problems, low-level vision is still a very active area. The reason is that, from a Bayesian perspective, solving them as statistical estimation problems does not only rely on models for the likelihood (i.e. the reconstruction task), but also on natural image priors as a key component.

A variety of models for natural image statistics have been explored in the past. Traditionally, models for gradient statistics [27, 17], including total-variation, have been a popular choice. Another line of works explores patch-based image statistics, either as per-patch sparse model [11, 35] or modeling non-local similarity between patches [9, 10, 13]. These prior models are general in the sense that they can be applied for various likelihoods, with the image formation and noise setting as parameters. However, the resulting optimization problems are prohibitively expensive, rendering them impractical for many real-time tasks especially on mobile platforms.

Recently, a number of works [29, 8] have addressed this issue by truncating the iterative optimization and learning discriminative image priors, tailored to a specific reconstruction task (likelihood) and optimization approach. While these methods allow to trade-off quality with the computational budget for a given application, the learned models are highly specialized to the image formation model and noise parameters, in contrast to optimization-based approaches. Since each individual problem instantiation requires costly learning and storing of the model coefficients, current proposals for learned models are impractical for vision applications with dynamically changing (often continuous) parameters. This is a common scenario in most real-world vision settings, as well as applications in engineering and scientific imaging that rely on the ability to rapidly prototype methods.

In this paper, we combine discriminative learning techniques with formal proximal optimization methods to learn generic models that can be truly transferred across problem domains while achieving comparable efficiency as previous discriminative approaches. Using proximal optimization methods [12, 23, 3] allows us to decouple the likelihood and prior which is key to learn such shared models. It also means that we can rely on well-researched physically-motivated models for the likelihood, while learning priors from example data. We verify our technique using the same model for a variety of diverse low-level image reconstruction tasks and problem conditions, demonstrating the effectiveness and versatility of our approach. After training, our approach benefits from the proximal splitting techniques, and can be naturally transferred to new likelihood terms for untrained restoration tasks, or it can be combined with existing state-of-the-art priors to further improve the reconstruction quality. This is impossible with previous discriminative methods. In particular, we make the following contributions:

• •

We propose a discriminative transfer learning technique for general image restoration. It requires a single-pass training and transfers across different restoration tasks and problem conditions.

• •

We show that our approach is general by demonstrating its robustness for diverse low-level problems, such as denoising, deconvolution, inpainting, and for varying noise settings.

• •

We show that, while being general, our method achieves comparable computational efficiency as previous discriminative approaches, making it suitable for processing high-resolution images on mobile imaging systems.

• •

We show that our method can naturally be combined with existing likelihood terms and priors after being trained. This allows our method to process untrained restoration tasks and take advantage of previous successful work on image priors (e.g., color and non-local similarity priors).

2 Related work

Image restoration aims at computationally enhancing the quality of images by undoing the adverse effects of image degradation such as noise and blur. As a key area of image and signal processing it is an extremely well studied problem and a plethora of methods exists, see for example [22] for a recent survey. Through the successful application of machine learning and data-driven approaches, image restoration has seen revived interest and much progress in recent years. Broadly speaking, recently proposed methods can be grouped into three classes: classical approaches that make no explicit use of machine learning, generative approaches that aim at probabilistic models of undegraded natural images and discriminative approaches that try to learn a direct mapping from degraded to clean images. Unlike classical methods, methods belonging to the latter two classes depend on the availability of training data.

Classical models focus on local image statistics and aim at maintaining edges. Examples include total variation [27], bilateral filtering [32] and anisotropic diffusion models [34]. More recent methods exploit the non-local statistics of images [1, 9, 21, 10, 13, 31]. In particular the highly successful BM3D method [9] searches for similar patches within the same image and combines them through a collaborative filtering step.

Generative learning models seek to learn probabilistic models of undegraded natural images. A simple, yet powerful subclass include models that approximate the sparse gradient distribution of natural images [19, 17, 18]. More expressive generative models include the fields of experts (FoE) model [26], KSVD [11] and the EPLL model [35]. While both FoE and KVSD learn a set of filters whose responses are assumed to be sparse, EPLL models natural images through Gaussian Mixture Models. All of these models have in common that they are agnostic to the image restoration task, i.e. they are transferable to any image degradation and can be combined in a modular fashion with any likelihood and additional priors at test time.

Discriminative learning models have recently become increasingly popular for image restoration due to their attractive tradeoff between high image restoration quality and efficiency at test time. Methods include trainable random field models such as cascaded shrinkage fields (CSF) [29], regression tree fields (RTF) [16], trainable nonlinear reaction diffusion (TRD) models [8], as well as deep convolutional networks [15] and other multi-layer perceptrons [4].

Discriminative approaches owe their computational efficiency at run-time to a particular feed-forward structure whose trainable parameters are optimized for a particular task during training. Those learned parameters are then kept fixed at test-time resulting in a fixed computational cost. On the downside, discriminative models do not generalize across tasks and typically necessitate separate feed-forward architectures and separate training for each restoration task (denoising, demosaicing, deblurring, etc.) as well as every possible image degradation (noise level, Bayer pattern, blur kernel, etc.).

In this work, we propose the discriminative transfer learning technique that is able to combine the strengths of both generative and discriminative models: it maintains the flexibility of generative models, but at the same time enjoys the computational efficiency of discriminative models. While in spirit our approach is akin to the recently proposed method of Rosenbaum and Weiss [25], who equipped the successful EPLL model with a discriminative prediction step, the key idea in our approach is to use proximal optimization techniques [12, 23, 3] that allow the decoupling of likelihood and prior and therewith share the full advantages of a Bayesian generative modeling approach.

Table 1 summarizes the properties of the most prominent state-of-the-art methods and puts our own proposed approach into perspective.

3 Proposed method

3.1 Diversity of data likelihood

The seminal work of fields-of-experts (FoE) [26] generalizes the form of filter response based regularizers in the objective function given in Eq. 1. The vectors ${\bf b}$ and ${\bf x}$ represent the observed and latent (desired) image respectively, the matrix ${\bf A}$ is the sensing operator, ${\bf F}_{i}$ represents 2D convolution with filter ${\bf f}_{i}$, and $\phi_{i}$ represents the penalty function on corresponding filter responses ${\bf F}_{i}{\bf x}$. The positive scalar $\lambda$ controls the relative weight between the data fidelity (likelihood) and the regularization term.

[TABLE]

The well-known anisotropic total-variation regularizer can be viewed as a special case of the FoE model where ${\bf f}_{i}$ is the derivative operator $\nabla$, and $\phi_{i}$ the $\ell_{1}$ norm.

While there are various types of restoration tasks (e.g., denoising, deblurring, demosaicing) and problem parameters (e.g., noise level of input images), each problem has its own sensing matrix ${\bf A}$ and optimal fidelity weight $\lambda$. For example, ${\bf A}$ is an identity matrix for denoising, a convolution operator for deblurring, a binary diagonal matrix for demosaicing, and a random matrix for compressive sensing [5]. $\lambda$ depends on both the task and its parameters in order to produce the best quality results.

The state-of-the-art discriminative learning methods (CSF[29], TRD[8]) derive an end-to-end feed-forward model from Eq. 1 for each specific restoration task, and train this model to map the degraded input images directly to the output. These methods have demonstrated a great trade-off between high-quality and time-efficiency, however, as an inherent problem of the discriminative learning procedure, they require separate training for each restoration task and problem condition. Given the diversity of data likelihood of image restoration, this fundamental drawback of discriminative models makes it time-consuming and difficult to encompass all tasks and conditions during training.

3.2 Decoupling likelihood and prior

It is difficult to directly minimize Eq. 1 when the penalty function $\phi_{i}$ is non-linear and/or non-smooth (e.g., $\ell_{p}$ norm, $0<p\leq 1$). Proximal algorithms [3, 12, 6] instead relax Eq. 1 and split the original problem into several easier subproblems that are solved alternately until convergence.

In this paper we employ the half-quadratic-splitting (HQS) algorithm [12] to relax Eq. 1, as it typically requires much fewer iterations to converge compared with other proximal methods such as ADMM [3] and PD [6]. The relaxed objective function is given in Eq. 2:

[TABLE]

where a slack variable ${\bf z}$ is introduced to approximate ${\bf x}$, and $\rho$ is a positive scalar.

With the HQS algorithm, Eq. 2 is iteratively minimized by solving for the slack variable ${\bf z}$ and the latent image ${\bf x}$ alternately as in Eq. 3 and 4 ($t=1,2,...,T$).

[TABLE]

where $\rho^{t}$ increases as the iteration continues. This forces ${\bf z}$ to become an increasingly good approximation of ${\bf x}$, thus making Eq. 2 an increasingly good proxy for Eq. 1.

Note that, while most related approaches including CSF [29] relax Eq. 1 by splitting on ${\bf F}_{i}{\bf x}$, we split on ${\bf x}$ instead. This is critical for deriving our approach. With this new splitting strategy, the prior term and the data likelihood term in the original objective Eq. 1 are now separated into two subproblems that we call the “prior proximal operator” (Eq. 3) and the “data proximal operator” (Eq. 4), respectively.

3.3 Discriminative transfer learning

We observed that, while the data proximal operator in Eq. 4 is task-dependent because both the sensing matrix ${\bf A}$ and fidelity weight $\lambda$ are problem-specific as explained in Sec. 3.1, the prior proximal-operator (i.e. ${\bf z}^{t}$-update step in Eq. 3) is independent of the original restoration tasks and problem conditions.

This leads to our main insight: Discriminative learned models can be made transferable by using them in place of the prior proximal operator, embedded in a proximal optimization algorithm. This allows us to generalize a single discriminative learned model to a very large class of problems, i.e. any linear inverse imaging problem, while simultaneously overcoming the need for problem-specific retraining. Moreover, it enables learning the task-dependent parameter $\lambda$ in the data proximal operator for each problem in a single training pass, eliminating tedious hand-tuning at test time.

We also observed that, benefiting from our new splitting strategy, the prior proximal operator in Eq. 3 can be interpreted as a Gaussian denoiser on the intermediate image ${\bf x}^{t-1}$, since the least-squares consensus term is equivalent to a Gaussian denoising term. This inspires us to utilize existing discriminative models that have been successfully used for denoising (e.g. CSF, TRD).

For convenience, we denote the prior proximal operator as ${\bf prox}_{\Theta}$, i.e.

[TABLE]

where the model parameter $\Theta$ includes a number of filters ${\bf f}_{i}$ and corresponding penalty functions $\phi_{i}$. Inspired by the state-of-the-art discriminative methods [29, 8], we propose to learn the model ${\bf prox}_{\Theta}$, and the fidelity weight scalar $\lambda$, from training data. Recall that with our new splitting strategy introduced in Sec. 3.2, the image prior and data-fidelity term in the original objective (Eq. 1) are contained in two separate subproblems (Eq. 3 and 4). This makes it possible to train together an ensemble of diverse tasks (e.g., denoising, deblurring, or with different noise levels) each of which has its own data proximal operator, while learning a single prior proximal operator ${\bf prox}_{\Theta}$ that is shared across tasks. This is in contrast to state-of-the-art discriminative methods such as CSF [29] and TRD [8] which train separate models for each task.

For clarity, in Fig. 1 we visualize the architecture of our method. The input images may represent various restoration tasks and problem conditions. At each HQS iteration, each image ${\bf x}^{t}_{p}$ from problem $p$ is updated by its own data proximal operator in Eq. 4 which contains separate trainable fidelity weight $\lambda_{p}$ and pre-defined sensing matrix ${\bf A}_{p}$; then each slack image ${\bf z}^{t}_{p}$ is updated by the same, shared prior proximal operator implemented by a learned, discriminative model.

Recurrent network. Note that in Fig. 1 each HQS iteration uses exactly the same model parameters, forming a recurrent network. This is in contrast to previous discriminative learning methods including CSF and TRD, which form feed-forward networks. Our recurrent network architecture maintains the convergence property of the proximal optimization algorithm (HQS), and is critical for our method to transfer between various tasks and problem conditions.

Shared prior proximal operator. While any discriminative Gaussian denoising model could be used as ${\bf prox}_{\Theta}$ in our framework, we specifically propose to use the multi-stage non-linear diffusion process that is modified from the TRD [8] model, for its efficiency. The model is given in Eq. LABEL:eq:tnrd_denoise.

[TABLE]

where $k$ is the stage index, filters ${\bf F}^{k}_{i}$, function $\psi^{k}_{i}$ are trainable model parameters at each stage, and ${\bf z}^{t}_{0}$ is the initial value of ${\bf z}^{t}_{k}$. Note that, different from TRD, our model does not contain the reaction term which would be $-\rho^{t}\alpha_{k}({\bf z}^{t}_{k-1}-{\bf x}^{t-1})$ with step size $\alpha_{k}$. The main reasons for this modification are:

• •

The data constraint is contained in ${\bf x}^{t}$ update in Eq. 4;

• •

More importantly, by dropping the reaction term our model gets rid of the weight $\rho^{t}$ which changes at each HQS iteration. Therefore, our proximal operator ${\bf prox}_{\Theta}({\bf x}^{t-1},\rho^{t})$ is simplified to be:

[TABLE]

The parameter $\Omega$ to learn in our method includes $\lambda$’s for each problem class $p$ (restoration task and problem condition), and $\Theta=\{{\bf F}^{k}_{i},\psi^{k}_{i}\}$ in the prior proximal operator shared across different classes, i.e. $\Omega=\{\lambda_{p},\Theta\}$. Even though the scalar parameters $\lambda_{p}$ are trained, our method allows users to override them at test time to handle non-trained problem classes or specific inputs as we will show in Sec. 4. This contrasts to previous discriminative approaches whose model parameters are all fixed at test time. The subscript $p$ indicating the problem class in $\lambda_{p}$ is omitted below for convenience. The values of $\rho^{t}$ are pre-selected: $\rho^{1}=1$ and $\rho^{t}=2\rho^{t-1}$ for $t>1$.

Note that a multi-stage model as in Eq. LABEL:eq:tnrd_denoise is not possible if we split on ${\bf F}_{i}{\bf x}$ instead of ${\bf x}$ in Eq. 1 and 2. For clarity, an overview of the proposed algorithm is given in Algorithm 1.

3.4 Training

We consider denoising and deconvolution tasks at training, where the sensing operator ${\bf A}$ is an identity matrix, or a block circulant matrix with circulant blocks that represents 2D convolution with randomly drawn blur kernels respectively. In denoising tasks, the ${\bf x}^{t}$ update in Eq. 4 has a closed-form solution:

[TABLE]

In deconvolution tasks, the ${\bf x}^{t}$ update in Eq. 4 has a closed-form solution in the Fourier domain:

[TABLE]

where $\mathcal{F}$ and $\mathcal{F}^{-1}$ represent Fourier and inverse Fourier transform respectively. Note that, compared to CSF [29], our method does not require FFT computations for denoising tasks. We use the L-BFGS solver [28] with analytic gradient computation for training. The training loss function $\ell$ is defined as the negative average Peak Signal-to-Noise Ratio (PSNR) of reconstructed images. The gradient of $\ell$ w.r.t. the model parameters $\Omega=\{\lambda_{p},\Theta\}$ is computed by accumulating gradients at all HQS iterations, i.e.

[TABLE]

The 1D functions $\psi^{k}_{i}$ in Eq. LABEL:eq:tnrd_denoise are parameterized as a linear combination of equidistant-positioned Gaussian kernels whose weights are trainable.

Progressive training. A progressive scheme is proposed to make the training more effective. First, we set the number of HQS iterations to be 1, and train $\lambda^{\prime}s$ and the model $\Theta$ of each stage in ${\bf prox}_{\Theta}$ in a greedy fashion. Then, we gradually increase the number of HQS iterations from 1 to $T$ where at each step the model $\Omega=\{\lambda,\Theta\}$ is refined from the result of the previous step. The L-BFGS iterations are set to be 200 for the greedy training steps, and 100 for the refining steps. Fig. 2 shows examples of learned filters in ${\bf prox}_{\Theta}$.

4 Results

Denoising and generality analysis. We compare the proposed discriminative transfer learning (DTL) method with state-of-the-art image denoising techniques, including KSVD [11], FoE [26], BM3D [9], LSSC [21], WNNM [13], EPLL [35], opt-MRF [7], ARF [2], CSF [29] and TRD [8]. The subscript in CSF5 and TRD5 indicates the number of cascaded stages (each stage has different model parameters). The subscript and superscript in our method DTL${}^{5}_{3}$ indicate the number of diffusion stages ($K=3$ in Algorithm 1) in our proximal operator ${\bf prox}_{\Theta}$, and the number of HQS iterations ($T=5$ in Alg. 1), respectively. Note that the complexity (size) of our model is linear in $K$, but independent of $T$. CSF, TRD and DTL use 24 filters of size 5$\times$5 pixels at all stages in this section.

The compared discriminative methods, CSF5 and TRD5 both are trained at single noise level $\sigma=15$ that is the same as the test images. In contrast, our model is trained on 400 images (100$\times$100 pixels) cropped from [26] with random and discrete noise levels (standard deviation $\sigma$) varying between 5 and 25. The images with the same noise level share the same data fidelity weight $\lambda$ at training.

To verify the generality of our method on varying noise levels, we test our model DTL${}^{3}_{3}$ (trained with varying noise levels in a single pass) and two TRD models (trained at specific noise levels 15 and 25) on 3 sets of 68 images with noise $\sigma=5,15,25$ respectively. The average PSNR values are shown in Fig. 3. Although performing slightly below the TRD model trained for the exact noise level used at test time, our method is more generic and works robustly for various noise levels. The performance of the discriminative TRD method drops down quickly as the problem condition (i.e. noise level) at test differs from its training data. In sharp contrast to discriminative methods (CSF, TRD, etc), which are inherently specialized for a given problem setting, i.e. noise level, the proposed approach transfers across different problem settings. More analysis can be found in the supplementary material.

All compared methods are evaluated on the 68 test images from [26] and the averaged PSNR values are reported in Table 2. The compared discriminative methods (CSF, TRD, etc) were trained for exactly the same noise level as the test images (i.e. the best case for them), while our model was trained with mixed noise levels and works robustly for arbitrary noise levels. Our results are comparable to generic methods such as KSVD, FoE and BM3D, and very close to discriminative methods such as CSF5, while at the same time being much more time-efficient.

Run-time comparison. In Table 3 we compare the run-time of our method and state-of-the-art methods. The experiments were performed on a laptop computer with Intel i7-4720HQ CPU and 16GB RAM. WNNM and EPLL ran out-of-memory for images over 4 megapixels in our experiments. CSF5, TRD5 and DTL${}^{3}_{3}$ all use “parfor” setting in Matlab. DTL${}^{3}_{3}$ is significantly faster than all compared generic methods (WNNM, EPLL, BM3D) and even the discriminative method CSF5. Run-time of DTL${}^{3}_{3}$ is about 1.5 times that of TRD5, which is expected as they use 5 versus 9 diffusion steps in total. In addition, we implement our method in Halide language [24], which has become popular recently for high-performance image processing applications, and report the run-time on the same CPU as mentioned above.

Deconvolution. In this experiment, we train a model with an ensemble of denoising and deconvolution tasks on 400 images (100$\times$100 pixels) cropped from [26], in which 250 images are generated for denoising tasks with random noise levels $\sigma$ varying between 5 and 25, and the other 150 images are generated by blurring the images with random 25$\times$25 kernels (PSFs) and then adding Gaussian noise with $\sigma$ ranging between 1 and 5. All images are quantized to 8 bits.

We compare our method with state-of-the-art non-blind deconvolution methods including Levin et al. [19], Schmidt et al. [30] and CSF [29]. Note that TRD [8] does not support non-blind deconvolution. We test the methods on the benchmark dataset from [20] which contains 32 real-captured images and report the average PSNR values in Table 4. The results of compared methods are quoted from [29].

As said in Sec. 3.3, while the scalar weight $\lambda$ is trained, our method allows users to override it at test time for untrained problem classes or specific inputs. Fig. 4 shows our results with different $\lambda$ on the experiments compared in Table 4. Within a fairly wide range of $\lambda$, our method outperforms all previous methods.

We further test the above model trained with ensemble tasks on the denoising experiment in Table 2. The result average PSNR is 30.98dB, which is comparable to the result with the model trained only on the denoising task.

Modularity with existing priors. As shown above, even though the fidelity weight $\lambda$ is trainable, our method allows users to override its value at test time. This property also makes it possible to combine our model (after being trained) with existing state-of-the-art priors at test time, in which case $\lambda$ typically needs to be adjusted. This allows our method to take advantage of previous successful work on image priors. Again, this is not possible with previous discriminative methods (CSF, TRD).

In Fig. 5 we show an example to incorporate a non-local patch similarity prior (BM3D [9]) with our method to further improve the denoising quality. BM3D performs well in removing noise especially in smooth regions but usually over-smoothes edges and textures. Our original model (DTL${}^{5}_{3}$) well preserves sharp edges however sometimes introduces artifacts in smooth regions when the input noise level is high. By combining those two methods, which is easy with our HQS framework, the result is improved both visually and quantitatively.

We give the derivation of the proposed hybrid method below. Let $\mathcal{S}({\bf x})$ represents the non-local patch similarity prior. The objective function is:

[TABLE]

Applying the HQS technique described in Sec. 3, we relax the objective to be:

[TABLE]

Then we minimize Eq. 12 by alternately solving the following 3 subproblems:

[TABLE]

where ${\bf prox}_{\Theta}$ is from our previous training, and the ${\bf v}^{t}$ subproblem is approximated by running BM3D software on ${\bf x}^{t-1}$ with noise parameter $\frac{\tau}{\rho_{s}^{t}}$ following [33, 14].

Similarly, our method can incorporate color image priors (e.g., cross-channel edge-concurrence prior [14]) to improve test results on color images, despite our model being trained on gray-scale images. An example is shown in Fig. 6. The hybrid method shares the advantages of our original model that effectively preserves edges and textures and the cross-channel prior that reduces color artifacts.

Transferability to unseen tasks. Our method allows for new data-fidelity terms that are not contained in training, with no need for re-training. We demonstrate this flexibility with an experiment on the joint denoising and inpainting task shown in Fig. 9. In this experiment, 60% pixels of the input image are missing, and the measured 40% pixels are corrupted with Gaussian noise with $\sigma=15$. Let vector ${\bf a}$ be the binary mask for measured pixels. The sensing matrix ${\bf A}$ in Eq. 1, assumed to be known, is a binary diagonal matrix (hence ${\bf A}={\bf A}^{\mathsf{T}}={\bf A}^{\mathsf{T}}{\bf A}$) with diagonal elements ${\bf a}$. To reuse our model trained on denoising/deconvolution tasks, we only need to specify ${\bf A}$ and $\lambda$. The subproblems of our HQS framework are given in Eq. 14.

[TABLE]

Analysis of convergence and model complexity. To better understand the convergence of our method, in Fig. 7 and 8 we show the results of each HQS iteration of our method on denoising and non-blind deconvolution.

To understand the effect of model complexity and the number of HQS iteration on results, in Table 5 we report test results of our method using models trained with different HQS iterations ($T$ in Algorithm 1), and with different stages in ${\bf prox}_{\Theta}$ ($K$ in Algorithm 1).

5 Conclusion

In this paper, we proposed the discriminative transfer learning framework for general image restoration. By combining advanced proximal optimization algorithms and discriminative learning techniques, a single training pass leads to a transferable model useful for a variety of image restoration tasks and problem conditions. Furthermore, our method is flexible and can be combined with existing priors and likelihood terms after being trained, allowing us to improve image quality on a task at hand. In spite of this generality, our method achieves comparable run-time efficiency as previous discriminative approaches, making it suitable for high-resolution image restoration and mobile vision applications.

We believe that in future work, our framework incorporating advanced optimization with discriminative learning techniques can be extended to deep learning, for training more compact and shareable models, and to solve high-level vision problems.