A D ESIGN S PACE S TUDY FOR LISTA AND B EYOND
Anonymous authors
Paper under double-blind review
A BSTRACT
In recent years, great success has been witnessed in building problem-speciﬁc deep networks from unrolling iterative algorithms, for solving inverse problems and beyond.
Unrolling is believed to incorporate the model-based prior with the learning capacity of deep learning.
This paper revisits the role of unrolling as a design approach for deep networks : to what extent its resulting special architecture is superior, and can we ﬁnd better?
Using LISTA for sparse recovery as a representative example, we conduct the ﬁrst thorough design space study for the unrolled models.
Among all possible variations, we focus on extensively varying the connectivity patterns and neuron types, leading to a gigantic design space arising from LISTA.
To efﬁciently explore this space and identify top performers, we leverage the emerging tool of neural architecture search (NAS).
We carefully examine the searched top architectures in a number of settings, and are able to discover networks that are consistently better than LISTA.
We further present more visualization and analysis to “open the black box”, and ﬁnd that the searched top architectures demonstrate highly consistent and potentially transferable patterns.
We hope our study to spark more reﬂections and explorations on how to better mingle model-based optimization prior and data-driven learning.
I NTRODUCTION
The signal processing and optimization realm has an everlasting research enthusiasm on addressing ill-conditioned inverse problems, that are often regularized by handcrafted model-based priors, such as sparse coding, low-rank matrix ﬁtting and conditional random ﬁelds.
Since closed-form solutions are typically unavailable for those model-based optimizations, many analytical iterative solvers arise to popularity.
More recently, deep learning based approaches provide an interesting alternative to inverse problems.
A learning-based inverse problem solver attempts to approximate the inverse mapping directly by optimizing network parameters, by ﬁtting “black box” regression from observed measurements to underlying signals, using synthetic or real-world sample pairs.
Being model-based and model-free respectively, the analytical iterative solvers and the learningbased regression make two extremes across the spectrum of inverse problem solutions.
A promising direction arising in-between them is called algorithm unrolling (Monga et al., 2019).
Starting from an analytical iterative solver designed for model-based optimization, its unrolled network architecture can be generated by cascading the iteration steps for a ﬁnite number of times, or equivalently, by running the iterative algorithm with early stopping.
The original algorithm parameters will also turn into network parameters.
Those parameters are then trained from end to end using standard deep network training, rather than being derived analytically or selected from cross-validation.
Unrolling was ﬁrst proposed to yield faster trainable regressors for approximating iterative sparse solvers (Gregor & LeCun, 2010), when one needs to solve sparse inverse problems on similar data repeatedly.
Later on, the unrolled architectures were believed to incorporate model-based priors while enjoying the learning capacity of deep networks empowered by training data, and therefore became a rising direction in designing principled and physics-informed deep architectures.
The growing popularity of unrolling lies in their demonstrated effectiveness in developing compact, data-efﬁcient, interpretable and high-performance architectures, when the underlying optimization model is assumed available.
Such approaches have witnessed prevailing success in applications such as compressive sensing (Zhang & Ghanem, 2018), computational imaging (Mardani et al., 2018), wireless communication (Cowen et al., 2019; Balatsoukas-Stimming & Studer, 2019), computer vision (Zheng et al., 2015; Peng et al., 2018), and other algorithms such as ADMM (Xie et al., 2019).
The empirical success of unrolling has sparkled many curiosities towards its deeper understanding.
A series of efforts (Moreau & Bruna, 2017; Giryes et al., 2018; Chen et al., 2018; Liu et al., 2019;
Ablin et al., 2019; Takabe & Wadayama, 2020) explored the theoretical underpinning of unrolling as a specially adapted iterative optimizer to minimizing the speciﬁc objective function, and proved the favorable convergence rates achieved over classical iterative solver, when the unrolled architectures are trained to (over)ﬁt particular data.
Orthogonally, this paper reﬂects on unrolling as a design approach for deep networks .
The core question we ask is:
For solving model-based inverse problems, what is the role of unrolling in designing deep architectures?
What can we learn from unrolling, and how to go beyond?
1. R ELATED WORKS : PRACTICES AND THEORIES OF UNROLLING
(Gregor & LeCun, 2010) pioneered to develop a learning-based model for solving spare coding, by unrolling the iterative shrinkage
thresholding algorithm (ISTA) algorithm
(Blumensath & Davies, 2008) as a recurrent neural network (RNN).
The unrolled network, called Learned ISTA (LISTA), treated the ISTA algorithm parameters as learnable and varied by iteration.
These were then ﬁnetuned to obtain optimal performance on the data for a small number of iterations.
Numerous works (Sprechmann et al., 2015; Wang et al., 2016; Zhang & Ghanem, 2018;
Zhou et al., 2018) followed this idea to unroll various iterative algorithms for sparse, low-rank, or other regularized models.
As the most classical unrolling example, a line of recent efforts have been made towards theoretically understanding LISTA.
Moreau & Bruna (2017) re-factorized the Gram matrix of the dictionary, and thus re-parameterized LISTA to show its acceleration gain, but still sublinearly.
Giryes et al.
(2018) interpreted LISTA as a projected gradient descent descent (PGD) where the projection step was inaccurate.
Chen et al.
(2018) and Liu et al.
(2019) for the ﬁrst time introduced necessary conditions for LISTA to converge linearly, and show the faster asymptotic rate can be achieved with only minimal learnable parameters (e.g., iteration-wise thresholds and step sizes).
Ablin et al.
(2019) further proved that learning only step sizes improves the LISTA convergence rate by leveraging the sparsity of the iterate.
Besides, several other works examined the theoretical properties of unrolling other iterative algorithms, such as iterative hard thresholding (Xin et al., 2016), approximated message passing (Borgerding & Schniter, 2016) , and linearized ADMM (Xie et al., 2019).
On a separate note, many empirical works (Wang et al., 2016; Gong et al., 2020) advocated that the unrolled architecture, when used as a building block for an end-to-end deep model, implicitly enforces some structural prior towards the model training (resulting from the original optimization objective)
(Dittmer et al., 2019).
That could be viewed as a special example of “architecture as prior” (Ulyanov et al., 2018).
A recent survey (Monga et al., 2019) presents a comprehensive discussion.
Speciﬁcally, the authors suggested that since iterative algorithms are grounded on domain-speciﬁc formulations, they embed reasonably accurate characterization of the target function.
The unrolled networks, by expanding the learnable capacity of iterative algorithms, become “tunable” to approximate the target function more accurately.
Meanwhile compared to generic networks, they span a relatively small subset in the function space and therefore can be trained more data-efﬁciently.
1. M OTIVATIONS AND CONTRIBUTIONS
This paper aims to quantitatively assess “how good the unrolled architectures actually are”, using
LISTA for sparse recovery as a representative example.
We present the ﬁrst design space ablation study 1 on LISTA: starting from the original unrolled architecture, we extensively vary the connectivity patterns and neuron types.
We seek and assess good architectures in a number of challenging settings, and hope to expose successful design patterns from those top performers.
As we enable layer-wise different skip connections and neuron types, the LISTA-oriented design space is dauntingly large (see Sections 2.1 and 2.2 for explanations).
As its manual exploration is infeasible, we introduce the tool of neural architecture search (NAS) into the unrolling ﬁeld.
NAS can explore a gigantic design space much more efﬁciently, and can quickly identify the subset of top-performing candidates for which we can focus our analysis on.
Our intention is not to innovate on NAS, but instead, to answer a novel question in LISTA using NAS as an experiment tool.
From our experiments, a beﬁtting quote may be: for designing neural networks to solve modelbased inverse problems, unrolling is a good answer, but usually not the best.
Indeed, we are able to discover consistently better networks in all explored settings.
From top candidates models, we observe highly consistent and potentially transferable patterns.
We conclude this paper with more discussions on how to better leverage and further advance the unrolling ﬁeld.
T ECHNICAL A PPROACH
Assume a sparse vector x ∗ =
[ x ∗ 1 , · · · , x ∗ M ] T ∈ R M , we observe its noisy linear measurements:
[Equation]
) where b ∈ R N , D =
[ d 1 , · · · , d M ] ∈ R N × M is the dictionary , and ε ∈ R N is additive Gaussian white noise.
For simplicity, each column of D is normalized, that is, (cid:107) d m (cid:107) 2 = (cid:107) D : ,m (cid:107) 2 = 1 , m = 1 , 2 , · · · , M .
Typically, we have N (cid:28) M .
A popular approach for solving the inverse problem: b → x , is to solve the LASSO below ( λ is a scalar):
[Equation]
using iterative algorithms such as the iterative shrinkage thresholding algorithm (ISTA):
[Equation]
where η θ is the soft-thresholding function 2 and L is a smoothness constant that decides the step size.
Inspired by ISTA, (Gregor & LeCun, 2010) proposed to learn the weight matrices in ISTA rather than ﬁxing them.
LISTA unrolls ISTA iterations as a recurrent neural network (RNN): if truncated to K iterations, LISTA becomes a K -layer feed-forward neural network with side connections:
[Equation]
If we set W b ≡ D T /L , W ( k ) x ≡ I − D T D
/L , α ( k ) ≡ 1 , θ ( k ) ≡ λ/L , then LISTA recovers ISTA 3 .
In practice, we start with x (0) set to zero.
As suggested by the seminal work (Liu et al., 2019), sharing W b and W x across layers does no hurt to the performance while reducing the parameter complexity.
We follow their weight-tying scheme in the paper, while learning layer-wise α ( k ) and θ ( k ) .
α ( k ) is separated from W x to preserve ﬂexibility.
We note one slight difference between our formulation (4), and the parameterization scheme suggested by Theorem 1 of (Chen et al., 2018).
The latter also restricted { W b , W x } to be coupled layerwise, and showed it yielded the same representation capability as the orginal LISTA in Gregor & LeCun (2010).
We instead have { W b , W x } as two independent learnable weight matrices.
The main reason we did not follow (Chen et al., 2018) is that W b - W x coupling was directly derived for unrolling ISTA, which is no longer applicable nor intuitive for other non-LISTA architecture variants.
Our empirical results also back that our parameterization in (4) is easier to train for varied architectures from search, and performs the same well when adopted in training the original LISTA.
Now, given pairs of sparse vector and its noisy measurement ( x ∗ , b ) , our goal is to learn the parameters Θ = { W b , W x , α ( k ) , θ ( k ) } K − 1 k =0 such that x ( K ) is close to x ∗ for all sparse x ∗ following some distribution P .
Therefore, all parameters in Θ are subject to end-to-end learning:
[Equation]
This problem is approximately solved over a training dataset { ( x ∗ i , b i ) }
Ni =1 sampled from P .
2. V ARYING THE C ONNECTIVITY IN LISTA
We ﬁrst present how we construct new LISTA-oriented architectures based on extensively varying the connectivity patterns, starting from the original unrolled LISTA.
The way we add connections to the original unrolled model is called Learnable Weighted Average (LWA).
Speciﬁcally, before the output of the k -th layer x ( k ) is fed into the next layer, we allow the outputs of previous layers
[Figure]
𝑊
[Figure]
𝛼
[Figure]
(𝑘) +
[Figure]
Original LISTA 𝑥 (𝑘) 𝑏 𝑊 𝑏 𝑊 1 𝛼 1 (𝑘) + 𝑔 𝑘,𝑘 𝑥 (𝑘−1) 𝑐 𝑘,𝑘 + 𝑔 𝑘−1,𝑘 𝑐 𝑘−1,𝑘 𝑥 (1) 𝑔 1,𝑘 𝑐 1,𝑘 ⋯⋯
[Figure]
to be possibly (not necessarily) connected to x ( k ) .
The input to the next layer is calculated with a learnable average of all the connections.
Formally,
[Equation]
6) where c i,k are (unbounded) model parameters that will be trained with data, and g i,k are gates whose values are chosen to decide speciﬁc connectivity.
Note that LWA introduces a few extra parameters (coefﬁcients
c i,k ) than LISTA, but those constitute only a very minor portion w.r.t.
the size of Θ .
We completely recognize that there are other possibilities to vary connections, and did our due diligence to make the informed choice of LWA.
First, we only focus on adding connections, motivated from both the deep learning and the empirical unrolling ﬁelds.
On one hand, prevailing success in deep learning (He et al., 2016; Huang et al., 2017) endorse that more densely connected networks are more likely to yield better performance.
Prior unrolling works (Wang et al., 2016) also empirically found that removing connections from an unrolled architectures would deteriorate its effectiveness.
We also provide additional experiments in the Appendix A.4 to demonstrate that pruning existing connections in LISTA will be consistently harmful.
Second , besides injecting new connections by
LWA, we have also experimented with several other options, including a naive averaging way (simply averaging the outputs of all chosen layers), and a momentum-inspired averaging way (which is inspired by placing a momentum term into the unrolled algorithm).
Through various experiments, we have conﬁrmed that both options are almost consistently inferior to LWA, and therefore focus on
LWA in the main paper.
The details of the two other options, as well as the experimental comparisons, can be founded in the Appendix A.3 too.
We by default keep all existing connections in LISTA itself, i.e. the gate function g k,k that controls the connection from the immediate preceding layer is always 1.
In this case, for a K -layer network, the total number of possible architectures is 2 1+2+ ··· +( K − 2) =
2 ( K − 1)( K − 2).
When K = 16 as commonly set in previous unrolling works (Chen et al., 2018; Borgerding & Schniter, 2016), only varying the connnecitivy will lead to a total of over 2 105 ≈ 4 × 10 31 different architectures.
2. V ARYING THE N EURON T YPE IN LISTA
Besides connection patterns, another important factor that might affect the model performance is the choice of the neuron (non-linear activation function) type.
The default neuron in LISTA is the soft-thresholding operator with learnable thresholds (bias term), naturally from unrolling ISTA.
To study whether they might be potential better neuron options or combinations, we augment the search space to allow each layer to select its neuron from three types: { soft-thresholding, ReLU,
LeakyReLU
( α = 0 . 1 ) } .
That increases the total architecture numbers further by 3 16 times, resulting in the ﬁnal gigantic space of ∼ 1 .
75 × 10 39 candidate architectures deriving from a 16-layer LISTA.
2. H OW TO E FFICIENTLY F IND AND A NALYZE G OOD A RCHITECTURES
The daunting size of our LISTA search space poses a severe challenge: how to effectively ﬁnd the best small subset of architectures, that we can observe and reliably summarize patterns from?
A manual exploration is infeasible, and hence we refer to the tool of neural architecture search (NAS) (Zoph & Le, 2016) to fulﬁll our goal.
Here NAS is leveraged as an off-the-shelf tool to explore the search space efﬁciently and to quickly focus on the most promising subset of top candidates.
Are NAS results stable and trustworthy?
Since our main goal is to analyze the found topperformer models and to identify design patterns, the effectiveness and stability of NAS are crucial to the trustworthiness of our conclusions.
We take several actions to address this potential concern.
First, to avoid the instability arising in training and evaluating the sampled candidate models, we do not adopt any parameter sharing (Pham et al., 2018; Liu et al., 2018).
Parameter sharing is currently a popular scheme in NAS to save search time, but can introduce search bias (Chu et al., 2019).
Despite the accompanied tremendous computation costs, we train each individual model from scratch to fully converge, to ensure the most stable, thorough and fully reproducible search results as possible.
Second, to avoid the possible artifact of any single NAS algorithm, we repeat our experiments using three different state-of-the-art, non-weight-sharing NAS algorithms, including reinforcement learning (RL) based method (Zoph & Le, 2016), regularized evolution based method (Real et al., 2019) and random sampling based method (Li & Talwalkar, 2019).
Those algorithms usually yield diverse search behaviors, and in general there is no clear winner among the three (Li & Talwalkar, 2019).
We use the negative validation NMSE as the optimization reward for them all.
Somehow surprisingly, we ﬁnd: (1) the best candidate architectures (e.g., top-50 in terms of NMSE) found by the three methods reach almost the same good NMSE values; and (2) those best subsets are highly overlapped (nearly identical) in the speciﬁc models found, and show quite consistent architecture patterns.
Hence due to the space limit, we focus on reporting all results and ﬁndings from RL-based
NAS, as the other two produce almost the same.
We hypothesize that our gigantic LISTA search space might have a smoother and more friendly landscape compared to typical NAS benchmarks, that facilitates state-of-the-art NAS algorithms to explore.
We leave that for future work to verify.
Third , the top-1 architecture found by NAS will inevitably be subject to sampling randomness and other search algorithm ﬂuctuations (e.g., hyperparameters).
Therefore, instead of overly focusing on analyzing the top-1 or few architectures, we take one additional step to create an “average architecutre” from the top-50 architectures, to further eliminate random factors that might affect our results’ credibility.
Speciﬁcally, for every possible connection between layers, we compute the percentage (between 0 and 1) of models from the top-50 set that have this connection.
We then use a threshold of 0.5 to decide whether the average architecture will have this connection (if the percentage ≥ 0.5) or not (if < 0.5).
As we will visualize and analyze in Section 4, the top-50 architectures usually have a high agreement level on connections, making most percentages naturally close to either 0 or 1.
The average architecture helps more robustly clearly perceive the “good patterns” shared by top models.
E XPERIMENTS AND ANALYSIS 3. E XPERIMENT S ETTING AND I MPLEMENTATION D ETAILS
Searching.
As discussed in Section 2.3, we focus on the RL-based NAS (Zoph & Le, 2016) without weight sharing, i.e., train every sampled model from scratch until full convergence.
Our experiments are performed on 400 Cloud TPUs, with more than 25,000 architectures sampled during each search.
We will publicly release our codes, and searched & trained models, after paper acceptance.
After the search is complete, we focus on the “average architecture” from the top-50 sampled set, as deﬁned in Section 2.3.
Typical NAS works that often ﬁnd the few best models to have notable performance gaps between each other; different from them, our top-50 searched models have only minor NMSE deviations, in addition to their consistent architecture patterns.
That endorses the stability of our search, and potentially reafﬁrms our smoothness assumption of LISTA search space.
To emphasize, we have only performed NAS in the problem setting of Section 3.2 (synthetic, noiseless).
For all remaining experiments, we directly re-train and evaluate the same average architecture, on corresponding training and testing sets.
The main motivation is to avoid the hassle of re-doing the tedious model search for every new problem .
Instead, we wish to show that by searching only once in a problem setting, the searched architecture not only performs best in the same setting, but also can be directly applied (with re-training) to other related settings with superior performance.
Such transferrability study has also been a popular subject in NAS literature (Zoph et al., 2018), where one searches on one dataset and re-trains the searched model on another similar dataset.
Training and evaluating sampled models.
We adopt the same stage-wise training strategy from (Chen et al., 2018; Liu et al., 2019) and use the same default K = 16 layers.
Every model is trained
[Table]
[Table] [Table]
with batch size 128 and learning rate 0.0005 on 2 Cloud TPUs.
In most experiments, the training, validation and testing sets are sampled i.i.d.
without overlap, except the training/testing mismatch experiments where we purposely have training and testing sets sampled from different distributions.
Also by default, we sample 102,400 samples for training, and 10,240 for validation and testing each, constituting data abundant settings.
A data limited setting will be discussed later.
For all synthetic data experiments, the model performance is evaluated with normalized MSE (NMSE) in decibel (Chen et al., 2018) averaged over the testing set: the lower the better.
Baselines.
We compare with three strong hand-crafted baselines: the original unrolling LISTA model; the densely connected variant of LISTA (every layer connected to all its preceding layers), denoted as Dense-LISTA ; and another improved unrolled sparse recovery model from FISTA (Moreau & Bruna, 2017) called LFISTA , which could also perceived as model-based injection of more connections into LISTA.
For the fair comparison with LFISTA, we adopt the same parameterization to (4) and (6) to model its Nesterov acceleration term.
We ﬁnd that to consistently improve the training stability and ﬁnal performance of LFISTA too.
More details are in the Appendix A.1.
3. S EARCH AND E VALUATION IN THE N OISELESS S ETTING
We perform our search on the simplest synthetic setting with no additive noise, i.e. ε = 0 in (1).
We generate all our data, with a random sampled dictionary D ∈ R m × n with m = 250 , n = 500 , following (Chen et al., 2018; Borgerding & Schniter, 2016).
In Appendix A.2, we also investigate the effect of problem size by varying m and n , as well as the effect of an ill-conditioned dictionary (instead of Gaussian).
We sample the entries of D i.i.d.
from the Gaussian distribution D ij ∼
N (0 , 1 /m ) and normalize its columns to unit (cid:96) 2 norm.
To sample the sparse vectors x , we decide each of its entry to be non-zero following the Bernoulli distribution with p = 0 . 1 , and then generate its non-zero entries from the N (0 , 1) .
We tried to adjust the dictionary formulation, sparsity level or nonzero magnitude distributions in some experiments, and observe the conclusions to be highly consistent.
We therefore only report in this setting due to the space limit.
The results are reported and in the row [ b ] in Table 1, where we compared our searched topaverage architecture with three baselines.
Our main observations are:
• The searched average model (signiﬁcantly) outperforms hand-crafted ones.
That proves our concept: much stronger models could be found by NAS in the LISTA-oriented design space.
Particularly, it surpasses the vanilla LISTA by a large gap of 10.9 dB. • The choices of neuron types are “embarrassingly uniform”.
Even we allow each of thelayers to independently select neurons, none shows to favor ReLU or leaky ReLU.
All our toparchitectures unanimously adopt soft-thresholding as the only neuron type for all their layers.
It is somehow a surprise, and hints the important role of model-based prior in unrolled networks.
• Non-trivial connectivity patterns are discovered, While LFISTA and Dense-LISTA both improve over LISTA, our search result indicates that “the denser the better” is not the true claim here (Table 1, row [ a ] ).
That differs from the general mind in computer vision models (Huang et al., 2017).
More discussions on the searched connectivity are in Section 4.
3. T RASNFERABILITY S TUDY FOR THE S EARCHED M ODEL
We now present the transferablity study, to re-train and evaluate the above-searched average architecture in more challenging settings.
The baselines are also re-trained and compared in fair settings
Noisy measurement [ c, d ] We ﬁrst add non-zero Gaussian noise ε (i.i.d.
across training and testing) to synthetic data.
We use two noise levels, corresponding to SNR = 40 dB and 20 dB, respectively.
Results in rows [ c ] , [ d ] of Table 1 show similar trends as [ a ] , though with small gaps.
Non-exactly sparse vector [ e ]
We next challenge LISTA to recover the non-exactly sparse x ∗ in (1).
We generate x ∗ from the Gamma distribution Γ( α, β ) with α = 1.0 and β = 0.1.
In this way, x ∗ has a handful of coordinates with large magnitudes, while the remaining coordinates have small yet nonzero values, making it approximately sparse or not exactly so.
In order to disentangle different factors’ inﬂuences, no additive noise is placed in this case (SNR = ∞ ).
From row [ e ] of
Table 1, the same trend is observed, so is the superiority of the searched model.
Training-testing mismatch [ f, g, h, i ]
We proceed to checking more challenging and practical scenarios of model robustness where the training and testing distribution may mismatch.
We consider three cases in Table 1: (1) Transfer-Noise (Gaussian): we train models on the noiseless training set (SNR= ∞ ), and directly apply to the two noisy testing sets, i.e., SNR = 40 dB as in row [ f ] and 20 dB as in [ g ] .
(2) Transfer-Noise (Gaussian → salt & pepper): the models trained with
SNR = ∞ are encountered with an unseen noise type in testing: salt & pepper noise with 1% density; (3) Perturbed dictionary 4 : we keep our training/validation sets generated by the same dictionary
D (SNR = ∞ ), while re-generating a testing set by the following way: we sample a small ∆ D from a Laplacian distribution L (0 , 10 − 3 ) , perturb ¯ D = D + ∆ D to create an unseen basis, normalize ¯ D to unit column and then use it to generate new testing samples (no additive noise).
Not surprisingly, all models see degraded performance under those challenging settings, but our searched architecture sticks as the best (or tie) among all.
LFISTA and Dense-LISTA are also able to outperform LISTA.
Besides, the gaps between Dense-LISTA and ours are generally reduced in all those mismatch cases, implying that denser connections might beneﬁt model robustness here.
Limited training data [ j, k ] We lastly verify a hypothesis raised in (Monga et al., 2019), that unrolling provides model-based prior and helps train more generalizable networks in the data-limited training regime.
To prove this concept, we reduce the training set to 10,240 and 5,120 samples, i.e., 10% and 5% of the default training size, respectively 5 .
The validation set is also reduced to 1, samples (10% default size), but the testing set size remains unchanged.
We observe two speciﬁcally interesting phenomenons: 1) for the ﬁrst time in our experiments,
LFISTA outperforms Dense-LISTA notably, and the margin seems to enlarge as the training size decreases; (2) the searched architecture largely outperforms the other three when the training size is 10,240; yet it becomes comparable to LFISTA at the smaller 5,120 training size, albeit still clearly surpassing LISTA and Dense-LISTA.
The lessons we learned seem to convey compound information: (i) LFISTA’s robust performance suggests that model-based unrolling indeed provides strong inductive prior, that is advantageous under data-limited training; (ii) compared to [ b − i ] where
Dense-LISTA has consistently strong performance, adding overly dense connections seems not to be favored in data-limited regimes; (3) our searched architecture appears to be the right blend of model-based prior and data-driven capacity, that also possesses robustness to data-limited training.
O PEN THE B OX : W HAT GOOD DESIGN PATTERNS ARE INSIDE ?
A closer look at the averaged connectivity.
Experiments in Section 3 indicate that signiﬁcantly improved architectures can be spotted by NAS, whose strong performance generalizes in many unseen settings.
We now dive into analyzing our searched average architecture.
Figure 2 plots the average of the top-50 architectures before threshold , i e., each value indicate the “percentage” of top-50 models having that connection (1 denotes all and 0 none).
The following patterns are observed:
• The consensus on connectivity is overall high.
There are nearly 50% connections that over 90% of the top-50 models agree not to use.
It seems the sparse recovery task obeys certain model-based prior and does not favor overly densely connections, as seen in Dense-LISTA.
• Besides the default connection (diagonal line), every layer uses at least one extra skip connection.
If counting in threshoding (0.5), all layers except layers 12, 14, 15 “activate” no less than 1/ possible input skips, and more than half layers “activate” no less than 1/2.
Those extra connections might represent or be interpreted as acceleration terms, potentially more sophisticated than studied in (Moreau & Bruna, 2017): we leave for future research.
[Figure]
architecture”.
The vertical and horizontal axes denote the “destination” and “origin” layer indices for a skip connection.
Note that every layer is always taken as input for its immediate preceding layer, therefore the diagonal line is all one.
• Early and late layers have better consensus on connections, while middle layers (9-11) display some diversity of choice.
The last layer (16) looks particularly special: it has connections from most preceding layers with high agreements.
Our further experiments ﬁnd that those connections to the last layer contribute a notable part to our searched performance: if we only pick those extra connections to layer-16 from our average architecture, and add them to LISTA/LFISTA, then we can immediately boost LISTA from -43.3 dB to -48.2 dB, and LFISTA from -47.3 dB to -48.4 dB (noiseless setting).
Understanding the role of ”densely connected last layer” in unrolling seems to be an interesting open question.
Transferring found pattern to other unrolling.
We brieﬂy explore whether the LISTA-found design patterns can generalize to other unrolling.
We choose differentiable linearized ADMM (DLADMM)
(Xie et al., 2019), another competitive unrolling model for sparse recovery, as the test bed.
We directly transplant the connectivity in Figure 2, by adding those extra (non-diagonal) connections to augment D-LADMM, with no other change.
We compare (a) the original D-LADMM; (b) our augmented one; and (c) a densely-connected D-LADMM constructed similar to Dense-LISTA.
The three models are trained and evaluated in three representative settings: (1) noiseless: the three models’ NMSEs are -54.2/-55.6/-55.4 dB, respectively; (2) noisy (SNR = 20): NMSEs -18.2/-19.1/- 19.0 dB; (3) data-limited (training size 10,240): NMSEs -24.8/-33.5/-22.6 dB. Results are generally very consistent with our LISTA cases: (1) the transferred connectivity pattern immediately boosts
D-LADMM in all three settings; (2) the dense variant performs comparably with our augmented one when training data is sufﬁcient, but degrades even behind the original in the data-limited regime.
D ISCUSSIONS AND F UTURE W ORK
While unrolling often yields reasonably good architectures, we seem to discover consistently better networks in this study.
But are we denying the advances and promise of the unrolling ﬁeld?
Absolutely Not .
As veterans studying unrolling, we hope this work to provide a reference and spark more reﬂections, on when unrolling is useful, how to improve it more, and (broadly) how to mingle modelbased optimization and data-driven learning better.
A good perspective was taken in (Monga et al., 2019), suggesting that with the original model-based optimization (Dittmer et al., 2019) induced as a prior, the unrolled architectures behave as “an intermediate state between generic networks and iterative algorithms”, possessing relatively low bias and variance simultaneously.
Unrolled networks might be more data-efﬁcient to learn (supposed by our data-limited training experiments).
Meanwhile if training data is sufﬁcient, or if linear sparse model is not the accurate prior, then unrolling does not have to be superior.
We then recommend data-driven model search or selection, perhaps leveraging the unrolled model as a robust starting point in the design space.
We shall further comment that this study is in a completely different track from theoretically understanding LISTA as a learned optimization algorithm (Moreau & Bruna, 2017; Giryes et al., 2018;
Chen et al., 2018; Liu et al., 2019; Aberdam et al., 2020).
Those works’ interests are mainly on interpreting the unrolled architecture as an early-truncated iterative optimizer (with adaptive weights): unrolling plays the central role in bridging this optimization-wise interpretability.
Unrolling can also connect to more desired stability or robustness results from the optimization ﬁeld (Aberdam et al., 2020; Heaton et al., 2020), which makes great research questions yet is beyond this paper’s scope.
A A PPENDIX A. H AND - CRAFTED C ONNECTIVITY B ASELINES
Other than the original unrolled LISTA, we also compare our searched connectivity pattern with two hand-crafted connectivity baselines LFISTA and Dense-LISTA.
Figure 3 shows a straightforward visualization for these two hard-crafted connectivity patterns.
x (1) x (2) x (3) x (4) x (5) x (6) x (7) x (8) x (9) x (10) x (11) x (12) x (13) x (14) x (15)
Layer (1) Layer (2) Layer (3)
Layer (4) Layer (5) Layer (6)
Layer (7) Layer (8) Layer (9)
Layer (10) Layer (11) Layer (12)
Layer (13) Layer (14) Layer (15)
Layer (16) 11 1  1   1    1     1      1       1        1         1          1           1
[Figure]
[Figure] [Figure]
4) x
[Figure]
5) x (6)
[Figure]
[Figure] [Figure]
9) x
[Figure] [Figure]
11) x
[Figure]
12) x
[Figure]
13) x (14) x
[Figure]
[Figure] [Figure]
A.1. LFISTA
The formulation of FISTA (Beck & Teboulle, 2009) is
[Equation]
It is natural to parameterize the second step as
[Equation]
which will reduce FISTA to LWA that is introduced in the main text.
Of course we can plug y ( k +1) into the third equation in (7), re-arrange and combine terms and parameterize the iteration as
[Equation]
However, this will introduce more parameters than LWA.
As we mainly consider LWA in the main text, we use the parameterization in (8) in this paper for fair comparison.
Introducing FISTA to our design space would add 14 extra skip connections.
A.1. D ENSE -LISTA
Inspired by successful application of dense connections in deep learning (Huang et al., 2017), we manually enable all possible skip connections in our connectivity design space.
For this denselyconnected LISTA instantiation, each layer will take the outputs of all its previous layers as input, leading to 105 extra skip connections comparing to the original LISTA.
[Table]
A. C OMPLEMENTARY E XPERIMENTS A.2. D IFFERENT PROBLEM DIMENSIONALITIES
We are also curious whether our ﬁnding is also applicable to sparse coding problems with different dimensionalities.
Here we follow the same noiseless setting as Sec.3.2 in the main text, but enlarge the dictionary to 512 × 1024, i.e. m = 512 , n = 1024 , and regenerate the training, valiation and test data with the new dictionary.
As shown in the ﬁrst row of Table 2, our results with new dimensionality setting are highly consistent to the noiseless measurement setting, indicating our searched pattern does not overﬁt to a speciﬁc problem size.
A.2. I LL - CONDITIONED DICTIONARIES
Most dictionaries used in the experiments are well-conditioned Gaussian random matrices that enjoy ideal properties such as natural incoherence.
However, in real-world applications, this is too ideal to be true.
Here we manually create ill-conditioned matrices and use them as the dictionaries.
We sample two Gaussian matrices U ∈ R 250 × 200 and V ∈ R 200 × 500 .
Then we use D = UV as the dictionary, which is inherently low-rank.
With an ill-conditioned dictionary, our searched architecture still outperforms other counterparts as demonstrated in the second row of Table 2.
A. C ONNECTION I NJECTION O PTIONS
Other than the Learnable Weighted Average (LWA) we use by default in the main text, we also consider two other options to inject new connections
. A.3. N AIVE A VERAGE (NA)
The most naive approach to fuse extra connectivity into the unrolled LISTA (4) is to simply take the average of the outputs of all chosen layers, before we apply the linear transform W ( k ) .
Formally, we replace x ( k ) in (4) with
[Equation]
where g i,k ∈ { 0 , 1 } is a gate that controls whether the i -th iterate is chosen as the input to the k -th layer (same hereinafter).
Denote the gates that connect the k -th layer, e.g. all gates in (NA) as vector g k
= ( g 1 ,k , . . .
, g k,k ) T , and then concatenate into a long ordered vector g = (cid:2) g T 1 , . .
. , g TK − 1 (cid:3) T .
All possible combination of the gate values of the gates form the search space H = { g } .
A.3. M OMENTUM (MM)
We also consider a more complicated way to manipulate extra skip connections, inspired by momentum approaches Sutskever et al.
(2013) in optimization.
Several acceleration approaches for ISTA leverage momentum terms, such as the renowned Fast ISTA (FISTA) Beck & Teboulle (2009);
Bioucas-Dias & Figueiredo (2007).
To allow for ﬂexibly learnable momentum, we illustrate our idea with a simple example, where only last two iterates are considered in the momentum.
In this way, the ISTA iteration (3) is extended as:
[Equation]
where β ( k ) is a hyperparameter that controls the strength of momentum.
Note that (9) reduces to
ISTA (3) if we set β ( k ) ≡ 1 /L , θ ( k ) ≡ λ/L and ν ( k ) =
∇ x f ( x ( k ) )
=
D T ( b − Dx ( k ) ) .
ν ( k ) is the
[Equation]
update direction that already takes momentum into consideration involving the last two iterations:
[Equation]
Substituting (12) in (9) gives
[Equation]
Denote ˆ W
[Equation]
b
[Equation]
and
[Equation]
D , and then we get the untied version (do not share ˆ W ( k ) b , ˆ W ( k ) 1 , ˆ W ( k ) 2 within layers) of (8):
[Equation]
As suggested by the seminal work (Liu et al., 2019), sharing the above weight matrices (corresponding to using layer-invariant β and γ ) across layers does no hurt to the performance while reducing the parameter complexity.
However, we introduce two scaling parameters α ( k )and α ( k )(will be trained using data) to loosen the constraint, yielding a simpler form of (14):
[Equation]
Extending (15) to the general case where all previous iterates could be chosen by a gate to get involved in the momentum term ν ( k ) , we get
[Equation]
where the subscript in ˆ W i means its input to come from the i -th last iterate and α ( k ) i is the step size parameter associated with ˆ W i in the k -th layer.
Step sizes α ( k ) i are initialized with 1.
A.3. E MPIRICAL RESULTS
Before launching our large-scale search experiment, we ﬁrst compare three of our connection injection methods (LWA, NA and MM) by randomly sampling about 500 architectures from our design space, and train each of them individually following our single model training protocol in Section 3.
Since we are focusing on evaluating injection methods, we use soft-thresholding as our neuron type by default.
Noiseless measurement and Gaussian noise.
We ﬁrst conduct experiments under three most common settings (noiseless measurement, noisy measurement with 40dB and 20dB additive Gaussian noise).
As shown in Figure 5, NA always hurts the performance in all three settings.
While
MM could perform slightly better than LWA in noiseless setting, LWA clearly beats MM in noisy
[Figure]
[Figure]
setting for both SNR=40dB and SNR=20dB. Based on our observation that NA constantly yields worse performance, we conduct the rest comparative experiments using only LWA and MM.
More complicated settings.
Other than additive Gaussian noise cases, we also evaluate and compare LWA and MM on the two more challenging and realistic settings: non-exactly sparse vector and limited training data.
We follow the same setting as described in Section 3.3.
We useas training data size in limited training data setting.
From NMSE distributions in both Figureand noisy setting in Figure 5, we can clearly tell that LWA has better robustness when the same connectivity pattern is applied.
At the same time, probably due to the larger number of parameters introduced, MM suffers more from limited training data.
Since MM only wins in the easiest noiseless case with sufﬁcient training data, we can safely tell that
LWA has more practical values as it is consistently the winner of settings where we have different noise in the data and where the data is limited.
Since we are not in an ideal world, we choose to conduct all our main experiments using LWA to fuse injected connections.
A. S IDE C ONNECTION P RUNING
Despite of our focus on adding skip connections in this paper, we also empirically evaluate what will happen if we also allow for pruning connections.
As shown in Figure 7(a), we insert gate functions to side connections to decide whether they will be removed from the original unrolled LISTA.
Speciﬁcally, we allow all the side connections except the ﬁrst one to have the probability being removed from the architecture.
As we use K = 16 LISTA layers as our default setting, our pruning search space is consisted of 2 15 potential candidates for original unrolled LISTA architecture.
For our joint injection and pruning setting, the search space is enlarged to 2 105 × 2 15 = 2 120 .
Note we also utilize the tool of random sampling here as in Section A.3.3 and choose to use soft-thresholding for all layers.
[Figure]
A.4. P RUNING FROM ORIGINAL LISTA
We ﬁrst apply the side connection pruning search space to the original LISTA.
As shown in the
Figure 7(b), removing side connections will always do harm to the unrolled LISTA model, which echos the ﬁnding observed by (Wang et al., 2016)
[Figure]
A.4. J OINT INJECTION AND PRUNING
Next we apply both connection injection and side connection removal at the same time.
Our results in Figure 8(a) illustrate that although connection injection can greatly reduce the inﬂuence of side connection pruning, most architectures from this design space suffer from the removal of side connections.
We also randomly select some models in this setting, keep their injected skip connections and reconnect all their side connections.
Figure 8(b) shows how NMSEs change for LWA models.
The arrows point from the “pruned side connections” case to “full side connections” case and the x axis means the total number of skip connections and side connections.
As we can see, we can almost always gain performance by reconnecting all side connections, which is aligned with our observations in A.4.1.
Our observation here indicates that we should always keep all side connections.
[Figure]
A. S TABILITY OF AVERAGED ARCHITECTURE
To further demonstrate the stability of our averaged architecture, we recalculate an averaged architecture from top-30 models, and compute its connection differences with the default top-50 average one.
The difference map is visualized in the Figure 9.
Clearly, they are highly aligned and echo our observations that those top architectures already have high agreement on connections to be activated.