[论文解读] 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 Astronomy被引用 3

一句话总结

PRISTA-Net 是一种深度展开网络，通过将迭代软阈值算法（ISTA）与可学习的非线性变换及注意力机制相结合，实现从编码衍射图样（CDPs）中进行相位恢复。通过端到端学习所有参数——包括阈值、步长和非线性映射——该方法在实现快速推理的同时，显著提升了对噪声的鲁棒性，并取得了当前最优的重建质量。

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}.

研究动机与目标

• 通过提升可解释性并减少对手动超参数调优的依赖，解决黑箱深度学习模型在相位恢复中的局限性。
• 通过设计一种结构化且可学习的网络架构，克服基于插件式（PnP）框架在相位恢复中计算复杂度高且泛化能力差的问题。
• 通过整合稀疏先验、注意力机制以及基于傅里叶变换（FFT）的全局特征学习，提升编码衍射图样（CDPs）的重建质量。
• 实现所有网络组件——包括非线性变换、阈值和步长——的端到端训练，无需手动初始化。
• 通过一种新型基于对数的损失函数，在保持高噪声鲁棒性的同时实现快速推理速度，尤其在低噪声水平下表现优异。

提出的方法

• 提出一种基于一阶 ISTA 算法的深度展开网络（DUN），用可学习的非线性变换替代固定的近端映射。
• 在展开框架中引入可学习的阈值机制与步长，实现对所有参数的端到端优化。
• 利用傅里叶变换（FFT）在空间域与频域同时使用卷积神经网络（CNN），以捕捉局部与全局图像特征。
• 引入 CBAM 注意力模块，聚焦于相位敏感特征，如边缘、纹理与结构细节。
• 设计一种基于对数的损失函数，通过强调幅度恢复中的相对误差，提升低噪声条件下的性能表现。
• 采用双分支模块：收缩梯度下降（SGD）模块用于细节优化，近端点映射（PPM）模块用于去噪与稀疏先验强制。

实验结果

研究问题

• RQ1 可学习且可解释的深度展开网络是否能在重建质量与推理速度方面超越现有的黑箱模型与基于 PnP 的相位恢复方法？
• RQ2 注意力机制与基于 FFT 的全局特征学习的结合，如何提升从 CDP 中进行相位恢复的性能？
• RQ3 与标准损失函数相比，所提出的基于对数的损失函数在低噪声条件下能多大程度上提升性能？
• RQ4 所学习的端到端参数（包括阈值、步长与非线性映射）是否能在不同 CDP 掩码与图像类型间实现良好泛化？
• RQ5 当掩码数量（即测量数）减少时，所提出的架构是否仍能保持高性能？

主要发现

• 在使用 J=4 个掩码时，PRISTA-Net 在 UNT-6 数据集上不同噪声水平下，平均 PSNR 分别较基线方法提升 0.7 dB、0.96 dB 与 2.24 dB。
• 在 256×256 测试数据集上，PRISTA-Net 显著优于当前最先进方法（包括 PrDeep、TFPnP 与 PrComplex PnP），尤其在保留精细纹理与边缘方面表现突出。
• 该方法在不同 CDP 掩码下均保持优异性能，展现出强大的泛化能力，且无需重新训练。
• 基于对数的损失函数在低噪声条件下带来显著性能提升，尤其在高信噪比（SNR）区域增强了重建保真度。
• 视觉结果表明，伪影更少，结构细节更丰富，即使在高噪声水平下，也能清晰重建复杂特征（如门框与船桅）。
• 网络中间输出显示，各阶段逐步减少伪影并增强细节，验证了 PPM（去噪）与 SGD（锐化）模块的有效性。