Source-linked AI summary

FISTA-Net: Learning A Fast Iterative Shrinkage Thresholding Network for Inverse Problems in Imaging

Jinxi Xiang, Yonggui Dong, Yunjie Yang

arXiv:2008.02683v3eess.IVphysics.app-ph

TL;DR

Inverse imaging requires recovering images from measurements while respecting forward-model and noise constraints. FISTA-Net unfolds FISTA into a learned model-based network with trainable gradient, proximal, and momentum components, and it outperforms comparative methods on EMT and sparse-view CT.

  • Problem

    Inverse imaging must recover unknown images from measurements, while fully data-driven approaches may omit forward-operator physics and image post-processing may lack data-consistency feedback.

  • Method

    FISTA-Net unfolds FISTA into cascaded gradient-descent, proximal-mapping, and momentum modules with learned gradient matrices, parameters, and nonlinear thresholding.

  • Results

    FISTA-Net outperformed the comparative methods across EMT metrics and improved PSNR over FBPConvNet by 3.340 dB at 60 views and 4.098 dB at 120 views.

  • Takeaways & Limitations

    FISTA-Net showed generalizability across EMT and sparse-view CT and can be extended to other inverse problems by changing the physical model A.

Abstract

from arXiv · show

Inverse problems are essential to imaging applications. In this paper, we propose a model-based deep learning network, named FISTA-Net, by combining the merits of interpretability and generality of the model-based Fast Iterative Shrinkage/Thresholding Algorithm (FISTA) and strong regularization and tuning-free advantages of the data-driven neural network. By unfolding the FISTA into a deep network, the architecture of FISTA-Net consists of multiple gradient descent, proximal mapping, and momentum modules in cascade. Different from FISTA, the gradient matrix in FISTA-Net can be updated during iteration and a proximal operator network is developed for nonlinear thresholding which can be learned through end-to-end training. Key parameters of FISTA-Net including the gradient step size, thresholding value and momentum scalar are tuning-free and learned from training data rather than hand-crafted. We further impose positive and monotonous constraints on these parameters to ensure they converge properly. The experimental results, evaluated both visually and quantitatively, show that the FISTA-Net can optimize parameters for different imaging tasks, i.e. Electromagnetic Tomography (EMT) and X-ray Computational Tomography (X-ray CT). It outperforms the state-of-the-art model-based and deep learning methods and exhibits good generalization ability over other competitive learning-based approaches under different noise levels.

I. INTRODUCTION

Imaging inverse problems recover unknown images from measurements under forward models, but conventional model-based and fully data-driven approaches each face practical limitations. Model-based deep learning combines physical modeling with learned representations to address these challenges.

  • Problem formulation: Inverse imaging estimates an unknown image x from measurements b linked through a forward operator.The operator may be linear for MRI and CT or nonlinear for ET and diffuse optical tomography.
  • Model-based methods: Optimization formulations combine data fidelity with a regularizer encoding priors such as sparsity, smoothness, low-rank structure, or self-similarity.The regularization coefficient satisfies λ > 0, and nonsmooth regularizers motivate proximal-gradient algorithms.
  • Model-based methods: ISTA, ADMM, PDHG, and FISTA provide computationally efficient iterative approaches for inverse problems with nonsmooth regularizers.Variants also address nonlinear forward operators by evaluating their gradients within frameworks such as FISTA.
  • Deep learning approaches: Fully data-driven networks learn mappings from measurements to images without explicitly modeling the forward operator, making them difficult to train for physics-dependent inversions.Their generalization has also not been thoroughly demonstrated, especially when mapping low-dimensional measurements to high-dimensional images.
  • Deep learning approaches: Image post-processing first forms a coarse image through analytical inversion and then uses a trained network to suppress noise and remove artifacts.The network produces a finer final image from the initial reconstruction.

1) Fully learned approach:

Deep learning approaches improve inverse-problem reconstruction but may lose structure or require substantial data, motivating FISTA-Net's model-based learned iterative design. The method unfolds FISTA while learning its parameters and proximal mapping.

  • Image post-processing: Image post-processing can improve metrics in low-dose CT but may cause over-smoothing and loss of critical structural details.Its performance is limited by the information content of the initial image and lacks a feedback mechanism for data consistency.
  • Model-based deep learning: Model-based deep learning combines a physical data-generation model, a statistical noise model, and a learned prior to constrain reconstruction.This architecture uses the forward operator and statistical model to enforce consistency with measurements.
  • FISTA-Net: FISTA-Net unfolds FISTA iterations into cascaded blocks and replaces nonlinear soft-thresholding with a learned deep proximal operator network.Its modules perform gradient descent, proximal mapping, and momentum operations in sequence.
  • FISTA-Net: FISTA-Net avoids forward-operator matrix inversion, which is desirable for large-scale or ill-conditioned inverse problems.A CNN handles proximal mapping, is shared across iterations, and uses different thresholding values for varying noise levels.
  • FISTA-Net: FISTA-Net learns gradient matrices and core parameters from data while imposing positivity and gradual variation constraints for proper convergence.A momentum module is added to damp oscillations and speed iterations.
  • Evaluation: The paper evaluates FISTA-Net on CT and Electromagnetic Tomography, covering linear and nonlinear imaging problems.The experiments compare performance across two distinct imaging modalities.

