{'fraction_non_alphanumeric': 0.03036884621246761, 'fraction_numerical': 0.008078036884621246, 'mean_word_length': 4.774477447744775, 'pattern_counts': {'":': 0, '<': 0, '<?xml version=': 0, '>': 0, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Loop Unrolled Shallow Equilibrium Regularizer (LUSER) -- A Memory-Efficient Inverse Problem Solver

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 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.


Introduction
In an inverse problems we face the task of reconstructing some data or parameters of an unknown signal from indirect observations. The forward process, or the mapping from the data to observations, is typically well known, but ill-posed or non-invertible. More formally, we consider the task of recovering some underlying signal x from measurements y taken via some forward operator A according to where η represents noise. The forward operator can be nonlinear, but to simplify the notation, we illustrate the idea in linear form throughout this paper. A common approach to recover the signal is via an iterative method based on the least squares loss: For many problems of interest, A is ill-posed and does not have full column rank. Thus, attempting to solve (2) does not yield a unique solution. To address this, we can extend (2) by including a regularizing term to bias the inversion towards solutions with favorable properties. Common examples of regularization include 2 , 1 , and Total Variation (TV). Each regularizer encourages certain properties on the estimated signal x (e.g., smoothness, sparsity, piece-wise constant, etc.) and is often chosen based on task-specific prior knowledge.
Recent works [1] attempt to tackle inverse problems using more data-driven methods. 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 incorporating 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 a fixed number of iterations (typically around 5-10) due to computational constraints. [2] proposes an interesting alternative that takes advantage of deep equilibrium (DEQ) models [3][4][5][6] that we refer to as DEQ4IP. Equilibrium models are designed to recursively iterate on their input until a "fixed point" is found (i.e., the input no longer changes after passing through the model). They extend this principle to 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 LU architectures while reducing its memory consumption by using a shallow (relative to existing feed-forward models) DEQ as the learned regularizer update. Unlike DEQ4IP that converts the entire LU architecture into a 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 a single architecture for solving inverse problems. Our proposed architecture (i ) reduces the number of forward/adjoint operations compared to the work proposed by [2], (ii ) reduces the memory footprint during training without loss of expressiveness as demonstrated by our experiments, and (iii ) is more stable to train in practice than DEQ4IP. We empirically demonstrate better reconstruction across multiple tasks than state-of-the-art LU alternatives with a similar number of parameters, with the ability to reduce computational memory costs during training by a factor of up to 8×.
The remainder of the paper is organized as follows. Section 2 reviews related works in solving inverse problems. Section 3 introduces the proposed LUSER, which we compare with other baseline methods in image deblurring, CT, and MRI tasks in Section 4. We conclude in Section 5 with a brief discussion.

Loop Unrolling
As noted above, a common approach to tackling an inverse problem is to cast it as an optimization problem consisting of the sum of a data consistency term and a regularization term where r is a regularization function mapping from the domain of the parameters of interest to a real number and γ ≥ 0 is a well-tuned coefficient. The regularization function is chosen for specific classes of signals to exploit any potential structure, e.g., x 2 for smooth signals and x 0 or x 1 for sparse signals. The totalvariation (TV) norm is another popular example of a regularizer that promotes smoothness while preserving edges and is often used for When r is differentiable, the solution of (3) can be obtained in an iterative fashion via gradient descent. For some step size λ at iteration k = 1, 2, . . . , K, we apply the update: For non-differentiable r, 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 k+1 = prox τ,r (x k + λA (y − Ax k )).
The loop unrolling (LU) method considers running the update in (4) or (5), but replaces λγ∇r or the proximal operator with a learned neural network instead. The overall architecture repeats the neural network based update for a pre-determined number of iterations, fixing the overall computational budget. Note that the network is only implicitly learning the regularizer. In practice, it is actually learning an update step, which can be thought of as de-noising or a projection onto the manifold of the data. LU is typically trained endto-end, i.e., when fed some initialization x 0 , the network will output the final estimate x K , compute a loss with respect to the ground truth, and then back-propagate across the entire computational graph to update network parameters. While end-to-end training is easier to perform and encourages faster convergence, it requires all intermediate activations to be stored in memory. Thus, the maximum number of iterations is always kept small compared to classical iterative inverse problem solvers. Due to the limitation in memory, there is a 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, [2] extends the LU model to potentially infinite number of iterations using an implicit network, and [7] allows deeper unrolling iterations using invertible networks, while requiring recalculation of the intermediate 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 expensive to apply. For example, the forward operator may involve solving differential equations such as the wave equation for seismic wave propagation [8] and the Lorenz equations for atmospheric modeling [9].
Alternatively, one can design a richer regularization network. For example [10] uses a transformer as the regularization network and achieves extremely competitive results in the fastMRI challenge [11], but requires multiple 24GB GPU for training with batch size of 1, 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 a rich regularization network without the additional computational memory costs during training.

Deep Equilibrium Models for Inverse Problems (DEQ4IP)
Deep equilibrium (DEQ) models introduce an alternative to traditional feed-forward networks [3][4][5][6]. Rather than feed an input through a fixed (relatively small) number of layers, DEQ models solve for the "fixed-point" given some input. More formally, given a network f θ and some input x (0) and y, we recursively apply the network via until convergence. 1 In this instance, y acts as an input injection that determines the final output. This is termed the forward process. The weights θ of the model can be trained via implicit differentiation, removing the need to store all intermediate activations from recursively applying the network. This allows for deeper, more expressive models without the associated memory footprint to train such models. [2] demonstrates an application of one such model, applying similar principles to a single iteration of an LU architecture. Such an idea is a natural extension as it allows the model to "iterate until convergence" rather than rely on a "fixed budget". In another word, looping over (4) and (5) many times (in practice, usually around 50 iterations) until x k 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 a learned regularizer/de-noiser). It is also important to note is that such a model would have to apply the forward operator (and potentially the adjoint) many times during the forward process. Although this can be accelerated to reduce the number of applications, it is still often more than the number of applications for an LU equivalent. This can be an issue if the forward operator is computationally expensive to apply, an issue LU methods somewhat mitigate by fixing the total number of iterations.

Alternative Approaches to Tackle Memory Issues
is a deep invertible network that address the memory issue by recalculating the intermediate results from the output. However it is not ideal when forward model is computationally expensive. Gradient checkpointing [12] is another practical technique 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 a weight-tied neural network.

Methodology
LU methods have been the state-of-the-art approach for solving inverse problems due to their stability when training, inclusion of the forward model, and their near instantaneous inference times. However, there is a noticeable trade-off in terms of the memory requirements when training these models and the accuracy for even medium scale problems, making them nearly impossible to use for much larger scale inverse problems. DEQ4IP offers an interesting alternative, drastically reducing the memory required during training and allowing the flexibility to adjust accuracy during reconstruction when performing inference. However, these models can potentially suffer when the forward/adjoint operators are computationally expensive, particularly when it takes more iterations to converge than a standard fixed LU method.
To address these concerns, we propose a novel architecture for solving inverse problems called Loop Unrolled Shallow Equilibrium Regularizer (LUSER) as an approach that limits the number of forward/adjoint calls for particularly complex inverse problems while 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 learned regularizer update in a standard LU architecture. The implicit DEQ models are smaller in size but just as expressive as typical convolutional models used as the learned regularizer update allowing for an accurate reconstruction with less computational memory costs. Furthermore, learning a proximal update is a far simpler task compared to solving the inverse problem as a whole, reducing the instability during training commonly faced by DEQ models.
We adopt a "proximal gradient descent" styled LU architecture. The network takes in measurements y and some initial estimate x 0 . The architecture consists of K stages, alternating between a gradient descent step d k = x k + λA (y − Ax k ) for k = 1, 2, . . . , K, followed by a feed-forward pass of the shallow equilibrium model acting as a proximal update block. Figure 1 illustrates a single block in LUSER as a learned proximal operator. The input d k from the previous stage is processed only once by a set of input injection layers to avoid redundant computation. The input injection will determine the final fixed point output [5]. 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 a fixed point is found. Although Figure 1 shows the simplest way of finding the fixed point, in practice acceleration or other fixed point solvers are applied to solve for the fixed point. We apply Anderson acceleration [13] for all of our models.  Figure 1: A proximal block in LUSER At each stage, the shallow equilibrium regularizer attempts to output its best estimate of the ground truth x * . We introduce a skip connection between the input d k and the final output so that the regularizer is learning the residual r between the input d k and the ground truth x * instead. When x k converges to x * , d k 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 d k in the hopes that fewer iterations will be needed at later stages as the current estimate x k converges towards the ground truth. Let r (i) denote the i th update of the residual, where r (0) = 0 and r (∞) 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: Because of the fixed point solver within each regularization block, LUSER uses fewer layers to achieve a similar level of performance, thus we call LUSER a shallow equilibrium regularizer. We also explore two variants of LUSER, dubbed LUSER-SW and LUSER-DW. LUSER-SW refers to a shared-weight version of the proposed algorithm, where the proximal operator in all loop unrolled 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 intermediate reconstructions can be vastly different. Thus, training a single model to handle all these instances can be a 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 to 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.

Experiments and Results
In this section, we compare our proposed networks to LU with DnCNN as proximal operator (LU-DnCNN) 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 [14], LoDoPaB-CT [15], and single-coil knee data from fastMRI [11] 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).

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 to all channels. For accelerated MRI tasks, measurements in k-space (or frequency domain) are often subsampled due to the cost in measurement. The goal of MRI reconstruction is to recover the underlying physical structure from subsampled noisy measurements. We simulate the forward operator A with a 2-dimensional Fourier transform with randomly selected rows. We consider two common subsampling scenarios: 4x and 8x acceleration, or subsampling the columns in full measurement by a factor of 4 and 8 respectively. For the CT task, the forward operator is a Radon transform, and we uniformly select 200 out of 1000 angles in measurements. The adjoint of measurement A y 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 [16], but we use the adjoint for the purpose of consistency.
We fix the budget of LU-DnCNN 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 [17] 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, we use 2 convolutional layers each for the input injection layer and data layer. The mixing layer contains 3 convolutional layers.
In order to stablize training for DEQ inspired models, we wrap all convolutional layers in DEQ4IP and LUSER with Spectral Norm (SN) [18]. We list more details for the learned proximal network for LUSER in Table 1 for the case when the input has a single channel. C in and C out denote the input and output channels, Injection Layer SN(conv(C in :1, C out :32, ks:3, pad:1)) + LeakyReLU SN(conv(C in :32, C out :32, ks:3, pad:1)) + LeakyReLU Data Layer SN(conv(C in :1, C out :32, ks:3, pad:1)) GN + LeakyReLU SN(conv(C in :1, C out :32, ks:3, pad:1)) GN + LeakyReLU Mix Layer SN(conv(C in :64, C out :64, ks:3, pad:1)) GN + LeakyReLU SN(conv(C in :64, C out :64, ks:3, pad:1)) GN + LeakyReLU SN(conv(C in :64, C out :1, ks:3, pad:1)) ks refers to the kernel size of a convolutional layer, pad denotes the padding in 2-dimension, and GN stands for GroupNorm. We use two metrics to evaluate the quality of reconstruction: Peak Signal-to-Noise Ratio (PSNR) in dB and the Structural Similarity Index (SSIM): where MSE(x, x) = E[ x − x 2 ] is the mean squared error. In SSIM, µ x and µ x are pixel means of x and x, σ 2 x and σ 2 x are variances of x and x, and σ x x is their covariance. C 1 and C 2 are small positive numbers to stabilize SSIM. Note that although we use the same models as [2], we train our 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 a total of 2 channels), the metrics we recorded are more aligned with those reported in [2].
All models are trained with a single RTX6000 24GB GPU. Table 2 compares the average PSNR and SSIM. The different weight version of LUSER outperforms LU-DnCNN with a similar number of network parameters. LUSER-SW achieves similar level of performance in most tasks with only 5 layers, versus 17 layers in LU-DnCNN. DEQ4IP attains the best performance in image delurring task, but it requires repetitive computation of forward/adjoint operators. Training shared-weight architectures with auxillary losses improves the reconstruction quality in most tasks.     Table 3 compares the network sizes (number of parameters) as 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 1 for stochastic gradient descent. Notice that in MRI, because the Fourier transformation is implemented in tensor form, the minimum batch size it can take is 2. We use the batch size as a proxy for the memory requirements during training. Since DEQ4IP is an extension of LU-DnCNN, their networks are of the same size, but implicit DEQ models support larger batch sizes making DEQ4IP far more memory-efficient during training. The advantages of using DEQ models for LUSER and DEQ4IP are particularly highlighted in the case of limited memory (smaller GPUs). LU-DnCNN is unable to even train for the CT and MRI 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 of parameters. For example, even though LUSER-DW and LUSER-SW have different numbers of parameters, they share the same architecture/depth and thus use roughly the same amount of memory during training.

