Unrolled Optimization with Deep Priors

Introduction

In inverse imaging problems, we seek to reconstruct a latent image from measurements taken under a known physical image formation. Such inverse problems arise throughout computational photography, computer vision, medical imaging, and scientific imaging. Residing in the early vision layers of every autonomous vision system, they are essential for all vision based autonomous agents. Recent years have seen tremendous progress in both classical and deep methods for solving inverse problems in imaging. Classical and deep approaches have relative advantages and disadvantages. Classical algorithms, based on formal optimization, exploit knowledge of the image formation model in a principled way, but struggle to incorporate sophisticated learned models of natural images. Deep methods easily learn complex statistics of natural images, but lack a systematic approach to incorporating prior knowledge of the image formation model. What is missing is a general framework for designing deep networks that incorporate prior information, as well as a clear understanding of when prior information is useful.

In this paper we propose unrolled optimization with deep priors (ODP): a principled, general purpose framework for integrating prior knowledge into deep networks. We focus on applications of the framework to inverse problems in imaging. Given an image formation model and a few generic, high-level design choices, the ODP framework provides an easy to train, high performance network architecture. The framework suggests novel network architectures that outperform prior work across a variety of imaging problems.

The ODP framework is based on unrolled optimization, in which we truncate a classical iterative optimization algorithm and interpret it as a deep network. Unrolling optimization has been a common practice among practitioners in imaging, and training unrolled optimization models has recently been explored for various imaging applications, all using variants of field-of-experts priors . We differ from existing approaches in that we propose a general framework for unrolling optimization methods along with deep convolutional prior architectures within the unrolled optimization. By training deep CNN priors within unrolled optimization architectures, instances of ODP outperform state-of-the-art results on a broad variety of inverse imaging problems.

Our empirical results clarify the benefits and limitations of encoding prior information for inverse problems in deep networks. Layers that (approximately) invert the image formation operator are useful because they simplify the reconstruction task to denoising and correcting artifacts introduced by the inversion layers. On the other hand, prior layers improve network generalization, boosting performance on unseen image formation operators. For deblurring and compressed sensing MRI, we found that a single ODP model trained on many image formation operators outperforms existing state-of-the-art methods where a specialized model was trained for each operator.

Moreover, we offer insight into the open question of what iterative algorithm is best for unrolled optimization, given a linear image formation model. Our main finding is that simple primal algorithms that (approximately) invert the image formation operator each iteration perform best.

In summary, our contributions are as follows:

We introduce ODP, a principled, general purpose framework for inverse problems in imaging, which incorporates prior knowledge of the image formation into deep networks.

We demonstrate that instances of the ODP framework for denoising, deblurring, and compressed sensing MRI outperform state-of-the-art results by a large margin.

We present empirically derived insights on how the ODP framework and related approaches are best used, such as when exploiting prior information is advantageous and which optimization algorithms are most suitable for unrolling.

Motivation

The proposed ODP framework is inspired by an extensive body of work on solving inverse problems in imaging via maximum-a-posteriori (MAP) estimation under a Bayesian model. In the Bayesian model, an unknown image $x$ is drawn from a prior distribution $\Omega(\theta)$ with parameters $\theta$. The imaging system applies a linear operator $A$ to this image, representing all optical processes in the capture, and then measures an image $y$ on the sensor, drawn from a noise distribution $\omega(Ax)$ that models sensor noise, e.g., read noise, and noise in the signal itself, e.g., photon-shot noise.

Let $P(y|Ax)$ be the probability of sampling $y$ from $\omega(Ax)$ and $P(x;\theta)$ be the probability of sampling $x$ from $\Omega(\theta)$. Then the probability of an unknown image $x$ yielding an observation $y$ is proportional to $P(y|Ax)P(x;\theta)$.

The MAP point-estimate of $x$ is given by $x=\mathop{\rm argmax}_{x}P(y|Ax)P(x;\theta)$, or equivalently