B. FISTA

FISTA accelerates ISTA through momentum, and FISTA-Net unfolds this procedure into cascaded gradient descent, proximal mapping, and momentum modules with learned operators and parameters.

  • B. FISTA: FISTA applies shrinkage at y^(k), a linear combination of x^(k) and x^(k−1), instead of directly at x^(k).This momentum choice accelerates convergence without requiring more than one gradient evaluation per iteration.
  • B. FISTA: FISTA-Net alternates data-consistency gradient descent, learned nonlinear proximal mapping, and a two-step update across cascaded iterations.The gradient step minimizes the measurement-consistency term and imposes a physics-based constraint.
  • B. FISTA: The gradient module uses a layer-dependent operator W^(k), decomposed into a learned step size µ^(k) and a precomputed matrix W̃.W̃ is obtained by minimizing a constrained Frobenius norm, after which µ^(k) is learned end-to-end.
  • B. FISTA: The proximal module learns an invertible sparse transform and nonlinear thresholding network rather than relying on a fixed handcrafted transform.Its convolutional representation uses four linear convolutional operators, ReLU activations, and a skip connection; the transform loss includes symmetry and sparsity terms.
  • B. FISTA: FISTA-Net shares the transform network across iterations while learning iteration-specific thresholds θ^(k) to accommodate changing noise and artifact levels.This avoids requiring separate pretrained denoisers for different noise levels.

B. Model-based Parameter Constraints

FISTA-Net constrains its learned iteration parameters to remain positive and evolve monotonically, supporting proper convergence behavior.

  • B. Model-based Parameter Constraints: The learned step size, thresholding value, and momentum scalar are constrained to be positive rather than manually tuned.The constraints address the possibility of non-positive step sizes and thresholds during iteration.
  • B. Model-based Parameter Constraints: The gradient step size µ^(k) should decay smoothly across iterations.
  • B. Model-based Parameter Constraints: The thresholding value θ^(k) should decrease iteratively as noise variances are progressively suppressed.
  • B. Model-based Parameter Constraints: The two-step update weight ρ^(k) should increase monotonically, matching the corresponding FISTA update weight.Softplus parameterizations enforce positivity, while ρ^(k) remains in (0, 1), consistent with FISTA.

C. Initialization

Initialization uses problem-dependent warm starts: Laplacian-regularized solutions for ill-conditioned EMT and unregularized inverse Radon transforms for well-conditioned CT.

  • C. Initialization: For ill-conditioned EMT, FISTA-Net is initialized with the solution of a Laplacian-regularized reconstruction problem.The initialization uses R(x) = ∥Lx∥2^2, where L is the Laplacian matrix.
  • C. Initialization: For well-conditioned CT, initialization requires no regularization and uses x0 = A^T b computed with Python or Matlab’s iradon-related functionality.
  • C. Initialization: Providing the CT initial guess marginally reduced training time but did not improve the final reconstruction result.
  • C. Initialization: The proximal convolutional network is initialized with Xavier initialization and fixed values for {w1, w2, w3, c1, c2, c3}.The listed initialization is {−0.5, −0.2, 1, −2, −1, 0}.

D. Implementation Details

FISTA-Net is trained end-to-end with shared networks, a fixed training iteration count, Adam optimization, and mini-batches of 64.

  • D. Implementation Details: The iteration step is fixed during end-to-end training, while shared networks permit choosing a different reconstruction iteration count.
  • D. Implementation Details: All networks are optimized with Adam using a mini-batch size of 64.
  • D. Implementation Details: The training procedure estimates two learning rates, lr1 and lr2, from the first and second moments of the gradients.

IV. EXPERIMENTS AND RESULTS

The experiments evaluate FISTA-Net on nonlinear EMT and linear sparse-view CT using standard image-quality metrics.

  • FISTA-Net is evaluated on nonlinear Electromagnetic Tomography and linear sparse-view Computed Tomography.The study uses EMT and sparse-view CT as representative imaging modalities.
  • RMSE, PSNR, and SSIM quantify the reconstructed image quality across the experiments.These are the three reported evaluation metrics.

A. Nonlinear Case: EMT

The EMT study uses simulated circular conductivity phantoms and compares FISTA-Net with model-based, post-processing, and model-based deep-learning baselines.

  • EMT phantoms vary in object location, conductivity, size, and count within the sensing region.Conductivity values range from (0.05, 0.5) S/m, and testing includes a separate three-object set.
  • The EMT forward problem is solved with COMSOL Multiphysics to generate measurement data from simulated conductivity distributions.The resulting inverse problem is nonlinear and ill-conditioned because of the dispersed electromagnetic field.
  • FISTA-Net is compared with Laplacian regularization, FISTA-TV, FBPConvNet, and ISTA-Net.The baselines include traditional model-based methods, an image post-processing network, and a model-based deep-learning method.

