Towards Neurally Augmented ALISTA
Abstract
It is well-established that many iterative sparse reconstruction algorithms such as
ISTA can be unrolled to yield a learnable neural network for improved empirical performance.
Recently, ALISTA has been introduced, combining the strong empirical performance of a fully learned approach like LISTA, while retaining theoretical guarantees of classical compressed sensing algorithms and signiﬁcantly reducing the number of parameters to learn.
However, these parameters are trained to work in expectation, often leading to suboptimal reconstruction of individual targets.
In this work we therefore introduce Neurally-Augmented-ALISTA, which computes step sizes and thresholds individually for each target vector during reconstruction.
This adaptive approach is theoretically motivated by revisiting the recovery guarantees of ALISTA and is able to outperform existing algorithms in sparse reconstruction.
Introduction
Compressed sensing is concerned with recovering a sparse vector from far fewer compressive linear observations.
Formally, consider the set of s -sparse vectors in R N , i.e. Σ Ns := (cid:8) x ∈ R N (cid:12)(cid:12) (cid:107) x (cid:107) 0 ≤ s (cid:9) .
Furthermore, let Φ ∈ R M × N be the measurement matrix, with typically M (cid:28) N .
For a given noiseless observation y =
Φ x ∗ of an unknown but s -sparse x ∗ ∈ Σ
Ns we therefore wish to solve: argmin
[Equation]
In [2] it has been shown, that under certain assumptions on Φ , the solution to the combinatorial problem in (1) can be instead obtained by solving the convex LASSO problem.
A very popular approach for solving LASSO is the iterative shrinkage thresholding algorithm (ISTA) [4], in which a reconstruction x ( k ) is obtained after k iterations from initial x (0) = 0 via the iteration:
[Equation]
where η θ is the soft thresholding function given by η θ ( x ) = sign ( x ) max(0 ,
| x | − θ ) (applied coordinate-wise) and L is the Lipschitz constant (i.e. the largest eigenvalue) of Φ T Φ .
Famously, the computational graph of ISTA with K iterations can be unrolled to yield Learned ISTA (LISTA) [7], a K -layer neural network in which all parameters involved can be trained using backpropagation and gradient descent.
LISTA achieves impressive empirical reconstruction performance for many sparse datasets but loses the theoretical guarantees of ISTA.
Bridging the gap between LISTA’s strong reconstruction quality and the theoretical guarantees for ISTA, ALISTA [11] was introduced.
In ALISTA, all matrices are excluded from the learning process to retain desirable properties for compressed sensing, and
W ∈ R M × N corresponding to Φ is instead computed by optimizing the generalized coherence:
[Equation]
For each layer of ALISTA, only a scalar step size parameter γ ( k ) and a scalar threshold θ ( k ) is learned from the data, yielding the iteration:
[Equation]
[Equation] As in LISTA, the parameters for ALISTA are learned end-to-end using backpropagation and stochastic gradient descent by empirically minimizing the mean squared error between the reconstruction and the target vector.
The authors rigorously upper-bound the reconstruction error of ALISTA in the noiseless case and demonstrate strong empirical reconstruction quality even in the noisy case.
The empirical performance similar to LISTA, the retained theoretical guarantees, and the reduction of number of parameters to train from either O ( KM 2 + NM ) in vanilla LISTA or O ( MNK ) in the variant of LISTA-CPSS [3] to just O ( K ) , make ALISTA an appealing algorithm to study and extend.
Thresholds that are adaptive to the current target vector have been explored in ALISTA-AT [9], where thresholds are computed depending on the absolute values of the components from the current reconstruction, for which the authors demonstrate superior recovery over ALISTA for a speciﬁc setting of M, N and s .
In a related approach [15] identify undershooting, meaning that reconstructed components are smaller than target components, as a shortcoming of LISTA and propose GatedLISTA to address these issues.
The authors show that adding their proposed gates to ALISTA, named
AGLISTA, improves its performance in the same setting of M, N and s as ALISTA-AT.
In this paper, motivated by essential proof steps of ALISTA’s recovery guarantee, we propose an alternative method for adaptively choosing thresholds and step sizes during reconstruction.
Our method directly extends ALISTA by using a recurrent neural network to predict thresholds and step sizes depending on an estimate of the (cid:96) 1 -error between the reconstruction and the unknown target vector after each iteration.
We refer to our method as Neurally Augmented ALISTA (NA-ALISTA), as the method falls into the general framework of neural augmentation of unrolled algorithms [14, 12, 5].
To summarize, our main contributions are:
We introduce Neurally Augmented ALISTA (NA-ALISTA), an algorithm which learns to adaptively compute thresholds and step-sizes for individual target vectors during recovery.
We provide theoretical motivation inspired by guarantees for sparse reconstruction which show that NA-ALISTA can achieve tighter error bounds depending on the target x ∗ .
We show that NA-ALISTA empirically outperforms ALISTA and other state-of-the-art algorithms in all evaluated settings and that the gains increase with decreasing M/N
Theoretical Motivation
The thresholds θ ( k ) in (4) play an important role in the analysis of ALISTA.
In this section we motivate the choice of adaptive thresholds - the key improvement in our proposed NA-ALISTA.
More speciﬁcally, we repeat the conditions under which ALISTA guarantees no false positives and highlight an intermediate step in the error bound from [11], which tightens when the thresholds can adapt to speciﬁc instances of x ∗ .
Assumption (adapted from Assumption 1 in [11]):
Let x ∗ ∈ Σ Ns be a ﬁxed s –sparse target vector.
Let W be such that it attains the inﬁmum (3) yielding ˜ µ = µ ( W, Φ) .
Let s < (1 + 1 / ˜ µ ) / 2 , { γ ( k ) } kk =1 ⊂ (0 , 2 2˜ µs − ˜ µ +1 ) and { θ ( k ) }
Kk =1 with:
[Equation]
Because in ALISTA, { θ ( k ) } Kk =1 and { γ ( k ) }
Kk =1 are optimized in expectation over the training data sampled from a compact set
X , (5) holds only if the thresholds are larger than the worst case (cid:96) 1 -error committed by the algorithm over all training vectors, i.e. θ ( k ) ≥ γ ( k ) ˜ µ sup x ∗ ∈X (cid:107) x ( k ) − x ∗ (cid:107) 1 , see eq.
(9) in [11].
Only in this case it is guaranteed that no false positives are in the support of x ( k ) , meaning that supp ( x ( k ) ) ⊆ supp ( x ∗ ) , see Lemma 1 in [11].
However, the threshold θ ( k ) also appears in the error upper bound.
We employ an intermediate step of the error upper bound in Theorem 1 from [11]:
Reconstruction error bound : Under the settings of the Assumption, it holds that:
[Equation]
This inequality is derived in detail in Appendix A of [11].
Hence it is desirable that θ ( k ) is as small as possible, but such that it still satisﬁes (5).
This means that ALISTA has to learn thresholds at least proportional to the largest possible committed (cid:96) 1 -error over all possible x ∗ in order to guarantee good reconstruction, for which it is in turn penalized in the error bound.
However, the thresholds that make the error bound tighter vary depending on the x ∗ that is to be recovered.
In fact, if an algorithm would have access to (cid:107) x ( k ) − x ∗ (cid:107) {+
+} and were allowed to choose thresholds adaptively, depending on this quantity, the more relaxed inequality (5) could be employed directly, without taking the supremum.
Neurally Augmented ALISTA
In order to tighten the error upper bound in (6), we introduce Neurally Augmented ALISTA (NAALISTA), in which we adaptively predict thresholds θ ( k,x ∗ ) depending on an estimate for the (cid:96) 1 -error between x ( k ) and the unknown x ∗ .
As can be observed from (6), such θ ( k,x ∗ ) must be proportional to (cid:107) x ( k ) − x ∗ (cid:107) 1 .
In theory, this true (cid:96) 1 -error could be recovered exactly.
This is because there are no false positives in x ( k ) , making it s -sparse and for ˜ µ < 1 / (2 s − 1) the column-normalized W T Φ is restricted-invertible for any 2 s -sparse input [6] [Corollary 5.4, p.113].
However, it is infeasible to solve such an inverse problem at every iteration k .
Furthermore, in practice the sparsity is often much larger than what is admissible via coherence bounds (see experiments of [7, 11, 15, 9], even assuming optimal admissible coherence by Welch bound [13]).
Therefore NA-ALISTA approximates the (cid:96) 1 error.
For this, consider the (cid:96) 1 -norms:
[Equation]
(7) of the residual and the iterative update quantity in (4).
Both are known to the algorithm even though x ∗ is unknown and do not incur additional computational costs other than taking the norm.
Since W T Φ has low generalized mutual coherence, it is a restricted isometry for sparse vectors and therefore u ( k ) will be correlated with the true error for sparse vectors.
This correlation is validated in
Figure 1 for u ( k ) , but also holds for r ( k ) .
Other useful quantities for approximating the (cid:96) 1 -error are (cid:107) x (0) − x ∗ (cid:107) 1 , . . .
, (cid:107) x
( k − 1) − x ∗ (cid:107) 1 (also shown in Figure 1.
The latter suggests the use of a recurrent neural network in NA-ALISTA, leveraging the past estimations.
We therefore propose to use an
LSTM [8] which has two input neurons, receiving u ( k ) and r ( k ) at each iteration k .
This is used to update the LSTM cell state c ∈ R H and hidden state h ∈ R H and produce the outputs θ ( k,x ∗ ) and γ ( k,x ∗ ) via matrix multiplication with a learned matrix U ∈ R 2 × H and then passing through the softsign function:
[Equation]
The threshold and stepsize are then used in the ﬁnal step of one NA-ALISTA iteration:
[Equation]
Experiments
In this section, we evaluate NA-ALISTA in a sparse reconstruction task and compare it against
ALISTA [11], ALISTA-AT [9], AGLISTA [15], as well as the classical ISTA [4] and FISTA [1].
To emphasize a fair and reproducible comparison between the models, the code for all experiments listed is available on Github 1 .
Experimental Setup.
Following the same experimental setup as [11, 15, 3, 9], the support of x ∗ ∈ R N is determined via i.i.d.
Bernoulli random variables with parameter S/N , leading to an [Figure] d [Figure]
[Figure]
expected sparsity of S .
Its non-zero components are sampled according to N (0 , 1) .
The entries of Φ are also sampled from N (0 , 1) and each column is normalized to unit (cid:96) 2 -norm.
Then, the generalized coherence in (3) between W and Φ is minimized via the Frobenius-Norm approximation using projected gradient descent, identical to [11, 15, 9].
Adam [10] is used to minimize the reconstruction error for all algorithms.
A test set of 10000 samples is ﬁxed before training and recovery performance is measured with the normalized mean squared error.
A support selection trick was introduced in [3] to speed up convergence and stabilize training and has been subsequently used extensively (see supplementary code of [11, 9, 15]).
For a fair comparison, we employ support selection in all learned models compared in this paper.
Our AGLISTA implementation follows the description from [15] and uses exponential gain gates and inverse-proportional-based overshoot gains.
The λ parameter in
ISTA and FISTA was tuned by hand, we found that λ
= 0 .
4 led to the best performance in our tasks.
When not otherwise indicated we use the following settings for experiments and algorithms: M =250,
N =1000, S =50, K =16, H =128, and y =
Φ x ∗ + z with additive white Gaussian noise z with a signal to noise ratio of 40 dB. We train all algorithms for 400 epochs, with each epoch containing 50, sparse vectors with a batch size of 512.
Comparison with Competitors.
As shown in Figure 6, we ﬁrst ﬁx N and observe the reconstruction error for a varying amount of iterations.
In Figure 3 we then decrease the compression ratio by increasing N and keeping everything else ﬁxed.
We observe that NA-ALISTA outperforms state-of-the-art adaptive methods in all evaluated scenarios, particularly by a larger margin as the compression ratio becomes more challenging.
Computational Cost.
In Figure 4, we examine the effect of changing the size H of the LSTM network used.
One can observe that an exponential increase in hidden neurons yields only a small error reduction for different N , suggesting that the size H =128 is a sufﬁcient default value regardless of N .
This veriﬁes that the added computation in NA-ALISTA (adding the H 2 term in O ( MN + H 2 ) for each iteration) is negligible in practice.
Veriﬁcation.
As an empirical veriﬁcation the Assumption 1 in (5) we need to check for every x ∗ , whether the ratio θ ( k,x ∗ ) /γ
( k,x ∗ ) is proportional to the (cid:96)
1 -error
||
x
∗ − x ( k ) || 1 .
Since it is infeasible to check the assumption for the whole space of sparse x ∗ , we empirically verify 5 for a sample of inputs from the training distribution.
In Figure 5 the means of both values are proportional to each other for such a test sample, suggesting that the reconstruction bound from [11] holds for NA-ALISTA as well.
[Figure]
in NA-ALISTA.
An exponential increase of the hidden layer size only yields a marginal improvement once H =64 is surpassed.
[Figure]
We report the mean and standard deviation over the test set.
A Supplementary Experiments
[Figure]
For an increased N =1000 under the same noise level in 6c, our algorithm outperforms the other methods clearly.
For a noise level of 20dB all algorithms perform similarly for N = 500 and N = and NA-ALISTA outperforms the others at N = 2000 . B Algorithm
[Table]