Deep Learning and Inverse Problems

Summer 2022
Reinhard Heckel

6 Unrolled neural networks

In the previous section, we discussed that a neural network can be trained end-to-end to denoise
an image. In this section, we discuss this learning-based approach in the context of a linear inverse
problem, where the goal is to reconstruct a signal from the measurement y = Ax + z, where
A ∈ Rm×n is typically a wide matrix, as before, or a matrix that is poorly conditioned, and z is
noise.
We discuss a technique for designing a neural network suitable for solving signal recovery al-
gorithms, based on modifying existing, iterative algorithms. This technique is called unrolling or
unfolding.

6.1 Posing reconstruction as a denoising problem

We first note that we can cast image reconstruction from a measurements y = Ax + z as a
denoising problem. Specifically, A† y is the least-squares reconstruction of the image x, which is
very coarse reconstruction that typically has a lot of artifacts. Given pairs of measurement and
corresponding clean images (y, x), we can train a neural network (e.g., a U-net) to denoise this
coarse reconstruction. For magnetic resonance imaging, this relatively simple approach already
outperforms sparsity based approaches [Zbo+18] in terms accuracy and computational cost for
reconstructing an image (here, the computational cost for training is not taken into account).

6.2 Unrolling the iterative soft thresholding algorithm

One of the earliest unrolled algorithms is Gregor and LeCun [GL10]’s work on learning a fast
algorithm for recovering a sparse signal. Recall that the iterative shrinkage thresholding algorithm
(ISTA) is an iterative algorithm for solving `1 -regularized least-squares:
1
min kAx − yk22 + λkxk1 .
x 2
ISTA is initialized with x0 = 0 and its iterations are given by

xt+1 = τλη (xt − ηAT (Axt − y)),

where τλ is the soft-thresholding operator defined earlier. The motivation of Gregor and LeCun
[GL10]’s work is to accelerate the existing ISTA algorithm through training. Thus, the goal is to
learn a faster algorithm, but not necessarily a better-performing algorithm.
We next formulate the ISTA algorithm as a neural network. The iterates of ISTA have the form

xt+1 = τθ (Qxt + By),

with Q = (I − ηAT A), B = ηAT , and θ = λη. Suppose we carry out K many ISTA iterations.
This corresponds to performing a forward-pass through a recurrent neural network, with fixed

1
weights. To see this, we view the computations in the t-th layer of the network as applying a linear
operation to the results from the previous layer, xt , followed by application of a non-linearity (i.e.,
the soft-thresholding function).
Gregor and LeCun [GL10]’s idea is to view the final result after k iterations, fθ (y) = xt , as
a function of the parameters θ = {Q, B, θ}, and train those parameters on a given dataset by
minimizing the loss in equation (??). This algorithm is called LISTA for Learned ISTA.
In their original paper, Gregor and LeCun [GL10] find that LISTA offers a significant speedup
for few iteration. For example, a LISTA network of depth k = 7 reaches prediction of similar
quality than ISTA with about ten times the number of iterations. However, the performance tends
to stagnate with the depth, and in the regime of many iterations, ISTA typically outperforms
LISTA.
Even in the regime where LISTA works well there is a caveat: For real-world problems the size
of the matrices that are learned is often prohibitively large. For example, suppose we apply this
approach to magnetic resonance imaging, where the images have size say 320×320 pixels. This gives
a parameter matrix Q of size 3202 × 3202 , which contains about 1010 entries. This matrix doesn’t
fit in the memory of modern GPUs, because it would require close to 300GB in high precision float
numbers. The computational and storage burden can be decreased by parameterizing Q = UVT ;
but Gregor and LeCun [GL10] found this to come at the cost of accuracy. Even if the computational
burden can be dealt with, note that this algorithm doesn’t take knowledge of the forward model, A,
into account, so the network has to learn both the signal and the forward model which is intuitively
wasteful because we know the forward model.
Gregor and LeCun [GL10]’s work pertains to sparse signal recovery, but of course the idea of
unrolling a proximal gradient descent algorithm can be applied to general iterative algorithms for
solving a variety of problems, and is a useful tool for designing neural networks for signal and image
recovery, as we will see in the next section.

6.3 A parameter efficient variational network

We next discuss the so-called variational network proposed by Chen et al. [Che+15] and Ham-
mernik et al. [Ham+18], which is an un-rolled neural network inspired by unrolling total-variation-
regularized least-squares. This network works well for accelerated magnetic resonance imaging, has
relatively few parameters, and is therefore a parameter-efficient approach to signal reconstruction.
A variant of this network, which has much more parameters [Sri+20] give state-of-the-art per-
formance in the 2020 fastMRI challenge, a competition for accelerated MRI reconstruction with
deep neural networks [Kno+20]. We’ll discuss this variant in the next subsection.
As discussed previously, it is common to reconstruct a signal by solving a regularized least-
squares problem
1
min kAx − yk22 + R(x), (1)
x 2

where R is a regularizer that often induces sparsity in an appropriate domain. A popular choice
for the regularizer for imaging problems is the total-variation norm, which works well if an images
consists of patches that are relatively constant, which turns out to be a reasonably good model for

2
natural images. The total variation norm for a vector is:
n−1
X
kxkT V = |xi − xi+1 |.
i=1

This can we written as kxkT V = kCxk1 , where C is a circular matrix with the first column equal
to c = [1, −1, 0, ..., 0]. In other words, the total variation norm can be viewed as convolution with
the filter [−1, 1] followed by applying the non-linearity | · | element-wise, followed by summing up
the entries:

kxkT V = h|Cx|, 1i ,

where 1 = [1, 1, . . . , 1] is the all-ones vector.

A generalization of the TV-norm regularizer is the so called Fields of Expert model [RB09]
defined as
k
X
R(x) = hψi (Ci x), 1i ,
i=1

where the operator Ci : Rn → Rn is a circulant matrix implementing a convolution with a filter

ci of size s. For example, the filter previously used for toal-variation minimization is [−1, 1],
and a filter of size s is simply [c1 , c2 , . . . , cs ], where the cj ’s are the filter parameters. Moreover,
ψi : R → R is a non-linear scalar function, applied element-wise to a vector. The original Fields of
Expert model [RB09] parameterizes the non-linearity, similarly as the soft-thresholding function in
LISTA is parameterized by the soft-thresholding parameter. Specifically, the activation function is
parameterized as a linear superposition of Gaussian radial basis functions (RBFs)
31 (z−µj )2
wij e−
X
ψi (z) = 2σ 2 ,
j=1

where µj are nodes distributed equidistant in some appropriate range, and the weights wij are
the parameters of the function ψi . To summarize, the Field-of-Experts model is parameterized
by the convolutional filters c1 , . . . , ck as well as the weights of the corresponding non-linearities
wij , i = 1, . . . , k, j = 1, . . . , 31. In its original form, the Fields-of-Experts model was used for
denoising by training the parameters of the model accordingly.
The variational network uses the Fields-of-Expert models in an un-rolled gradient algorithm.
Gradient descent applied to the Field-of-Expert regularized least squares problem (1) is
k
!
X
xt+1 = xt − η AT (Axt − y) + CTi ψi0 (Ci xt ) ,
i=1

where ψi0 is the derivative of ψi . Chen et al. [Che+15] proposed to unfold those iterations for a few
iterates and make all parameters independently trainable (i.e, the filters in the first, second, etc
iterations) are independent of each other. Chen et al. [Che+15] applied this approach successfully
to denoising and JPEG deblocking. Analogously, in the context of magnetic resonance imaging,

3
 x0 x1 xT −1
y AT −η0 AT (Ax0 − y) + ... −ηT AT (Ax0 − y) + x̂

neural network neural network

Figure 1: A variational network for image reconstruction consisting of T cascades. [Ham+18]

proposed such a network with T = 10 cascades where the neural network is a simple CNN with
a single hidden layer. This results in a small (about 1.3e5 parameters) and effective network for
image reconstruction. Substituting the CNN with a large U-net significantly improves performance
and gives image quality.

Hammernik et al. [Ham+18] proposed to unfold those iterations. The iterates are initialized as
x0 = AT y, and the remaining iterates are given by

xt+1 = xt − ηt AT (Axt − y) + CNNt (xt ), (2)

where
k
X
0
CNNt (x) = CTt,i ψt,i (Ct,i x) (3)
i=1

can be viewed as a two-layer CNN, where the trainable parameters are the convolutional filters
ct,i and the parameters of the activation function ψt,i0 . Note that the parameters in each of the

unrolled iterations are trainable, and the stepsize parameter ηt is trainable as well. The network is
illustrated in Figure 1.
Note that the variational network takes the forward model explicitly into account, and it also
takes prior intuition about the data into account, by expressing the regularizer in a particular form
inspired by the total-variation norm. The paper [Ham+18] trained a variational network consisting
of T = 10 steps. In each step, k = 48 filters of size 11 × 11 were learned. The activation function
was parameterized by 31 RBFs. This resulted in a network with about 1.3e5 parameters, easily
manageable by a computer.

6.4 A large variational network

The variational network discussed in the previous section outperforms a traditional un-trained
methods like total variation norm minimization, and is relatively parameter efficient. However,
substituting the simple two-layer CNN in equation (3) with a large U-net significantly improves
performance, and gives near state-of-the-art performance for MRI reconstruction (see the results for
VNU in [Sri+20] for performance evaluations). The corresponding network has T = 12 cascades and
about 30M parameters, which is by a factor 230 larger than the variational network in the previous

4
section. Tweaking the network further with MRI specific steps gives another slight performance
boost (see [Sri+20]) and gives state-of-the-art performance as measured by in the 2020 fastMRI
challenge, a competition for accelerated MRI reconstruction with deep neural networks [Kno+20].

References
[Che+15] Y. Chen, W. Yu, and T. Pock. “On Learning Optimized Reaction Diffusion Processes
for Effective Image Restoration”. In: Proceedings of the IEEE Conference on Computer
Vision and Pattern Recognition. 2015, pp. 5261–5269.
[GL10] K. Gregor and Y. LeCun. “Learning Fast Approximations of Sparse Coding”. In: Pro-
ceedings of the 27th International Conference on International Conference on Machine
Learning. ICML’10. USA: Omnipress, 2010, pp. 399–406.
[Ham+18] K. Hammernik, T. Klatzer, E. Kobler, M. P. Recht, D. K. Sodickson, T. Pock, and F.
Knoll. “Learning a Variational Network for Reconstruction of Accelerated MRI Data”.
In: Magnetic Resonance in Medicine 79.6 (2018), pp. 3055–3071.
[Kno+20] F. Knoll et al. “Advancing Machine Learning for MR Image Reconstruction with an
Open Competition: Overview of the 2019 fastMRI Challenge”. In: Magnetic Resonance
in Medicine 84.6 (2020), pp. 3054–3070.
[RB09] S. Roth and M. J. Black. “Fields of Experts”. In: International Journal of Computer
Vision 82.2 (2009), p. 205.
[Sri+20] A. Sriram, J. Zbontar, T. Murrell, A. Defazio, C. L. Zitnick, N. Yakubova, F. Knoll, and
P. Johnson. “End-to-End Variational Networks for Accelerated MRI Reconstruction”.
In: arXiv:2004.06688 [cs, eess] (2020).
[Zbo+18] J. Zbontar et al. “fastMRI: An Open Dataset and Benchmarks for Accelerated MRI”.
In: arXiv:1811.08839 [physics, stat] (2018).