where the data term $f(y,Ax)=-\log P(y|Ax)$ and prior term $r(x,\theta)=-\log P(x;\theta)$ are negative log-likelihoods. Computing $x$ thus involves solving an optimization problem [3, Chap. 7].

Unrolled iterative methods

A large variety of algorithms have been developed for solving problem (1) efficiently for different convex data terms and priors (e.g., FISTA , Chambolle-Pock , ADMM ). The majority of these algorithms are iterative methods, in which a mapping $\Gamma(x^{k},A,y,\theta)\to x^{k+1}$ is applied repeatedly to generate a series of iterates that converge to solution $x^{\star}$, starting with an initial point $x^{0}$.

Iterative methods are usually terminated based on a stopping condition that ensures theoretical convergence properties. An alternative approach is to execute a pre-determined number of iterations $N$, in other words unrolling the optimization algorithm. This approach is motivated by the fact that for many imaging applications very high accuracy, e.g., convergence below tolerance of $10^{-6}$ for every local pixel state, is not needed in practice, as opposed to optimization problems in, for instance, control. Fixing the number of iterations allows us to view the iterative method as an explicit function $\Gamma^{N}(\cdot,A,y,\theta)\to x^{N}$ of the initial point $x^{0}$. Parameters such as $\theta$ may be fixed across all iterations or vary by iteration. The unrolled iterative algorithm can be interpreted as a deep network .

Parameterization

The parameters $\theta$ in an unrolled iterative algorithm are the algorithm hyperparameters, such as step sizes, and model parameters defining the prior. Generally the number of algorithm hyperparameters is small (1-5 per iteration), so the model capacity of the unrolled algorithm is primarily determined by the representation of the prior.

Many efficient iterative optimization methods do not interact with the prior term $r$ directly, but instead minimize $r$ via its (sub)gradient or proximal operator $\mathbf{prox}_{r(\cdot,\theta)}$, defined as

The proximal operator is a generalization of Euclidean projection. In the ODP framework, we propose to parameterize the gradient or proximal operator of $r$ directly and define $r$ implicitly.

Unrolled optimization with deep priors

We propose the ODP framework to incorporate knowledge of the image formation into deep convolutional networks. The framework factors networks into data steps, which are functions of the measurements $y$ encoding prior information about the image formation model, and CNN steps, which represent statistical image priors. The factorization follows a principled approach inspired by classical optimization methods, thereby combining the best of deep models and classical algorithms.

The ODP framework is summarized by the network template in Algorithm 1. The design choices in the template are the optimization algorithm, which defines the data step $\Gamma$ and algorithm state $z^{k}$, the number of iterations $N$ that the algorithm is unrolled, the function $\phi$ to initialize the algorithm from the measurements $y$, and the CNN used in the prior step, whose output $x^{k+1/2}$ represents either $\nabla r(x^{k},\theta^{k})$ or $\mathbf{prox}_{r(\cdot,\theta^{k})}(x^{k})$, depending on the optimization algorithm. Figure 1 shows an example ODP instance for deblurring under Gaussian noise.

Instances of the ODP framework have two complementary interpretations. From the perspective of classical optimization based methods, an ODP architecture applies a standard optimization algorithm but learns a prior defined by a CNN. From the perspective of deep learning, the network is a CNN with layers tailored to the image formation model.

ODP networks are motivated by minimizing the objective in problem (1), but they are trained to minimize a higher-level loss, which is defined on a metric between the network output and the ground-truth latent image over a training set of image/measurement pairs. Classical metrics for images are mean-squared error, PSNR, or SSIM. Let $\Gamma(y,\theta)$ be the network output given measurements $y$ and parameters $\theta$. Then we train the network by (approximately) solving the optimization problem

