{'fraction_non_alphanumeric': 0.043677431059506534, 'fraction_numerical': 0.011565674891146589, 'mean_word_length': 4.734460338101431, 'pattern_counts': {'":': 0, '<': 1, '<?xml version=': 0, '>': 0, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Memory-Augmented Deep Unfolding Network for Compressive Sensing

Mapping a truncated optimization method into a deep neural network, deep unfolding network (DUN) has attracted growing attention in compressive sensing (CS) due to its good interpretability and high performance. Each stage in DUNs corresponds to one iteration in optimization. By understanding DUNs from the perspective of the human brain's memory processing, we find there exists two issues in existing DUNs. One is the information between every two adjacent stages, which can be regarded as short-term memory, is usually lost seriously. The other is no explicit mechanism to ensure that the previous stages affect the current stage, which means memory is easily forgotten. To solve these issues, in this paper, a novel DUN with persistent memory for CS is proposed, dubbed Memory-Augmented Deep Unfolding Network (MADUN). We design a memory-augmented proximal mapping module (MAPMM) by combining two types of memory augmentation mechanisms, namely High-throughput Short-term Memory (HSM) and Cross-stage Long-term Memory (CLM). HSM is exploited to allow DUNs to transmit multi-channel short-term memory, which greatly reduces information loss between adjacent stages. CLM is utilized to develop the dependency of deep information across cascading stages, which greatly enhances network representation capability. Extensive CS experiments on natural and MR images show that with the strong ability to maintain and balance information our MADUN outperforms existing state-of-the-art methods by a large margin. The source code is available at https://github.com/jianzhangcs/MADUN/.


INTRODUCTION
Compressive sensing (CS) is a novel methodology of acquisition and reconstruction. Signal is first sampled and compressed with linear random transformations. Then, the original signal can be reconstructed from far fewer measurements than required by the sub-Nyquist sampling rate [25] [20]. Due to its attractive merits of the simple/fast sampling process and the low demand for data transmission and storage, the CS method has spawned many applications, including but not limited to single-pixel imaging [6] [24], accelerated magnetic resonance imaging (MRI) [21], wireless remote monitoring [47], and snapshot compressive imaging [36] [35].
Mathematically, a random linear measurements y ∈ R can be formulated as y = Φx, where x ∈ R is the original signal and Φ ∈ R × is the measurement matrix with ≪ .
is the CS ratio. Obviously, CS reconstruction is an ill-posed inverse problem. To obtain a reliable reconstruction, the conventional CS methods commonly solve an energy function, arg min where F (x) denotes a prior term with regularization parameter . For the traditional CS methods [12][18] [43][42] [7][22] [52], the prior term can be the sparsifying operator corresponding to some pre-defined transform basis, such as discrete cosine transform (DCT) and wavelet [49] [50]. They enjoy the advantages of strong convergence and theoretical analysis in most cases, but usually inevitably suffer from high computational complexity and face the trouble of choosing optimal transforms and parameters [51]. Recently, fueled by the powerful learning ability of deep networks, several deep network-based image CS reconstruction algorithms have been proposed [16] [32][8] [2], which can be generally grouped into two categories: deep non-unfolding networks (DNUNs) and deep unfolding networks (DUNs). For DNUN, it aims to directly learn the inverse mapping from the CS measurement domain to the original signal domain [23] [11]. For DUN, it combines deep network with optimization and trains a truncated unfolding Figure 1: Architecture of our proposed MADUN, which consists of stages. x denotes the full-sampled image for training, Φ denotes measurement matrix, y is the under-sampled data, x (0) denotes the initialization, x ( ) denotes the recovered result of MADUN and x ( ) stands for the output images of ( )-th stage. The dotted red line denotes High-throughput Short-term Memory (HSM) transmission and the dotted green line denotes Cross-stage Long-term Memory (CLM) transmission.
DUN is usually composed of a fixed number of stages where each stage is mainly influenced by its direct former one, and the mechanism can be considered as short-term memory, e.g., ISTA-Net [40] and DPDNN [44]. In each stage, the information transmission is usually hampered by the feature transformation with channel number reduction from multiple to one. Besides, the information in the prior stages is easily forgotten and the long-term dependency problem is rarely realized, which causes inadequate recoveries.
To address the above issues, in this paper, we propose a Memory-Augmented Deep Unfolding Network (MADUN), focusing on CS reconstruction. In our MADUN, we design a Memory-Augmented Proximal Mapping Module (MAPMM) which contains two different memory augmentation mechanisms, as shown in Figure 1. To reduce information loss between each two adjacent stages, we design a High-throughput Short-term Memory mechanism (HSM) to adaptively add multi-channel information and well ensure maximum signal flow. Also, a Cross-stage Long-term Memory mechanism (CLM) is exploited to explicitly construct the deep and adjustable long-term information path across stages. Our MADUN is a comprehensive framework, considering the merits of previous optimization-inspired DUNs and meanwhile, improving the network representation ability by augmenting entire information transmission among all stages. It enjoys the advantages of both the satisfaction of interpretability and the ensurance of information abundance. The major contributions are summarized as follows: • We propose a novel Memory-Augmented Deep Unfolding Network (MADUN) with persistent memory for CS reconstruction, which is able to adaptively capture the adequate features and recover more details and textures.  [32] design a dual-path network to learn the structure-texture decomposition in a data-driven manner. Residual learning and U-Net structure are adopted in CS-MRI to successfully learn the aliasing artifacts [10]. Some works jointly learn the measurement by a sampling subnetwork and the recovery by a reconstruction sub-network from the training data. CSNet [27] can avoid blocking artifacts by learning an end-to-end mapping between measurements and the whole reconstructed images. Shi et al. [28] exploit a convolutional neural network for scalable sampling and quality scalable reconstruction. Apparently compared with traditional methods, DNUNs can represent image information flexibly with the fast inferencing efficiency. However, in DNUNs, the learning performance seriously depends on the careful tuning and the sampling matrix is not wellembedded in the reconstruction process, which not only result in the very tricky training schemes but also drag down the network performance due to the difficulty of directly learning the recovery mapping without explicit matrix guidance.
Although existing DUNs inherit a good structure from optimization and exhibit friendly interpretability, we find that, from the perspective of neural networks, DUNs' inherent design of taking one-channel images as inputs and outputs for each stage will greatly hinder network information transmission and lose much more image details. Chen et al. [2] propose a contextual memory (CM) module which can augment the sequential links across stages. However, the CM module is still hard to maintain abundant details due to the poor storage capacity of one-channel stage outputs.

PROPOSED METHOD
In this section, we will elaborate on the design of our proposed Memory-Augmented Deep Unfolding Network (MADUN) for CS.

Basic Network Architecture
As we all know, DUN is a kind of CNN method that combines an efficient iterative algorithm. Iterative shrinkage-thresholding algorithm (ISTA) is well suited for solving many large-scale linear inverse problems [40], e.g., general CS and CS-MRI. Traditional ISTA solves the CS reconstruction problem in Eq. (1) by iterating between the following update steps: x ( ) = arg min where is the step size. Inspired by ISTA-Net + [40] which is a DUN unfolded by ISTA, Eq. (3) and Eq. (4) can be expressed as two modules in the ( )-th stage (0 < ≤ ) as shown in Figure 2 (a): Eq. (5) is called gradient descent module (GDM) where ( ) is a learnable parameter, and Eq. (6) is called proximal mapping module (PMM) which is actually a CNN-based denoiser.
To keep a simple structure with the high recovery accuracy, PMM uses basic convolution layers (Conv) and residual blocks (RB) which generate residual outputs by the structure of Conv-ReLU-Conv, and adopts an efficient reconstruction network [45] as shown in Figure 2 (b), which can be divided into four parts: • A convolution layer receives one-channel network inputs r ( ) and generates multi-channel outputs, yielding Conv 1 (r ( ) ). • A mapping layer extracts deep representation, which consists of two residual blocks (RB 1 (·) and RB 2 (·)) in PMM due to optimizing easily [9]. Here, s ( ) is the intermediate result produced between the two residual blocks. • A convolution layer Conv 2 (·) outputs residual result by the feature conversions from multi-channel to one-channel. • A residual learning strategy produces final module outputs.

Memory-Augmented DUN
Motivated by the viewpoint that the human brain achieves memory by preserving and storing what they acquire or are informed previously [4], the information flow across stages can also be regarded as short-term memory transmission in DUNs, where x ( ) in Figure 2 is as the link between each two adjacent stages. However, Conv 2 (·) in Eq. (7) is a lossy conversion process which would cause memory loss. Here, we propose a Memory-Augmented Deep Unfolding Network (MADUN) which augments memory across stages, each of which shares a unified two-step recovering scheme as follows: Here, GDM is trivial [40] and is the same as Eq. (3), and Memory-Augmented Proximal-Mapping Module (MAPMM) is introduced with two types of memory augmentation mechanisms which can highly strengthen the information transmission among all stages.

Memory-Augmented Proximal Mapping
On the basis of PMM, MAPMM introduces High-throughput Shortterm Memory (HSM) and Cross-stage Long-term Memory (CLM), which greatly enhance the network representation capability.

High-throughput Short-term Memory (HSM).
To address the information-lossy intra-stage transmission, we introduce HSM with multi-channel z ( −1) (of size × × ) to bridge a high-capacity path between adjacent stages. As shown in Figure 2 (c) and as the output of RB 2 (·) at ( − 1)-th stage, z ( −1) is: Concatenating z ( −1) with r ( ) followed by Conv 3 (·) generates -channel features. The remaining process is the same with Eq. (7). Thus, PMM with HSM, denoted by x ( ) = H ( ) PMM-HSM (r ( ) ), is: where ∥ denotes concatenating the feature maps. Obviously, z ( −1) transmits high-throughput information from the previous stage to the current stage, achieving multi-channel short-term memory which ensures maximum flow and reduces information loss caused by the conventional channel-shrinking transformation.

Cross-stage Long-term Memory (CLM).
For modeling explicitly the long-range dependencies among all cascading stages, we utilize a ConvLSTM layer [29], which has been well-validated to be a stable and powerful way to balance the past and current states, to develop the cross-stage long-term memory (CLM) and further enhance the signal conversion and transmission.
(15) Comparing Eq. (15) with Eq. (7), it is obvious to observe that z ( −1) is able to effectively add high-throughput information between each two adjacent stages, while [h ( −1) , c ( −1) ] can augment mid/ high-frequency information by applying ConvLSTM layers across cascading stages. Therefore, our proposed MAPMM enhances information transmission among all stages by adaptively learning these two types of memory mechanisms.

Initialization of HSM and CLM
Given the measurements y as the known information, HSM z ( ) is initialized by applying a light-weight one-convolution layer Conv 0 (·) on the image Φ ⊤ y to generate feature maps, yielding: Considering that CLM [h ( ) , c ( ) ] has its physical meaning and it is practical to assume no prior knowledge is used at the beginning, therefore, we simply initialize [h ( ) , c ( ) ] to be zero.

Network Parameters and Loss Function
The learnable parameter set in MADUN, denoted by , can be expressed as is a one-convolution layer with one input channel and output channels. H ( ) MAPMM (·) consists of Conv 3 (·), Conv 2 (·), RB 1 (·), RB 2 (·) and ConvLSTM(·) whose parameters are different in each stage. Conv 3 (·) with + 1 input channels and output channels receives  networks inputs, and Conv 2 (·) with input channels and one output channel outputs residual results. Also, the parameters of RB 1 (·) and RB 2 (·) are both from two convolution layers without the change of the channel number (with input channels and output channels), while ConvLSTM(·) is same with [19]. Here, all the convolutions in MADUN adopt 3 × 3 filter kernels. Given a set of full-sampled images x =1 and some sampling patterns with CS ratio , the under-sampled -space data is obtained by y = Φx , producing the train data pairs {(y , x )} =1 . Our MADUN takes y as inputs and generates the reconstruction result x ( ) as outputs with the initialization x (0) = Φ ⊤ y . The loss function is designed to use 1 loss between x and x ( ) as: where , and represent the number of training images, the size of each image and the stage number of MADUN respectively.

EXPERIMENT 4.1 Implementation Details
We use the 400 training images of size 180 × 180 [3], generating the training data pairs {(y , x )} =1 by extracting the luminance component of each image block of size 33 × 33, i.e. = 1, 089. Meanwhile, we apply in the data augmentation technique to increase the data diversity. For a given CS ratio, the corresponding measurement matrix Φ ∈ R × is constructed by generating a random Gaussian matrix or a jointly learned matrix and then orthogonalizing its rows, i.e. ΦΦ ⊤ = I, where I is the identity matrix. Applying y = Φx yields the set of CS measurements, where x is the vectorized version of an image block.  Table 2.  To train the network, we use Adam optimization [13] with a learning rate of 0.0001, and a batch size of 64. We also use momentum of 0.9 and weight decay of 0.999. For MADUN, the models are trained with 410 epochs separately for each CS ratio. Each image block of size 33 × 33 is sampled and reconstructed independently for the first 400 epochs, and for the last ten epochs, we adopt larger image blocks of size 99 × 99 as inputs to further fine-tune the model. To alleviate blocking artifacts, we firstly unfold the blocks of size 99×99 into overlapping blocks of size 33×33 while sampling process Φx and then fold the blocks of size 33 × 33 into larger blocks while initialization Φ ⊤ y [30]. We also unfold the whole image with this approach during testing. The default stage number is set to be 25 and the default number of feature maps is set to be 32. And the learnable parameters ( ) is initialized to 1. The CS reconstruction accuracies on the all datasets are evaluated with PSNR and SSIM.  Figure 3 is the results when ratio is 25% on Set11 dataset. To get a better trade-off between model performance and complexity, we choose = 25 as the default setting in all experiments.

Performance Effect of MADUN Components.
Based on the PMM structure, we do ablation experiments on the HSM and CLM at ratio=10% on Set11 and Urban100 dataset, as shown in Table 1. To demonstrate the effectiveness of HSM and CLM, we removes them from MADUN to form two simplified versions, corresponding to Figure 2 (c) and Figure 2 (d) respectively. We conclude that they are all effective ways to compensate the signal flow among stages, and HSM plays a more important role by bridging each two adjacent stages with a high-throughput short-term transmission.

Analysis of Memory Augmentation
We now separately illustrate how our different memory augmentation mechanisms affect information transmission. For HSM, high-throughput information which is added to the original one-channel short-memory between each two adjacent stages significantly reduces information loss and improves performance, as shown in Table 1. For further analysis of the effect on different positions of linking memory in HSM, we make two comparative experiments (c) and (d), where * HSM represents that the outputs of Conv 3 (·) are the points of memory transmission and •HSM denotes the outputs of RB 1 (·). As we can see, our default version of HSM transmits more refined memory than others.
And for CLM, we do experiments in different memory methods to prove the superiority of ConvLSTM, as shown in Table 2. In Figure 4, on the basis of PMM, (a) is the "Plus" strategy which takes a direct addition operation, (b) is the "Concat" strategy which adopts a single convolution layer to merge the all memory into the main intermediate features, and (c) is the implementation of CLM. CLM achieves better performance than others, which reveals that information is selectively retained by applying LSTM is useful for CS. To get deeper insights of the long-term information flow established by our CLM, as Figure 5 illustrates, we give the average kernel weight norms of the three convolution gates in each CLM, and plot the 1-D spectral density curves of memory features h ( ) in different CLM variants by integrating the 2-D power spectrums along each concentric circle [33] (i.e. by averaging the feature components with the same frequency in the DCT domain and normalizing the frequency range into [0,1]). From the upper subplot, we observe that the weight norms of the three convolution gates increase as the stage index increases, and the average norm of the forget gate is larger than the other two in most stages, which means that the cell  outputs c ( ) get close attention in the main trunk of MADUN and CLM plays a more important role in later stages. From the spectral density curves, we can see that, without long-term memory makes the network hard to keep the mid/high-frequency information; the "Concat" strategy also causes the information loss due to its weak adaptability; the "Plus" strategy introduces much feature space noise and weakens the proximal mapping process. However, our CLM with a flexible gated memory mechanism, achieves a better balance in the long-term information flow, which obtains the higher performances compared with others.

Qualitative Evaluation
We conduct two types of experiments on fixed random Gaussian sampling matrix and jointly learned sampling matrix, and compare them with the corresponding methods respectively.   Table 3. The models of ReconNet [16], ISTA-Net + [40], DPA-Net [32], IRCNN [46], MAC-Net [2] and iPiano-Net [30] are trained by the same methods and training datasets with the corresponding works, and DPDNN [5] and GDN [8] utilize the same training dataset with our method due to no CS reconstruction task in their original works. One can observe that our MADUN outperforms all the other competing methods in PSNR and SSIM across all the cases. Figure 6 further show the visual comparisons on challenging images on Set11 dataset and Urban100 dataset, respectively. Our MADUN generates images that are visually pleasant and faithful to the groundtruth.

Jointly Learned Matrix.
We mimic the CS sampling process y = Φx ∈ R using a convolution layer, and utilize another convolution layer with filters from Φ ⊤ ∈ R × to obtain initialization Φ ⊤ y, like [41]. We compare MADUN with some methods which jointly learn the sampling matrix and the reconstruction process. Table 4 presents quantitative results on Set11 and CBSD68 dataset. As we can see, our MADUN achieves the best performance on all ratios. Figure 7 further shows the visual comparisons on challenging images, one can see that MADUN recovers richer image structures and texture details than all the other methods.

Application to Compressive Sensing MRI
To demonstrate the generality of MADUN, we directly extend MA-DUN to the specific problem of CS-MRI reconstruction, which aims at reconstructing MR images from a small number of under-sampled data in k-space. In this application, we set the sampling process Φ in Eq. (5) to Φ = BF, where B is an under-sampling matrix and F is the discrete Fourier transform. In this case, we compare against seven recent methods for the CS-MRI domain. Utilizing the same training and testing brain medical images as ADMM-Net [37], the CS-MRI results of MADUNs are summarized in Table 5 for different CS ratios. From Table 5, we can see that our proposed method outperforms the state-of-the-art methods on testing brain dataset with all tested CS ratios, especially when ratios are smaller.

CONCLUSION
We propose a novel Memory-Augmented Deep Unfolding Network (MADUN) for CS, which addresses the problem of the informationlossy short-term transmission between each two adjacent stages in traditional DUNs by establishing and strengthening both the short-term and long-term memory flows with large capacity. Our two types of memory augmentation mechanisms, dubbed Highthroughput Short-term Memory (HSM) and Cross-stage Long-term Memory (CLM), are developed and integrated into our Memory-Augmented Proximal Mapping Module (MAPMM). HSM and CLM collaboratively achieve persistent and adaptive memory by bridging adjacent stages with the multi-channel signal path and developing the dependency of deep information across all cascading stages, respectively. Extensive CS experiments on both natural and MR images exhibit that the proposed MADUN achieves better signal balance based on both memory mechanisms and outperforms existing state-of-the-art deep network-based methods with large margins.
In the future, we will further extend our MADUN to other image inverse problems and video applications.