Variational Learning ISTA
Authors: Fabio Valerio Massoli
Christos Louizos
Arash Behboodi
Venue: Accepted by TMLR ()

Abstract: Compressed sensing combines the power of convex optimization techniques with a sparsity-inducing prior on the signal space to solve an underdetermined system of equations. For many problems, the sparsifying dictionary is not directly given, nor its existence can be assumed. Besides, the sensing matrix can change across different scenarios. Addressing these issues requires solving a sparse representation learning problem, namely dictionary learning, taking into account the epistemic uncertainty of the learned dictionaries and, finally, jointly learning sparse representations and reconstructions under varying sensing matrix conditions. We address both concerns by proposing a variant of the LISTA architecture. First, we introduce Augmented Dictionary Learning ISTA (A-DLISTA), which incorporates an augmentation module to adapt parameters to the current measurement setup. Then, we propose to learn a distribution over dictionaries via a variational approach, dubbed Variational Learning ISTA (VLISTA). VLISTA exploits A-DLISTA as the likelihood model and approximates a posterior distribution over the dictionaries as part of an unfolded LISTA-based recovery algorithm. As a result, VLISTA provides a probabilistic way to jointly learn the dictionary distribution and the reconstruction algorithm with varying sensing matrices. We provide theoretical and experimental support for our architecture and show that our model learns calibrated uncertainties.

Published in Transactions on Machine Learning Research (07 (224))
Variational Learning ISTA
Fabio Valerio Massoli
fmassoli@Btii qualcomom.com com
Qualcomm Al Research
Christos Louizos Qualcomm Al Research
clouizos www.:oos@@...ggggoooo.ooomm qww.cqura.ccoooommmommm
Arash Behboodi
dichboodi@qtk,
www.::/awwttggalliomm.omm qualcouram com
Qualcomm Al Research
Reviewed on OpenReview: https: // openreview. net/ forun? id= AQkOUsituu tuG
//
Abstract
Compressed sensing combines the power of convex optimization techniques with a sparsity- inducing prior on the signal space to solve an underdetermined system of equations. For many problems, the sparsifying dictionary is not directly given, nor its existence can be assumed. Besides, the semsing matrix can change across different scenarios. Addressing these issues requires solving a sparse representation learning problem, namcly dictionary learning, taking into account the epistemic uncertainty of the learned dictionaries and, finally, jointly learning sparse representations and reconstructions under varying scnsing matrix conditions. We address both concerns by proposing variant of the LISTA architecture. First, introduce Augmented Dictionary Learning ISTA (A-DLISTA), which incorporates augmentation module to adapt parameters to the current measurement setup. Then, we propose to learn distribution over dictionaries via a variational approach, dubbed Variational Learning ISTA (VLISTA). VLISTA exploits A-DLISTA as the likelihood model and approximates a posterior distribution over the dictionaries as part of an unfolded LISTA- recovery algorithm. As a result, VLISTA provides a probabilistic way to jointly learn the dictionary distribution and the reconst.uuetio algorithm with varying sensing matrices. We provide theoretical and experimental support for architecture and show that our model learns calibrated uncertainties.
1 Introduction
By imposing prior on the signal structure, compressed sensing solves underdetermined inverse prob- lems. Canonical examples of signal structure and sensing medium sparsity and linear inverse problems. unknown signal of interest, E R", from set of linear
Compressed sensing aims at reconstructing measurements, E Rm acquired by means of a linear transformation, 4 E Ru where m < n. Due to the underdetermined nature the problem, s is typically assumed to be sparse in a given basis. Hence, s Vxx, where E Raxb is a matrix whose columns represent the sparsifying basis vectors, and æ E R&dh is the sparse representation of s. Therefore, given noiscless observations y $s, of an unknown signal, Lx, we sook to solve the LASSO problem:
argminlly Dv ø|
(1)
where P is a constant scalar controlling the sparsifying penalty. Iterative algorithms, such as Iterative Soft-Thresholding Algorithm (ISTA) (Daubechies al. 2004 represent a popular approach to solving
"Qualcomm AT Research is initiative Qualcomm Technologies, Inc.

---

Published in Transactions on Machine Learning Research (07/2024)
such problems. A number of studies have been conducted to improve compressed sensing solvers. A typical approach involves unfolding iterative algorithms as layers of neural networks and learning parameters end-to- end (Gregor & LeCun 2010). Such ML algorithms are typically trained by minimizing the reconstruction objective:
C(x, Bu-Dll 1T
(2)
where the expected value taken over data sampled from D, and the subscript "T" refers to the last step, or layer, of the unfolded model. Variable sensing matrices and unknown sparsifying dictionaries are some of the main challenges of data-driven approaches. By learning dictionary and including it in optimization iterations, the work in Aberdam et al. (2021); D; Schnoor et al. (2022) aims to address these issues. However, data samples might not have exact sparse representations, so no ground truth dictionary is available. The issue can be more severe for heterogeneous datasets where the dictionary choice might vary from one sample to another. A real-world example would be the channel estimation problem in a Multi-Input Multi-Output (MIMO) mmwave wireless communication system (Rodriguez-Fernández Fernández et al. 2018). Such a problem can be cast as an inverse problem of the form y = a and solved using compressive sensing techniques. The sensing matrix, 3, represents the so-called beamforming matric while the dictionary, V represents the sparsifying basis for the wireless channel itself. Typically, & changes from one set of measurements to the next and the channel model might require different basis vectors across time, Adaptive acquisition is another example of application: in MRI image reconstruction, the acquisition step can be adaptive. Here, the sensing matrix is sampled from a known distribution to reconstruct the signal. Therefore, given the adaptive nature of the process, each data sample characterized by different. (Bakker al. 2020 Yin et al. 2021).
Our Contribution A principled approach to this problem would be to leverage Bayesian framework and define distribution over the dictionaries with proper uncertainly quantification. We follow two steps to accomplish this goal. First, we introduce Augmented Dictionary Learning ISTA (A-DLISTA), an augmented version of the Learning Iterative Soft- Toft-Thresholding Algorithm (LISTA)-like model, capable of adapting some of its parameters to the current measurement setup. We theoretically motivate its design and empirically prove its advantages compared to other non-adaptive L.STA-like models in a non-static measurement scenario, i.e., considering varying sensing matrices. Finally, to learn a distribution over dictionaries, we introduce Variational Learning ISTA (VLISTA), variational formulation that leverages A-DLISTA as the likelihood model. VLISTA refines the dictionary iteratively after each iteration based on the outcome of the previous layer. Intuitively, our model can be understood as form of a recurrent variational autoencoder, c.g., Chung et al. 2015)), where at each iteration of the algorithm we have an approximate posterior distribution over the dictionaries conditioned on the outcome of the previous iteration. Moreover, VLISTA provides uncertainty estimation to detect Out-Of Of-Distribution (OOD) samples. We train A-DLISTA using the same objective as in Equation 2 while for VLISTA we maximize the ELBO (Equation 15). We refer the reader to Appendix D for the detailed derivation of the ELBO. [Behrens et al. (2021) proposed an augmented version of LISTA, termed Neurally Augmented ALISTA (NALISTA). However, there are key differences with A-DLISTA. In contrast to NALISTA, our model adapts some of its parameters to the current scnsing matrix and learned dictionary. Hypothetically, NALISTA could handle varying sensing matrices. However, that comes at the price of solving, for each datum, the inner optimization step to evaluate the W matrix. Finally, while NALISTA uses an LSTM as augmentation network, A-DLISTA employs a convolutional neural network (shared across all layers). Such difference reflects the type of dependencies between layers and input data that the networks try to model. We report in Appendix B and Appendix C detailed discussions about the theoretical motivation and architectural design for A-DLISTA.
Our work's main contributions can be summarized as follows:
We design an augmented version of a LISTA-like type of model, dubbed A-DLISTA, that can handle non-static measurement setups, i.e. per-sample sensing matrices, and adapt parameters to the current. data instance.
We propose VLISTA that learns a distribution over sparsifying dictionaries. The model can be interpreted Bayesian LISTA model that leverages A-DLISTA the likelihood model.
2

---

Published in Transactions on Machine Learning Research (07/2024)
VLISTA adapts the dictionary to optimization dynamics and therefore can be interpreted hierarchical representation learning approach, where the dictionary atoms gradually permit more refined signal recovery.
The dictionary distributions can be used successfully for out-of-distriuutio sample detection.
The remaining part of the paper is organized as follows. section 2 we brielly report related works relevant. to the current rescarch. In section 3) and |section 4 we introduce some background notions and details of our model formulations, respectively. Datasets, baselines, and experimental results are described in seccionn 5)
Finally, we draw our conclusion in section 6
2 Related Works
In compressed sensing, recovery algorithms have been extensively analyzed theoretically and numerically FFoucart& Rauhut, 2013). One of the most prominent approaches is using iterative algorithms, such as: ISTA
Daubechies et al. 2004), Approximate message passing (AMP) (Donoho al.) 2009) Orthogonal Matching Pursuit (OMP) (Pati et al. 1993 Davis et al. 1994). and Iterative Hard- Thresholding Algorithm (IHTA) Blumensath & Davies 2009). These algorithms have associated hyperparameters, including the number of iterations and soft threshold, which can be adjusted to better balance performance and complexity. With unfolding iterative algorithus as layers of neural networks, these parametens can be learned in an end-to-end fashion from a dataset see, for instance, some variants Zhang & Ghanem (2018); Metzler et al. (2017): p yang ct al. (2016): Borgerding et al. (2017); (Sprechmann et al. (2015). In previous studies by Zhoou et al. (2009) 2012 non- non-parametric Bayesian method for dictionary learning was presented. The authors focused on a fully Bayesian joint, compressed sensing inversion and dictionary learning, where the dictionary atoms were drawn and fixed beforehand. Bayesian compressive sensing (BCS) (Ji et al. 2008)) uses relevance vector machines (RVMs) (Tipping, 2001) and a hicrarchical prior to model distributions of each entry. This line of work quantifies the uncertainty of recovered entries while assuming a fixed dictionary. Our current work differs by accounting for uncertainty in the unknown dictionary by defining a distribution over it. Learming ISTA was initially introduced by Gregor & LeCun (2010). Since then, many works have followed, including those by [ehrenn et al.] (2021): Liu et al. (2019); Chen et al.) (2021); Wu et al. (2020). These subsequent works provide guidelines for improving LISTA, for example, in convergence, parameter efficiency, step size and threshold adaptation, and overshooting. However, they assume fixed and known sparsifying dictionaries and sensing matrices. Researches by Aberdam et al. (2021): Behboodi et al. (2022); Schnoor et. al. (2022) have explored ways to relax these assumptions, including developing models that can handle varying sensing matrices and learn dictionarics. The authors in [Schnoor ct al. (2022): ): Behboodi et al. (2022) provide an architecture that can both incorporate varying sensing matrices and learn dictionaries. However, their focus is on the theoretical analysis of the model. Furthermore, there are theoretical studies on the convergence and generalization of unfolded networks, see for example: Giryes et al. (2018); Pu et al. (2022); Aberdam et al. (2021) Pu et al. (2022): Chen et al. (2018): Behboodi et al. (2022): : Schnoor et al. (2022). Our paper builds these ideas by modelling a distribution over dictionaries and accounting for epistemic uncertainly. Previous studies have explored theoretical aspects of unfolded networks, such as convergence and generalization, and we contribute to this body of research by considering the impact of varying sensing matrices and dictionaries. The framework of variational autoencoders (VAEs) enables the learning of a generative model through latent variables (Kingma & Welling 2013 Rezende et al., 2014). When there are data-sample-specffc dictionaries in our proposed model, reminisces extensions of VAls to the recurrent setting (Chung et al.) 2015 2016), which assumes a sequential structure in the data and imposes temporal correlations between the latent variables. Additionally, there are connections and similarities to Markov state-space models, such as the ones described in Krishnan et al. (2017). By using global dictionaries in VLISTA, the model becomes variational Bayesian Recurrent Neural work. Variational Bayesian neural networks were first introduced in [lundel] et al. (2015), with independent priors and variational posteriors for each layer. This work has been extended to recurrent settings in Fortunato et al. (2019). The main difference between these works and our sctting is the prior and variational posterior. At each step, the prior and variational posterior are conditioned on previous steps instead of being fixed across steps.
3

