Understanding approximate and unrolled dictionary learning for pattern recovery

Dictionary learning consists of finding a sparse representation from noisy data and is a common way to encode data-driven prior knowledge on signals . Alternating minimization ( AM ) is standard for the underlying optimization , where gradient descent steps alternate with sparse coding procedures . The major drawback of this method is its prohibitive computational cost , making it unpractical on large real-world data sets . This work studies an approximate formulation of dictionary learning based on unrolling and compares it to alternating minimization to find the best trade-off between speed and precision . We analyze the asymptotic behavior and convergence rate of gradients estimates in both methods . We show that unrolling performs better on the support of the inner problem solution and during the first iterations . Finally , we apply unrolling on pattern learning in magnetoencephalography ( MEG ) with the help of a stochastic algorithm and compare the performance to a state-of-the-art method . 1 INTRODUCTION . Pattern learning provides insightful information on the data in various biomedical applications . Typical examples include the study of magnetoencephalography ( MEG ) recordings , where one aims to analyze the electrical activity in the brain from measurements of the magnetic field around the scalp of the patient ( Dupré la Tour et al. , 2018 ) . One may also mention neural oscillations study in the local field potential ( Cole & Voytek , 2017 ) or QRS complex detection in electrocardiograms ( Xiang et al. , 2018 ) among others . Dictionary learning ( Olshausen & Field , 1997 ; Aharon et al. , 2006 ; Mairal et al. , 2009 ) is particularly efficient on pattern learning tasks , such as blood cells detection ( Yellin et al. , 2017 ) and MEG signals analysis ( Dupré la Tour et al. , 2018 ) . This framework assumes that the signal can be decomposed into a sparse representation in a redundant basis of patterns – also called atoms . In other words , the goal is to recover a sparse code Z ∈ Rn×T and a dictionary D ∈ Rm×n from noisy measurements Y ∈ Rm×T which are obtained as the linear transformation DZ , corrupted with noise B ∈ Rm×T : Y = DZ + B . Theoretical elements on identifiability and local convergence have been proven in several studies ( Gribonval et al. , 2015 ; Haeffele & Vidal , 2015 ; Agarwal et al. , 2016 ; Sun et al. , 2016 ) . Sparsity-based optimization problems related to dictionary learning generally rely on the usage of the ` 0 or ` 1 regularizations . In this paper , we study Lasso-based ( Tibshirani , 1996 ) dictionary learning where the dictionary D is learned in a set of constraints C by solving min Z∈Rn×T , D∈C F ( Z , D ) , 1 2 ‖DZ − Y ‖22 + λ ‖Z‖1 . ( 1 ) Dictionary learning can be written as a bi-level optimization problem to minimize the cost function with respect to the dictionary only , as mentioned in Mairal et al . ( 2009 ) , min D∈C G ( D ) , F ( Z∗ ( D ) , D ) with Z∗ ( D ) = arg min Z∈Rn×T F ( Z , D ) . ( 2 ) Computing the data representation Z∗ ( D ) is often referred to as the inner problem , while the global minimization is the outer problem . Classical constraint sets include the unit norm , where each atom is normalized to avoid scale-invariant issues , and normalized convolutional kernels to perform Convolutional Dictionary Learning ( Grosse et al. , 2007 ) . Classical dictionary learning methods solve this bi-convex optimization problem through Alternating Minimization ( AM ) ( Mairal et al. , 2009 ) . It consists in minimizing the cost function F over Z with a fixed dictionary D and then performing projected gradient descent to optimize the dictionary with a fixed Z . While AM provides a simple strategy to perform dictionary learning , it can be inefficient on large-scale data sets due to the need to resolve the inner problems precisely for all samples . In recent years , many studies have focused on algorithm unrolling ( Tolooshams et al. , 2020 ; Scetbon et al. , 2021 ) to overcome this issue . The core idea consists of unrolling the algorithm , which solves the inner problem , and then computing the gradient with respect to the dictionary with the help of back-propagation through the iterates of this algorithm . Gregor & LeCun ( 2010 ) popularized this method and first proposed to unroll ISTA ( Daubechies et al. , 2004 ) – a proximal gradient descent algorithm designed for the Lasso – to speed up the computation of Z∗ ( D ) . The N + 1-th layer of this network – called LISTA – is obtained as ZN+1 = ST λ L ( W 1Y + W 2ZN ) , with ST be- ing the soft-thresholding operator . This work has led to many contributions aiming at improving this method and providing theoretical justifications in a supervised ( Chen et al. , 2018 ; Liu & Chen , 2019 ) or unsupervised ( Moreau & Bruna , 2017 ; Ablin et al. , 2019 ) setting . For such unrolled algorithms , the weights W 1 and W 2 can be re-parameterized as functions of D – as illustrated in Figure A in appendix – such that the output ZN ( D ) matches the result of N iterations of ISTA , i.e . W 1D = 1 L D > and W 2D = ( I − 1 L D > D ) , where L = ‖D‖2 . ( 3 ) Then , the dictionary can be learned by minimizing the loss F ( ZN ( D ) , D ) over D with backpropagation . This approach is generally referred to as Deep Dictionary Learning ( DDL ) . DDL and variants with different kinds of regularization ( Tolooshams et al. , 2020 ; Lecouat et al. , 2020 ; Scetbon et al. , 2021 ) , image processing based on metric learning ( Tang et al. , 2020 ) , and classification tasks with scattering ( Zarka et al. , 2019 ) have been proposed in the literature , among others . While these techniques have achieved good performance levels on several signal processing tasks , the reasons they speed up the learning process are still unclear . In this work , we study unrolling in Lasso-based dictionary learning as an approximate bi-level optimization problem . What makes this work different from Bertrand et al . ( 2020 ) , Ablin et al . ( 2020 ) and Tolooshams & Ba ( 2021 ) is that we study the instability of non-smooth bi-level optimization and unrolled sparse coding out of the support , which is of major interest in practice with a small number of layers . In Section 2 , we analyze the convergence of the Jacobian computed with automatic differentiation and find out that its stability is guaranteed on the support of the sparse codes only . De facto , numerical instabilities in its estimation make unrolling inefficient after a few dozen iterations . In Section 3 , we empirically show that unrolling leads to better results than AM only with a small number of iterations of sparse coding , making it possible to learn a good dictionary in this setting . Then we adapt a stochastic approach to make this method usable on large data sets , and we apply it to pattern learning in magnetoencephalography ( MEG ) in Section 4 . We do so by adapting unrolling to rank one convolutional dictionary learning on multivariate time series ( Dupré la Tour et al. , 2018 ) . We show that there is no need to unroll more than a few dozen iterations to obtain satisfying results , leading to a significant gain of time compared to a state-of-the-art algorithm . 2 BI-LEVEL OPTIMIZATION FOR APPROXIMATE DICTIONARY LEARNING . As Z∗ ( D ) does not have a closed-form expression , G can not be computed directly . A solution is to replace the inner problem Z∗ ( D ) by an approximation ZN ( D ) obtained through N iterations of a numerical optimization algorithm or its unrolled version . This reduces the problem to minimizing GN ( D ) , F ( ZN ( D ) , D ) . The first question is how sub-optimal global solutions of GN are compared to the ones of G. Proposition 2.1 shows that the global minima of GN converge as fast as the numerical approximation ZN in function value . Proposition 2.1 Let D∗ = arg minD∈C G ( D ) and D∗N = arg minD∈C GN ( D ) , where N is the number of unrolled iterations . We denote by K ( D∗ ) a constant depending on D∗ , and by C ( N ) the convergence speed of the algorithm , which approximates the inner problem solution . We have GN ( D ∗ N ) −G ( D∗ ) ≤ K ( D∗ ) C ( N ) . The proofs of all theoretical results are deferred to Appendix C. Proposition 2.1 implies that when ZN is computed with FISTA ( Beck & Teboulle , 2009 ) , the function value for global minima of GN converges with speed C ( N ) = 1N2 towards the value of the global minima of F . Therefore , solving the inner problem approximately leads to suitable solutions for equation 2 , given that the optimization procedure is efficient enough to find a proper minimum of GN . As the computational cost of zN increases with N , the choice of N results in a trade-off between the precision of the solution and the computational efficiency , which is critical for processing large data sets . Moreover , learning the dictionary and computing the sparse codes are two different tasks . The loss GN takes into account the dictionary and the corresponding approximation ZN ( D ) to evaluate the quality of the solution . However , the dictionary evaluation should reflect its ability to generate the same signals as the ground truth data and not consider an approximate sparse code that can be recomputed afterward . Therefore , we should distinguish the ability of the algorithm to recover a good dictionary from its ability to learn the dictionary and the sparse codes at the same time . In this work , we use the metric proposed in Moreau & Gramfort ( 2020 ) for convolutions to evaluate the quality of the dictionary . We compare the atoms using their correlation and denote as C the cost matrix whose entry i , j compare the atom i of the first dictionary and j of the second . We define a sign and permutation invariant metric S ( C ) = maxσ∈Sn 1 n ∑n i=1 |Cσ ( i ) , i| , where Sn is the group of permutations of [ 1 , n ] . This metric corresponds to the best linear sum assignment on the cost matrix C , and it can be computed with the Hungarian algorithm . Note that doing so has several limitations and that evaluating the dictionary is still an open problem . Without loss of generality , let T = 1 and thus z ∈ Rn in the rest of this section . Gradient estimation in dictionary learning . Approximate dictionary learning is a non-convex problem , meaning that good or poor local minima of GN may be reached depending on the initialization , the optimization path , and the structure of the problem . Therefore , a gradient descent onGN has no guarantee to find an adequate minimizer of G. While complete theoretical analysis of these problems is arduous , we propose to study the correlation between the gradient obtained with GN and the actual gradient of G , as a way to ensure that the optimization dynamics are similar . Once z∗ ( D ) is known , Danskin ( 1967 , Thm 1 ) states that g∗ ( D ) = ∇G ( D ) is equal to∇2F ( z∗ ( D ) , D ) , where∇2 indicates that the gradient is computed relatively to the second variable in F . Even though the inner problem is non-smooth , this result holds as long as the solution z∗ ( D ) is unique . In the following , we will assume that D > D is invertible on the support of z∗ ( D ) , which implies the uniqueness of z∗ ( D ) . This occurs with probability one if D is sampled from a continuous distribution ( Tibshirani , 2013 ) . AM and DDL differ in how they estimate the gradient of G. AM relies on the analytical formula of g∗ and uses an approximation zN of z∗ , leading to the approximate gradient g1N ( D ) = ∇2F ( zN ( D ) , D ) . We evaluate how well g1N approximates g∗ in Proposition 2.2 . Proposition 2.2 Let D ∈ Rm×n . Then , there exists a constant L1 > 0 such that for every number of iterations N ∥∥g1N − g∗∥∥ ≤ L1 ‖zN ( D ) − z∗ ( D ) ‖ . Proposition 2.2 shows that g1N converges as fast as the iterates of ISTA converge . DDL computes the gradient automatically through zN ( D ) . As opposed to AM , this directly minimizes the loss GN ( D ) . Automatic differentiation yields a sub-gradient g2N ( D ) such that g2N ( D ) ∈ ∇2F ( zN ( D ) , D ) + J + N ( ∂1F ( zN ( D ) , D ) ) , ( 4 ) where JN : Rm×n → Rn is the weak Jacobian of zN ( D ) with respect to D and J+N denotes its adjoint . The product between J+N and ∂1F ( zN ( D ) , D ) is computed via automatic differentiation . Proposition 2.3 Let D ∈ Rm×n . Let S∗ be the support of z∗ ( D ) , SN be the support of zN and S̃N = SN ∪ S∗ . Let f ( z , D ) = 12 ‖Dz − y‖ 2 2 be the data-fitting term in F . Let R ( J , S̃ ) = J+ ( ∇21,1f ( z∗ , D ) 1S̃ ) +∇22,1f ( z∗ , D ) 1S̃ . Then there exists a constant L2 > 0 and a subsequence of ( F ) ISTA iterates zφ ( N ) such that for all N ∈ N : ∃ g2φ ( N ) ∈ ∇2f ( zφ ( N ) , D ) + J + φ ( N ) ( ∇1f ( zφ ( N ) , D ) + λ∂‖·‖1 ( zφ ( N ) ) ) s.t . : ∥∥∥g2φ ( N ) − g∗∥∥∥ ≤ ∥∥∥R ( Jφ ( N ) , S̃φ ( N ) ) ∥∥∥∥∥zφ ( N ) − z∗∥∥+ L22 ∥∥zφ ( N ) − z∗∥∥2 . This sub-sequence zφ ( N ) corresponds to iterates on the support of z∗ . Proposition 2.3 shows that g2N may converge faster than g 1 N once the support is reached . Ablin et al . ( 2020 ) and Tolooshams & Ba ( 2021 ) have studied the behavior of strongly convex functions , as it is the case on the support , and found similar results . This allowed Tolooshams & Ba ( 2021 ) to focus on support identification and show that automatic differentiation leads to a better gradient estimation in dictionary learning on the support under minor assumptions . However , we are also interested in characterizing the behavior outside of the support , where the gradient estimation is difficult because of the sub-differential . In practice , automatic differentiation uses the sign operator as a sub-gradient of ‖·‖1 . The convergence behavior of g2N is also driven by R ( JN , S̃N ) and thus by the weak Jacobian computed via back-propagation . We first compute a closed-form expression of the weak Jacobian of z∗ ( D ) and zN ( D ) . We then show that R ( JN , S̃N ) ≤ L ‖JN − J∗‖ and we analyze the convergence of JN towards J∗ . Study of the Jacobian . The computation of the Jacobian can be done by differentiating through ISTA . In Theorem 2.4 , we show that JN+1 depends on JN and the past iterate zN , and converges towards a fixed point . This formula can be used to compute the Jacobian during the forward pass , avoiding the computational cost of back-propagation and saving memory . Theorem 2.4 At iteration N + 1 of ISTA , the weak Jacobian of zN+1 relatively to Dl , where Dl is the l-th row of D , is given by induction : ∂ ( zN+1 ) ∂Dl = 1|zN+1| > 0 ( ∂ ( zN ) ∂Dl − 1 L ( Dlz > N + ( D > l zN − yl ) In + D > D ∂ ( zN ) ∂Dl ) ) . ∂ ( zN ) ∂Dl will be denoted by JNl . It converges towards the weak Jacobian J ∗ l of z ∗ relatively to Dl , whose values are J∗l S∗ = − ( D > : ,S∗D : ,S∗ ) −1 ( Dlz ∗ > + ( D > l z ∗ − yl ) In ) S∗ , on the support S∗ of z∗ , and 0 elsewhere . Moreover , R ( J∗ , S∗ ) = 0 . This result is similar to Bertrand et al . ( 2020 ) where the Jacobian of z is computed over λ to perform hyper-parameter optimization in Lasso-type models . Using R ( J∗ , S∗ ) = 0 , we can write∥∥∥R ( JN , S̃N ) ∥∥∥ ≤ ∥∥∥R ( JN , S̃N ) −R ( J∗ , S∗ ) ∥∥∥ ≤ L ‖JN − J∗‖ , ( 5 ) as ∥∥∇21,1f ( z∗ , D ) ∥∥2 = L. If the back-propagation were to output an accurate estimate JN of the weak Jacobian J∗ , ∥∥∥R ( JN , S̃N ) ∥∥∥ would be 0 , and the convergence rate of g2N could be twice as fast as the one of g1N . To quantify this , we now analyze the convergence of JN towards J ∗ . In Proposition 2.5 , we compute an upper bound of ∥∥JNl − J∗l ∥∥ with possible usage of truncated backpropagation ( Shaban et al. , 2019 ) . Truncated back-propagation of depth K corresponds to an initial estimate of the Jacobian JN−K = 0 and iterating the induction in Theorem 2.4 . Proposition 2.5 Let N be the number of iterations and K be the back-propagation depth . We assume that ∀n ≥ N−K , S∗ ⊂ Sn . Let ĒN = Sn\S∗ , let L be the largest eigenvalue ofD > : ,S∗D : ,S∗ , and let µn be the smallest eigenvalue of D > : ,SnD : ,Sn−1 . Let Bn = ∥∥∥PEn −D > : ,ĒnD† > : ,S∗PS∗∥∥∥ , where PS is the projection on RS and D† is the pseudo-inverse of D. We have∥∥JNl − J∗l ∥∥ ≤ K∏ k=1 ( 1− µN−k L ) ‖J∗l ‖+ 2 L ‖Dl‖ K−1∑ k=0 k∏ i=1 ( 1−µN−i L ) ( ∥∥zN−kl − z∗l ∥∥+BN−k ‖z∗l ‖ ) . Proposition 2.5 reveals multiple stages in the Jacobian estimation . First , one can see that if all iterates used for the back-propagation lie on the support S∗ , the Jacobian estimate has a quasi-linear convergence , as shown in the following corollary . Corollary 2.6 Let µ > 0 be the smallest eigenvalue of D > : ,S∗D : ,S∗ . Let K ≤ N be the backpropagation depth and let ∆N = F ( zN , D ) − F ( z∗ , D ) + L2 ‖zN − z ∗‖ . Suppose that ∀n ∈ [ N −K , N ] ; Sn ⊂ S∗ . Then , we have∥∥J∗l − JNl ∥∥ ≤ ( 1− µL ) K ‖J∗l ‖+K ( 1− µL ) K−1 ‖Dl‖ 4∆N−KL2 . Once the support is reached , ISTA also converges with the same linear rate ( 1 − µL ) . Thus the gradient estimate g2N converges almost twice as fast as g 1 N in the best case – with optimal subgradient – as O ( K ( 1− µL ) 2K ) . This is similar to Ablin et al . ( 2020 , Proposition.5 ) and Tolooshams & Ba ( 2021 ) . Second , Proposition 2.5 shows that ∥∥J∗l − JNl ∥∥ may increase when the support is not well-estimated , leading to a deterioration of the gradient estimate . This is due to an accumulation of errors materialized by the sum in the right-hand side of the inequality , as the term BN ‖z∗‖ may not vanish to 0 as long as SN 6⊂ S∗ . Interestingly , once the support is reached at iteration S < N , the errors converge linearly towards 0 , and we recover the fast estimation of g∗ with g2 . Therefore , Lasso-based DDL should either be used with a low number of steps or truncated back-propagation to ensure stability . These results apply for all linear dictionaries , including convolutions . Numerical illustrations . We now illustrate these theoretical results depending on the number N of unrolled iterations . The data are generated from a random Gaussian dictionary D of size 30×50 , with Bernoulli-Gaussian sparse codes z ( sparsity 0.3 , σ2z = 1 ) , and Gaussian noise ( σ2noise = 0.1 ) – more details in Appendix A . Figure 1 confirms the linear convergence of JNl once the support is reached . However , the convergence might be unstable when the number of iteration grows , leading to exploding gradient , as illustrated in the second case . When this happens , using a small number of iterations or truncated back-propagation becomes necessary to prevent accumulating errors . It is also of interest to look at the proportion of unstable Jacobians ( see Figure 2 ) . We recover behaviors observed in the first and second case in Figure 1 . 40 % samples suffer from numerical instabilities in this example . This has a negative impact on the gradient estimation outside of the support . We display the convergence behavior of the gradients estimated by AM and by DDL with different back-propagation depths ( 20 , 50 , full ) for simulated data and images in Figure 3 . We unroll FISTA instead of ISTA to make the convergence faster . We observed similar

Summary: The authors investigate the asymptotic behavior of unrolling applied to dictionary learning. The applicability of unrolling to dictionary learning results from the circumstance that dictionary learning can be reformulated in terms of bilevel optimization, where the lower-level (or inner) problem is a sparse coding problem (in the LASSO case considered here). Unrolling means to replace the argmin in the lower-level problem with the $N$-th iterate of a suitable optimization algorithm. Gradients of the upper-level loss can then be computed by means of backpropagation through algorithmic iterates. The authors study the behavior of resulting gradients depending on $N$ and draw a comparison with alternating minimization. They find that unrolling constitutes a scalable alternative to alternating minimization, where unrolling a relatively small number of iterations or using truncated backpropagation is favorable to ensure stable approximate gradients.