[Paper Review] PRISTA-Net: Deep Iterative Shrinkage Thresholding Network for Coded Diffraction Patterns Phase Retrieval

Aoxu Liu, Xiaohong Fan+2 more|arXiv (Cornell University)|Sep 8, 2023

Advanced X-ray Imaging TechniquesPhysics and Astronomy3 citations

TL;DR

PRISTA-Net is a deep unfolding network that integrates the iterative shrinkage thresholding algorithm (ISTA) with learnable nonlinear transformations and attention mechanisms to improve phase retrieval from coded diffraction patterns (CDPs). By end-to-end learning of all parameters—including thresholding, step size, and nonlinear mappings—it achieves state-of-the-art reconstruction quality with fast inference and enhanced robustness to noise.

ABSTRACT

The problem of phase retrieval (PR) involves recovering an unknown image from limited amplitude measurement data and is a challenge nonlinear inverse problem in computational imaging and image processing. However, many of the PR methods are based on black-box network models that lack interpretability and plug-and-play (PnP) frameworks that are computationally complex and require careful parameter tuning. To address this, we have developed PRISTA-Net, a deep unfolding network (DUN) based on the first-order iterative shrinkage thresholding algorithm (ISTA). This network utilizes a learnable nonlinear transformation to address the proximal-point mapping sub-problem associated with the sparse priors, and an attention mechanism to focus on phase information containing image edges, textures, and structures. Additionally, the fast Fourier transform (FFT) is used to learn global features to enhance local information, and the designed logarithmic-based loss function leads to significant improvements when the noise level is low. All parameters in the proposed PRISTA-Net framework, including the nonlinear transformation, threshold parameters, and step size, are learned end-to-end instead of being manually set. This method combines the interpretability of traditional methods with the fast inference ability of deep learning and is able to handle noise at each iteration during the unfolding stage, thus improving recovery quality. Experiments on Coded Diffraction Patterns (CDPs) measurements demonstrate that our approach outperforms the existing state-of-the-art methods in terms of qualitative and quantitative evaluations. Our source codes are available at \emph{https://github.com/liuaxou/PRISTA-Net}.

Motivation & Objective

• Address the limitations of black-box deep learning models in phase retrieval by enhancing interpretability and reducing reliance on manual hyperparameter tuning.
• Overcome the computational complexity and poor generalization of plug-and-play (PnP) frameworks in phase retrieval by designing a structured, learnable architecture.
• Improve reconstruction quality of coded diffraction patterns (CDPs) by integrating sparse priors, attention mechanisms, and global feature learning via FFT.
• Enable end-to-end training of all network components—including nonlinear transformations, thresholding, and step size—without manual initialization.
• Achieve fast inference speed while maintaining high robustness to noise, especially at low noise levels, through a novel logarithmic-based loss function.

Proposed method

• Proposes a deep unfolding network (DUN) based on the first-order ISTA algorithm, replacing fixed proximal mappings with learnable nonlinear transformations.
• Introduces a learnable thresholding mechanism and step size within the unfolding framework, enabling end-to-end optimization of all parameters.
• Employs a convolutional neural network (CNN) in both spatial and frequency domains using FFT to capture local and global image features.
• Incorporates a CBAM attention module to focus on phase-sensitive features such as edges, textures, and structural details.
• Designs a logarithmic-based loss function that improves performance under low-noise conditions by emphasizing relative error in amplitude recovery.
• Uses a two-branch module: a shrinkage gradient descent (SGD) module for detail refinement and a proximal-point mapping (PPM) module for denoising and sparse prior enforcement.

Experimental results

Research questions

• RQ1 Can a learnable, interpretable deep unfolding network outperform existing black-box and PnP-based phase retrieval methods in terms of reconstruction quality and inference speed?
• RQ2 How does the integration of attention mechanisms and global feature learning via FFT improve phase retrieval from CDPs?
• RQ3 To what extent does the proposed logarithmic-based loss function enhance performance under low-noise conditions compared to standard loss functions?
• RQ4 Can the end-to-end learned parameters (thresholding, step size, nonlinear mapping) generalize across different CDP masks and image types?
• RQ5 Does the proposed architecture maintain high performance when the number of masks (and thus measurements) is reduced?

Key findings

• PRISTA-Net achieves an average PSNR improvement of 0.7 dB, 0.96 dB, and 2.24 dB over baseline methods on UNT-6 at different noise levels when using J=4 masks.
• On the 256×256 test dataset, PRISTA-Net outperforms state-of-the-art methods including PrDeep, TFPnP, and PrComplex PnP, especially in preserving fine textures and edges.
• The method maintains superior performance across different CDP masks, demonstrating strong generalization capability without requiring retraining.
• The logarithmic-based loss function leads to significant performance gains under low-noise conditions, particularly improving reconstruction fidelity in high-SNR regimes.
• Visual results show fewer artifacts and richer structural details, with clear reconstruction of complex features like door frames and boat masts even at high noise levels.
• The intermediate outputs of the network show progressive artifact reduction and detail enhancement across stages, confirming the effectiveness of the PPM (denoising) and SGD (sharpening) modules.