Source-linked AI summary
Trainable ISTA for Sparse Signal Recovery
Daisuke Ito, Satoshi Takabe, Tadashi Wadayama
TL;DR
Sparse signal recovery must estimate sparse signals from noisy, undersampled measurements across varied sensing matrices. TISTA unfolds ISTA with MMSE shrinkage, estimated error variance, and trainable parameters. Experiments report applicability across sensing-matrix classes and faster convergence than AMP and LISTA in many cases.
Problem
Sparse recovery methods need to estimate sparse signals under varied sensing matrices, while AMP’s performance is limited when its Gaussian-matrix assumptions are violated.
Method
TISTA combines a linear estimator, MMSE-based shrinkage, an estimated error variance, and trainable variables controlling the step size and shrinkage variance.
Results
TISTA applies to Gaussian, binary, and large-condition-number sensing matrices and converges significantly faster than AMP or LISTA in many cases.
Takeaways & Limitations
TISTA uses a number of trainable variables equal to the iteration rounds, supporting highly stable and fast training processes.
Abstract
from arXiv · showhide
In this paper, we propose a novel sparse signal recovery algorithm called Trainable ISTA (TISTA). The proposed algorithm consists of two estimation units such as a linear estimation unit and a minimum mean squared error (MMSE) estimator-based shrinkage unit. The estimated error variance required in the MMSE shrinkage unit is precisely estimated from a tentative estimate of the original signal. The remarkable feature of the proposed scheme is that TISTA includes adjustable variables controlling a step size and the error variance for the MMSE shrinkage. The variables are adjusted by standard deep learning techniques. The number of trainable variables of TISTA is equal to the number of iteration rounds and it is much smaller than that of known learnable sparse signal recovery algorithms. This feature leads to highly stable and fast training processes of TISTA. Computer experiments show that TISTA is applicable to various classes of sensing matrices such as Gaussian matrices, binary matrices and matrices with large condition numbers. Numerical results also demonstrate that TISTA shows significantly faster convergence than those of AMP and LISTA in many cases.
I. INTRODUCTION
Compressed sensing estimates a sparse signal from undersampled noisy measurements, with ISTA providing a standard iterative linear-estimation and soft-shrinkage baseline. TISTA extends this structure with trainable variables and an MMSE-based shrinkage design.
- Problem setup: Compressed sensing estimates sparse x from noisy undersampled observations y = Ax + w, where A has M < N rows.The recovery goal is to estimate x as accurately as possible from y.
- Lasso and ISTA: Lasso promotes sparse reconstructions by minimizing a quadratic data-fit term plus λ||x||1.The regularization constant λ controls the sparsity-promoting penalty.
- Lasso and ISTA: ISTA solves the Lasso problem through repeated linear estimation followed by soft-thresholding shrinkage.ISTA is also interpretable as proximal gradient descent.
- Related algorithms: AMP generally converges faster than ISTA, but its derivation assumes i.i.d. Gaussian sensing-matrix components.Related methods such as OAMP and VAMP broaden the sensing-matrix settings they address.
- Learnable recovery: Unfolding iterative sparse-recovery algorithms produces signal-flow graphs that can be equipped with learnable parameters.This deep-learning perspective motivates trainable algorithmic variants.
- TISTA: TISTA combines ISTA’s structure with an MMSE shrinkage function, an error-variance estimator, and trainable linear-estimator variables.Its trainable variables are adjusted using deep-learning techniques, while the error-variance estimator is adopted from OAMP.
B. AMP
AMP uses an Onsager correction and threshold estimation to obtain state-evolution tracking and fast convergence under Gaussian sensing assumptions. The section also introduces MMSE shrinkage functions and contrasts AMP’s matrix restrictions with OAMP’s broader design.
- AMP: AMP uses an Onsager correction so linear-estimator outputs become statistically decoupled in large systems.The resulting errors behave as white Gaussian variables, enabling state-evolution analysis of error variances.
- AMP: AMP’s threshold τ_t is estimated for its shrinkage function, and AMP was reported to converge much faster than ISTA under its sensing-matrix condition.The threshold estimator is a key distinction between AMP and ISTA.
- AMP limitations: AMP may not provide excellent recovery for non-Gaussian, nonzero-mean, high-variance, or large-condition-number sensing matrices.Its favorable behavior is tied to the stated Gaussian sensing assumptions.
- OAMP: OAMP uses a divergence-free nonlinear estimator and permits transpose, pseudo-inverse, or LMMSE choices for its linear-estimator matrix W.For i.i.d. Gaussian or unitary-invariant sensing matrices, the two estimator errors are proved statistically orthogonal.
- MMSE estimator: MMSE shrinkage is derived for additive Gaussian channels and Bernoulli-Gaussian or finite-discrete sparse priors, then used as a TISTA building block.For the Bernoulli-Gaussian prior, the nonzero component is Gaussian and p denotes its occurrence probability.
- MMSE estimator: Fig. 1 shows η_MMSE as a function of received signal y for σ^2 = 0.2 and 0.8, with α^2 = 1 and p = 0.1.The resulting curves resemble soft thresholding while remaining differentiable everywhere with respect to y.
B. Recursive formula for TISTA
TISTA combines a linear estimator, MMSE-based shrinkage, and an error-variance estimator under a full-rank sensing-matrix assumption. Learnable step sizes control the linear update and shrinkage variance, while propositions justify variance estimates under residual-error assumptions that are later acknowledged as unverifiable.
- Recursive structure: TISTA assumes a full-rank sensing matrix and uses the pseudo-inverse W = A^T(AA^T)^−1 in its recursive update.The initial state is s0 = 0, and the final estimate is ŝ = sT.
- Learnable variables: TISTA's learnable scalars γt provide step sizes and control the variance used by the MMSE shrinkage.Parameters p and α can also be optimized, adding two trainable parameters when used.
- Variance estimation: The MMSE shrinkage requires an accurate estimate of the residual error variance, derived under Gaussian and independence assumptions on residual errors.The assumptions motivate viewing the linear-estimator output as a virtual Gaussian-noise observation.
- Variance estimation: Propositions establish variance-estimation identities when residual errors have i.i.d. zero-mean entries and specified cross terms vanish.The derivation rewrites the residual as (I − γtZ)(st − x) + γtWw and uses trace identities for full-rank matrices.
- Scope and limitations: TISTA lacks OAMP's Onsager correction or divergence-free function, so stochastic orthogonality and state-evolution analysis are not guaranteed.The validity of the residual-error assumptions is instead to be checked experimentally.
C. Time complexity and number of trainable variables
TISTA has quadratic per-iteration complexity and a small trainable-variable count that does not depend on system size. Its computation is dominated by matrix-vector products, while the pseudo-inverse can be precomputed once.
- Time complexity: O(N^2) time is required per TISTA iteration, matching the stated per-iteration complexity of ISTA and AMP.Matrix-vector products and trace evaluations dominate the iteration cost.
- Time complexity: O(N^3) pseudo-inverse computation is performed only once in advance rather than at every iteration.The pseudo-inverse is used in the recursive computation after precomputation.
- Trainable variables: TISTA uses T trainable variables, or T + 2 when α and p are also optimized, for T iteration rounds.Its trainable-variable count is independent of N and M.
- Trainable variables: TISTA requires fewer trainable variables than LISTA and LAMP in the T-round process.The cited comparison gives LISTA and LAMP larger per-round parameter counts and identifies TISTA as the least-parameter method.
D. Incremental training for TISTA
TISTA is trained by unfolding its iterations into a feedforward signal-flow graph and optimizing its variables incrementally. Each generation adds one layer, updates the objective, and retains previous variables as initialization.
- Training procedure: Unfolding TISTA's recursion produces a signal-flow graph analogous to a multilayer feedforward neural network.Mini-batch training uses stochastic gradient descent methods such as SGD, RMSprop, or Adam.
- Incremental training: Incremental training minimizes E[||st − x||^2] at generation t by tuning the variables associated with the currently trained layers.The objective advances from the estimate at one layer to the next after processing D mini-batches.
- Incremental training: After training through layer t, incremental training appends layer t + 1 and retrains the entire network for D mini-batches.The optimization objective changes to the newly extended estimate.
- Incremental training: Variables learned in earlier generations initialize the optimization when a new layer is appended.The procedure updates γt sequentially from the first layer to the last.
IV. PERFORMANCE EVALUATION
This section evaluates TISTA's sparse-recovery performance through computer experiments.
- Evaluation: TISTA's sparse-recovery performance is evaluated by computer experiments.
A. Details of experiments
The experiments evaluate TISTA under Gaussian sensing with sparse Bernoulli-Gaussian signals, comparing NMSE, variance estimation, learned parameters, and convergence against established algorithms.
- Experimental setup: The experiments use Bernoulli-Gaussian sparse signals with p = 0.1 and α2 = 1, Gaussian noise, and Gaussian sensing matrices with variance 1/M.The main Gaussian-matrix setting uses N = 500 and M = 250 at SNR = 40 dB.
- Error-variance estimation: The estimated error variance τ2 accurately tracks the empirically estimated true error variance ¯τ2 despite non-monotonic variance behavior.The residual error nevertheless decreases rapidly, indicating successful signal recovery.
- NMSE comparison: TISTA has the steepest NMSE curve during the first 12 rounds among TISTA, ISTA, LISTA, AMP, and OAMP.In the N = 500, M = 250, SNR = 40 dB experiment, the comparison uses average normalized MSE versus iteration.
- NMSE comparison: 5 rounds are sufficient for TISTA to reach NMSE = −30 dB, compared with 6 rounds for OAMP and 10 rounds for LISTA.TISTA saturates around −42 dB and has gains of about 5.8 dB at T = 5 and 4.0 dB at T = 7 relative to OAMP.
- Learned variables: Three independently trained γt sequences approximately coincide and exhibit a zigzag pattern with values ranging from 1 to 10.Training α2 and p also improves TISTA's NMSE, suggesting that the MMSE estimator's true prior parameters are not always optimal for TISTA.
2) Large-scale problem:
The experiments test TISTA's scalability and robustness beyond the conventional Gaussian sensing ensemble. TISTA maintains strong NMSE behavior for large systems, altered Gaussian variance, and binary sensing matrices.
- Large-scale problem: TISTA achieves good NMSE performance for the large system (N, M) = (5000, 2500), saturating around −43 dB, 0.6 dB lower than OAMP.LISTA is omitted because it is computationally intractable to execute in the experimental environment.
- Running time: 37 times faster than LISTA, TISTA is in the N = 500 setting while also providing better NMSE performance.For T = 7, TISTA's CPU time for N = 10^4 signals is nearly equal to LISTA's for N = 500.
- Binary sensing matrices: TISTA's NMSE curves approximately coincide between binary and Gaussian sensing matrices, supporting robustness to non-Gaussian sensing matrices.The binary matrix entries are sampled uniformly from {±1}.
E. Sensing matrices with a large condition number
The paper evaluates TISTA on sensing matrices with large condition numbers and finds robust noiseless recovery, while noisy performance degrades as conditioning worsens. A trainable gradient-descent analogy is used to motivate learned zigzag-like acceleration.
- Large-condition-number sensing matrices: κ = 5000 causes almost no TISTA NMSE degradation without observation noise, whereas AMP diverges for κ ≥5.The comparison is made in the noiseless setting σ2 = 0.
- Large-condition-number sensing matrices: At SNR = 60 dB, TISTA has smaller NMSE than LISTA for κ = 1, 15, and 100, but its NMSE severely degrades as κ increases.The degradation is attributed to noise enhancement from the pseudo inverse linear estimator W.
- Hypothesis on zigzag shapes: Small fixed GD step sizes converge slowly, while large step sizes induce oscillation, slow convergence, or divergence.This illustrates why step-size selection is difficult for ill-conditioned quadratic problems.
- Hypothesis on zigzag shapes: TGD significantly outperforms fixed-step GD in averaged error and provides much faster convergence.TGD learns an appropriate step-size strategy through trainable parameters.
- Hypothesis on zigzag shapes: The learned γ_i values form a zigzag pattern, which the paper hypothesizes may also accelerate TISTA on ill-conditioned problems.The authors state that this hypothesis still requires confirmation in future work.
VI. SPARSE SIGNAL RECOVERY FOR MNIST IMAGES
The paper tests TISTA on non-i.i.d. sparse handwritten-digit images from MNIST rather than only synthetic sparse signals. TISTA produces higher-quality reconstructions and faster convergence than OAMP in the reported experiment.
- Experimental setup: MNIST provides 784-dimensional rasterized handwritten-digit signals, with M = 392 Gaussian measurements and observation-noise variance 4 × 10−4.TISTA is compared with OAMP using the Bernoulli-Gaussian MMSE estimator without assuming knowledge of the image prior PDF.
- Results: At t = 8, TISTA reconstructs an image with MSE = 0.0091, while OAMP reaches MSE = 0.0148 even at t = 100.The reported image comparison uses the same noisy observation for TISTA and OAMP.
- Results: TISTA shows much faster convergence than OAMP when comparing reconstructions after t = 1, 4, and 8 iterations.The paper reports that this agrees with the convergence results from Section IV.
- Conclusion: The MNIST experiment strongly suggests that TISTA can handle real non-i.i.d. sparse signals when enough training data are available.This conclusion is explicitly framed as a suggestion from the reported experiment.
VII. EXTENSIONS
The paper extends TISTA for sensing matrices with nonzero-mean components and large condition numbers. Mean removal substantially improves recovery for nonzero-mean matrices without increasing time complexity.
- Extensions: The proposed extensions target sensing matrices with nonzero-mean components or large condition numbers without additional computational costs in learning.The extensions are presented as outperforming original TISTA in their respective settings.
- Nonzero-mean sensing matrices: Mean removal modifies the sensing matrix and residual calculations so TISTA operates with sufficiently small-mean components.The resulting method is called TISTA-MR.
- Nonzero-mean sensing matrices: For A_i,j ∼ N(1, 1/M), TISTA-MR outperforms original TISTA, whose NMSE saturates around −10 dB in noiseless and SNR = 60 dB cases.The original TISTA saturation is reported for both evaluated settings.
- Nonzero-mean sensing matrices: At SNR = 60 dB and T = 10, TISTA-MR reaches −38 dB NMSE, a gain of about 28 dB over TISTA.The paper reports this improvement without increasing time complexity.
B. Sensing matrices with a large condition number
Large sensing-matrix condition numbers degrade original TISTA, so the paper replaces its pseudo-inverse linear estimator with an LMMSE-based alternative. TISTA-LMMSE substantially improves NMSE for κ = 1000 without increasing recovery-time complexity.
- Large condition numbers cause nonnegligible performance degradation for TISTA except in the noiseless case.
- TISTA-LMMSE replaces the original TISTA linear estimator with an LMMSE matrix to improve recovery for ill-conditioned sensing matrices.
- The regularized matrix W = A^T(A A^T + βI)^−1 reduces W's condition number and prevents noise enhancement.Here β is a real constant selected to minimize post-training NMSE.
- Matrix inversion is required only once at the beginning, giving TISTA-LMMSE the same time complexity as original TISTA.The approach avoids an M × M inversion at every iteration.
- For κ = 1000 and SNR = 60 dB, TISTA-LMMSE with β = 5.0 × 10−4 shows much better NMSE than original TISTA.Figure 17 compares LISTA, original TISTA, and TISTA-LMMSE using the regularized matrix.
- TISTA's trainable-variable count equals the number of iterations, supporting stable and fast training across Gaussian, binary, and large-condition-number matrices.The paper also reports faster convergence than AMP or LISTA in many cases and applicability to non-i.i.d. MNIST signals.