Loop Unrolled Shallow Equilibrium Regularizer (LUSER) - A Memory-Efficient Inverse Problem Solver
Authors: Peimeng Guan
Jihui Jin
Justin Romberg
Mark A. Davenport
Venue: Submitted to ICLR 2023 (2023)
Keywords: Deep Equilibrium Models
Inverse Problems
Deep Learning
Loop Unrolled

Abstract: In inverse problems we aim to reconstruct some underlying signal of interest from potentially corrupted and often ill-posed measurements. Classical optimization-based techniques proceed by optimizing a data consistency metric together with a regularizer. Current state-of-the-art machine learning approaches draw inspiration from such techniques by unrolling the iterative updates for an optimization-based solver and then learning a regularizer from data. This \emph{loop unrolling} (LU) method has shown tremendous success, but often requires a deep model for the best performance leading to high memory costs during training. Thus, to address the balance between computation cost and network expressiveness, we propose an LU algorithm with shallow equilibrium regularizers (LUSER). These implicit models are as expressive as deeper convolutional networks, but far more memory efficient during training. The proposed method is evaluated on image deblurring, computed tomography (CT), as well as single-coil Magnetic Resonance Imaging (MRI) tasks and shows similar, or even better, performance while requiring up to $8 \times$ less computational resources during training when compared against a more typical LU architecture with feedforward convolutional regularizers.

