[論文レビュー] 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）を、学習可能な非線形変換とアテンションメカニズムと統合した深層アンフォールディングネットワークであり、符号化された回折パターン（CDP）からの位相再構成を改善する。すべてのパラメータ—しきい値、ステップサイズ、非線形写像—をエンド・トゥ・エンドで学習することで、高速な推論とノイズに対する高い耐性を実現し、最先端の再構成品質を達成する。

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

研究の動機と目的

• ブラックボックス型の深層学習モデルの制限を克服し、解釈可能性を高め、手動によるハイパーパrameterチューニングへの依存を減らす。
• 位相再構成におけるプラグアンドプレイ（PnP）フレームワークの計算コストの高さと一般化性能の低さを克服するため、構造的かつ学習可能なアーキテクチャを設計する。
• スパースプライア、アテンションメカニズム、およびFFTによるグローバル特徴学習を統合することで、符号化された回折パターン（CDP）の再構成品質を向上させる。
• 非線形変換、しきい値、ステップサイズといったすべてのネットワーク部品を、手動の初期化なしにエンド・トゥ・エンドで学習可能にする。
• 新規に設計された対数損失関数により、低ノイズ条件下でも高いノイズ耐性を維持しながら、高速な推論速度を実現する。

提案手法

• 1次ISTAアルゴリズムに基づく深層アンフォールディングネットワーク（DUN）を提案し、固定されたプロキシマレーションマッピングの代わりに学習可能な非線形変換を導入する。
• アンフォールディングフレームワーク内に学習可能なしきい値とステップサイズを導入し、すべてのパラメータをエンド・トゥ・エンドで最適化可能にする。
• 空間領域および周波数領域の両方で畳み込みニューラルネットワーク（CNN）を用いて、局所的およびグローバルな画像特徴を捉える。
• エッジ、テクスチャ、構造的詳細などの位相に敏感な特徴に注目するため、CBAMアテンションモジュールを統合する。
• 低ノイズ条件下での性能向上を図るため、相対誤差の重みを強調する対数に基づく損失関数を設計する。
• 2ブランチモジュールを採用：詳細の精緻化に用いるスリッピング勾配降下（SGD）モジュールと、ノイズ除去とスパースプライアの強制に用いるプロキシマルポイントマッピング（PPM）モジュール。

実験結果

リサーチクエスチョン

• RQ1 学習可能で解釈可能な深層アンフォールディングネットワークは、従来のブラックボックス型およびPnPベースの位相再構成手法に比べ、再構成品質と推論速度の両面で優れるか？
• RQ2 アテンションメカニズムとFFTによるグローバル特徴学習の統合が、CDPからの位相再構成にどのように寄与するか？
• RQ3 提案された対数損失関数は、標準損失関数と比較して、低ノイズ条件下での性能向上にどの程度寄与するか？
• RQ4 エンド・トゥ・エンドで学習されたパラメータ（しきい値、ステップサイズ、非線形写像）は、異なるCDPマスクや画像タイプに一般化可能か？
• RQ5 測定数（マスク数）が減少した場合でも、提案アーキテクチャは高い性能を維持できるか？

主な発見

• J=4マスクを用いた場合、UNT-6データセットにおいて、さまざまなノイズレベルでベースライン手法よりも平均して0.7 dB、0.96 dB、2.24 dBのPSNR向上を達成した。
• 256×256のテストデータセットにおいて、PrDeep、TFPnP、PrComplex PnPといった最先端手法を上回り、特に繊細なテクスチャやエッジの再現性に優れた性能を示した。
• さまざまなCDPマスクに対して優れた一般化能力を示し、再トレーニングなしに優れた性能を維持した。
• 対数損失関数により、低ノイズ条件下での性能向上が顕著に見られ、特に高SNR領域での再構成忠実度が向上した。
• 視覚的結果では、アーティファクトが少なく、構造的詳細が豊かに再構成されており、高ノイズ条件下でもドアフレームやボートマストのような複雑な特徴が明確に再構成された。
• ネットワークの中間出力は、段階を経るごとにアーティファクトが徐々に減少し、詳細が強化されていることが確認され、PPM（ノイズ除去）とSGD（鋭化）モジュールの有効性が裏付けられた。