---

Published in Transactions on Machine Learning Research (07 2024)
3 Background
3.1 Sparse linear inverse problems
We consider linear inverse problems of the form: y = y where we have access to a set of linear measurements y E _ of an unknown signal E R", acquired through the forward operator & E R=x=" Typically, in compressed sensing literature, & is called the sensing, or measurements, matrix and it represents an underdetermined system of equations for m n. The problem of reconstructing from (y, 4) is ill-posed due the shape of the forward operator. To uniquely solve for 8, the signal is assumed to admit sparse representation, E R,, in a given basis, {e, R") The vectors are called atoms and are collected the columns of a matrix & E Raxb termed the sparsifying dictionary. Therefore, the problem of estimating s, given a limited number of observations through the operator Φ, is translated to sparse recovery problem: x" = arg ming || æ0lo s.t. y = æ Given that the lo pseudo-norm requires solving an NP-hard problem, the I norm is used instead as a convex relaxation of the problem. A proximal gradient descenttbbsse approach for solving the problem yields the ISTA algorithm (Daubechies et al., 2004 Beck & Teboulle, 2009):
3t = me, (x1-1 (( -V2x,-1)
where t is the index of the current iteration, De (x1-1) is the reconstructed sparse vector at the current (previous) layer, and Og and are the sofl-threshold and slep size hyperparametens, respectively. Specifically, 6r characterizes the soft- seft-threshold function given by: no, (x) = sign (x) ([z) 0,)+. In the ISTA formulation, those two parameters are shared across all the iterations: Te, 02 + 7, 0. In what follows, we use the terms "layers" and "iterations" interchangeably when describing ISTA and its variations.
3.2 LISTA
LISTA Gregor & LeCun 2010 an unfolded version of the ISTA algorithm in which each iteration is parametrized by learnable matrices. Specifically, LISTA reinterprets Equation 3 defining the layer of a feed-forward neural network implemented as So, (V,æ,-1 Wyy) where V,, W. learnt from a dataset. Tn that way, those weights implicitly contain information about o and y that assumed to be fixed. As LISTA, also its variations, e.g., Analytic LISTA (ALISTA) (Liu et al. 2019), NALISTA (Behrens et al. 2021 and HyperLISTA (Chen et al. 2021), require similar constraints such as fixed dictionary and sensing matrix to reach the best performance. However, there are situations where one or none of the conditions are met (see examples in section I).
4 Method
4.1
Augmented Dictionary Learning ISTA (A-DLISTA)
To deal with situations where T is unknown and + changing across samples, one can unfold the ISTA algorithm and re-parametrize the dictionary a learnable matrix. Such an algorithm is termed Dictionary
Learning ISTA (DLISTA) (Pezeshki et al. 2022: Behboodi et al. 2022 Aberdam et al. 2021 and, similarly to Equution 3 each layer is formulated as:
170, (I + Ye ( Φ p..))))
(4)
with one last linear layer mapping x to the reconstructed signal 8. The model can be trained end-to-end to learn all lz, 7t. vt Differently from ISTA (Equation 3), DLISTA (Equation 4) learns a dictionary specific for each layer, indicated by the subscript "t". The model can be trained end-to-end to learn all A, %t V,- The base model is similar (Behboodi et al.) [022 Aberdam al., (2021). However, requires additional changes. Consider the t-th layer of DLISTA with the varying sensing matrix & and define the following
4

---

Published in Transactions on Machine Learning Research (07 2024)
parameters:
ptt, Φ*)
max ISI/ISN
(( ())) ( (1)
(5)
(t, Φ*)= ==
max 1S1JSN
(( ( (T"( vo)l
(6)
5(y, Φ²)
max 7 (1 v 1.):ll2
(7)
where ii is the mutual coherence of & Foucart Rauhut 2013 Chapter 5) and P2 closely connected to
generalized mutual coherence (Liu al. 2019p. However, in contrast to the generalized mutual coherence, jo2
includes the diagonal inner product for j. Finally, 8(-) is reminiscent of the restricted isometry property (RIP) constant {Foucart & Rauhut 2013); key condition for many recovery guarantees in compressed sensing. When the columns of the matrix & are normalized, the choice of r yields 8(y, t, x*) 0.
The following proposition provides conditions cach layer to improve the reconstruction error.
Proposition 4.1. Suppose that yd 4 where is the ground truth sparse vector with support supp(ar.) = S, and vo is the ground truth dictionary. For DLISTA ilerations given as
2t 10, (xt-1 + v ( "(y² (aat)))
(8)
have:
If for all t, the pairs (0e, Yt, 1) satisfy
nt (ā |x, 2,-1lh u2 [2212=11 |1=x111)= S O..
(9)
then there is no false posilive in each iteration. In other words, for all L, we have supp(a.) C supp(æ.)).
2. Assuming that the conditions of the last step hold, then we get the following bound on the error:
| æ.111 S (0(7)) + -n((S|) 1)) ||at--1 xx SS ||x.|1 + 1810,-
We provide the derivation of Proposition 4.1 together with additional thcoretical results Appendix Proposition 4.1 provides insights about the choice of 7t. and O,, and also suggests that (6(7)) + ( 1)) needs to be smaller than one to reduce the error each step.
(
[0%,vl3
[0%
(W.)
(45)
(66 rt)
*;
V
Figure 1: Models architectures. Left: A-DLISTA architecture. Each blue block represents a single ISTA-like iteration parametrized by the dictionary & the threshold and step size (0),nE) 7.). The red blocks represent the augmentation network (with shared parameters across layers) that adapts {0), 7 for layer based on the dictionary V, and the current measuremnntt setup o for the i-th data sample. Right: VLISTA (inference) architecture. The red and blue blocks correspond the same operations for A-DLISTA. The pink blocks represent the posterior model used to refine the dictionary based input data {y), 4' and the sparse vector reconstructed at layer t, Xx.
Upon examining Proposition 4.1 it becomes evident that % and 6, play key role the convergence of the algorithm. However, there is trade-off to consider when making these choices. For instance, suppose set.
5

---

Published in Transactions on Machine Learning Research (07/2024)
Algorithm 1 Augmented Dictionary Learning ISTA (A-DLISTA)- Inference Algorithm
Algorithm 2 Variational Learning ISTA (VLISTA)-
Inference Algorithm
Require:
"") the sensing matrix changes across samples; augmentation model fe ab +0
Require: = {(y' +') the sensing matrix
changes across samples; augmentation model fe:
posterior model fo
for t= 1, ..... T do
(0), 11) +- ( v.
> Augmentation step
for 1 1 1...,,,, do
x- y' - - ) o
(Mc, 02t + fo (2) Φ¹)
b Posterior
u - x + Vg
paramcters estimation
xi= no (u)
V.~N(Velle o²
Dictionary sampling
end for return wt;;
(0% 18) - ( v.
p Augmentation step
g + y' (
t x + thg
x nas (u)
end for return w/i
Og and decrease "/- In that case, we may ensure good support selection, but it could also increase d(7))- Tn situations where the sensing matrix remains fixed, the network can possibly learn optimal choices through end-to-end training. However, when the sensing matrix differs across various data samples (i.c., $ - D'), it is no longer guaranteed that there exists unique choice of 7/ and 0, for all o.. Since these parameters can be determined when & and v, are fixed, we suggest utilizing augmentation network to determine Yt and Og from each pair of d and v, For more thorough theoretical analysis, please refer to Appenddix B We show in Figure (left plot) the resulting A-DLISTA model. At each layer, A-DLISTA performs two basic operations ions, namely, soft-threshold (blue blocks in Figure 1) and augmentation (red blocks in Figure 1}. The former represents an ISTA-like iteration parametrized by the set of weights {v,,, 02 78,, whilst the latter is implemented using a convolutional neural network. As shown in the figure, the augmentation network takes as input the sensing matrix for the given data sample, & together with the dictionary learned at the layer for which the augmentation model will generate the 0; and 7f parameters: (0), 7/) fe ( I,,,, where O are the augmentation models' trainable parameters. Through such an operation, A-DLISTA adapts the soft- fett-hheshhod and step size of each layer to the current data sample. The inference algorithmic description of A-DLISTA is given in Algorithm [] We report more details about the augmentation network in the supplementary materials Appendix Cp.
4.2 Variational Learning ISTA
Although A-DLISTA possesses adaptivity to data samples, it still assumes the existence of a ground truth dictionary. We relax such a hypothesis by defining a probability distribution over vi, and formulating a variational approach, titled VLISTA, to solve the dictionary learning and sparse recovery problems jointly. To forge our variational framework whilst retaining the helpful adaptivity property of A-DLISTA, we re-interpret the soft-threshokding layers of the latter as part of likelihood model defining the output mean for the reconstructed signal. Given its recurrent-lin structure (Chung et al. 2015); , we equip VLISTA with a conditional trainable prior where the condition is given by the dictionary sampled at the previous iteration. Therefore, the full model comprises three components, namely, the conditional prior P:(+), the variational posterior (-), and the likelihood model, (.).). All components are parametrized by neural networks whose outputs represent the parameters for the underlying probability distribution. In what follows, we describe more in detail the various building blocks of the VLISTA model.
4.2.1 Prior Model
The conditional prior, Pe( Pe (.,),,,
1),,, is modelled as a Gaussian distribution whose parameters are conditioned the previously sampled dictionary. We implement (-) neural network, fe(-) [8], o (.), ogs.

---

Published in Transactions on Machine Learning Research (07/2024)
with trainable parameters & (Su. (2) The model's architecture comprises shared convolutional block followed by two branches generating the Gaussian distribution's mean and standard deviation, respectively. Therefore, at layer t, the prior conditional distribution is given by: Pe ( .--1) = II,, ((Priil Pl = f (Aso (i-i)) Øei.j fg. (9so ())))))s)), where the indices i, run over the rows and columns of v To
simplify our expressions, we will abuse notation and refer to distributions like the former one as:
Pe 1 ( where o²= = fE. (050 (V.-)))
(10)
F4 fg (05o 40 (-1))
We will use the same type of notation throughout the rest of the manuscript to simplify formulas. The prior design allows for enforcing a dependence of the dictionary at iteration f to the one sampled at the previous iteration. Thus allowing
refine L the iterations proceed. The only exception is the prior imposed over the dictionary at t= 1, where there is no previously sampled dictionary. To handle this exception, we assume standard Gaussian distributed L, The joint prior distribution over the dictionaries for VLISTA is given by:
Pe 4 1: = (1 1) II N N(W.lp
(11)
where Me and o², are defined in Equation 10
4.2.2 Posterior Model
Similarly to the prior model, the variational posterior too is modelled as Gaussian distribution parametrized by a neural network fa(-) = 18 h-o f o haw (-) that outputs the mean and variance for the underlying probability distribution
do (v(V.l11111)) y' x0) -N(V:lur;o')) where Fe ln (h (I.-1. y', Φ')); ))) a²: 12 (hop (1; (b))
(12)
The posterior distribution for the dictionary, Ve, at layer / is conditioned on the data, {v', D'), as well the reconstructed signal at the previous layer, F4-1- Therefore, the joint posterior probability over the dictionaries is given by:
go 1:T L.T y' x)) to V (
II
,,, D')
(13)
When considering our selection of Gaussian distributions for our prior and posterior models, we prioritized computational and implementation convenience. However, it's important to note that our framework is not limited to this distribution family. As long as the distributions used are reparametrizable (Kingma & Welling 2013), meaning that
obtain gradients of random samples with respect to their parameters and can evaluate and differentiate their density, VL.STA can support any flexible distribution family. This includes mixtures of Gaussians to incorporate heavier tails and distributions resulting from normalizing flows (Rezende & Mohamed 2015).
4.2.3 Likelihood Model
The soft-threshokding block of A-DLISTA is at the heart of the reconstruction module. Similarly to the prior and posterior, the likelihood distribution is modelled as a Gaussian parametrized by the output of an A-DLISTA block. In particular, the network generates the mean vector for the Gaussian distribution while we treat the standard deviation a tunable hyperparameter. By combining these elements, we can formulate the joint ligeliiel distríbution as:
7

---

Published in Transactions on Machine Learning Research (07 (274))
9Φ xxclxx-1110 )0
Ps (
4
4t
x
8 Yt-1
Y
Figure 2: VLISTA graphical model. Dependencies on y' and & are factored out for simplicity. The sampling done only based on the posterior qo (f/ y', Φ²). Dashed lines represent variational approximations.
T
log (a = o V Wi.Ty', σb)
Σ
log N (argel a²t), where
(14)
H. A-DLISTA L Tt 21 O): a², =
where is a hyperparameter of the network, xiya represents the ground truth value for the underlying unknown sparse signal for the i-th data sample, θ is the set of A-DLISTA's parameters, and pa( x (-) represents the likelihood for the ground truth ziy, at each time step & [1,T], under the likelihood model given the current reconstruction. Note that in Equation 14 we use the same æ through the entire sequence t E []. We report Figure (right plot) the inference architecture for VLISTA. It crucial to note that during inference VLISTA does not require the prior model which used for training only. Additionally, its graphical model and algorithmic description are shown in Figure 2 and Algorithm 2 respectively. We report further details concerning the architecture of the models and the objective function in the supplementary material Appendix C and Appendix D
4.2.4 Training Objective
Our approach involves training all VLISTA components in end-to-end fashion. To accomplish that we maximize the Evidence Lower Bound (ELBO):
E B ( loo- D:) [og pe(a: = ar 11., Dল)
ELBO =
(15)
Σ E Eyan (B1 w/2 ..D:) ..D)) DKY (V. D() (M 2
DKL go ( 1a0o Dt) p( (T)
As we can see from Equation 15 the ELBO comprises three terms. The first term is the sum of expected log-likeli ligliikli of the target signal at each time step. The second term the sum KL divergences between the approximate posterior and the prior at each time step. The third term the KL divergence between the approximate posterior at the initial time step and a prior. In our implementation, we set to "T the number of layers and initialize the input signal to
8

---

Published in Transactions on Machine Learning Research (07/2024)
To evaluate the likelihood contribution in Equation 15 posterior 1 D In contrast, the last two terms in the equation represent the KT. divergence contribution between the prior and posterior distributions. It's worth noting that the prior in the last term is not conditioned on the previously sampled dictionary, given that P + p( ) = ( |0: 1) (refer to Equation and Equation 11 We refer the reader to Appendix for the derivstion of the ELBO.
marginalize over dictionaries sampled from the
5 Experiments
5.1 Datasets and Baselines
We evaluate our models' performance by comparing them against classical and ML-based baselines on three datasets: MNIST, CIFARIO, and synthetic dataset. Concerning the synthetic dataset, we follow a similar approach as in Chen et al. 2018 Lin & Chen (2019); Behrens et al. (2021). However, in contrast to the mentioned works, we generate a different & matrix for each datum by sampling i.idd. entrics from a standard Gaussian distribution. We generale the ground truth sparse signals by sampling the entries from a standard Gaussian and setting each entry to be non-zero with a probability of 0.1. We generate 5K samples and use train 3K for training, 1K for model selection, and 1K for testing. Concerning the MNIST and CIFAR1O, the models using the full images, without applying any crop. For CIFAR10, gray-scale and normalize the images. We generate the corresponding observations, y', by multiplying each sensing matrix with the ground truth image: y¹ = ='ss'. We compare the A-DLISTA and VLISTA models against classical and ML baselines. Our classical baselines use the ISTA algorithm, and we pre-compute the dictionary by either considering the canonical or the wavelet basis or using the SPCA algorithm. Our MT, baselines use different unfolded learning versions of ISTA, such as LISTA. To demonstrate the benefits of adaptivity, we perform an ablation study on A-DLISTA by removing its augmentation network and making the parameters 0z, Ytt learnable only through backpropagation. We refer to the non-augmented version of A-DLISTA as DLISTA. Therefore, for DLISTA. 0 and cannot be adapted to the specific input sensing matrix. Morcover, we consider BCS (Ji et al. 2008) as a specific Bayesian baseline for VLISTA. Finally, we conduct Out-Of-Distribution (OOD) detection experiments. We fixed the number of layers to three for all ML models to compare their performance. The classical baselines do not possess learnable parameters. Therefore, we performed an extensive grid search to find the best hyperparameters for them. More details concerning the training procodure and ablation studies can be found in Appendix D and Appendix F
5.2 Synthetic Dataset
Regarding the synthetic dataset, we evaluate models performance by computing the median of the c.d.f. for the reconstruction NMSE (Fig ure 3). A-DLISTA'S adaptivity appears to offer an advantage over other models. However, con- cerning VLISTA, we observe drop in perfor- mance. Such behaviour is consistent across experiments and can be attributed to a few fac- tors. One possible reason for the drop in perfor- mance is the noise introduced during training due to the random sampling procedure used to generate the dictionary. Additionally, the amor- tization gap that affects all models based on amortized variational inference (Cremer et al. 2018) can also contribute to this effect. Despite this, VLISTA still performs comparably to BCS. Lastly, we note that ALISTA and NALISTA do not perform as well as other models. This likely due to the optimization procedure these two models require to evaluate the weight matrix W. The computation of the W matrix requires fixed
0.0
-0.5
-1.0
NMSE
-1.5
USTA
ALISTA
ADUSTA (our)
-2.0
BCS
VLISTA (our)
Measurements
Figure 3: NMSE's median. The y-axes is in dB (the lower the better) for different number of measurements (I-axes).
9

---

Published in Transactions on Machine Learning Research (07 (024))
sensing matrix, condition not satisfied in the current setup, Regarding non-static measurements, we averaged across multiple Φ thus obtaining non-optimal W matrix. To support our hypothesis, we report in Appendix F results considering a static measurement scenario for which ALISTA and NALISTA report very high performance.
5.3 Image Reconstruction MNIST & CIFAR10
When evaluating different models on the MNIST and CIFAR10 datasets, we use the Structural Similarity Index Measure (SSTM) to measure their performance. As for the synthetic dataset, we experience strong instabilities in ALISTA and NALISTA training due to non-static measurement setup. Therefore, we do not provide results for these models. is important note that the poor performance of ALISTA and NALISTA is a result of our specific experiment setup, which differs from the static case considered in the formulation of these models. We refer to Appendix F for results using static measurements scenario. By looking at the results in Table ]] and Table 2, we draw similar conclusions as for the synthetic dataset. Additionally, we report results from three classical baselines (subsection 5.1p. Among non- -an-asasia models, A-DLISTA shows the best results. Furthermore, by comparing A-DLISTA with its non-augmented version, DLISTA, one can notice the benefits of using an augmentation network to make the model adaptive. Concerning Bayesian approaches, VLISTA outperforms BCS. However, is important to note that BC does not have trainable parameters, unlike VLISTA. Therefore, the higher performance of VLISTA comes the price of an expensive training procedure. Similar to the synthetic dataset, VLISTA exhibits drop in performance compared to A-DLISTA for MNIST and CIFAR1O.
Table MNIST SSIM (the higher the better) for different number of mcasurements. First three rows correspond "classical" baselines. We highlight in bold the best performance for Bayes and Non-Bayes models.
SSTM +
number of measurements
10
100
300
500
(xe-1)
(xe-1)
(xe-1)
(xe=l)
(xe-1)
Canonical
0.39
0.5640.04
2.20
3.75 3.75+0.05
4.94
Non- Non-Bayes Bayes
Wavelet
0.40
0.5610.06
2.30
3.90
5.05.0.01
SPCA
0.45+0.11
0.65+0.06
2.72+0.00
3.52+0.08
4.9840.08
LISTA
0.9640.01
1.1140.01
3.70
5.3640.01
6.31
DLISTA
0.9640.01
1.0940.01
4.01
5.57+0.01
6.2640.01
A-DLISTA (our)
0.961
1.17 10.01
4.79
6.15.10.01
6.70
Bayes
BCS
0.05
0.60
1.1040.01
4.48
6.23
VLISTA (our)
0.80
0.94
3.2940.01
4.73
6.02
5.4 Out Of Distribution Detection
This section focuses on a crucial distinction between non-Bayesian models and VLISTA for solving inverse linear problems. Unlike any non-Bayesian approach compressed sensing, VLISTA allows quantifying uncertainties the reconstructed signals. This means that it can detect out-of-disttribation samples without requiring ground truth data during inference. In contrast to other Bayesian techniques that design specific priors to meet the sparsity constraints after marginalization (Ji et al.) 2008 Zhou et al. 2014), VLISTA completely overcomes such an issue as the thresholding operations are not allected by the marginalization over dictionaries. To prove that VLISTA can detect OOD samples, we employ the MNIST dataset. First, we split the datasct into two distinct subsets the In-Distribution (ID) set and the OOD. The ID set comprises images from three randomly chosen digits, while the OOD set includes images of the remaining digits. Then, partitioned the ID set into training, validation, and test sets for VLISTA. Once the model was trained,
10

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Table 2: CIFAR1O SSIM (the higher the better) for different number of measurements. First three rows correspond "classical" baselines. We highlight in bold the best performance for Bayes and Non-Bayes models.
SSIM t
number of mcasurements
10
100
300
500
(xe-1)
(xe=))
(xe
(xe-l)
(xe
Canonical
0.17=0.10
0.21+0.02
0.33
0.47+0.02
0.5840.03
Non-Bayos
Wavelet
0.23
0.4240.02
1.44
2.52+0.09
3.43
SPCA
0.31
0.43+0.02
1.53
2.6640.08
3.5840.07
LISTA
1.34
1.67
3.10
4.20
4.71
DLISTA
1.16
1.9640.02
4.50=0.01
5.1540.01
5.42+0.01
A-DLISTA (our)
1.34+0.02
1.77
4.74+0.01
5.26=0.01
5.83+0.01
Bayes
BCS
0.04
0.48+0.01
0.59
1.29
1.91
VLISTA (our)
0.8640.08
1.2550.03
3.5940.02
4.01
4.3640.01
it was tasked with reconstructing images from the ID test and OOD sets. To assess the model's ability detect OODs, utilized a huo-sample t-test. We accomplished that by leveraging the per-pixel variance of the reconstructed ID, f (var!Draas)) Ji=0 and OOD, (varOODa) Ji=0 images (with P being the number of pixels). To compute the per-pixel variance, we reconstruct each image 100 times by sampling a different dictionary for each of Irial. We then construct. the empirical c.d.f. of the per-pixel variance for each image. By using the mean of the c.d.f. as a summary statistics, we can apply the two-sample t-test to detect OOD samples. We report the results in Figure 4 As reference p-value for rejecting the null hypothesis about the two variance distributions being the same, we consider a significance level equal 0.05 (green solid line). We conducted multiple tests at different noise levels to assess the robustness of OOD detection to measure noise. For the current task, used BCS as baseline. However, due to the different nature of the BCS framework, utilized slightly different evaluation procedure to determine its p-values. We employed the same ID and OOD splits as VLISTA but considered the c.d.f. of the reconstruction error the model evaluates. The remainder of the process was identical to that of VLISTA.
6 Conclusion
10a
Our study introduces a novel approach called VLISTA, which combines dictionary learning and sparse recovery into single variational framework. Traditional compressed sensing methods rely on known ground truth dictio- nary to reconstruct signals. Moreover state-of- the-art LISTA-type of models, typically assume a fixed measurement setup. In our work, we relax both assumptions. First, we propose soft-thresholdig algorithm, termed A-DLISTA, that can handle different sensing matrices. We thcoretically justify the use of an augmentation network to adapt the threshold and step size for each layer based on the current input and the learned dictionary. Finally, we propose a probabilistic assumption about the existence ground truth dictionary and use it to create
10
p-value
10-6
8
BCS
VLISTA
10-10
p-value
0.05
10-1
10
101
102
SNR (dB)
Figure 4: p-value for OOD rejection as function of the noise level. The green line represents reference p-value equal to 0.05.
11

---

Published in Transactions on Machine Learning Research (07/2024)
the VLISTA framework. Our empirical results show that A-DLISTA improves upon performances of classical and ML. baselines in non-static measurement scenario. Although VLISTA does not outperform A-DLISTA, it allows for uncertainty evaluation in reconstructed signals, a valuable feature for detecting out-of-distributioo data. In contrast, none of the non-Bayesian models can perform such a task. Unlike other Bayesian approaches, VLISTA does not require specific priors to preserve sparsity after marginalization. Instead, the averaging operation applies to the sparsifying dictionary, not the sparse signal itself.
Impact Statement
This work proposes two new models to jointly solve the dictionary learning and sparse recovery problems, especially concerning scenarios characterized by a varying sensing matrix. We believe the potential societal consoquences of our work being chiefly positive, since it might contribute to a larger adoption of LISTA-type of models to applications requiring fast solutions to underdetermined inverse problems, especially concerning varying forward operators. Nonetheless, it crucial to exercise caution and thoroughly comprehend the behavior of A-DLISTA and VLISTA, with any other LISTA model, in order to obtain reliable predictions.
References
Aviad Aberdam, Alona Golts, and Michael Elad. Ada-LISTA: Learned Solvers Adaptive to Varying Models. IEEE Transactions on Pullern Analysis and Machine Intelligence, pp- 1-1, 2021.
Tim Bakker, Herke van Hoof, and Max Welling. Experimental design for mri by greedy policy search. Advances in Newral Information Processing Systems, 33:18954 18966, 2020.
Amir Bock and Mare Teboulle. A fast iterative shrinkage-thresholding algorithm for linear inverse problems.
SIAM journal on imaging sciences, 2(1):183-200 2002 2009.
Arash Behboodi, Holger Rauhut, and Flkeehrdd Schnoor. Compressive Sensing and Neural Networks from Statistical Learning Perspective. In Gitta Kutyniok, Holger Rauhut, and Robert J. Kunsch (eds.), Compressed Sensing in Information Processing, Applied and Numerical Harmonic Analysis, pp. 247 277. Springer International Publishing, Cham, 2022.
Freya Behrens, Jonathan Sauder, and Peter Jung. Neurally Augmented ALISTA, In International Conference on Learning Representations, 2021.
Thomas Blumensath and Mike E. Davics. Iterative hard thresholding for compressed sonsing. Applied and Computational Harmonic Analysis, 27(3):265 274, November 2009.
Charles Blundell, Julien Cornebise, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International conference on machine learning, pp. 1613 1622. PMLR, 2015.
Mark Borgerding, Philip Schniter, and Sundeep Rangan. AMP-inspired deep networks for sparse linear inverse problems. IEEE Transactions Signal Processing, 42(16) -3088, 2017. Publisher: IEEE.
Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Theoretical linear convergence of unfolded ISTA and its practical weights and thresholds. Advances in Neural Information Processing Systems, 31, 2018.
Xiaohan Chen, Jialin Liu, Zhangyang Wang, and Wotao Yin. Hyperparameter tuning is all you need for lista. Advances in Neural Information Processing Systems, 34, 2021.
Junyoung Chung, Kyle Kastner, Laurent Dinh, Kratarth Goel, Aaron C Courville, and Yoshua Bengio. A recurrent latent variable model for sequential data. Advances in neural information processing systems, 28, 2015.
Junyoung Chung, Sungjin Ahn, and Yoshua Bengio. Hierarchical multiscale recurrent neural networks. arXiv preprint ar Xrx:::6609.011 014407 2016.
12

---

Published in Transactions on Machine Learning Research (07/2024)
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders. In International Conference on Machine Learning, pp-
PMLR, 2018.
Ingrid Daubechies, Michel Defrise, and Christine De Mol. An iterative thresholding algorithm for linear inverse problems with a sparsity constraint. Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institule of Mathemalicul Sciences, 57(11):1413 1457, 2004.
Geoffrey M. Davis, Stephane G. Mallat, and Zhifeng Zhang. Adaptive time-fremuent freeurruy decompositions. Optical Engineering, 33(7):2183-221) 2191, July 1994.
David L. Donoho, Arian Maleki, and Andrea Montanari. Message-passing algorithms for compressed sensing. Proceedings of the National Academy of Sciences, 106(4)):18914-1891 November 2009.
Meire Fortunato, Charles Blundell, and Oriol Vinyals. Bayesian Recurrent Neural Networks. arXiv: 004.02798 arXiv: 1704.02798.
[es, statj, May 2019. URL http:://axtii orggabs////a.aa...
http:://axxx..orrraa////OO 02798
Simon Foucart and Holger Rauhut. Mathematicul Introduction to Compressive Sensing. Applied and Numerical Harmonic Analysis. Springer New York, New York, NY, 2013.
Raja Giryes, Yonina C. Eldar, Alex M. Bronstcin, and Guillermo Sapiro. Tradeoffs between convergence speed and reconstruction accuracy in inverse problems. IEEE Transactions Signal Processing, 66(7): Publisher: IEFF.
1676-1690, 2018.
Karol Gregor and Yann LeCun. Learning fast approsimations of sparse coding. In 27th International Conference on Machine Learning, ICML 2010, 2010.
Shihao Ji, Ya Xue, and Lawrence Carin. Bayesian Compressive Sensing. IEEE Trunsactions on Signal Processing, 56(6):23446 22666, June 2008. Conference Name: IEEE Transactions on Signal Processing.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint ar arXiv:1312.6114 2013.
Rahul Krishnan, Uri Shalit, and David Sontag. Structured inference networks for nonlincar state space models. Proceedings of the AAAI Conference on Artificial Intelligence, 31(1), Feb. 2017.
Jialin Liu and Xiaohan Chen. Alista: Analytic weights are as good as learned weights in lista. In International Conference on Learning Representations (TCLR), 2019.
Jialin Liu, Xiaohan Chen, Zhangyang Wang, and Wotao Yin. ALISTA: Analytic Weights Are As Good As Learned Weights in LISTA. In International Conference Learning Representations, 2019.
Chris Metzler, Ali Mousavi, and Richard Baraniuk. Learned D-AMP: Principled Neural Network based Com- pressive Image Recovery. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems 30, pp. 1770 1781. Curran Associates, Inc., 2017.
Y.C. Pati, R. Rezaiifar, and P.S. Krishnaprasad. Orthogonal matching pursuit: recursive function approxima- tion with applications wavelet decomposition. In Proceedings of 27th Asilomar Conference on Signals, Systems and Computers, pp. 40-44 vol.1, November 1993. doi: 10.1109/ACSSC.1993
Hamed Pezeshki, Fabio Valerio Massoli, Arash Behboodi, Taesang Yoo, Arumugam Kannan, Mah- moud Taherzadeh Boroujeni, Qiaoyu Li, Tao Luo, and Joseph B Soriaga. Beyond codebook-based analog beamforming at mmwave: Compressed sensing and machine learning methods. In GLOBECOM 2022-2022 IEEE Global Communications Conference, pp. 776-781. IEEE, 2022.
Wei Pu, Yonina C. Eldar, and Miguel R. D. Rodrigues. Optimization Guarantecs for ISTA and ADMM Based Unfolded Networks. In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp.
May 2022.
13

---

Published in Transactions on Machine Learning Research (07/2024)
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International conference machine learning, pp. 1530-1538. PMLR, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International conference machine learning, pp- 1278 1286.
PMLR, 2014.
Javier Rodriguez- Rodriguezz-Frnsändee Fernández, Nuria González-Prelcic, Kiran Venugopal, and Robert W Heath. Frequency- domain compressive channel estimation for frequency-sclecthe hybrid millimeter wave mimo systems. IEEE Transactions Wireless Communications,, 17(5):2946 2960, 2018.
Ekkehard Schnoor, Arash Behboodi, and Holger Rauhut. Generalization Error Bounds for Iterative Recovery Algorithms Unfolded as Neural Networks. arXiv.2112.04364, January 2022.
P. Sprechmann, A. M. Bronstein, and G. Sapiro. Learning Elficient Sparse and Low Rank Models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 97(9):1821-188)) September 2015. ISSN 0162-8828.
Michael E. Tipping. Sparse Bayesian learning and the relevance vector machine. Journal of machine learning research, 1(Jun):211 244, 2001.
Kailun Wu, Yiwen Guo, Ziang Li, and Changshui Zhang. Sparse coding with gated learned ista. International Conference on Learning Representations, 2020.
yan yang, Jian Sun, Huibin Li, and Zongben Xu. Deep ADMM-Net for Compressive Sensing MRL In D. Lee, M. Sugiyama, U. Luxburg, I. Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016.
Tianwei Yin, Zihui Wu, He Sun, Adrian V Dalca, Yisong Yue, and Katherine I. Bouman. End-to-end sequential sampling and reconstruction for mri. arXiv preprint orXiv:2105.064 2021.
Jian Zhang and Bernard Ghanem. ISTA-Net: Interpretable optimizationninspires deep network for image compressive sensing. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 1828-1837, 2018.
Mingyuan Zhou, Haojun Chen, Lu Ren, Guillermo Sapiro, Lawrence Carin, and John Paisley. Non-Parametric Bayesian Dictionary Learning for Sparse Image Representations. In Advances in Neural Information Processing Systems, volume 22. Curran Associates, Inc., 2009.
Mingyuan Zhou, Haojun Chen, John Paisley, Lu Ren, Lingbo Li, Zhengmming Xing, David Dunson, Guillermo Sapiro, and Lawrence Carin. Nonparametric Bayesian Dictionary Learning for Analysis of Noisy and Incomplete Images. IEEE Transactions on Image Processing, 21 21(1):130-11 144, January 2012.
Zhou Zhou, Kaihui Liu, and Jun Fang. Bayesian compressive sensing using normal product priors. IEEE Signal Processing Letters, 22((5)5883555 587, 2014.
14

---

Published in Transactions on Machine Learning Research (07 /024)
A Appendix
B Theoretical Analysis
In this section, we report about theoretical motivations for the A-DLISTA design choices. The design is molivated by considering the convergence analysis of LISTA method. We start by recalling a result from Chen et al. (2018), upon which our analysis relies. The authors of Chen al. (2018) consider the inverse problem Ax., with æ, as the ground truth sparse vector, and use the model with each layer given by:
Zt = no, (+--1+ + W. Ax(-1)),
(16)
where (We, (t) are learnable paramcters.
The following result from Chen et al. 20018 Theorem 2) is adapted to noiseless setting.
Theorem B.1. Suppose that the ileralions of LISTA are given by equalion 16 and assume |2c. S B and |æ+1lo S s. There ezists sequence of porameters {W,, 0,} such that
| x=112 S sB exp(( Vt = 1,2,
for constant e> that depend only on the sensing matrir A and the sporsity level 8.
It important note that the above convergence result only assures the eristence of the paramcters that are good for convergence but does not guarantee that the training would necessarily find them. The latter result is in general difficult to obtain.
The proof in Chen ct al. 2018) has two main steps:
1. No false positive: the thresholds are chosen such that the entries outside the support of æ. remain zero. The choice of threshold, among others, depend on the coherence value of W. and A. We will provide more details below.
2. Error bounds for æ,; assuming proper choice of thresholds, the authors derive bounds on the recovery error.
We focus on adapting these sleps to our setup. Note that to assure there is no false positive, it is common in classical ISTA literature to start from large thresholds, so the soft thresholding function aggressively maps many entries to zero, and then gradually reduce the threshold value as the iterations progress.
B.1 Analysis with known ground-truth dictionary
Let's consider the extension of Theorem B.1 to our sctup:
2t = 11.. (x,-1 + (I" (gy xx " v,)r))))
(17)
Note that in our case, the weight W, replaced with n ($)))) with learnable V, and 7t. Besides, the matrix A is replaced by t* and the forward model is given by o $ ,... The sensing matrix 4 can change across samples, hence the dependence on the sample index k.
If the learned dictionary V is equal to V.. the layers of our model are equal to classical iterative soft- thresholding algorithms with learnable step-size and threshold 0z.
There are many convergence results in the literature, for example see Daubechies et al. 2204hh. We can use convergence analysis of iterative soft thresholding algorithms using the mutual coherence similar to Chen
al. 2018): Behrens al. (2021) As reminder, the mutual coherence of the matrix M is defined as:
((M) = max ISiFiSN
MM [TMM]
(18)
where M; is the i'th column of M.
15

---

Published in Transactions on Machine Learning Research (07/2024)
The convergence result requires that the mutual coherence (( be sufficiently small, for example in order of (2s))) with the sparsity, and the matrix $ is column normalized, i.e., [[ = 1. Then the step size can be chosen equal to one, Le., > = 1. The thresholds 0r are chosen to avoid false positive using a similar schedule mentioned above, that is, first starting with a large threshold d and then gradually decrcasing it to
certain limit. We do not repeat the derivations, and interested readers can refer to Daubechies et al. (2004); Behrens al. ); Chen et al. (2018) and references therein.
Remark B.2. When the dictionary V, is known, we can adapt the algorithm to the varying sensing matrix $ by first normalizing the column $ T. What is important to note is that the threshold choice is function the mutual coherence of the sensing matrix. So with each new sensing matrix, the thresholds should be adapted following the mutual coherence value. This observation partially justifies the choice of thresholds as function of the dictionary and the sensing matrix, hence the augmentation network.
B.2 Analysis with unknown dictionary
We now move the scenario where the dictionary is itself learned, and not known in advance. Consider the layer t of DLISTA with the sensing matrix s and define the following parameters:
jutt, Φh) == max
(( ()))) Le)s
(19)
i21 Φ²)== = max
((Tb">,)) (T*(IV, v.),,
(20)
151JE5N
8(7, 1, 2b)) = max - 7 [| v.)
(21)
Some comments are in order:
The term ji is the mutual coherence of the matrix $
The term #2 is closely connected to generalized mutual coherence, however, it differs in that unlike generalized mutual coherence, it includes the diagonal inner product for i j. It captures the effect of mismatch with ground-truth dictionary.
Finally, the term 8(-) is reminiscent of restricted isometry property (RIP) constant (Foucart & Rauhut 2013), key condition for many recovery guarantees in compressed sensing. When the columns of the matrix & v, is normalized, the choice of 7 =1 yield 8(y, Φ*) 0.
For the rest of the paper, for simplicity, we only kept the dependence on in the notation and dropped the dependence of Jt, u2 and on i, t and v,,
Proposition B.3. Suppose that n L. with the support supp(æ>>) S. For DLISTA iterations give as
at = 10, (x=-1 + 7e ($*,, $b" 4b' (,
(22)
have:
1. If for all t, the pairs (0,, ns, vi) satisfy
th (ū lx. æt-1ll, A2 (1x= < On,
(23)
then there is no false positive in each iteration. In other words, for all we have supp(æ:) S supp(x.).
Assuming that the conditions of the last step hold, then we get the following bound the error:
|x0, x.l11 S (6(7)) + SP - 1)) || - æ. + T ||c.||, + 1S00,
16

---

Published in Transactions on Machine Learning Research (07 (224))
B.2.1 Guidelines from Proposition.
We remark on some of the guidelines we can get from the above result.
Thresholds. Similar to the discussion in previous sections, there are thresholds such that, there is no false positive at each layer. The choice of 0 is function of % and, through coherence terms, & and v Since & changes for cach sample k we learn a neural network that yiclds this paramcter as a function of & and V.
Step size. The step size " can be chosen to control the error decay. Ideally, we would like to have the term (ō(1)) + [H][S)] 1)) to be strictly smaller than one. In particular, Yt directly impacts 8(1)), also a function of & and Vr. We can therefore consider Tt as a function of & and L,, which hints at the augmentation neural network we introduced for giving " as function of those parameters.
Remarks on Convergence. One might wonder the convergence possible given the bound on the error. We try to sketch a scenario where this can happen. First, note that once we have chosen Yt, and given & and Vt, we can select 0g using the condition 23 Also, if the network gradually learns the ground truth dictionary at later stages, the term jā2 vanishes. We need to choose the term Yt carefully such that the term (8(1)) Mu[[S| 1)) is smaller than one. Similar to ISTA analysis, we would need to assumee bounds on the mutual coherence jú and the column norm for T With standard assumptions, sketched above as well, the error gradually decreases per iteration, and we can reuse the convergence results of ISTA. We would like to emphasize that this is a heuristic argument, and there is no guarantee that the training yields a model with the parameters in accordance with these guidelines. Although we show experimentally that the proposed methods provide the promised improvements.
B.3 Proof of Proposition B.3
In what follows, we provide the derivations for Proposition B.3
Convergence proofs of ISTA type models involve two steps in general. First, it is investigated how the support is found and locked in, and second how the error shrinks at each step. We focus on these two steps, which matter mainly for our architecture design. Our analysis is similar in nature to Chen al. (2018); Aberdam ct al. (2021), however it differs from Aberdam et al. (2021) in considering unknown dictionaries and from Chen et al. (2018) in both considered architecture and varying sensing matrix. In what follows, consider noiseless setting. However, the results can be extended to noisy setups by adding additional terms containing noise norm similar to Chen et al. (2018). We make following assumptions:
1. There is ground-truth (unknown) dictionary vo such that 8, V V..
2. As a consequence, y T....
3. We assume that a, is sparse with its support contained in S. In other words: zi,, = 0 for i i S.
To simplify the notation, we drop the index k, which indicates the varying sensing matrix, from & and y and use & and y for the rest. We break the proof two lemma, each proving one part of Proposition B.3
B.3.1 Proof step 1: no false positive condition
The following lemma focuses assuring that we do not have false posilive in support recovery after each iteration of our model. In other words, the models continues updating only the entries in the support and keep the entries outside the support zero.
Lemma B.4. Suppose that the support of x. is given as supp(x.)) = S. Consider iterations given by:
x² 11, (x + 7e ( @V(0)-1)),
17

---

Published in Transactions on Machine Learning Research (07 (2/22)
with Zo O. If we have for all 1,
To (it |10, - 21-1ll1 + 22 1| s Os,
then there will no false posilive, i.e, Fe,s = for Vi $ S, VL.
Proof. We prove this by induction. Since wo = O, the induction base trivial. Suppose that the support of #t-1 already included in that of v. namely supp( 11) S supp(æ.) = Consider E S We have
Tti no, (7.( Ty L -
(24)
To avoid false positives, we need to guarantee that for i i S:
no, (T((()))))) (y o (V;=)))) =0 (
(V S 8,,
(25)
which means that the soft tote-thresholing function will have zero output for these entries. First note that:
( ))) 2
P a. v. $ (( (( 4 (v vx
(( ( v.x.)))
(26)
Σ ((bV)).) (#NV)))( (sWW:)) - 21.1))) + +((()))))) (VV... v,xx,)
(27)
jES
We can bound the first term by:
Σ
#1-15)
Σ ((
(a 2(-15)1
jES
jES
S 1. B1-1ll1>
where we use the definition of mutual coherence for the upper bound. The last term is bounded by
(( (())))) Φ
vxx.)]
Σ (( o ( ,,)),a},
(28)
jES
((DV)):) (.))s
(29)
jES
$ Øg |x.111.
(30)
Therefore, we get ()((())))))) ( (y Tb,,,.1)) no (ī ||a. æ2-1l11 uia 11a.111) 11
The following choice guarantees that there is no false positive:
Ta (īt ||c. - x1-1ll1 + g2 ||x: s O.
(31)
B.3.2 Proof step 2: controlling the recovery error
The previous lemma provided the conditions such that there is no false positive. We see under which conditions the model can reduce the error inside the support S.
Lemma B.5. Suppose that the threshold parameter 0g has been chosen such that there is no false positive after each éleration. We have:
|1. æ.lli (5(%)) n 1)) || æ..11 + 71 || 18
18

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Proof. For i S. we have:
[r,s Ra,il S [4-1, + ( (vaa v.,a O2.
(32)
At the iteration for i E S, we can separate the dictionary mismatch and the rest of the error as follows: 21-14 + 7o ())) Ty -V)x=-1)=
It-15 + 7el Σ (( ( (DV1))(a.s x1-1,)) + ((+ @ v,x))).
We can decompose the first part further as:
Σ (<v.)))
Tt
JES
(6 1-1 + ( (
7
Σ (
1 ( 1.j)
JES
Using triangle inequality for the previous decomposition we get:
[--11 + (c(()))))) (y- 2v,20-11 x.,o - (c(V)))) (LV-
Σ ((+U))) ( II-1
+ %
JESJ7i
+ 71 ((+))))) T D(V.o. v,x.))
$ 6(7) ( z-,))
Σ [xx_ x1-1 + Y ||x,||
+ 7
jES,j
It suffices sum up the errors and combine previous inequalities get:
| x.111
Σ æ1, x.,1 I iES
$ (6(m)) + n(((S) - 1)) | x=11, + S ||x,|| + |S|0.-
Since we assumed there is no false positive, we get the final result:
|2 a2.111 = Σ
|mr, xn,il = (5(1)) + mi([S) - 1)) | - æ..11 S | + |S)0
||
iES
C Implementation Details
In this section we report details concerning the architecture of A-DLISTA and VLISTA.
C.1 A-DLISTA (Augmentation Network)
As previously stated in the main paper (subsection 4.1), A-DLISTA consists of two architectures: the DLISTA model (blue blocks in Figure 1) representing the unfolded version of the ISTA algorithm with parametrized V and the augmentation (or adaptation) network (red blocks in Figure I). At given reconstruction layer t, the augmentation model takes the measurement matrix E' and the dictionary I as input and generates the parameters {7e, 0.) for the current iteration. The architect ure for the augmentation network is illustrated in Figure 5) which shows a feature extraction section and two output branches, one for each generated parameter. To ensure that the estimated (1t, 0,) paramcters are positive, each branch equipped with a softplus function. As noted the main paper, the weights of the augmentation model are shared across all A-DLISTI layers.
19