Since ODPs are close to conventional CNNs, we can approximately solve problem (2) using the many effective stochastic gradient based methods developed for CNNs (e.g., Adam ). Similarly, we can initialize the portions of $\theta$ corresponding to the CNN prior steps using standard CNN initialization schemes (e.g., Xavier initialization ). The remaining challenge in training ODPs is initializing the portions of $\theta$ corresponding to algorithm parameters in the data step $\Gamma$. Most optimizaton algorithms only have one or two parameters per data step, however, so an effective initialization can be found through standard grid search.

2 Design choices

The ODP framework makes it straightforward to design a state-of-the-art network for solving an inverse problem in imaging. The design choices are the choice of the optimization algorithm to unroll, the CNN parameterization of the prior, and the initialization scheme. In this section we discuss these design choices in detail and present defaults guided by the empirical results in Section 5.

The choice of optimization algorithm to unroll plays an important but poorly understood role in the performance of unrolled optimization networks. The only formal requirement is that each iteration of the unrolled algorithm be almost everywhere differentiable. Prior work has unrolled the proximal gradient method , the half-quadratic splitting (HQS) algorithm , the alternating direction method of multipliers (ADMM) , the Chambolle-Pock algorithm , ISTA , and a primal-dual algorithm with Bregman distances . No clear consensus has emerged as to which methods perform best in general or even for specific problems.

In the context of solving problem (1), we propose the proximal gradient method as a good default choice. The method requires that the proximal operator of $g(x)=f(Ax,y)$ and its Jacobian can be computed efficiently. Algorithm 2 lists the ODP framework for the proximal gradient method. We interpret the CNN prior as $-\alpha_{k}\nabla r(x,\theta^{k})$. Note that for the proximal gradient network, the CNN prior is naturally a residual network because its output $x^{k+1/2}$ is summed with its input $x^{k}$ in Step 4.

The algorithm parameters $\alpha_{0},\ldots,\alpha_{N-1}$ represent the gradient step sizes. The proposed initialization $\alpha_{k}=C_{0}C^{-k}$ is based on an alternate interpretation of Algorithm 2 as an unrolled HQS method. Adopting the aggressively decreasing $\alpha_{k}$ from HQS minimizes the number of iterations needed .

In Section 5.5, we compare deblurring and compressed sensing MRI results for ODP with proximal gradient, ADMM, linearized ADMM (LADMM), and gradient descent. The ODP formulations of ADMM, LADMM, and gradient descent can be found in the supplement. We find that all algorithms that approximately invert the image formation operator each iteration perform on par. Algorithms such as ADMM and LADMM that incorporate Lagrange multipliers were at best slightly better than simple primal algorithms like proximal gradient and gradient descent for the low number of iterations typical for unrolled optimization methods.

CNN prior

The choice of parameterizing each prior step as a separate CNN offers tremendous flexibility, even allowing the learning of a specialized function for each step. Algorithm 2 naturally introduces a residual connection to the CNN prior, so a standard residual CNN is a reasonable default architecture choice. The experiments in Section 5 show this architecture achieves state-of-the-art results, while being easy to train with random initialization.

Choosing a CNN prior presents a trade-off between increasing the number of algorithm iterations $N$, which adds alternation between data and prior steps, and making the CNN deeper. For example, in our experiments we found that for denoising, where the data step is trivial, larger CNN priors with fewer algorithm iterations gave better results, while for deconvolution and MRI, where the data step is a complicated global operation, smaller priors and more iterations gave better results.

Initialization

The initialization function $(x^{0},z^{0})=\phi(f,A,y,\theta^{0})$ could in theory be an arbitrarily complicated algorithm or neural network. We found that the simple initialization $x^{0}=A^{H}y$, which is known as backprojection, was sufficient for our applications [23, Ch. 25].

Related work

The ODP framework generalizes and improves upon previous work on unrolled optimization and deep models for inverse imaging problems.

Field-of-experts