Under review as a conference paper at ICLR 2023
LOOP UNROLLED SHALLOW EQUILIBRIUM REGULAR- IZER (LUSER) - A MEMORY-EFFICIENT INVERSE PROBLEM SOLVER
Anonymous authors Paper under double- blind review
ABSTRACT
In inverse problems we aim to reconstruct some underlying signal of interest from potentially corrupted and often ill-posed measurements. Classical optimization- based techniques proceed by optimizing data consistency metric together with a regularizer. Current state-of-the-ard machine learning approaches draw inspira- tion from such techniques by unrolling the iterative updates for an optimization- based solver and then learning a regularizer from data. This loop unrolling (LUU method has shown tremendous success, but often requires a deep model for the best performance leading to high memory costs during training. Thus, to address the balance between computation cost and network expressiveness, we propose an LU algorithm with shallow equilibrium regularizers (LUSER). These implicit models are as expressive as deeper convolutional networks, but far more memory efficient during training. The proposed method is evaluated on image deblurring, computed tomography (CT), as well as single-coil Magnetic Resonance Imaging (MRI) tasks and shows similar, even better, performance while requiring up 8x less computational resources during training when compared against more typical LU architecture with feedforward convolutional regularizers.
INTRODUCTION
In an inverse problems we face the task of reconstructing some data or parameters an unknown signal from indirect observations. The forward process, or the mapping from the data to observa- tions, is typically well known, but ill-posed or non-invertible. More formally, we consider the task of recovering some underlying signal from measurements taken via forward operator A according to
y Aa + 7,
where represents noise. The forward operator can be nonlinear, but to simplify the notation, we illustrate the idea in linear form throughout this paper. common approach recover the signal via an iterative method based the least squares loss:
arg min ly Arl
For many problems of interest, A is ill-posed and does not have full column rank. Thus, attempting solve (2) does yield unique solution. To address this, we can extend 2 by including regularizing term to bias the inversion towards solutions with favorable properties. Common examples of regularization include l2, L1, and Total Variation (TV). Each regularizer encourages certain properties on the estimated signal & (e.g., smoothness, sparsity, piece-wise constant, etc.) and is often chosen based on task-specific prior knowledge.
Recent works (Ongie et al. 2020) attempt to tackle inverse problems using more data-driven meth- ods. Unlike typical supervised learning tasks that attempt to learn a mapping purely from examples, deep learning for inverse problems have access to the forward operator and thus should be able to guide the learning process for more accurate reconstructions. One popular approach to incorporat- ing knowledge of the forward operator is termed loop unrolling (LU). These methods are heavily inspired by standard iterative inverse problem solvers, but rather than use a hand tuned regularizer, they instead learn the update with some parameterized model. They tend to have fixed number

---

Under review as a conference paper at ICLR 2023
of iterations (typically around 5-10) due to computational constraints. Gilton al. (2021) proposes an interesting alternative that takes advantage of deep equilibrium (DEQ) models (Bai et al. 2019 2020 Fung et al. 2021, El Ghaoui et al., (021) that refer to DEQ4IP. Equilibrium models are designed to recursively iterate on their input until a "fixed point" found (i.e., the input no longer changes after passing through the model). They extend this principle the LU method, choosing to iterate until convergence rather than for a "fixed budget".
Our Contributions. We propose an alternative novel architecture for solving inverse problems called Loop Unrolled Shallow Equilibrium Regularizer (LUSER). It incorporates knowledge of the forward model by adopting the principles of L.U architectures while reducing its memory consump- tion by using shallow (relative to existing feed-forward models) DEQ as the learned regularizer update. Unlike DEQ4IP that converts the entire LU architecture into DEQ model, we only convert the learned regularizer at each stage. This has the advantage of simplifying the learning task for DEQ models, which can be unstable to train in practice. To our knowledge, this is the first use of multiple sequential DEQ models within single architecture for solving inverse problems. Our proposed architecture (i) reduces the number of forward/adjoint operations compared to the work proposed by Gilton et al. (2021), and (ii) reduces the memory footprint during training without loss of ex- pressiveness as demonstrated by our experiments. We empirically demonstrate better reconstruction across multiple tasks than LU alternatives with comparable number of parameters, with the ability to reduce computational memory costs during training by a factor of up to 8x.
The remainder of the paper is organized follows. Section reviews related works in solving inverse problems. ( the proposed LUSER, which we compare with other baseline methods in image deblurring, CT, and MRI tasks in Section 4 We conclude in Section with a brief discussion.
2
RELATED WORK
2.1 LOOP UNROLLING
As noted above, a common approach to tackling an inverse problem to cast it as an optimization problem consisting of the sum of data consistency term and regularization term
min y
Ar 3 +r(x),
where is a regularization function mapping from the domain of the parameters of interest to a real number and y 2 0 is a well-tuned coefficient. The regularization function chosen for specific classes of signals to exploit any potential structure, e.g., ||]|| for smooth signals and ||æ:|o ||>||| for sparse signals.
When is differentiable, the solution of (3) be obtained in iterative fashion via gradient descent. For some step size at iteration A: 1,2,.... N..,,, K, we apply the update:
Fb+1 = Bg XA
- T (y- -
For non-differentiable the more generalized proximal gradient algorithm can be applied with the following update, where is a well-tuned hyperparameter related to the proximal operator:
x++1 = prox prox,,,
(ax + XA (y Arx)).
(5)
The loop unrolling (LU) method performs the update in 4 or (5). but replaces Vr or the proximal operator with learned neural network instead. The overall architecture repeats the neural network based update for pre-determined number of iterations, fixing the overall computational budget. Note that the network is only implicitly learning the regularizer. In practice, it actually learning an update step, which can be thought of de-noising projection onto the manifold of the data. LU is typically trained end-to-end. While end-to-end training casier to perform and encourages faster convergence, it requires all intermediate activations to be stored in memory. Thus, the maximum number iterations is always kept small compared to classical iterative inverse problem solvers.
Due to the limitation in memory, there is trade-off between the depth of a LU and the richness of each regularization network. Intuitively, one can raise the network performance by increasing the number of loop unrolled iterations. For example, Gilton al. (2021) extends the LU model to
2

---

Under review as a conference paper at ICLR 2023
potentially infinite number of iterations using implicit network, and (Putzky & Welling 2019)
allows deeper unrolling iterations using invertible networks, while requiring recalculation of the in- termediate results from the output in training phase. This approach can be computationally intensive for large-scale inverse problems or when the forward operator is nonlinear and computationally ex- pensive apply. For example, the forward operator may involve solving differential equations such as the wave equation for seismic wave propagation (Chapman 2004) and the Lorenz equations for atmospheric modeling(Destreiccee 2022).
Alternatively, one can design a richer regularization network. For example Fabian & Soltanolkotabi (2022) uses transformer as the regularization network and achieves extremely competitive results in the fastMRI challenge (Zbontar et al. 2018)) but requires multiple 24GB GPÚ for training, which is often impractical, especially for large systems. Our design strikes a balance in the expressiveness in regularization networks and memory efficiency during training. Our proposed work is an alternative method to achieve rich regularization network without the additional computational memory costs during training.
2.2 DEEP EQUILIBRIUM MODELS FOR INVERSE PROBLEMS (DEQ4IP)
Deep equilibrium (DEQ) models introduce an alternative to traditional feed-forward networks (Bai
al., 2219: 2020 Fung et al., 2021; El Ghaoui et al., 2021). Rather than feed an input through
fixed (relatively small) number of layers, DEQ models solve for the "fixed-point" given some input.
More formally, given a network fo and some input x (0) and y. we recursively apply the network via
x (k+1) = fe( fæ (x) y),
(6)
until convergence| this instance, y acts as an input injection that determines the final output. This is termed the forward process. The weights 0 of the model can be trained via implicit differentiation, removing the need to store all intermediate activations from recursively applying the network (Bai
ct al. 2019 Fung et al. (222I). In particular, adopt the "Jacobian-fire" backpropagation strategy outlined in Fung et al. (2021). This allows for deeper, more expressive models without the associated memory footprint train such models.
Gilton et al. (2021) demonstrates an application of one such model, applying similar principles to a single iteration of an LU architecture. Such idea is a natural extension as it allows the model to "iterate until convergence" rather than rely on a "fixed budget". More specifically, the model repeats 4 and (5) many times (in practice, usually around 50 iterations) until 24 converges. However, such a model can be unstable to train and often performs best with pre-training of the learned portion of the model (typically acting as learned regularizer/de-noiseer) It is also important note that such a model would have to apply the forward operator (and its adjoint) many times during the forward process. Although this can be accelerated to reduce the number of applications, often signficiantly more than the number of applications for an LU equivalent which can be an issue the forward operator computationally expensive to apply.
still
2.3 ALTERNATIVE APPROACUEES TO TACKLE MEMORY ISSUES
I-RIM (Putzky& & Welling 2019) is deep invertible network that address the memory issue by
calculating the intermediate results from the output. However it is not ideal when forward model is computationally expensive. Gradient checkpointing (Sohoni et al.) (019)) is another practical tech- nique to reduce memory costs for deep neural networks. It saves intermediate activations of some checkpoint nodes, and recomputes the forward pass between two checkpoints for backpropagation. However, it is not an easy and efficient technique to implement for weight-ticd neural network.
3 METHODOLOGY
LU methods currently dominate the state-of-the-arr approaches for solving inverse problems due to their stability when training, inclusion of the forward model, and their near instantaneous inference
Note that, since our approach will ultimately use both methods. to aid clearer presentation we use subscript, i.e., æa to denote the LU iterations, and superscript with parenthesis, i.c., r()) to denote the iterations in the deep equilihrium model.
3

---

Under review as a conference paper at ICLR 2023
times. However, there is noticeable trade-off in terms of the memory requirements accuracy when training these models. Even for medium scale problems, some of the proposed architectures require multiple GPUs just to train, making this approach infeasible to use for much larger scale in- verse problems. DEQ4IP offers an interesting alternative, drastically reducing the memory required during training and allowing the flexibility adjust accuracy during reconstruction when perform- ing inference. However, these models can potentially suffer when the forward/adjoint operators are computationally expensive, particularly since it takes more iterations to converge than standard fixed LU method.
To address these concerss, we propose a novel architecture for solving inverse problems called Loop Unrolled Shallow Equilibrium Regularizer (LUSER), We adopt a LU approach to limit the number of forward/adjoint calls for particularly complex inverse problems while also drastically reducing the memory requirements for training, allowing us to scale up to much larger inverse problems (or require less GPU memory for existing problems). LUSER achieves this by adopting DEQ models as the trainable regularizer update in standard LU architecture. The implicit DEQ models are smaller in size but just as expressive (Bai et al. 2019) as typical convolutional models used as the learned regularizer update allowing for an accurate reconstruction with less computational memory costs. Furthermore, learning proximal update is a far simpler task compared to solving the inverse problem whole.
We adopt a "proximal gradient descent" styled LU architecture. The network takes in measurements y and some initial estimate øg. The architecture consists of K stages, alternating between gradient descent step dy 2% + AA (y A for k = 2, K, followed by a feed-forward pass of the shallow equilibrium model acting proximal update block.
de
Injection Layer
Initialize with weros
Yes
Data Layer
concatenate
Mix Layer
converged
#b-1
(oo)
p0)
No
Figure 1: A proximal block in LUSER
Figure [ illustrates a single block in LUSER as a learned proximal operator. The input de from the previous stage is processed only once by set of input injection layers to avoid redundant com- putation. The input injection will determine the final fixed point output. The recursive portion of the proximal block consists of two sets of layers: the data layer and the mixing layer. The current estimate of the fixed point is first passed through the data layer before being concatenated with the input injection and processed by the mixing layer. This process repeats until fixed point is found. In practice, convergence to fixed point is achieved when the difference between two time steps (+(+1)) - [A)] is within a small 6 or when the maximum number of iterations is reached. A more thorough summary can be found in Fung et al. (2021).
At each stage, the shallow equilibrium regularizer attempts to output its best estimate of the ground truth x.. We introduce skip connection between the input dy and the final output so that the regularizer is learning the residual between the input dy and the ground truth x² instead, similar the approach adopted by DnCNNs Ryu et al.] (2019). When 2& converges x", de will also converge to x²" and we expect the residual to be closer and closer to the zero vector. Therefore, we initialize the input of data layer with a zero vector with the same dimensions as the input de in the hopes that fewer iterations will be needed at later stages as the current estimate æg converges towards the ground truth. Let (6) denote the 1th update of the residual, where r (0) = 0 and ((o) be the fixed point solution of the residual. Let ① denote the concatenation operator. The process in
one loop unrolled block can be formulated as the following:
# = d + MixLayer (DataLayer(n(xx) @ InjectLayer((ee)
Although Figure D shows the simplest way of finding the fixed point, in practice acceleration or other fixed point solvers applied to solve for the fixed point. Similar to Gilton al. (2021),
4

---

Under review as a conference paper at ICLR 2023
also apply Anderson acceleration (Walker & Ni 2011 for all of our models when searching for the fixed point.
Using fixed point solver based architecture for the regularization block allows LUSER achieve a similar level of performance with a much shallower number of layers than traditional feedforward models. For the scope of this article, we restrict ourselves total of 5 layers to demonstrate memory savings, but initial experiments have shown that increasing the number of layers can yield higher performance.
We also explore two variants of LUSER, dubbed LUSER-SW and LUSER-DW. LUSER-SW refers to whaghd-weight version of the proposed algorithm, where the proximal operator in all loop un- rolled iterations are weight-tied and thus identical. Theoretically, the regularizer should be able to handle any input regardless of the iteration step. However, in practice, the distribution of inter- mediate reconstructions can be vastly different. Thus, training single model to handle all these instances can be daunting task, leading to poor generalization across the different stages. Since the total computational budget is fixed, one approach can be to use different weights (DW) for the learned proximal operator at each stage to handle the potentially different distributions. This will increase the number of parameters that need to be trained and stored, but since LUSER already has so few parameters to begin with, expanding to the different weight variant is still feasible. This paper will compare two variants of LUSER with other architectures in different tasks.
DEQ4IP relies on the learned proximal operator to have the same weights, thus we cannot include a different weight variant for comparison. On the other hand, the same principles can be applied the LU variant of DEQ4IP. However, we aim to only compare models of similar number of parameters (or less in the case of LUSER-SW), and thus restrict our attentions to the shared weight variant of LU only.
4 EXPERIMENTS AND RESULTS
In this section, we compare our proposed networks to LU with DrCNN as proximal operator (LU- DrCNN) and DEQ4IP on three different tasks: image deblurring, computed tomography (CT), and single-coil accelerated Magnetic Resonance Imaging (MRI). The datasets we use are RGB CelebA (Liu et al. 2015), LoDoPaB-CT (Leuschner et al. 202I), and single-coil knee data from fastMRI Zbontar et al. 2018) respectively. We also experiment with incorporating an auxiliary loss of including an MSE loss on intermediate reconstructions with the ground-truth instead of just the final output. This is done for the LU models (LU-DnCNN, LUSER).
4.1 EXPERIMENTAL SETUP
In the image deblurring task, the blurry image is obtained by applying a (9 9) Gaussian kernel with variance of 5 to an image with additive white Gaussian noise with a standard deviation of 0.01. If the image is RGB, the same kernel is applied all channels. For accelerated MRI tasks, mea- surements in s-space (or frequency domain) are often subsampled due to the cost in measurement. The goal of MRI reconstruction is recover the underlying physical structure from subsampled noisy measurements. We simulate the forward operator A with 2-dimensional Fourier transform with randomly selected rows. We consider two common subsampling scenarios: 4x and 8x accel- eration, or subsampling the columns in full measurement by factor of 4 and 8 respectively. For the CT task, the forward operator angles in measurements. The adjoint of measurement A is used as initialization for MRI and CT tasks, which brings the measurement back to the signal domain. However, in the deblurring task, since the measurement lies in a same domain as the underlying clean image, y is used as the initial guess. Notice that some works use the filtered backprojection as the initialization for CT, such as (Khorashadizadeh et al. 2022). but use the adjoint for the purpose of consistency.
Radon transform, and we uniformly select 200 out of 1000
We fix the budget of LU-DDCNN and LUSER to be a total of 8 iterations, while as we allow DEQ4IP to iterate until it reaches a fixed point. We use a DnCNN adopted from Ryu et al. (2019) with 17 convolutional layers with 64 channels, followed by BatchNorm and ReLU activations for the regularizer for LU-DnCNN and DEQ4IP, For the learned regularizer update in LUSER, use convolutional layers each for the input injection layer and data layer. The mixing layer contains convolutional layers.
5

---

Under review as a conference paper at ICLR 2023
Table
Architecture of Proximal Network in LUSER
Layer Details
SN(conv(Cinll, 32n ks:3, pad:1) LeakyReLU SN(:onv(Cn.:32, Cout:2, ks:3, pad: 1)) + LeakyReI.U
Injection Layer
SN(conv(C.nn) i, :3, ks:3, pad:1)) GN + ReLU
Data Layer
SN(conv(Cin :1 Cou 32, ks:3, pad:1)) GN + LeakyReLU
SN(:onv(Cin:64, Cour Cour144 16, ks:3, pad: ()) GN + LeakyReLU
Mix Layer
SN(conv(Cm::64 C4, ks:3, pad:1) 1)) GN + LeakyReLU
SN(conv(Cinn) Couil, ks:3, pad:1))
In order to stablize training for DEQ inspired models, we wrap all convolutional layers in DEQ4IP and LUSER with Spectral Norm (SN) [Miyato et al. 2018). We list more details for the learned proximal network for LUSER Table for the case when the input has a single channel. Cin and Cou denote the input and output channels, ks refers the kernel size of convolutional layer, pad denotes the padding in 2-dimension, and GN stands for GroupNorm.
We use two metrics evaluate the quality of reconstruction: Peak Signal-to-Non -oise Ratio (PSNR) in dB and the Structural Similarity Index (SSIM). Note that although we use the same models as Gilton et al. (2021) ). train models from scratch and report lower values on the MRI task. We suspect that this is due to evaluating with a single channel only. When we include the imaginary channel (for total of 2 channels), the metrics we recorded are more aligned with those reported in Gilton et al. (20TT. All models are trained with single RT'60000 24GB GPU.
4.2 RECONSTRUCTION QUALITY
Table@compares the average testing PSNR and SSIM. The different weight version of LUSER out- performs LU-DnCNN with similar number of network parameters specified in Table LUSER- SW achieves similar level of performance most tasks with only 5 layers, versus 17 layers in LU- DrCNN. DEQ4IP attains the best performance in image delurring task and CT, but LUSER achieves comparable quality. Training shared-weight architectures with auxiliary losses improves the recon- struction quality in most tasks. Note that restricted ourselves only 5 layers for LUSER-SW and LUSER-DW for all reported experiments to highlight the memory savings without degradation in performance. In initial exploration, noticed that increasing the number of layers can lead to boosts in performance. For example, in the CT task, LUSER-SV SW with 6 layers increases the PSNR and SSIM by 1.04 dB and 0.022 respectively.
Figure 2 shows representative reconstruction results. LUSER-DY achieves higher PSNR and SSIM some cases, and attains better qualities detailed structures, especially compared LU-DmCNN. The areas with improvements are emphasized with red boxes in the ground truth images.
Table 2: Average PSNR and SSIM for test set, the best two performances are in bold.
PSNR SSIM
LU-DnCNN
DEQ4IP Final loss
LUSER-SW
LUSER-DW
Final loss
Aux loss
Final loss
Aux loss
Final loss
Aux loss
Deblurring
29.93
30.39
31.57
30.30
30.65
31.40
31.15
(3, 218, 178)
0.862
0.871
0.895
0.869
0.878
0.891
0.888
CT (00, 300)
30.59
31.59
32.19
28.82
28.04
31.83
31.66
0.844
0.859
0.871
0.801
0.797
0.860
0.859
4x MRI
29.01
29.02
29.01
28.82
29.18
29.86
29.37
(2, 320, 320)
0.668
0.671
0.678
0.662
0.685
0.740
0.713
8x MRI (2, 320, 320)
27.50
27.65
27.51
27.42
27.42
28.06
27.55
0.576
0.572
0.570
0.562
0.560
0.630
0.596

---

Under review as a conference paper at ICLR 2023
Initialization
LU-DnCNN
DEQ4IP
LUSER-DW
Ground Truth
PSNR: 25.46
PENK: 27.95 SSIM 0867
PENR. 28.87 SSIM 0884
SSIM 0833
A
Deblurring
II
D
D
SSIM:
SSIM:
SSIM:
CT
SSIM: 0.803
SSIM: 0.797
SSIM: 0334
4x MRI
PSNR SSIM: 696
PSNR: SSIM: 0.744
PSNR:
8x MRI
Figure 2: Representative reconstructions, where each row represents one task. The left-most column shows the initialization or input to the networks; the middle three columns show the reconstructions for LU-DnCNN, DEQ4IP and LUSER-DW; and the right-most column shows the underlying true images. Regions corresponding to qualitative improvements are emphasized in red boxes in the last column.
Table 3: Comparison of network sizes and maximum possible batch sizes during training. Entries with '-' denote that the architecture with a batch size of cannot fit into particular GPU RAM capacity.
Maximum Batch Size
GPU RAM
LU-DnCNN
DEQ4IP (17 layers)
LUSER-SW (5 layers)
LUSER-DW (5 layers)
Capacity
(17 layers)
#Params
558,580
558,580
96,503
770,073
Deblurring
8 GB 10 GB
16
10
10
(3, 218, 178)
16
14
14
24 GB
4
68
34
#Params
556,033
556,033
93,954
751,625
CT 30, 300)
8 GB
4
2
2
10 GB
24 GB
2
20
10
10
#Params
557,185
557,185
95,107
760,849
MRI
8 GB
4
4
(2, 320, 320)
10 GB
2
4
24 GB
4
16
12
12
7