---

Published in Transactions on Machine Learning Research (07/202) /204)
AdsptvolugPoo2D
DonvED
mntcc1,out.m1 km3, s72)
LaskyRetlU))
ConvzD (ftee1, (rt-e11,od_oo put_or1,_11 k-3, a=2)
LaakyRelUU)
ConvzD (ince11 out 6F1, k=3, s=2)
LoakyReLJ))
LeakyReLUU)
Lineer
(in [n_feat=122 fcat-25)
LeskyReLU0
Linear (n_ fut_feat-1)
Lineer (nate22, foat
Softplun)
Softplue
Figure 5: Augmentation model's architecture for A-DLISTA.
20

---

Published in Transactions on Machine Learning Research (07 (07/2024) 2024)
F2 k-3,
Tarki
Comvg)) (n ox32 mul 0404, and]
Convio
Conv10 (A, poet out pw10, L+784)
(n.cel] out 0=10, jo nb, meso)
CenviD
Tarny]
(8, 10.
iPLombb_mee ov, kr784)
c-64, 128, 1-3, s=2)
(8. 10, 11
(B, 100,
Tarih))
ConVTZo
OoMWTZD (nt.cn o0r,c7644 6:2)
cro4, ke7), en2)
Tomi)
Tanhi!
(1n30, Out-420)
Com/720
CoMV120
In c-044 Out p-52, I-) and)
(n. 5-64, P33 am3, 2*2)
Taml
Tanh!!
ConYT2D
ConVT2D
ConViZD om0)) 1- Imd,
Don/TZD 0nt.0330 olt o:2)
[n_or1,oi__o02 km? sn2)
[i_eet) out__r22, [nt_ert,ovi__002 kn1, a12)
Tashi)
Tanho
ConVT2D
kn5
/
Tanho)
Tanho
ConvTZD
Gonv
(h.cr1, WLF51, knS, 512) Taihy)
0n.ch1, oMt krG sx2)
Tanto
Cane12D
Cony
Figure 6: Left: prior network architecture. Right: posterior network architecture. For the posterior model, we show the output shape from each of the three input heads in the figure. Such a structure is necessary since the posterior model accepts three quantities as input: observations, the sensing matrix, and the reconstruction from the previous layer. Different shapes characterize these quantities. The letter "B" indicates batch size.
C.2 VLISTA
As described in (ubsectton 4.2 of the meain paper, VLISTA comprises three different. components: the likelihood, and the prior and posterior models.
C.2.1 VLISTA Likelihood model
The likelihood model (subsection 4.2) represents a Gaussian distribution with mean value parametrized using the A-DLISTA model. There is. However, a fundamental difference between the likelihood model and the A-DLISTA architecture is presented in subsection 4.1 Indeed, unlike the latter, the likelihood model of VLISTA does not learn the dictionary using backpropagation. Instead, it uses the dictionary sampled from the posterior distribution.
C.2.2 VLISTA Posterior & Prior models
We report Figure 6
the prior (left image) and the posterior (right image) architectures. We implement both models using an encoder-decoder scheme based on convolutional layers. The prior network comprises two convolutional layers followed by two separate branches dedicated to generating the mean and variance of (uubbbtttiisuustttttonno for the prior. In contrast to the prior, the posterior network accepts three different quantitics as input: the sensing matrix, the observations, and the reconstructed sparse vector from the previous ileration. To process the three inputs together, the posterior accounts for three separated "input" layers followed by an aggregation step. Subsoquently, two branches are used to generate the mean and the standard deviation of the Gaussian distribution of the dictionary (subsection 4.2.
the Gaussian distribution (subsection 4.2) We use the dictionary sampled at the previous iteration as input
We offer the reader unified overview of our variational model in Figure 7
21

---

Published in Transactions on Machine Learning Research (07/202) 2024)
Prior Model Pt (V; (-1)
6y'
Posterior Model 2x +!)
X--1
Likellhood Model
A-DLISTA 0 (o
T
Figure 7: VLISTA iterations' schematic view.
D
Training Details
This section provides details regarding the training of the A-DLISTA and VLISTA models. Using the Adam optimizer, we train the reconstruction and augmentation models for A-DLISTA jointly. We set the initial learning rate to 1.e-2 and 1.e-3 for the reconstruction and augmentation network, respectively, and we drop value by factor of 10 every time the loss stops improving. Additionally, we set the weight decay to 5.e and the batch size 128. We use Mean Squared Error (MSE) as the objective function for all datascts. We train all the components of VLISTA using the Adam optimizer, which similar to A-DLISTA. We the learning rate to 1.e-3 and drop its value by a factor of 10 every time the loss stops improving. Regarding the objective function, we maximize the ELBO and set the weight for KL divergence to I.e We report in Equation 33 the ELBO derivation.
logp(æ11 = T'rly', "lll',"") 4²) log
p( o Im, y" Φ d')p(( ( 1 W
(33)
p(x11T æ I. y )p 4 )a
B y x*)
log
d 1:T
1.7 miT: y¹ 4
( m
1.T. + )')pp ( (1:T)
4³) log
d WTT
1.r T. y¹ Φ0)
( 1:T E1.7 wi Φ') log p(æ1.
L 1:T. 3 & B')deir
pl 1 1
+
q( ((P17 (ILT. Φ²) log
d 1:T
g( 1:7 æI.T y Φi)
E.
B W log p(ar. sra D 1:t, y' ΦD))
Σ
E 1:1- 1 DKI. (1(VV x¹) (V,
DKL ( ((v1))
Note that in Equation 33
consider the same ground truth, ayy for each iteration / E [1.7].
22

---

Published in Transactions on Machine Learning Research (07/2024)
E
Computational Complexity
This section provides complexity analysis of the models utilized in our research. Table 3 displays the number of trainable parameters and average inference time for each model, while Table showcases the MACs count. To better understand the quantities appearing in Table 4 we have summarized their meaning in [Table 5) The average inference time was estimated by testing over 1000 batches containing 32 data points using a GeForce RTX 2080 Ti.
Trainable Parameters and Average Inference Time. To compute the values in Table considered the architectures used in the main corpus of the paper, e.g. same number of layers. From Table 3 it's worth noting that although ISTA appears to have the longest inference time, that can be attributed to the cost of computing the spectral norm of the matrix = L Such an operation, can consume up to 98% of the total inference time. Interestingly, neither NALISTA nor A-DLISTA require the computation of the spectral norm as they dynamically generate the step size. Additionally, LISTA does not require it all. NALISTA and A-DLISTA have comparable inference times due to the similarity of their operations, whereas LISTA is the fastest model, whilst VLISTA has higher average inference time given the use of the posterior model and the sampling procedure. Interestingly, LISTA and A-DLISTA have comparable nunber of trainable parameters, while NALISTA has significantly fewer. However, it's essential to emphasize that the number of trainable parameters depends on the problem setup, such as the number of measurements and atoms. We use the same experimental setup described the main paper, which includes 500 measurements, 1024 atoms, and three layers for each MI. model. As outlinod in the main paper, the likelihood model for VLISTA is similar in architecture to A-DLISTA, as reflected in the MACs count shown in Table 4 However, the likelihood model of VLISTA has different number of trainable parameters compared to A-DLISTA. Such dufference is due to VLISTA sampling its dictionary from the posterior rather than training it like A-DLISTA. Despite this difference, the time required for the likelihood model (shown in Table 3) is comparable to that of A-DLISTA. It's important to note that the inference time for the likelihood reported "per iteration", so we must multiply it by the number of layers A-DLISTA uses to make a fair comparison.
Macs count. Our attention now turns to the MACs count for the A-DLISTA augmentation network. As shown in Table the count is upper bounded by HWK² + BP. Note that the height and width of the input are halved after each convolutional layer, while the input and output channels are always one, and the kernel size equals three for each layer (see details in Figure 5). To obtain the upper bound for the
Table 3: Number of trainable parameters (Millions) and Average inference time (milliseconds) for different models. Concerning the inference time, we report the average value with its error considering 10 and 50 measurements setups.
Parameters (M)
Average Inference Time (ms)
mcas. 10
mcas. 500
ISTA
0.00
54.0 (norm: 41.540..)
(1.55 ) (norm.:
(1.53+0.02)e²)
NALISTA!
3.33e
5.8
7.040.3
LISTA
3.15
1.1
1.5
A-DLISTA²
3.15 (Aug.
NN: 3.11e
8.2
9.1
VLISTA
3.131
19.7
21.3
VLISTA
Prior Model
1.10
Posterior Model
2.08
3.6 3.6
1.6
4.1 1..2
Likelihood Model
1.05
2.7 10.2
3.05
LSTM hidden size equal to 256;
2 Each layer learns its own dictionary;
Full model at inference Prior
model NOT used;
Single iteration.
23

---

Published in Transactions on Machine Learning Research (07/2024)
MACs count, set H = max(H;) Himput/2 and W = max(W.) = Wrinput /2', where H; and W, are the height and width at the output of the i-th convolutional layer, respectively. With that in mind, we upper bounds the MACs count for the convolutional part of the network by HWK* The convonlutional backbone is followed by two linear layers (see details in Figure 5). The first linear layer takes a vector of size B E R R /16 wimpse/16 as imput and outputs vector of length P= 25. Finally, this vector is fed into two heads, each generating scalar. Therefore, the overall upper bound for the MACs count for the augmentation network is O(HW BP + P) O(H P(B + 1)) O(HW + BP), with the factor +1 dropped. Similar reasoning applics to the prior and posterior models of VLISTA, where we estimate the MACs count by multiplying the MACs for the most expensive layers by the total number of layers of the same type.
+
F Additional Results
In this section we report additional experimental results. In subsection F.1 we report results concerning a fix measurement setup, ie. &' + ④, while in subsection F.2 we show reconstructed images for different classical basclines.
Table 4: MACs count. Concerning VLISTA, we report the MACs count for each of its component: the Prior, the Posterior, and the Likelihood models.
MACs
ISTA
O(DM(2L + N) + D²M)
NALISTA
O(DM(2L + N) + [4h(d + h) + 12)t)
LISTA
O(LD(N + +) LMN)
A-DLISTA
O(LDMN + [HWK2 + BBIt)
VILISAA
Prior Model
O W Cr C C + +T2 OLL M DC pa Cp. + +Li 110Hp WollpoWpo22 To)) O(LDMN + [HWK2 + BPP))
Posterior Model
Likelihood Model
Contribution from the augmentation network.
Table 5: Description of quantities appearing Table 4
M N
Number of measurements
Dimensionality of dictionary's atoms
D L h; d
Number of atoms
Number of layers
Hidden and input size for the LSTM
H; W
Height and width of A-DLISTA augmentation network's input
Kernel size for the Conv layers A-DLISTA augmentation network
B; P
Input and output size of the lincar layer of A-DLISTA augmentation network Input and output channels for the "æ-input." head of the posterior model Posterior model bottleneck input and output sizes Height and width of posterior model's transposed convolutions input Kernel size of the posterior model's transposed convolutions
Cpoi C
I S
pos Wpoi
T po
Hori W pr
Input and output sizes of convolutional (and transposed conv.) layers of the prior model Input and output channels convolutional (and transposed conv.) layers of the prior model kernel size for convolutions and transpose convolutions the prior model
Ciri pr
Kpri Typ
24

---

Published in Transactions on Machine Learning Research (07 2024)
F.1 Fixed Sensing Matrix
We provide in Table 6 and Table 7
results considering & fixed measurement scenario, i.e. using a single sensing matrix Φ. Comparing these results to (Table 1 and Table 2] we notice the following. To begin with, LISTA and A-DLISTA perform better compared to the set up in which we use a varying sensing matrix (scctonn 5y. We should expect such behaviour given that we simplified the problem by fixing the & matrix. Additionally, as we mentioned in the main paper, ALISTA and NALISTA exhibit high performances (superior other models when 300 and 500 measurements are considered). Such a result expected, given that these two models were designed for solving inverse problems in fixed measurement scenario. Furthermore, the results in Table 6 and [Table 7 support our hypothesis that the convergence issues observe in the varying sensing matrix sctup are likely related to the "inner" optimization that ALISTA and NALISTA require to evaluate the "W" matrix.
Table 6: MNIST SSIM (the higher the better) for different number of measurements with fixed sensing matrix, i.e., T - $. We highlight in bold the best performance. Note that whenever there is agreement within the error for the best performances, we highlight all of them.
SSIM 1
number of measurements
10
100
300
500
(xe=1)
(xe=1)
(xe=1)
(xe=1)
(xe-1)
LISTA
1.34
3.1220.02
5.98
6.7440.01
6.9640.01
ALISTA
0.84+0.01
0.94
1.70+0.01
5.71
6.65
NALISTA
0.91
1.124u.01
2.46.0.01
7.03
8.22:0.02
A-DLISTA (our)
1..02
3.58_..01
5.6640.01
6.47
6.84
Table 7: CIFAR10 SSIM (the higher the better) for different number of measurements with fixed sensing
matrix, i.., o' + Φ. Note that whenever there is agreement within the error for the best performances, highlight all of them.
SSIM 1
number of measurements
100
300
500
(xe-1)
(xe-1)
(xe-1)
(xe-1)
(xe-1)
LISTA
2.5210.01
3.19.0.01
4.48
6.29
6.74
ALISTA
0.21+0.03
0.54+0.02
0.88+0.01
3.54
5.52+0.01
NALISTA
1.32
1.3240.02
1.06
4.59
6.88
A-DLISTA (our)
2.91
3.07
4.26++001
5.89
6.56
F.2 Classical baselines
We report additional results concerning classical dictionary learning methods tested on the MNIST and CIFARIO datasets. It is worth noting that classical baselines can reconstruct images with high quality if it is assumed that there are neither computational nor time constraints (although this would correspond to an unrealistic scenario concerning real-world applications). Therefore, while tuning hyperparameters, we consider a number of iterations up to a several thousand.
Figure 8 to Figure 13 showcase examples of reconstructed images for different baselines.
25

---

Published in Transactions on Machine Learning Research (07 (07/2024)
2
3
8
7
?
2
3
8
7
7
Figure 8: Example of reconstructed MNIST images using the canonical basis. Top row: reconstructed
images. Bottom row: ground truth images. To reconstruct images we use 500 measurements and the number of layers optimized to get the best reconstruction possible.
8
8
7
i
7
Figure 9: Example of reconstructed MNIST images using the wavelet basis. Top row: reconstructed images.
Bottom row: ground truth images. To reconstruct images we use 500 measurements and the number of layers optimized to get the best reconstruction possible.
26

---

Published in Transactions on Machine Learning Research (07 (07/2024)
4
3
3
4
3
3
Figure 10: Example of reconstructed MNIST images using SPCA. Top row: reconstructed images. Bottom row: ground truth images. To reconstruct images we use 500 measurements and the number of layers optimized to get the best reconstruction possible.
Figure 11: Example of reconstructed CIFAR10 images using the canonical basis. Top row: reconstructed
images. Bottom row: ground truth images. To reconstruct images use 500 measurements and the nunber of layers optimized to get the best reconstruction possible.
27

---

Published in Transactions on Machine Learning Research (07 (07/202) /204)
Figure 12: Example reconstructed CIFAR1O images using the wavelet basis. Top row: reconstructed
images. Bottom row: ground truth images. To reconstruct images we use 500 measurements and the number of layers optimized to get the best reconstruction possible.
Example of reconstructed CIFAR10 images using SPCA. Top row: reconstructed images. Bottom row: ground truth images. To reconstruct images we use 500 measurements and the number of layers optimized to get the best reconstruction possible.
Figure 13:
28


Reviews:
[{"review_id": "rUkq3tIldT", "reviewer": "Reviewer_8VsP", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: The proposed methods have certain novelty with ample compressive sensing applications, and numerical experiments seem to be convincing. \nWeaknesses: The assumptions for the prior and posterior models could be further justified for real applications. Their extensions to other probability distributions could also be illustrated in more detail.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes two major algorithms: (1) Augmented Dictionary Learning ISTA (A-DLISTA), which uses an augmentation\nmodule to adaptively learn the parameters according to the sensing matrix; (2) Variational Learning ISTA (VLISTA) which can learn dictionary distributions in the variational framework for compressive sensing recovery with varying sensing matrices. Numerical experiments on the three data sets have illustrated the proposed effectiveness and show that the proposed dictionary distribution learning can be applied successfully for out-of-distribution sample detection.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Substantial language editing and proofreading should be done. \n2. Theoretical and numerical justifications of prior/posterior models could be provided.", "broader_impact_concerns": "The paper lacks broader impact discussions or statements, which could be further added.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "O8MLqbQnqq", "reviewer": "Reviewer_YHuk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "# Strengths:\n(1) The introduction and related works sections are pretty clear and well written. My only suggestion there is to add the cost function sooner and give more motivation and intuition to what you mean by sample-adaptive algorithm.\n(2) I think the A-DLISTA is a good contribution. But it is not sufficiently explained (see Weakness#2)\n\n# Weaknesses\n(1) The results section is based off of pretty unrealistic toy data. The VLISTA seems like a lot of complicated machinery with a staggeringly huge number of parameters. I would never use this without a huge analysis of how prone it is to overfitting and how it performs encoding noisy /realistic data with non-Gaussian sensing matrices. Random zeros are not possible to implement in physical systems (to my knowledge) so it's unclear how well this can translate to e.g. MRI data. I'm also having a hard time understanding the presentation of OOD results (but that may be my fault).\n(2) You need to spell out better what an inference looks like for ADLISTA. I imagine that first you train ADLISTA to predict ground truth (which does not exist in sparse coding...). Then, during inference, you also perform something that must still look a bit like training, i.e. joint optimization of codes and matrix coefficients. See the reference at the bottom of Requested Changes as this represents my understanding of your algorithm (something that is mixing forward-inference of the [formerly] convex algorithm with backpropagation updates of your matrices). You could hugely benefit by an enumerated algorithm box that really spells it out for us. \n\nMost critically of all: What is the cost function for training your sample-specific $\\Phi$? How many iterations of code optimization (T)? How many updates applied to the matrix coefficients in between code inferences? These details aren't addressed at all unless I missed something. These are the core reasons I have said you haven't backed up your claims but please also see Necessary Changes in the Requested Changes section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose two extensions of learnable iterative soft thresholding algorithm (ISTA): (1) Augmented Dictionary Learning ISTA (A-DLISTA) that is designed to adapt with the current sample being processed, and (2) Variational Learning ISTA (VLISTA), which implements a Bayesian framework for learning a dictionary distribution that leverages A-DLISTA as the corresponding likelihood model.\n\nMy understanding of A-DLISTA in the context of its purpose is as follows. Normally, (F)ISTA, a well established convex optimization algorithm family, is applied on a sample-to-sample basis with a provided/fixed sensing and sparsification representation (everything is sample-specific, and nothing is learned/ no assumptions are made about any other samples). LISTA is the same algorithm, but with learned parameters that are intended to encode a dataset faster than ISTA. LISTA is a truncated ISTA trained to predict a fully converged ISTA that can readily be applied across a dataset. LISTA is a recurrent neural network with shared layers.\n\nDLISTA as presented in Eqn (2) breaks the recurrence by allowing the learned matrix to be optimized for each iteration (subscript t). The learned matrix is however decomposed into two multiplicative terms via the signal model $y=\\Phi\\Psi x^*$. At first --because $\\Phi\\Psi$ always occur together-- it appears that it doesn't matter whether the $\\Phi$ term is fixed since $\\Psi$ is allowed to change every iteration, however I believe the fixed sensing matrix helps with convergence analysis when $\\Phi$ has convenient properties.\n\nThe proposed ADLISTA has the exact same operations, but removes the remaining recurrent structure (ie it is now a fully feedforward neural network) by letting $\\Phi$ also be updated by error backpropagation across samples and iterations. So the signal model has sample-specific sequence of matrices $\\Phi_t$ and dataset-specific sequence of matrices $\\Psi_t$. This is the meaning of the authors' term \"augmented network\". and \"adaptivity to samples\". This is central to the paper's thesis of its contributions.\n\nVLISTA uses a variational bayesian framework to predict the dataset-specific sparsification dictionary $\\Psi_t$, given $\\Phi_t, x_{t-1}, y_t$, and implicitly given $\\Psi_{t-1}$ (via the training procedure).\n\nEmpirical results are composed from synthetic compressed sensing scenario. MNIST, CIFAR10, and a (seemingly?) unspecified synthetic dataset are sampled by optimal compressed sensing conditions (Gaussian matrices).", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "# (a) Necessary Suggestions\n1) Show the cost function you wish to optimize before presenting the recursive solution in Eqn (1). I suggest putting this on the first page because the notions of sensing matrix and sparsification dictionary/transform can easily be conflated depending on the specific field a reader is coming from. For a while I thought you were considering the case with separated synthesis/sensing and analysis/sparsifying terms in the cost function.\n2) formatting:  (A) cite the (F)ista paper at Eqn(1) instead of just stating ISTA; (B) some of your quotations marks are imbalanced, I can't make it look right in MarkDown but e.g. after Eqn(1) they are curly on the left and straight on the right. (C) From Eqn 1-2 you add a subscript to sparse dictionary atoms $\\Psi\\rightarrow\\Psi_t$ but also change the Hermetian-transpose term to regular transpose. This change should be either remedied or explained in the paper. (D) The ADLISTA and VLISTA lines look identical in Fig 3, you need to use some differentiating factor (markers, dashes, or one could be black).\n3) I don't think BCS is defined?... Also, you don't explain what $x^*$ are for the synthetic case. What are these, just noise vectors? Sinusoids?\n4) It's suspicious that LISTA gets better performance than ISTA, since LISTA is trained to predict converged ISTA. It shold only get the same result, faster. My guess is that you are not allowing ISTA to properly converge. The only cost functions in the paper are for VLISTA so I might just be confused. Speaking of which, I was expecting a discussion for T (number of iterations). How many are required for your method to function?\n5) I apologize, but I don't really understand Figure 4 and the OOD experiments in general. Are you computing 100 p values for each noise level and each every test image and then averaging across all images and trials? I feel like a ROC curve would be more useful to understand how your algorithms fare across the test set.\n6) I am having trouble understanding your generation of data. You describe picking a sample, then generating a Gaussian $\\Phi_i$. Then you randomly set ~10% of $\\Phi$'s entries to zero. Then your plots and tables have a \"number of measurements\" parameter, so here I believe you are keeping the first N rows. But what is with the 10% zeroes-- that's not enough to be considered sparse, right? 90% dense?\n7) If you have analytic results that are in the appendix, you really should put the statement of the theorems in the paper at least.\n\n# (b) Medium Suggestions\n1) Motivation: You may wish to give more specific examples or higher level motivation in the introduction, i.e. to clarify what problem you are trying to solve. The purpose of a dictionary, traditionally, is to provide common representation across many samples in a dataset. In compressive sensing the role of the \"dictionary\" is somewhat different in that it usually is a system model or, as you put it, the sensing matrix. By giving a specific example (such as how multi-channel MRI machines have instance-specific calibration matrices as part of the system model) it could help clarify what scenario you would need a linear operator to play the role of a dictionary for a specific sample instance. [later during review]: I see you have an example on page 4; this should be in the introduction IMO. I.e. using examples to explain the problem you want to solve, not just to say what LISTA cannot do.\n2) The statement \"...[in LISTA] those weights [V, W] implicitly contain information about Φ and Ψ that are assumed to be known and fixed.\" from Sec 3.2. This is not really accuracy because they are updated during learning, I think you're talking about ISTA? Or are you referring to the fact that LISTA is trained to predict ISTA? But I don't think you did that in your experiments.\n3) I think it would be really helpful to put VLISTA in the same diagram as ADLISTA. I am having a hard time understanding what a forward inference looks like, exactly, since forward inference includes backpropagation updates of $\\Phi$ but also a new set of neural networks outputing $Psi$. I'm not 100% convinced this should be considered a \"LISTA family\" algorithm due to the latter operations.\n4) Compressed sensing on MNIST in 2024 seems inappropriate. There's no written argument that can compensate for lack of realistic experimental results.\n\n# (c) Light suggestions\n1) I am satisfied with the Related Works section but [1] is another unrolled approach with sample adaptivity, but with an ADMM-family algorithm. Please consider putting an algorithm enumeration so spell out the steps. If [1] is unrelated then I do not understand your ADLISTA algorithm.\n2) I would be curious to see the wallclock time during inference. \n\n\n[1] Goehle, Geoff, and Benjamin Cowen. \"Joint Sparse Coding and Frame Optimization.\" IEEE Signal Processing Letters (2024).", "broader_impact_concerns": "I don't see any ethical concerns .", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "lALR9dYeRA", "reviewer": "Reviewer_3H4v", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Novel Approach: The paper introduces novel methods, A-DLISTA and VLISTA, to address the issues of dictionary learning and uncertainty modeling in compressed sensing, which contribute to the existing literature in this field.\n\n2. Theoretical and Experimental Support: The authors provide both theoretical and experimental evidence to support their proposed methods, demonstrating the effectiveness and feasibility of their approaches.\n\n3. Calibrated Uncertainties: The paper highlights the ability of the proposed model to learn calibrated uncertainties, which is a valuable aspect in real-world applications where understanding model confidence is crucial.\n\nWeaknesses:\n1. Lack of Practical Application Context: The paper lacks concrete examples or practical application scenarios where the proposed methods could be applied. Providing such context would enhance the understanding of the practical usefulness of these techniques.\n\n2. Limited Discussion on Computational Complexity: The paper does not extensively discuss the computational complexity of the proposed methods. Understanding the computational requirements and scalability is important for practical implementation.\n\n3. Generalization to Wider Machine Learning Field: The paper could benefit from a discussion on how the proposed methods may generalize and contribute to the broader machine learning field beyond compressed sensing.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The author proposed a variant of the LISTA architecture. Firstly, they introduced Augmented Dictionary Learning ISTA (A-DLISTA), which includes an augmentation module to adapt to the current measurement setup. Then, they proposed a method for learning the dictionary distribution through a variational approach, named Variational Learning ISTA (VLISTA). VLISTA utilizes A-DLISTA as a likelihood model and approximates the posterior distribution of dictionaries as part of an unfolded LISTA-based recovery algorithm. Consequently, VLISTA offers a probabilistic way to jointly learn the dictionary distribution and the reconstruction algorithm under varying sensing matrix conditions. The authors provided both theoretical and experimental support, demonstrating that their model learns calibrated uncertainties.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to the weaknesses.", "broader_impact_concerns": "NA.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The problem studied in this papers is a classic problem in signal processing: joint estimation of a signal and its sparsifying basis. That the authors also allow for changing sensing matrix in time is a nice twist adding a layer of complexity to the problem, with potential value to practitioners. The theoretical guarantees in a teacher-student set up are strengthening the  soundeness of the approches. 
Many of the main reserves by reviewers were easily answered (lack of discussion on the applicability and real-case uses, on the algorithmic complexity, or, mainly, on presentation issues). The discussion has been fruitful and the revised version of the paper seem to have improved clarity. One reviewer still recommended rejection. However, the two others seem pleased by the latest paper version. My own quick reading indicates paper of sufficient quality, both in originality and clarity of exposition, to deem publication in TMLR. Also because this problem is anything but niche and may trigger the interest of part of TMLR' audience.
My main suggestion to the authors before publication would be to add a link to a working code able to reproduce the experiments.
Decision: Accept (poster)