A Design Space Study for LISTA and Beyond

1 INTRODUCTION . 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 fitting and conditional random fields . 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 fitting “ 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 finite 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 first 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-efficient , interpretable and high-performance architectures , when the underlying optimization model is assumed available . Such approaches have witnessed prevailing success in applications such as compressive ∗The authors Tianjian Meng and Xiaohan Chen contribute equally to the work . 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 specific objective function , and proved the favorable convergence rates achieved over classical iterative solver , when the unrolled architectures are trained to ( over ) fit particular data . Orthogonally , this paper reflects 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.1 RELATED 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 finetuned to obtain optimal performance on the data for a small number of iterations . Numerous works ( Sprechmann et al. , 2015 ; Wang et al. , 2016a ; 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 first 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 ) . Besides , the safeguarding mechanism was also introduced for guiding learned updates to ensure convergence , even when the test problem shifts from the training distribution ( Heaton et al. , 2020 ) . On a separate note , many empirical works ( Wang et al. , 2016a ; b ; 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 . Specifically , the authors suggested that since iterative algorithms are grounded on domain-specific 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-efficiently . 1.2 MOTIVATIONS 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 first design space ablation study1 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 field . NAS 1We define a design space as a family of models derived from the same set of architecture varying rules . can explore a gigantic design space much more efficiently , 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 befitting 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 field . 2 TECHNICAL APPROACH . Assume a sparse vector x∗ = [ x∗1 , · · · , x∗M ] T ∈ RM , we observe its noisy linear measurements : b = ∑M m=1 dmx ∗ m + ε = Dx ∗ + ε , ( 1 ) where b ∈ RN , D = [ d1 , · · · , dM ] ∈ RN×M is the dictionary , and ε ∈ RN is additive Gaussian white noise . For simplicity , each column of D is normalized , that is , ‖dm‖2 = ‖D : ,m‖2 = 1 , m = 1 , 2 , · · · , M . Typically , we have N M . A popular approach for solving the inverse problem : b→ x , is to solve the LASSO below ( λ is a scalar ) : minx 1 2 ‖b−Dx‖22 + λ‖x‖1 ( 2 ) using iterative algorithms such as the iterative shrinkage thresholding algorithm ( ISTA ) : x ( k+1 ) = ηλ/L ( x ( k ) + DT ( b−Dx ( k ) ) /L ) , k = 0 , 1 , 2 , . . . ( 3 ) where ηθ is the soft-thresholding function2 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 fixing 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 : x ( k+1 ) = ηθ ( k ) ( Wbb + α ( k ) W ( k ) x x ( k ) ) , k = 0 , 1 , · · · , K − 1 . ( 4 ) If we set Wb ≡ DT /L , W ( k ) x ≡ I−DTD/L , α ( k ) ≡ 1 , θ ( k ) ≡ λ/L , then LISTA recovers ISTA3 . In practice , we start with x ( 0 ) set to zero . As suggested by the seminal work ( Liu et al. , 2019 ) , sharing Wb and Wx 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 Wx to preserve flexibility . 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 { Wb , Wx } to be coupled layerwise , and showed it yielded the same representation capability as the orginal LISTA in Gregor & LeCun ( 2010 ) . We instead have { Wb , Wx } as two independent learnable weight matrices . The main reason we did not follow ( Chen et al. , 2018 ) is that Wb-Wx 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 Θ = { Wb , Wx , α ( k ) , θ ( k ) } K−1k=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 : min Θ Ex∗ , b∼P‖x ( K ) ( Θ , b , x ( 0 ) ) − x∗‖22 . ( 5 ) This problem is approximately solved over a training dataset { ( x∗i , bi ) } Ni=1 sampled from P . 2Soft-thresholding function is defined in a component-wise way : ηθ ( x ) = sign ( x ) max ( 0 , |x| − θ ) 3Those can be parameter initialization in LISTA , which we follow ( Chen et al. , 2018 ) to use by default .

Summary: This paper studies a very interesting new problem of assessing unrolled models in a broader context using NAS methods. LISTA-style unrolling has been popular for deep learning-based inverse problems. But it is quantitatively unclear how good the unrolled models are, among all possible model variations. To fill in this gap, the authors first define a proper search space based on the varying connections and neurons from the unrolled LISTA backbone architecture. NAS is then exploited as the tool to find the best subsect of architecture from the large space.