---

Under review as a conference paper at ICLR 2023
4.3 MEMORY TRAINING
GPU capacity is a major bottleneck for training large-scale loop unrolled networks as discussed earlier. Table compares the network sizes (number of parameters) well as the maximum training batch sizes for three commonly seen GPU RAM capacities: 8 GB, 10 GB and 24 GB. Batch sizes are recorded with maximum even numbers, except when it is for stochastic gradient descent. Notice that in MRI, because the Fourier transformation is implemented in tensor form, the minimum batch size it can take Since DEQ4IP is an extension of LU-DaCNN, their networks are of the same size, but implicit DEQ models support larger batch sizes making DEQ4IP far more memory- advantages of using DEQ models for LUSER and DEQ4IP are particularly highlighted in the case of limited memory (smaller GPUs). LU-DnCNN is unable even train for the CT and MRT task with limited memory constraints, while LUSER and DEQ4IP can. This pattern is expected to repeat for more large scale tasks where standard LU architectures will be unable to train at all due to memory requirements. It is important to note that memory requirements depend more on the depth of the network than the number parameters. For example, even though LUSER-DW and LUSER SW have different numbers of parameters, they share the same architecture/depth and thus roughly the same amount memory during training.
We the batch size as proxy for the memory requirements during training.
during training. The
4.4 INFERENCE TIME
The evaluation time of these models depend on complexity of forward model. For the ex- periments performed, taking gradient step average took 5.3e-5, 6.9e-4, and 6.6e-3 seconds for deblurring, MRI and CT respectively. In this regime, the time to evaluate the DEQ models dominates the inference time. However, the forward model time increases, it has the potential to dominate.
In order to test this, we introduce artificial delays of 5e-2 and 5c-1 seconds to the CT forward model and evaluate the models the test set. Because the performance does not significantly vary (as the trained models didn't change), we only report the run times in Figure In these regimes, the choice of adopting Loop Unrolled structure presents an advantage over DEQ4IP. In our experiments, DEQ4IP took approximately 30-50 iterations to reach a reasonable level of performance, leading 30-50 calls of the forward model. The inference time shown in the shaded region in green of Figure 3al while as LUSER is shown in red for a max iteration of 10-30. It interesting to note that as the forward model run time increases, does the overlap between LUSER and DEQ4IP, with LUSER evaluating faster on average somewhere between le-2 and 1e-1 seconds.
The reason becomes more evident in Figure 3b which show cases the breakdown in timing for the solid line example from Figure 3a The blue bar represents the approximate portion of time spent evaluating the forward model while as the green and red represent the learned model portion (as well as any overhead) associated with DEQ4IP and LUSER respectively. As the forward modeling time
Inference Time Forward Model Time
Evaluation Times
10a
DEQ4IP
DEQ4IP
Inference Time (sec)
WSER
10
LUSER
Run Time (sec)
Forward Model
10a
ColobA 5.3e-55
MRI
CT
Simulated Simulated
10
10-1
10-2
10-1
6 6.9e-4s
6.6e-35
5,5e-25
5.De-1s
Forward Model Run Time (sec)
Task/Forward Model Time
(a)
(b)
Figure 3: (a) Plot of inference time across range of "max iterations" and (b) with approximate forward model time overlaid total time.
8

---

Under review as a conference paper at ICLR 2023
becomes non-negligible, the evaluation time suffers accordingly with LUSER evaluating faster for the 5e-2 second delay experiment.
4.5 EFFECT OF WEIGHT SHARING
As shown in Table LUSER-DY DD has better performance than LUSER-SW with equivalent block structure, but has larger number of parameters. We also explore the possibility of reusing multiple LUSER proximal blocks across loop unrolled iterations. For example, in MRI tasks, we repeat 4 different proximal blocks over 8 iterations. Let 9i denote the 8t proximal block where (1,2,3, 4), and we form the LUSER network with the sequence of proximal blocks in order of 91.91. 92. 92: 93. 91- Now, the number of parameters is only half of that in LUSER- -WW but maintains a similar level of performance. We denote this variant LUSER-PSWY -SWW which stands for partially shared-wer weight. Table 4 compares the average PSNR and SSIM of LUSER-SV LUSER- PSW and LUSER-DE with the same block structure, which are trained using final loss only.
94-
Table 4: Average PSNR and SSIM for LUSER with different weight-sharing schematics
LUSER-SW
LUSER-PSW
LUSER-DW
PSNR
28.82
28.85
29.86
4x MRI
SSIM
0.66
0.74
0.74
PSNR
27.42
28.04
28.06
8x MRI
SSIM
0.56
0.63
0.63
Table summarizes various prorperties among the three architectures. In particular, refer to the expressiveness the network loosely as the performance relative to the depth. LUSER is able to achieve comparable performance to both DEQ4IP and LU with much shallower network. We denote the rows related to timing with the asterisk to include the simulated regimes as well, though for negligible forward models, DEQ4IP can be faster than LUSER.
Table 5: Method Comparisons
LU
DEQ4IP
LUSER
Training Time
Fast
Slow
Moderate
Inference Time*
Fast
Slow
Moderate
Network Size
Large
Large
Small
Training Memory Usage
Large
Small
Moderate
Expressiveness
Low
Moderate
High
5 CONCLUSION
Loop unrolling architectures with deep convolutional layers the learned regularizer update a popular approach for solving inverse problems. Although its variants achieve state-of-the-art re- sults across a variety of tasks, LU algorithms incur a huge memory cost when training due the requirement of saving all intermediate activations, sometimes even requiring multiple GPUs to train on complex tasks. DEQ4IP offers an interesting alternative via extending L.U to infinitely many layers by finding a fixed point solution, but can be impractical when the forward/adjoint operators are nonlinear or larger in scale. To address the memory issue, we proposed two variants loop un- rolling architectures with deep equilibrium models the learned regularizer updates, LUSER-SW and LUSER-DY -DW. We verify the memory savings (by comparing batch sizes) relative to loop un- rolling algorithms with DnCNN model. Across all tasks, LUSER-DY outperforms LU-DnCNN with a similar number of network parameters, while reducing the memory requirements by a factor of 5 or more.

---

Under review as a conference paper at ICLR 2023
REFERENCES
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32, 2019.
Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Multiscale deep equilibrium models. Advances in Neural Information Processing Systems, 33:5238-5250, 2020.
Chris Chapman. Fundamentals of seismic wave propagation. Cambridge university press, 2004.
Laurent El Ghaoui, Fangda Gu, Bertrand Travacca, Armin Askari, and Alicia Tsai. Implicit deep learning. SIAM Journal on Mathematics of Data Science, 3(3):930-958, 2021.
Zalan Fabian and Mahdi Soltanolkotabi. Humus-net: Hybrid unrolled multi-scale network architec- ture for accelerated mri reconstruction. arXiv preprint arXiv:2203.082/3 2022.
Samy Wu Fung, Howard Heaton, Qiuwei Li, Daniel McKenzie, Stanley Osher, and Wotao Yin. Fixed point networks: Implicit depth models with jacobian-fiee backprop. 2021.
Davis Gilton, Gregory Ongie, and Rebecca Willett. Deep equilibrium architectures for inverse problems in imaging. IEEE Transactions on Computational Imaging, 7:1123-1133, 2021.
AmirEhsan Khorashadizadeh, Konik Kothari, Leonardo Salsi, Ali Aghababaei Harandi, Maarten de Hoop, and Ivan Dokmani'c c. Conditional injective flows for bayesian imaging. arXiv preprint 0rXiv:2204.07664, 2022.
Johannes Leuschner, Maximilian Schmidt, Daniel Otero Baguer, and Peter Maass. Lodopab-ct, benchmark dataset for low-dose computed tomography reconstruction. Scientific Data, 8(1):1-12, 2021.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild.
In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. arXiv preprint arXiv: (tt.:/8802.99900 2018
Christian Oestreicher. A history of chaos theory. Dialogues clinical neuroscience, 2022.
Gregory Ongie, Ajil Jalal, Christopher A Metzler, Richard G Baraniuk, Alexandros G Dimakis, and Rebecca Willett. Deep learning techniques for inverse problems in imaging. IEEE Journal Selected Areas in Information Theory, 1(1):39-56, 2020.
Patrick Putzky and Max Welling. Invert to learn to invert, 2019, URL https://axxiv.orr
abs/1911. 10914
Ernest Ryu, Jialin Liu, Sicheng Wang, Xiaohan Chen, Zhangyang Wang, and Wotao Yin. Plug-and- play methods provably converge with properly trained denoisers. In International Conference Machine Learning, pp. 5546-5557. PMLR, 2019.
Nimit Sharad Sohoni, Christopher Richard Aberger, Megan Leszezynski, Jian Zhang, and
Christopher Ré. Low-memory neural network training: A technical report. arXiv preprint
arXiv:1904.1066 2019.
Homer F Walker and Peng Ni. Anderson acceleration for exent-point iterations. SIAM Journal on Numerical Analysis, 49(4):1715-1735, 2011.
Jure Zbontar, Florian Knoll, Anuroop Sriram, Tullie Murrell, Zhengnan Huang, Matthew Muck- ley, Aaron Defazio, Ruben Stern, Patricia Johnson, Mary Bruno, Mare Parente, Krzysztof J. Geras, Joe Katsnelson, Hersh Chandarana. Zizhao Zhang, Michal Drozdzal, Adriana Romero, Michael Rabbat, Pascal Vincent, Nañissa Yakubova, James Pinkerton, Duo Wang, Erich Owens, C. Lawrence Zitnick, Michael P. Recht, Daniel K. Sodickson, and Yvonne W. Lui. fastmri: An open dataset and benchmarks for accelerated mri, 2018. URL https: /ettps//azzivv http//xxxi..rgaas//// org/abs/
1811. 08839
10


Reviews:
[{"review_id": "zDXOCWsM9nm", "reviewer": "Reviewer_ko7X", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper aims to propose a memory-efficient deep equilibrium model (DEQ) for solving inverse problems. However, after reading the paper, I am not sure if I have a clear understanding of what are the key contributions of the proposed LUSER model. It seems to me that the key difference between LUSER and prior DEQ models are:\n\n1. LUSER uses pre-learned DEQ models in its architecture (*\"...LUSER achieves this by adopting DEQ models as the learned\nregularizer update in a standard LU architecture....\"*)\n2. Somehow a shallow network can be used so that LUSER only needs to update a small number of parameters so that the memory requirement can be reduced.", "summary_of_the_review": "Not a good paper.", "strength_and_weaknesses": "*Weakness\n\n1. The paper is not well written. The methodology section lacks enough technical details as well as explanations for readers to understand the proposed model. For example, Fig. 1 is not illustrative, and the caption does not provide any explanation of the figure. Due to the above weakness, I am not equipped with enough knowledge of the model to evaluate the contribution of the model and interpret the numerical results.\n\n2. Please explain what is key difference between LUSER and prior models, such as DEQ4IP. I think DEQ4IP also only learns the regularizer network, isn't it? (Data-consistency layer is formed using the fixed forward model, while the regularizer layer is a neural network).\n\n3. The resolution of Fig. 2. is too low to see the visual difference.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "UE6GD22fw4", "reviewer": "Reviewer_3hmh", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The proposed work, LUSER, aims to improve the computational efficiency of deep unfolding network at training for solving  general image inverse problems.  To do so,  the proposed LUSER incorporated the fixed-point training of the neural network itself within each unrolling iterations by using the deep equilibrium models (DEQ). DEQ is a recently proposed framework for memory efficient learning of an infinite-depth unrolled network as implicitly defined by a fixed point of an operator. The authors have also conducted experiments to show that their approach is able to get similar quality reconstructions with reduced memory, compared to normal deep unfolding.", "summary_of_the_review": "Consequently, given the pros and cons on balance, I feel this is a very borderline paper, and I vote for borderline reject tentatively", "strength_and_weaknesses": "Strengths:\n\nThe studied problem is interesting and worth devoting to. The main advantages of the presented technique are the reduced execution time of the data-consistency layer compared to the one based on deep equilibrium models (DEQ4IP) and improved network capacity compared to traditional deep unfolding (LU-DnCNN). \n\nWeakness:\n\nHowever, there are some issues that I would like to highlight:\n\n1. One major concern is that the proposed method seems only improve the neural networks learning capacity by running it multiple times within each unrolling iteration. When using the same number of parameters, the memory costs suppose to be the same with normal deep unfolding networks. Unlike DEQ, it not fundamentally solves the memory burden issue when unfolding more iterations (layers) or using more sophisticated networks such as transformers. \n\n2. The conducted experiments cannot sufficiently support the authors’ claim and the numerical results are somehow vague to read. For example, the various properties of different approach in Table 4 are not well supported by the explicit numerical validation. One better way is to present the training/validation loss against time and epoch, compared to other baseline methods. Besides, it is unclear why DEQ4IP has the same PSNR/SSIM results to LU-DnCNN for MRI reconstruction. Since DEQ4IP can run more iterations during training, it should result in better performance, as reported in [Gilton et al. (2021)].\n\n3. While the proposed method aims to improve computational efficiency of the forward model, no large-scale imaging tasks are conducted in this paper, which also diminishes the contribution.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "L53bT4Gnq_", "reviewer": "Reviewer_4Y5M", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper combines algorithm unrolling, i.e. loop unrolling, with the deep equilibrium models for reduced memory and improved results. ", "summary_of_the_review": "The paper combines loop unrolling with deep equilibrium models. The experimental results show that the proposed approach does not achieve better results than DEQ. Its benefit over LU is also questionable as the paper does not show results for a large-scale problem such as processing 3D datasets.", "strength_and_weaknesses": "Strengths:\n\n1. The paper reduces the memory cost for loop unrolling by combining it with the deep equilibrium models. \n\nWeaknesses:\n1) The novelty of this paper is quite limited. Authors combine LU and DEQ, but results seem to be worse than DEQ. Author claims their approach to be better than LU and DEQ in several aspects such as stability. However, extensive experiments have not been conducted to backup these claims.\n\n2) Authors claim for even medium scale problems there is a noticeable tradeoff for accuracy for LU models. Authors have focused on 2D data to show that their approach requires less memory and can deal with large scale complex inverse problems. However, for many applications involving complex operations such as MRI,  LU have achieved a state-of-the-art results (see fastMRI challenge, Muckley et. al, TMI 2021) for 2D datasets.  Thus, memory cost has not been a main challenge for processing 2D data. This issue arises when processing 3D datasets. However, authors have not included any experiments focusing on processing 3D datasets.  \n\n3) Authors mention that their approach can tackle stability issues seen in DEQ4IP.  However, there is no experiment backing up this claim.\n\n4) Experimental results from Table 2 shows that, DEQ4IP outperforms LUSER. Luser can *quantitatively* achieve slightly better performance when parameters are not shared.  However, quantitative metrics may not be reflective of the performance.  In figure 2, MRI images for LUSER approach seem to be more blurry than DEQ4IP. Similarly, DEQ4IP achieves sharper results for CT. Thus, in practice, LUSER provides worse results than the counterpart approach. Also, why LUSER-SW is not included in figure 2?\n \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "H0z5Qvfu6O", "reviewer": "Reviewer_VVvv", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work proposes a new unrolled network architecture for inverse problems. It's main contribution is to replace the large feed-forward architecture (e.g., DnCNN) used in typical unrolled architectures with a small recurrent architecture, known as a Deep Equilibrium (DEQ) model. The DEQ model iterates until it converges to a fixed point and only the gradients associated with this fixed point need to be stored, thus making the proposed architecture more efficient to train.\n\nThe proposed architecture is applied to simulated image deblurring, CT, and MRI where it performs incrementally better than a standard unrolled projected gradient descent baseline. The proposed architecture is far more memory efficient to train than the baseline, but actual training times are not provided.\n", "summary_of_the_review": "The paper proposes a simple and novel technique for significantly reducing the memory costs associated with solving large scale imaging inverse problems with unrolled neural networks. While the current submission is missing important info (e.g., inference times), I remain mostly positive about this paper.", "strength_and_weaknesses": "# Strengths\nSolving large scale imaging inverse problems with unrolled network architectures is an important task and memory efficient training is one of the largest open problems in this field. This work represents a meaningful step in solving this problem. \n\nThe paper is generally well-written and easy to follow. \n\nThe proposed network architecture is well-motivated and simple -- this is a good thing.\n\n# Weaknesses\nWhile the paper is not claiming to be state-of-the-art, the comparisons in the paper are somewhat lacking and the unrolled DnCNN baseline is far from state-of-the-art. Even removing biases from the network may have significantly increased performance [A] so it's unclear if there are really any performance (rather than memory) benefits associated with the proposed method.\n\n[A] Mohan, Sreyas, et al. \"Robust And Interpretable Blind Image Denoising Via Bias-Free Convolutional Neural Networks.\" International Conference on Learning Representations. 2019.\n\nThe training procedure description states only that \"The weights \\theta of the model can be trained via implicit differentiation, removing the need to store all the intermediate activations\". This could use a citation or additional explanation. \n\nA number of statements lack evidence or citations. E.g., \"The implicit DEQ models are smaller in size but just as expressive as typical convolutional models\".\n\nWhile Table 4 is a useful summary of the relative difference in training times, evaluation times, etc., the actual training times, evaluation times, etc. need to be in the paper. \"Fast\" vs \"slow\" is almost meaningless without context and this imprecision makes it nearly impossible to fully evaluation the submission -- the proposed method could have 1 day inference times.\n\nIn table 5, two comparisons isn't enough to show any trends. Would 7 layers further improve performance? When does performance plateau?\n\nI would avoid stating LUSER provides a path forward for non-linear inverse problems. While this may be true, there's zero evidence for this in the paper.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Decision: Reject