Inverse problems are typically solved by first defining a model and then choosing an inference procedure. With this separation of modeling from inference, inverse problems can be framed in a modular way. For example, variational inference can be applied to a broad class of models. The modularity, however, typically goes away after model parameters have been trained under a chosen inference procedure. During training, model and inference often interact in a way that the model parameters will ultimately be adapted to the chosen inference procedure, posing the two components inseparable after training. But if model and inference become inseperable after training, why separate them in the first place? We propose a novel learning framework which abandons the dichotomy between model and inference. Instead, we introduce Recurrent Inference Machines (RIM), a class of recurrent neural networks (RNN), that directly learn to solve inverse problems. We demonstrate the effectiveness of RIMs in experiments on various image reconstruction tasks. We show empirically that RIMs exhibit the desirable convergence behavior of classical inference procedures, and that they can outperform state-ofthe-art methods when trained on specialized inference tasks. Our approach bridges the gap between inverse problems and deep learning, providing a framework for fast progression in the field of inverse problems.
1 INTRODUCTION
Inverse Problems are a broad class of problems which can be encountered in all scientific disciplines, from the natural sciences to engineering. The task in inverse problems is to reconstruct a signal from observations that are subject to a known (or inferred) corruption process known as the forward model. A typical example of an inverse problem is the linear measurement problem
y = Ax + n, (1)
where x is the signal of interest, A is an m × d corruption matrix, n is an additive noise vector, and y is the actual measurement. If A is a wide matrix such that m d, this problem is typically ill-posed. Many signal reconstruction problems can be phrased in terms of the linear measurement problem such as image denoising, super-resolution, deconvolution and so on. The general form of A typically defines the problem class. If A is an identity matrix the problem is a denoising problem, while in tomography A represents a Fourier transform and a consecutive sub-sampling of the Fourier coefficients.
Inverse problems are often formulated as an optimization problem of the form
min x d(y,Ax) + λR(x), (2)
where d(y,Ax) is the data fidelity term that enforces x to satisfy the observations y, and R(x) is a regularization term which restricts the solution to comply with a predefined model over x.
The difficulties that arise in this framework are two-fold: (1) it is difficult to choose R(x) such that it is an appropriate model for complex signals such as natural images, and (2) even under a well chosen R(x) the optimization procedure might become difficult.
Compressed sensing approaches give up on a versatileR(x) in order to define a convex optimization procedure. The idea is that the signal x has a sparse representation in some basis Ψ such that x = Ψu and that the optimization problem can be rephrased as
min u d(y,AΨu) + λ ‖u‖1 , (3)
where ‖·‖1 is the sparsity inducing L1-norm (Donoho, 2006a). Under certain classes of d(y,AΨu) such as quadratic errors the optimization problem becomes convex. Results from the compressed sensing literature offer provable bounds on the reconstruction performance for sparse signals of this form (Candès et al., 2006; Donoho, 2006b). The basis Ψ can also be learned from data (Aharon et al., 2006; Elad & Aharon, 2006).
Other approaches interpret equation (2) in terms of probabilities such that finding the solution is a matter of performing maximum a posteriori (MAP) estimation (Figueiredo et al., 2007). In those cases d(y,AΨu) takes the form of a log-likelihood and R(x) takes the form of a parametric logprior log pθ(x) over variable x such that the minimization becomes:
max x
log p(y|A,x) + log pθ(x). (4)
This allows for more expressiveness of R(x) and for the possibility of learning the prior pθ(x) from data. However, with more expressive priors optimization will become more difficult as well. In fact, only for a few trivial prior-likelihood pairs will inference remain convex. In practice one often has to resort to approximations of the objective and to approximate double-loop algorithms in order to allow for scalable inference (Nickisch & Seeger, 2009; Zoran & Weiss, 2011).
In this work we take a radically different approach to solving inverse problems. We move away from the idea that it is beneficial to separate learning a prior (regularizer) from the optimization to do the reconstruction. The usual thinking is that this separation allows for greater modularity and the possibility to interchange one of these two complementary components in order to build new algorithms. In practice however, we observe that the optimization procedure almost always has to be adapted to the model choice to achieve good performance (Aharon et al., 2006; Elad & Aharon, 2006; Nickisch & Seeger, 2009; Zoran & Weiss, 2011). In fact, it is well known that the optimization procedure used for training should match the one used during testing because the model has adapted itself to perform well under that optimization procedure (Kumar et al., 2005; Wainwright, 2006).
What we need is a single framework which allows us to backpropagate through the optimization procedure when we learn the free parameters. Hence, We propose to look at inverse problems as a direct mapping from observations to estimated signal,
x̂ = fφ(A,y) (5)
where x̂ is an estimate of signal x from observations (A,y). Here we define φ as a set of learnable parameters which define the inference algorithm as well as constraints on x. The goal is thus to define map whose parameters are directly optimized for solving the inverse problem itself. It has the benefits of both having high expressive power (if the map fφ is complex enough) as well as being fast at inference time.
This paradigm shift allows us to learn and combine the effect of a prior, the reconstruction fidelity and an inference method without the need to explicitly define the functional form of all components. The whole procedure is simply interpreted as a single RNN. As a result, there is no need for sparsity assumptions, the introduction of model constraints to allow for convexity, or even for double-loop algorithms (Gregor & LeCun, 2010). In fact the proposed framework allows for use of current deep learning approaches which have high expressive power without trading off scalability. It further allows us to move all the manual parameter tuning - which is still common in traditional approaches (Zoran & Weiss, 2011) - away from the inference phase and into the learning phase. We believe this framework can be an important asset to introduce deep learning into the domain of inverse problems.
2 RECURRENT INFERENCE MACHINES
The goal of this work is to find an inverse model as described in equation (5). Often, however, it will be intractable to find (5) directly, even with modern non-linear function approximators. For high-dimensional y and x, which are typically considered in inverse problems, it will simply not be possible to fit matrix A into memory explicitly, but instead matrix A will be replaced by an operator that acts on x. An example is the Discrete Fourier Transform (DFT). Instead of using a Fourier matrix which is quadratic in the size of x, DFTs are typically performed using the Fast Fourier Transform (FFT) algorithm which reduces computational cost and memory consumption significantly. The use of operators, however, does not allow us to feed A into (5) anymore, but instead we will have to resort to an iterative approach that alternates between updates of x and
evaluation of Ax. This is precisely what is typically done in gradient-based inference methods, and we will motivate our framework from there.
2.1 GRADIENT-BASED INFERENCE
Recall from equation (4) that inverse problems can be interpreted in terms of probability such that optimization is an iterative approach to MAP inference. In its most simple form each consecutive estimate of x is then computed through a recursive function of the form
xt+1 = xt + γt∇ ( log p (y|A,x) + log pθ (x) ) (xt) (6)
where we make use of the fact that p(x|A,y) ∝ p(y|A,x)pθ(x) and γt is the step size or learning rate at iteration t. Further, A is a (partially-)observable covariate, p(y|A,x) is the likelihood function for a given inference problem, and pθ (x) is a prior over signal x. In many cases where either the likelihood term or the prior term deviate from standard models, optimization will not be convex. In constrast, the approach presented in this work is completely freed from ideas about convexity, as will be shown in the next section.
2.2 RECURRENT FUNCTION DEFINITION
The central insight of this work is that update equation (6) can be generalized such that xt+1 = xt + gφ(∇y|x,xt) (7)
where we denote∇ log p(y|A,x)(xt) by∇y|x for readability and φ is a set of learnable parameters that govern the updates of x. In this representation, prior parameters θ and learning rate parameters γ have been merged into one set of trainable parameters φ.
To recover the original update equation (6), gφ(∇y|x,xt) is written as gφ(∇y|x,xt) = γt ( ∇y|x +∇x ) (8)
where we make use of ∇x to denote ∇ log pθ(x)(xt). It will be useful to dissect the terms on the right-hand side of (8) to make sense of the usefulness of the modification.
First notice, that in equation (6) we never explicitly evaluate the prior, but only evaluate its gradient in order to perform updates. If never used, learning a prior appears to be unnecessary, and instead it appears more reasonable to directly learn a gradient function ∇x = fθ(xt) ∈ Rd. The advantage of working solely with gradients is that they do no require the evaluation of an (often) intractable normalization constant of pθ(x).
A second observation is that the step sizes γt are usually subject to either a chosen schedule or chosen through a deterministic algorithm such as a line search. That means the step sizes are always chosen according to a predefined model Γ. Interestingly, this model is usually not learned. In order to make inference faster and improve performance we suggest to learn the model Γ as well.
In (7) we have made the prior pθ(x) and the the step size model Γ implicit in function gφ(∇y|x,ηt). We explicitly keep ∇y|x as an input to (7) because - as opposed to Λ and pθ(x) - it represents extrinsic information that is injected into the model. It allows for changes in the likelihood model p(y|x) without the need to retrain parameters φ of the inference model gφ. Figure 1 gives a visual summary of the insights from this section.
2.3 OUTPUT CONSTRAINTS
In many problem domains the range of values for variable x is naturally constraint. For example, images typically have pixels with strictly positive values. In order to model this constraint we make use of nonlinear link functions as they are typically used in neural networks, such that
x = Ψ(η) (9) where Ψ(·) is any differentiable link function and η is the space in which RIMs iterate such that update equation (7) is replaced by
ηt+1 = ηt + gφ(∇y|η,ηt) (10) As a result x can be constraint to a certain range of values through Ψ(·), whereas iterations are performed in the unconstrained space of η
2.4 RECURRENT NETWORKS
A useful extension of (7) is to introduce a latent state variable st into the procedure. This latent variable is typically used as a utility in recurrent neural networks to learn temporal dependencies in data processing. With an additional latent variable the update equations become
ηt+1 = ηt + hφ ( ∇y|η,ηt, st+1 ) (11) st+1 = h ∗ φ ( ∇y|η,ηt, st ) (12)
where h∗φ(·) is the update model for state variable s. The variable s will allow the procedure to have memory in order to track progression, curvature, approximate a preconditioning matrix Tt (such as in BFGS) and determine a stopping criterion among other things. The concept of a temporal memory is quite limited in classical inference methods, which will allow RIMs to have a potential advantage over these methods.
2.5 TRAINING
In order to learn a step-wise inference procedure it will be necessary to simulate the inference steps during training. I.e. during training, an RIM will perform a number of inference steps T . At each step the model will produce a prediction as depicted in figure Figure 1. Each of those predictions is then subject to a loss, which encourages the model to produce predictions that improve over time. In it’s simplest form we can define a loss which is simply a weighted sum of the individual prediction losses at each time step such that
Ltotal(φ) = T∑ t=1 wtL(xt(φ),x) (13)
is the total loss. Here, L(·) is a base loss function such as the mean square error, wt is a positive scalar and xt(φ) is a prediction at time t. In this work we follow Andrychowicz et al. (2016) in setting wt = 1 for all time steps.
3 RELATED WORK
The RIM framework can be seen as an auto-encoder framework in which only the decoder is trained, whereas the encoder is given by a known corruption process. In terms of the training procedure this makes RIMs very similar to denoising auto-encoders (Vincent et al., 2008). Though initially with the objective of regularization in mind, denoising auto-encoders have been shown to be effectively used as generative models (Vincent et al., 2010). The difference of RIMs to denoising auto-encoders and also more recently developed auto-encoders such as Kingma & Welling (2014); Rezende et al. (2014) is that RIMs enforce coupling between encoder and decoder both, during training and test time. In it’s typical form, decoder and encoder of an auto-encoder are only coupled during training time, while there is no information flow during test time (Kingma & Welling, 2014; Rezende et al., 2014; Vincent et al., 2008; 2010). An exception is the work from Gregor et al. (2016) which is conceptually strongly related to RIMs. There, an RNN model is used to generate static data by drawing on a fixed canvas. An error signal is propagated throughout the generation process.
There have been approaches in the past which aim to formulate a framework in which an inference procedure is learned. One of the best known frameworks is LISTA (Gregor & LeCun, 2010) which aims to learn a model that reconstructs sparse codes from data. LISTA models try to fit into the classical framework of doing inference as described in 1, whereas RIMs are completely removed from assumptions about sparsity. A recent paper by Andrychowicz et al. (2016) aims to train RNNs as optimizers for non-convex optimization problems. Though introduced with a different intention, RIMs can be seen as a generalization of this approach, in which the model - in addition to the gradient information - is aware about the absolute position of a prediction in variable space(see equation (7)).
4 EXPERIMENTAL RESULTS
We evaluate our method on various kinds of image restoration tasks which can each be formulated in terms of linear measurement problems as described in equation (1). We first analyze the properties
of our proposed method on a set of restoration tasks from random projections. Later we compare our model on two well known image restoration tasks: image denoising and image super-resolution.
4.1 MODELS
If not specified otherwise we use the same RNN architecture for all experiments presented in this work. The chosen RNN consists of three convolutional hidden layers and a final convolutional output layer. All convolutional filters were chosen to be of size 3 x 3 pixels. The first hidden layer consists of convolutions with stride 2 (64 features), subsequent batch normalization and a tanh nonlinearity. The second hidden layer represents the RNN part of the model. We chose a gated recurrent unit (GRU) (Chung et al., 2014) with 256 features. The third hidden layer is a transpose convolution layer with 64 features which aims to recover the original image dimensions of the signal, followed again by a batch normalization layer and a tanh nonlinearity. All models have been trained on a fixed number of iterations of 20 steps. All methods were implemented in Tensorflow1.
4.2 DATA
All experiments were run on the BSD-300 data set (Martin et al., 2001)2. For training we extracted patches of size 32 x 32 pixels with stride 4 from the 200 training images available in the data set. In total this amounts to a data set of about 400 thousand image patches with highly redundant information. All models were trained over only two epochs, i.e. each unique image patch was seen by a model only twice during training. Validation was performed on a held-out data set of 1000 image patches.
For testing we either used the whole test set of 100 images from BSDS-300 or we used only a subset of 68 images which was introduced by Roth & Black (2005) and which is commonly used in the image restoration community 3.
4.3 IMAGE RESTORATION
All tasks addressed in this work assume a linear measurement problem of the form as described in equation (1) with additive (isotropic) Gaussian noise. In this case the gradient of the likelihood takes the form
∇y|x = 1
σ2 AT (y −Ax) (14)
where σ2 is the noise variance. For very small σ this gradient diverges. In order to make the gradient more stable also for small σ we chose to rewrite it as
∇y|x = 1
σ2 + AT (y −Ax) (15)
where = softplus(φ ) and φ is a trainable parameter. As a link function Ψ (see (9)) we chose the logistic sigmoid nonlinearity4 and we used the mean square error as training loss.
4.4 MULTI-TASK LEARNING WITH RANDOM PROJECTIONS
To analyze the properties of our proposed framework in terms of convergence and to test whether all components of the model are useful, we first trained the model to reconstruct image patches from noisy random projections of grayscale image patches. We consider three types of random projection matrices: (1) Gaussian ensembles with elements drawn from a standard normal distribution, (2) binary ensembles with entries of values {−1, 1} drawn from a Bernulli distribution with p = 0.5, and (3) Fourier ensembles with randomly sampled rows from a Fourier matrix (see Donoho (2006b)).
We trained three models on these tasks: (1) a Recurrent Inference Machine (RIM) as described in 2, (2) a gradient-descent network (GDN) which does not use the current estimate as an input (compare
1https://www.tensorflow.org 2https://www2.eecs.berkeley.edu/Research/Projects/CS/vision/bsds/ 3http://www.visinf.tu-darmstadt.de/vi research/code/foe.en.jsp 4All training data was rescaled to be in the range [0, 1]
Andrychowicz et al. (2016)), and (3) a feed-forward network (FFN) which uses the same inputs as the RIM but where we replaced the GRU unit with a ReLu layer in order to remove state-dependence. Model (2) and (3) are simplifications of RIM in order to test the influence of each of the removed model components on prediction performance.
Figure 2 shows the reconstruction performance of all three models on random projections. In all tasks the RIM clearly outperforms both other models, showing overall consistent convergence behavior. The FFN performs well on easier tasks but starts to show degrading performance over time on more difficult tasks. This suggests that the state information of RIM plays an important role on the convergence behavior as well as overall performance. The GDN shows worst performance among all three models. For all tasks, the performance of GDN starts to degrade clearly after the 20 time steps that were used during training. We hypothesize that the model is able to compensate some of the missing information about the current estimate of x through state variable s during training, but the model is not able to transfer this ability to episodes with more iterations.
These results suggests that both the current estimate as well as the recurrent state carry useful information for performing inference. We will therefor only consider fully fledged RIMs from here on.
4.5 IMAGE DENOISING
After evaluating our model on 32 x 32 pixel image patches we wanted to see how reconstruction performance generalizes to full sized images and to an out of domain problem. We chose to reuse the RIM that was trained on the random projections task to perform image denoising. In this section we will call this model RIM-3task. To test the hypothesis that inference should be trained task specific, we further trained a model RIM-denoise solely on the denoising task. Table 2 shows the denoising performance through the mean PSNR on the BSD-300 test set for both models as compared to state-of-the-art methods in image denoising. The RIM-3task model shows very competitive results with other methods on all noise levels. This exemplifies that the model indeed has learned something reminiscent of a prior, as it was never directly trained on this task. The RIM-denoise model further improves upon the performance of RIM-3task and it outperforms most other methods on all noise levels. This is to say that the same RIM was used to perform denoising on different noise levels, and this model does not require any hand tuning after training.
Table 2 shows denoising perfomance on image that have been 8-bit quantized after adding noise(see Schmidt et al. (2016)). In this case performance slightly deteriorates for both models, though still making competitive with stateof-the-art methods. This effect could possibly be accommodated through further training, or by adjusting the forward model. Figure 3 gives some qualitative results on the denoising performance for one of the test images from BSD-300 as compared to the method from Zoran & Weiss (2011). RIM is able to produce more naturalistic images with less visible artifacts. The state variable in our RIM model allows for a growing receptive field size over time, which could explain the good long range interactions that the model shows.
Many denoising algorithms are solely tested on gray-scale images. Sometimes this is due to additional difficulties that multi-channel problems bring for some inference approaches. To show that it is straightforward to apply RIMs to multi-channel problems we trained a model to denoise RGB images. The denoising performance can be seen in table 1. The model is able to exploit correlations across color channels which allows for an additional boost in reconstruction performance.
4.6 IMAGE SUPER-RESOLUTION
We further tested our approach on the well known image super-resolution task. We trained a single RIM 5 on 36 x 36 pixel image patches from the BSD-300 training set to perform image super-
5The architecture of this model was slightly simplified in comparison to the previous problems. Instead of strided convolutions, we chose a trous convolutions. This model is more flexible and used only about 500.000 parameters. Previous experiments will be updated with the same model architecture.
resolution for factors 2, 3, and 46. We followed the same testing protocol as in Huang et al. (2015), and we used the test images that were retrieved from their website 7. Table 3 shows a comparison with some state-of-the-art methods on super-resolution for the BSD-300 test set. Figure 4 shows a qualitative example of super-resolution performance. The other deep learning method in this comparison, SRCNN Dong et al. (2014), is outperformed by RIM on all scales. Interestingly SRCNN was trained for each scale independently whereas we only trained one RIM for all scales. The chosen RIM has only about 500.000 parameters which amounts to about 2MB of disk space, which makes this architecture very attractive also for mobile computing.
6We reimplemented MATLABs bicubic interpolation kernel in order to apply a forward model (subsampling) in TensorFlow which agrees with the forward model in Huang et al. (2015).
7https://sites.google.com/site/jbhuang0604/publications/struct sr
Metric Scale Bicubic SRCNN A+ SelfExSR RIM (Ours)
PSNR 2x 29.55± 0.35 31.11± 0.39 31.22± 0.40 31.18± 0.39 31.39± 0.39 3x 27.20± 0.33 28.20± 0.36 28.30± 0.37 28.30± 0.37 28.51± 0.37 4x 25.96± 0.33 26.70± 0.34 26.82± 0.35 26.85± 0.36 27.01± 0.35
SSIM 2x 0.8425± 0.0078 0.8835± 0.0062 0.8862± 0.0063 0.8855± 0.0064 0.8885± 0.0062 3x 0.7382± 0.0114 0.7794± 0.0102 0.7836± 0.0104 0.7843± 0.0104 0.7888± 0.0101 4x 0.6672± 0.0131 0.7018± 0.0125 0.7089± 0.0125 0.7108± 0.0124 0.7156± 0.0125
Table 3: Image super-resolution performance on RGB images from BSD-300 test set. Mean and standard deviation (of the mean) of Peak Signal-to-Noise Ratio (PSNR) and Structural Similarity Index (SSIM) Wan (2004). Standard deviation of the mean was estimated from 10.000 boostrap samples. Test protocol and images taken from Huang et al. (2015). Only the three best performing methods from Huang et al. (2015) were chosen for comparison: SRCNN Dong et al. (2014), A+ Timofte et al. (2015), SelfExSR Huang et al. (2015). Best mean values in bold.
5 DISCUSSION
In this work, we introduce a general learning framework for solving inverse problems with deep learning approaches. We establish this framework by abandoning the traditional separation between model and inference. Instead, we propose to learn both components jointly without the need to define their explicit functional form. This paradigm shift enables us to bridge the gap between the fields of deep learning and inverse problems. We believe that this framework can have a major impact on many inverse problems, for example in medical imaging and radio astronomy. Although we have focused on linear image reconstruction tasks in this work, the framework can be applied to inverse problems of all kinds, such as non-linear inverse problems.
ACKNOWLEDGMENTS
The research was funded by the DOME project (Astron & IBM) and the Netherlands Organization for Scientific Research (NWO). The authors are greatful for helpful comments from Thomas Kipf, Mijung Park, Rajat Thomas, and Karen Ullrich.
Unfortunately, even after reading the authors' response to my pre-review question, I feel this paper in its current form lacks sufficient novelty to be accepted to ICLR.
Fundamentally, the paper suggests that traditional iterative algorithms for specific class of problems (ill-posed image inverse problems) can be replaced by discriminatively trained recurrent networks. As R3 also notes, un-rolled networks for iterative inference aren't new: they've been used to replace CRF-type inference, and _also_ to solve image inverse problems (my refs [1-3]). Therefore, I'd argue that the fundamental idea proposed by the paper isn't new---it is just that the paper seeks to 'formalize' it as an approach for inverse problems (although, there is nothing specific about the analysis that ties it to inverse problems: the paper only shows that the RIM can express gradient descent over prior + likelihood objective).
I also did not find the claims about benefits over prior approaches very compelling. The comment about parameter sharing works both ways---it is possible that untying the parameters leads to better performance over a fewer number of 'iterations', and given that the 'training set' is synthetically generated, learning a larger number of parameters doesn't seem to be an issue. Also, I'd argue that sharing the parameters is the 'obvious' approach, and the prior methods choose to not tie the parameters to get better accuracy.
The same holds for being able to handle different noise levels / scale sizes. A single model can always be trained to handle multiple forms of degradation---its just that its likely to do better when it's trained for specific degradation model/level. But more importantly, there is no evidence in the current set of experiments that shows that this is a property of the RIM architecture. (Moreover, this claim goes against one of the motivations of the paper of not training a single prior for different observation models ... but to train the entire inference architecture end-to-end).
It is possible that the proposed method does offer practical benefits beyond prior work---but these benefits don't come from the idea of simply unrolling iterations, which is not novel. I would strongly recommend that the authors consider a significant re-write of the paper---with a detailed discussion of prior work mentioned in the comments that highlights, with experiments, the specific aspects of their recurrent architecture that enables better recovery for inverse problems. I would also suggest that to claim the mantle of 'solving inverse problems', the paper consider a broader set of inverse tasks---in-painting, deconvolution, different noise models, and possibly working with multiple observations (like for HDR).