2) Comparison Study:

FISTA-Net outperforms competing methods on EMT reconstruction while retaining interpretable, progressively learned parameters and robustness across noise levels. Its performance improves with depth before stabilizing, but absolute conductivity remains inaccurate because EMT is nonlinear.

  • EMT comparison: FISTA-Net outperforms other methods on both EMT test sets using PSNR, SSIM, and RMSE, although FBPConvNet can match or slightly exceed its shape preservation.FBPConvNet uses 482,449 learnable parameters, whereas FISTA-Net uses 74,599; FBPConvNet nevertheless has moderate PSNR and RMSE because it lacks a data fidelity term.
  • EMT comparison: All methods fail to accurately reconstruct the absolute conductivity of different EMT objects.The sensor response is nonlinear, and linearizing the forward operator introduces nonlinear errors that affect absolute conductivity.
  • Layer and parameter behavior: FISTA-Net’s seven-layer configuration balances image quality and computational cost after performance stabilizes beyond seven layers.The evaluation varied depth from five to nine layers; larger configurations increase GPU cost, while performance tends to stabilize above seven.
  • Layer and parameter behavior: During iteration, FISTA-Net produces progressively clearer images and monotonically decreasing θ(k) and µ(k), unlike the less interpretable and sometimes negative parameters of ISTA-Net.Smaller θ(k) and µ(k) support more accurate solutions but slower convergence, consistent with decreasing reconstructed-image noise variance.
  • Noise robustness: FISTA-Net outperforms competing methods across the 22dB–40dB SNR range, with larger advantages at low SNR.ISTA-Net and FISTA-Net are more robust to additive noise, whereas FBPConvNet degrades faster; FISTA-TV shows similar robustness to the model-based deep learning methods.

B. Linear Case: Sparse-view CT

The sparse-view CT experiments evaluate FISTA-Net on Mayo Clinic data with 60 and 120 projection views, comparing it with model-based and learning-based baselines. FISTA-Net progressively recovers details while achieving the strongest reported quantitative results.

  • Clinical CT dataset: Sparse-view CT reduces projection views to 60 and 120 to lower radiation dose and scan time.The Mayo Clinic dataset contains 2,378 full-dose CT images from ten patients; eight patients were used for training and one for validation.
  • Comparison Study: FISTA-Net was compared with FBP, FISTA-TV, FBPConvNet, and ISTA-Net using seven network layers.FISTA-TV used 100 iterations with regularization parameter 0.001, while ISTA-Net used six cascaded stages.
  • Comparison Study: FISTA-Net preserved most fine structures while providing superior noise reduction on 60-view CT images.FBP showed strip artifacts, FISTA-TV introduced blocky artifacts and oversmoothing, and FBPConvNet blurred or distorted some details relative to ground truth.
  • Quantitative evaluation: FISTA-Net achieved the best results across all reported metrics on the Mayo Clinic CT dataset.The comparison covers both 60-view and 120-view reconstructions.
  • Quantitative evaluation: 3.340 dB and 4.098 dB PSNR improvements over FBPConvNet were obtained for 60-view and 120-view CT, respectively.FBPConvNet provided the strongest quantitative results among the comparative methods before these improvements were measured.
  • Iteration results: During 60-view CT reconstruction, noise removal and detail recovery progressed across iterations as ρ(k) increased and θ(k), µ(k) decreased.The learned parameter configuration resembled that of the conventional model-based reconstruction method.

3) Iteration results:

The convergence analysis compares FISTA-Net with FISTA-TV and ISTA-Net using the same initial guess and seven layers. FISTA-Net converges faster, while rigorous theoretical convergence analysis remains outside the paper’s scope.

  • C. Convergence Analysis: The convergence experiment compares FISTA-TV, ISTA-Net, and FISTA-Net from the same initial guess.FISTA-TV was hand-tuned for its best reconstruction, and both learned networks used seven layers.
  • C. Convergence Analysis: FISTA-Net converges faster than ISTA-Net and FISTA-TV.The comparison uses the evolution of RMSE across iterations.
  • C. Convergence Analysis: ISTA-Net converges stage-wise for the nonlinear EMT problem, possibly because its parameters are freely learned.The proposed explanation is attributed to the discussion of ISTA-Net’s learned parameters.
  • C. Convergence Analysis: Rigorous theoretical convergence analysis is outside the scope of this paper.The authors refer readers to convergence analysis for Learned ISTA as a potentially extendable basis.
  • V. CONCLUSION: FISTA-Net is reported to generalize across nonlinear EMT and linear sparse-view CT by decoupling data fidelity from the proximal operator.The physical model A can be changed to extend the network to different imaging inverse problems.
Loading 2008.02683v3…