Learning Structured Sparse Matrices for Signal Recovery via Unrolled Optimization
Abstract
Countless signal processing applications include the reconstruction of an unknown signal from very few indirect linear measurements.
Because the measurement operator is commonly constrained by the hardware or the physics of the observation process, ﬁnding measurement matrices that enable accurate signal recovery poses a challenging discrete optimization task.
Meanwhile, recent advances in the ﬁeld of machine learning have highlighted the effectiveness of gradient-based optimization methods applied to large computational graphs such as those arising naturally when unrolling iterative algorithms for signal recovery.
However, it has remained unclear how to leverage this technique when the set of admissible measurement matrices is both discrete and sparse.
In this paper, we tackle this problem and propose an efﬁcient and ﬂexible method for learning structured sparse measurement matrices.
Our approach uses unrolled optimization in conjunction with Gumbel reparametrizations.
We empirically demonstrate the effectiveness of our method in two prototypical compressed sensing situations.
Introduction
Linear measurement operators following a structure that is constrained by the physics of the measurement process or the hardware of the measuring device are ubiquitous in signal processing.
In particular, structured sparse matrices appear in many signal processing applications [20, 31, 9, 25].
The optimal design of such measurement operators—which reside in a discrete subset—to improve the performance of downstream tasks poses great computational challenges.
Classical approaches commonly employ discrete optimization, as no gradients can be directly computed.
On the other hand, in the ﬁeld of machine learning, gradient-based optimization via backpropagation through massive nonlinear computational graphs has shown impressive performance in many applications.
A promising concept to enable gradient-based learning in non-differentiable settings is given by Gumbel reparametrizations [16, 21], which allow estimating the gradients of categorical random variables.
In this work, we employ this technique to fuse gradient-based learning with the design of sparse measurement operators that are constrained to a discrete set.
Signal recovery problems are often solved by convex optimization methods.
It is well established that the computational graph of many iterative convex optimization schemes can be unrolled to obtain a neural network that can be readily backpropagated through [11], enabling the computation of gradients with respect to the measurement operator and other involved parameters.
However, it has remained unclear how this technique can be leveraged when the set of admissible measurement operators is sparse and discrete.
We present a novel approach to tackle this problem.
Our main contributions can be summarized as follows
: • We propose an efﬁcient and easy-to-implement method for learning structured sparse measurement matrices for signal recovery by using unrolled optimization and Gumbel reparametrizations.
Due to its compatibility with automatic differentiation, our method provides a novel framework for the data-driven design of practicable measurement operators. • We successfully apply our algorithmic approach in two prototypical situations, namely compressed sensing with left- d -regular graphs and single-pixel imaging, outperforming conventional setups based on randomization.
Background & Related Work 2.1 Linear Inverse Problems & Compressed Sensing
In linear inverse problems, the basic task is to recover an unknown target signal x ∈ R n from indirect observations of the form y =
Φ x + e , where Φ ∈ R m × n is a known measurement matrix and e ∈ R m noise.
The number of measurements m is usually much smaller than the signal dimension n , making the inverse problem ill-posed and only solvable under prior knowledge about the underlying signals.
A prominent example is given by compressed sensing, in which the signal is assumed to be sparse [6, 8].
In particular, when x is s -sparse, it is possible to robustly reconstruct x from m ∈
O ( s log( n/s )) measurements via convex optimization, given that Φ fulﬁlls the null space property [10].
The convex program that is to be solved corresponds to a LASSO problem [27] with hyperparameter λ :
[Equation]
Many randomized measurement operators fulﬁll the null space property with high probability [10], allowing theoretical guarantees on the reconstruction of generic sparse signals.
However, for realworld applications, randomization is not a necessary prerequisite.
A widely-used class of algorithms for solving LASSO-type problems in compressed sensing are gradient-based methods.
For example, an iterative proximal scheme, known as the Iterative Soft
Thresholding Algorithm (ISTA) [7], can be used to solve (1):
[Equation]
where prox λ (cid:107)·(cid:107) 1 ( v ) =
sign ( v ) max( | v | − λ, 0) .
Another popular method is based on Iterative Hard
Thresholding (IHT) [4], in which the proximal operator in is replaced by
2.2 Unrolled Optimization in Linear Inverse Problems
Many iterative optimization schemes can be viewed as a neural network via unrolling the computational graph of a ﬁnite number of iterations T .
Once unrolled, the parameters θ of a given recovery algorithm f θ : R m → R n as well as the measurement operator Φ can be ﬁt to a (training) dataset by minimizing a loss function L : R n × R n → R via (stochastic) gradient-based optimization:
[Equation]
In the context of linear inverse problems, a prominent example of unrolled optimization is Learned
ISTA (LISTA) [11], where ISTA is unrolled and the involved matrices, step-sizes, and thresholds are learned in an end-to-end fashion.
Such a data-driven approach can reduce the required number of iterations to solve recovery problems by orders of magnitude.
Further improvements to LISTA are made by Analytic LISTA [19], and Neurally Augemented LISTA (NA-ALISTA) [2].
Unrolled optimization provides a natural approach to learning data-driven Φ by including it in an end-to-end training procedure.
For unconstrained, dense Φ this has been explored in [30, 1].
However, in real-world applications, the measurement matrix must often follow constraints imposed by hardware or the underlying physics, limiting the practicability of such approaches.
2.3 Gumbel Reparametrizations
Gradient-based learning has been shown to be effective and scalable [26, 5].
However, when discrete nodes are included in the computational graphs, gradients cannot be computed directly and have to be estimated.
Formally, we consider a computational graph including a discrete random variable v taking one of the values 1 , . .
, a (in one-hot encoding), where its unnormalized log-probabilities are denoted by ϕ =
[ ϕ 1 , . .
ϕ a ] T ∈ R a .
The value of v is then passed through a deterministic, differentiable function f .
As v is discrete, it is not possible to directly backpropagate through v .
Therefore, an estimate of the gradient ∇ ϕ E v [ f ( v )] must be computed.
The Gumbel-softmax trick [16, 21] allows computing a differentiable relaxation of v by adding component-wise i.i.d.
Gumbel noise to ϕ before applying a softmax.
This modiﬁcation of the Gumbel-max trick [12, 22] enables the backpropagation through discrete random variables outperforming previous approaches to estimating the gradient of discrete nodes [3, 29, 23].
In cases where true discreteness is needed, it is possible to use the argmax operator in the forward pass and the softmax in the backward pass of backpropagation.
This is known as the straight-through Gumbel softmax estimator.
It can be directly extended to taking multiple samples without replacement from a categorical distribution, namely by selecting the top- K values instead of only the largest one [28].
Method
We propose a method that enables gradient-based learning of structured sparse matrices for signal recovery by using unrolled optimization in conjunction with Gumbel reparametrizations.
To this end, we model the measurement matrix Φ ∈ R m × n as being governed by a learnable parameter ϕ ∈ R m × n , employing Gumbel top-
k operators on a partition of the indices of ϕ .
Formally, consider the index set I := { 1 . . .
m } × { 1 . . .
n } .
Let P ( I ) denote a partition of I , i.e., a set of non-empty pairwise disjoint subsets I 1 , . .
, I l ⊂ I such that (cid:83) i ∈{ 1 ...l } I
i =
I .
We introduce the notation ϕ
[ I i ] for indexing the elements of ϕ in the index sets I i .
Then, for each i ∈ { 1 . . .
l } , we use a Gumbel reparametrization to select d i elements from the i -th set in P ( I ) .
This is done by adding element-wise i.i.d.
Gumbel noise to each element of ϕ
[ I i ] , and then taking the top- d i values in the forward pass.
In the backward-pass of the backpropagation, the gradient of the softmax is used instead of the hard top- d i .
A pseudo-code implementation of this procedure, assuming a software framework capable of automatic differentiation, is provided in Algorithm 1.
The binary matrix resulting from Algorithm 1 can be used like any other matrix in a framework with automatic differentiation, as the gradient with respect to ϕ is well-deﬁned.
This means that Φ can be constructed from this binary matrix by combination with other constants or learned parameters.
In that way, our method can be used in an out-of-the-box fashion to learn a masking pattern for ﬁxed matrices such as Fourier matrices.
A simpler special case is obtained by choosing the trivial partition, where all entries of Φ are governed by a single Gumbel reparametrization; in other words, a single top- K is used to select elements of the entire matrix.
This can be readily applied whenever learning a single sampling mask is required, for example, in the context of compressive magnetic resonance imaging [15, 14], where a ﬁxed number of rows from a Fourier matrix is selected.
However, the ﬂexibility of our method comes from the fact that the partition can be freely chosen to take any ﬁne-grained structure of the matrices into account.
In many applications, such as single-pixel imaging or compressed sensing with left- d -regular graphs, the partition is naturally given by the rows or columns of ϕ , while the number of elements d to be selected from each partition component is equal.
In this case, vectorized, and therefore computationally efﬁcient softmax top- K operations over an axis can be applied.
In principle, our approach enables learning any type of measurement operator that can be constructed by differentiable transforms of a sparse binary matrix (which is structured in the sense that entries are subselected from a speciﬁc partition).
With this in mind, Algorithm 1 should be seen as an instance of a more general framework for the data-driven design of measurement operators.
A full evaluation of this methodology and potential generalizations of Algorithm 1 are beyond the scope of this paper. Algorithm 1 Learning a
Input Learnable Parameters: Parameters of measurement ϕ
, parameters of reconstruction algorithm θ
Application: Single Pixel Imaging
A remarkable application of compressed sensing is single pixel imaging [9], which makes imaging possible using compressive measurements acquired by a spatial light modulator or a digital micromirror device that collects light onto a single pixel using a series of m masks.
Mathematically, this process can be modeled as a compressed sensing problem using a binary measurement matrix Φ .
While theoretical results highlight that random masks have favorable reconstruction properties for generic sparse signals, this may not be the case when the physics of the measurement process or additional image structures are taken into account.
Our approach is well suited to incorporate such aspects.
For our case study [Figure] we partition the indices of ϕ into the set of row vectors and use Algorithm 1 to learn adaptive masks for the MNIST dataset ( n = 784 ) [18].
We unroll IHT, which has no additional learnable parameters beyond the measurement operator [Figure] and NA-ALISTA, which adaptively predicts step-sizes and thresholds by an LSTM-network [13].
The results in Figure 1 show that our method reduces the number of measurements required for good reconstruction and greatly speeds up convergence; see Appendix for implementation details.
Figure 2 visualizes learned and random masks as well as some reconstructed
MNIST digits: the learned masks have clear structure and lead to better reconstruction.
Additional samples and masks can be found in the Appendix
[Figure] [Figure]
Application: Compressed Sensing with Left-d-Regular Graphs
Another class of structured random matrices suitable for compressed sensing is given by adjacency matrices left- d -regular graphs.
In particular, for large n , adjacency matrices of random left- d -regular graphs with d ∈ O ( n/s ) and m ∈ O ( s log( n/s )) allow stable signal recovery [10].
However, for smaller values of m and n , such random graphs are unlikely to enjoy the required properties and therefore may not be usable for compressed sensing.
This regime is again well suited for our approach.
Invoking Algorithm 1, we partition the entries of the measurement matrix into its row vectors, and select d = 7 ones per row.
As unrolled method, we consider Iterative Hard Thresholding for expanders (E-IHT):
[Equation]
Here, H s denotes the hard thresholding operator and M : R m → R n the median operator, taking the median of the d values connected to each left vertex by the graph Φ .
Similarly to E-IHT, we obtain E-NA-ALISTA by replacing the adjoint Φ T in each iteration by the median operator M .
The experimental results in Figure 3 demonstrate that our approach signiﬁcantly improves signal recovery of generic sparse vectors as well as the convergence speed; see Appendix for implementation details.
[Figure]
A Implementation Details
In our experimental setup, 1 we optimize Φ and θ in Eq.
(3) using the Adam optimizer [17] with β 1 = 0 .
9 , β 2 = 0 .
999 (default in PyTorch [24]) and a ﬁxed learning rate of 0.0002.
We use minibatches of size 512 training samples.
The parameters ϕ are initialized using a standard Gumbel distribution.
We ﬁnd that when learning both Φ and θ , training is more stable when rescaling the
Gumbel noise by a factor of 0.01, both for initialization as well as for the Gumbel reparametrizations during training.
We keep the softmax temperature in all Gumbel reparametrizations ﬁxed at τ = 1 , which has been demonstrated to work well in practice [16].
A.1 Single Pixel Imaging
As the MNIST digits are not very sparse in the image domain, we use bi-orthogonal 2.2 wavelets with one level as sparsifying transform.
For this, we denote the wavelet transform as Ψ and reparametrize x = Ψˆ x , such that the LASSO problem in Eq.
(1) yields the synthesis formulation:
[Equation]
[Equation]
The target loss function during training in Eq.
(3) is set to be the squared loss L (ˆ x, x ) = (cid:107) ˆ x − x (cid:107) 22 .
Additive i.i.d.
Gaussian noise with a signal-to-noise ratio of 40dB is added to the measurement vector y .
As a metric, we use the normalized mean squared error (NMSE), deﬁned as:
[Equation]
A.2 Compressed Sensing with Left-d-Regular Graphs
In the experiments on compressed sensing with left- d -regular graphs, the performance is evaluated on synthetic generic sparse vectors.
The support of the synthetic data is generated using i.i.d.
Bernoulli random variables, and the non-zero components are drawn from a normal distribution.
We use
L ( x, ˆ x ) = (cid:107) x − y (cid:107) 1 as the loss function, as the mean absolute error is the usual choice in expander theory.
The experiments in this setup are conducted with heavy tailed noise (student t-distributed with 1 degree of freedom) and a signal-to-noise ratio of 40dB. As a metric, we employ the normalized mean absolute error (NMAE), deﬁned as:
[Equation]
B Computational Cost
In Algorithm 1, ϕ has m · n learnable parameters, which must be stored in memory during training time.
This is identical to the cost of learning a dense sensing matrix without constraints.
At test time, a single mask is sampled after adding Gumbel noise and performing the top- K , and then kept ﬁxed.
This means that there is no additional computational cost compared to using a conventional (randomly chosen, but ﬁxed) matrix satisfying the constraints.
Hence, our method is feasible for training and comes at no additional expenses during testing or deployment.
C Supplementary Experimental Results
[Figure]
[Figure] [Figure]