[논문 리뷰] PRISTA-Net: Deep Iterative Shrinkage Thresholding Network for Coded Diffraction Patterns Phase Retrieval

Aoxu Liu, Xiaohong Fan외 2명|arXiv (Cornell University)|2023. 09. 08.

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) 프레임워크의 계산 복잡도와 낮은 일반화 능력 문제를 해결하기 위해 구조적이고 학습 가능한 아키텍처 설계.
• 스parser priors, 주의 메커니즘, FFT를 통한 전역 특징 학습을 통합하여 코딩된 회절 패턴(CDP)의 재구성 품질 향상.
• 비선형 변환, 임계값, 스텝 사이즈 등 네트워크의 모든 구성 요소를 수동 초기화 없이 엔드 투 엔드 학습 가능하게 설계.
• 새로운 로그 기반 손실 함수를 통해 저노이즈 조건에서 높은 노이즈 내성과 함께 고속 추론을 유지함으로써 성능 확보.

제안 방법

• 일阶 ISTA 알고리즘에 기반한 딥 언폴딩 네트워크(DUN)를 제안하며, 고정된 프록시 매핑을 학습 가능한 비선형 변환으로 대체.
• 언폴딩 프레임워크 내에서 학습 가능한 임계값 메커니즘과 스텝 사이즈를 도입하여 모든 파라미터의 엔드 투 엔드 최적화 가능.
• 공간 도메인과 주파수 도메인 양쪽에서 FFT를 활용한 컨볼루션 신경망(CNN)을 사용하여 국소적 및 전역적 이미지 특징 캡처.
• 모서리, 질감, 구조적 세부 정보와 같은 위상 민감한 특징에 집중하기 위해 CBAM 주의 모듈 통합.
• 저노이즈 조건에서 성능 향상을 위해 상대 오차 강조를 통한 로그 기반 손실 함수 설계.
• 이중 브런치 모듈 설계: 세부 사항 정밀화를 위한 스파이크 경사 하강(Shrinkage Gradient Descent, SGD) 모듈과 노이즈 제거 및 스퍼스 피어리어스 강화를 위한 프록시-포인트 매핑(Proximal-Point Mapping, PPM) 모듈.

실험 결과

연구 질문

• RQ1 학습 가능하고 해석 가능한 딥 언폴딩 네트워크가 기존의 블랙박스 및 PnP 기반 위상 복원 방법보다 재구성 품질과 추론 속도 측면에서 뛰어나게 성능을 냈는가?
• RQ2 주요 메커니즘과 FFT를 통한 전역 특징 학습의 통합이 CDP에서의 위상 복원에 어떤 영향을 미치는가?
• RQ3 제안된 로그 기반 손실 함수는 기존 표준 손실 함수에 비해 저노이즈 조건에서 성능 향상에 얼마나 기여하는가?
• RQ4 학습된 파라미터(임계값, 스텝 사이즈, 비선형 매핑)가 다양한 CDP 마스크와 이미지 유형 간에 일반화 가능한가?
• RQ5 측정 수(마스크 수)가 감소할 경우 제안된 아키텍처가 여전히 높은 성능을 유지하는가?

주요 결과

• J=4 마스크를 사용할 때, UNT-6 데이터셋에서 다양한 노이즈 수준에서 기준 방법 대비 평균 PSNR 향상 0.7 dB, 0.96 dB, 2.24 dB 달성.
• 256×256 테스트 데이터셋에서 PrDeep, TFPnP, PrComplex PnP와 같은 최신 기술 수준의 방법들을 초월하며, 특히 미세한 질감과 모서리를 잘 유지함.
• 다양한 CDP 마스크에서 우수한 성능 유지를 보이며, 재학습 없이도 강력한 일반화 능력을 입증.
• 로그 기반 손실 함수는 저노이즈 조건에서 뚜렷한 성능 향상을 가져오며, 특히 고SNR 영역에서 재구성 정밀도 향상.
• 시각적 결과에서는 더 적은 아티팩트와 richer한 구조적 세부 정보를 보이며, 고노이즈 조건에서도 문, 배의 마스트와 같은 복잡한 특징을 명확히 재구성.
• 네트워크의 중간 출력 결과는 단계별로 점진적인 아티팩트 감소와 세부 정보 강화를 보여주며, PPM(노이즈 제거) 및 SGD(선명화) 모듈의 효과성을 확인.