A more sophisticated approach than learned sparsity priors is to parameterize the prior gradient or proximal operator as a field-of-experts (FoE) $g(Cx,\theta)$, where $C$ is again a filterbank and $g$ is a separable nonlinearity parameterized by $\theta$, such as a sum of radial basis functions . The ODP framework improves upon the FoE approaches both empirically, as shown in the model comparisons in Section 5, and theoretically, as the FoE model is essentially a 2-layer CNN and so has less representational power than deeper CNN priors.

Deep models for direct inversion

Several recent approaches propose CNNs that directly solve specific imaging problems. These architectures resemble instances of the ODP framework, though with far different motivations for the design. Schuler et al. propose a network for deblurring that applies a single, fixed deconvolution step followed by a learned CNN, akin to a prior step in ODP with a different initial iterate . Xu et al. propose a network for deblurring that applies a single, learned deconvolution step followed by a CNN, similar to a one iteration ODP network . Wang et al. propose a CNN for MRI whose output is averaged with observations in $k$-space, similar to a one iteration ODP network but without jointly learning the prior and data steps . We improve on these deep models by recognizing the connection to classical optimization based methods via the ODP framework and using the framework to design more powerful, multi-iteration architectures.

Experiments

In this section we present results and analysis of ODP networks for denoising, deblurring, and compressed sensing MRI. Figure 2 shows a qualitative overview for this variety of inverse imaging problems. Please see the supplement for details of the training procedure for each experiment.

We consider the Gaussian denoising problem with image formation $y=x+z$, with $z\sim\mathcal{N}(0,\sigma^{2})$. The corresponding Bayesian estimation problem (1) is

We trained a 4 iteration proximal gradient ODP network with a 10 layer, 64 channel residual CNN prior on the 400 image training set from . Table 2 shows that the ODP network outperforms all state-of-the-art methods on the 68 image test set evaluated in .

2 Deblurring

We consider the problem of joint Gaussian denoising and deblurring, in which the latent image $x$ is convolved with a known blur kernel in addition to being corrupted by Gaussian noise. The image formation model is $y=k*x+z$, where $k$ is the blur kernel and $z\sim\mathcal{N}(0,\sigma^{2})$. The corresponding Bayesian estimation problem (1) is

To demonstrate the benefit of ODP networks specialized to specific problem instances, we first train models per kernel, as proposed in . Specifically, we train 8 iteration proximal gradient ODP networks with residual 5 layer, 64 channel CNN priors for the out-of-focus disk kernel and motion blur kernel from . Following Xu et al., we train one model per kernel on ImageNet , including clipping and JPEG artifacts as required by the authors. Table 2 shows that the ODP networks outperform prior work slightly on the low-pass disk kernel, which completely removes high frequency content, while a substantial gain is achieved for the motion blur kernel, which preserves more frequency content and thus benefits from the inverse image formation steps in ODP.

Next, we show that ODP networks can generalize across image formation models. Table 3 compares ODP networks (same architecture as above) on the test scenarios from . We trained one ODP network for the four out-of-focus kernels and associated noise levels in . The out-of-focus model is on par with of Schuler et al., even though Schuler et al. train a specialized model for each kernel and associated noise level. We trained a second ODP network on randomly generated motion blur kernels. This model outperforms Schuler et al. on the unseen test set motion blur kernel, even though Schuler et al. trained specifically on the motion blur kernel from their test set.

3 Compressed sensing MRI

In compressed sensing (CS) MRI, a latent image $x$ is measured in the Fourier domain with subsampling. Following , we assume noise free measurements. The image formation model is $y=P\mathcal{F}x$, where $\mathcal{F}$ is the DFT and $P$ is a diagonal binary sampling matrix for a given subsampling pattern. The corresponding Bayesian estimation problem (1) is

We trained an 8 iteration proximal gradient ODP network with a residual 7 layer, 64 channel CNN prior on the 100 training images and pseudo-radial sampling pattern from , range from sampling 20% to 50% of the Fourier domain. We evaluate on the 50 test images from the same work.