Main Results
We observe that LUSER takes longer to train in each epoch than LU-DnCNN, but the training converges in fewer iterations than LU-DnCNN. The total training time for all methods are roughly the same for the image deblurring and MRI tasks, but DEQ4IP takes longer to train on the CT problem as it is finding a fixed-point with a more complicated forward operator. During evaluation, LU-DnCNN is the fastest, and the DEQ4IP and LUSER are in similar speed. When the forward operator is computationally intensive, LUSER evaluates faster than DEQ4IP. Table 4 compares the various properties among the three architectures.

Effect of number of Layers in LUSER-SW
The experiments above show the results of LUSER-SW with only 5 layers as extreme case to highlight the effectiveness of even a vastly smaller network. The number of layers typically correlate with the expressiveness of a deep feedforward neural network. We also explore the effect of number of layers in the reconstruction quality for the learned proximal operators of LUSER. Table 5 shows the average PSNR and SSIM with one extra layer in the Mix Layer. It raises the PSNR by 1 dB and SSIM by 0.2, but sacrifices some memory in training. This suggests that one can design the network in LUSER according to their need in reconstruction quality and memory constraints.

Effect of weight sharing
As shown in Table 2, LUSER-DW has better performance than LUSER-SW with equivalent block structure, but has a 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 g i denote the i th proximal block where i = {1, 2, 3, 4}, and we form the LUSER network with the sequence of proximal blocks in order of g 1 , g 1 , g 2 , g 2 , g 3 , g 3 , g 4 , g 4 . Now, the number of parameters is only half of that in LUSER-DW, but maintains a similar level of performance. We denote this variant LUSER-PSW, which stands for partially shared-weight. Table 6 compares the average PSNR and SSIM of LUSER-SW, LUSER-PSW and LUSER-DW with the same block structure, which are trained using final loss only.

Conclusion
Loop unrolling architectures with deep convolutional layers as the learned regularizer update are a popular approach for solving inverse problems. Although its variants achieve state-of-the-art results across a variety of tasks, LU algorithms incur a huge memory cost when training due to the requirement of saving all intermediate activations, sometimes even requiring multiple GPUs to train on complex tasks. DEQ4IP offers an interesting alternative via extending LU 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 of loop unrolling architectures with deep equilibrium models as the learned regularizer updates, LUSER-SW and LUSER-DW. We verify the memory savings (by comparing batch sizes) relative to loop unrolling algorithms with a DnCNN model. Across all tasks, LUSER-DW outperforms LU-DnCNN with a similar number of network parameters, while reducing the memory requirements by a factor of 5 or more. LUSER offers a path forward for large-scale, complex, non-linear inverse problems that are currently infeasible to train with existing methods.