Table 4 shows that the ODP network outperforms even the the best alternative method, BM3D-MRI, for all sampling patterns. The improvement over BM3D-MRI is larger for sparser (and thus more challenging) sampling patterns. The fact that a single ODP model performed so well on all four sampling patterns, particularly in comparison to the ADMM-Net models, which were trained separately per sampling pattern, again demonstrates that ODP generalizes across image formation models.

4 Contribution of prior information

In this section, we analyze how much of the reconstruction performance is due to incorporating the image formation model in the data step and how much is due to the CNN prior steps. To this end, we performed an ablation study where ODP models were trained without the data steps. The resulting models are pure residual networks.

Table 6 shows that for denoising the residual network performed as well as the proximal gradient ODP network, while for deblurring and CS MRI the proximal gradient network performed substantially better. The difference between denoising and the other two inverse problems is that in the case of denoising the ODP data step is trivial: inverting $A=I$. However, for deblurring the data step applies $B=(A^{H}A+\gamma I)^{-1}A^{H}$, a complicated global operation that for motion blur satisfies $BA\approx I$. Similarly, for CS MRI the ODP proximal gradient network applies $A^{\dagger}$ in the data step, a complicated global operation (involving a DFT) that satisfies $BA\approx I$.

The results suggest an interpretation of the ODP proximal gradient networks as alternating between applying local corrections in the CNN prior step and applying a global operator $B$ such that $BA\approx I$ in the data step. The CNNs in the ODP networks conceptually learn to denoise and correct errors introduced by the approximate inversion of $A$ with $B$, whereas the pure residual networks must learn to denoise and invert $A$ directly. When approximately inverting $A$ is a complicated global operation, direct inversion using residual networks poses an extreme challenge, overcome by the indirect approach taken by the ODP networks.

5 Comparing algorithms

The existing work on unrolled optimization has done little to clarify which optimization algorithms perform best when unrolled. We investigate the relative performance of unrolled optimization algorithms in the context of ODP networks. We have compared ODP networks for ADMM, LADMM, proximal gradient and gradient descent in Table 6, using the same initialization and network parameters as for deblurring and CS MRI.

For deblurring, the proximal gradient and ADMM models, which apply the regularized pseudoinverse $(A^{H}A+\gamma I)^{-1}A^{H}$ in the data step, outperform the LADMM and gradient descent models, which only apply $A$ and $A^{H}$. The results suggest that taking more aggressive steps to approximately invert $A$ in the data step improves performance. For CS MRI, all algorithms apply the pseudoinverse $A^{\dagger}$ (because $A^{H}=A^{\dagger}$) and have similar performance, which matches the observations for deblurring.

Our algorithm comparison shows minimal benefits to incorporating Lagrange multipliers, which is expected for the relatively low number of iterations $N=8$ in our models. The ODP ADMM networks for deblurring and CS MRI are identical to the respective proximal gradient networks except that ADMM includes Lagrange multipliers, and the performance is on par. For deblurring, LADMM and gradient descent are similar architectures, but LADMM incorporates Lagrange multipliers and shows a small performance gain. Note that Gradient descent cannot be applied to CS MRI as problem (5) is constrained.

Conclusion

The proposed ODP framework offers a principled approach to incorporating prior knowledge of the image formation into deep networks for solving inverse problems in imaging, yielding state-of-the-art results for denoising, deblurring, and CS MRI. The framework generalizes and outperforms previous approaches to unrolled optimization and deep networks for direct inversion. The presented ablation studies offer general insights into the benefits of prior information and what algorithms are most suitable for unrolled optimization.

Although the class of imaging problems considered in this work lies at the core of imaging and sensing, it is only a small fraction of the potential applications for ODP. In future work we will explore ODP for blind inverse problems, in which the image formation operator $A$ is not fully known, as well as nonlinear image formation models. Outside of imaging, control is a promising field in which to apply the ODP framework because deep networks may potentially benefit from prior knowledge of physical dynamics.

References