[논문 리뷰] Online Deep Equilibrium Learning for Regularization by Denoising

Jiaming Liu, Xiaojian Xu외 3명|arXiv (Cornell University)|2022. 05. 25.

Numerical methods in inverse problems인용 수 6

한 줄 요약

이 논문은 영상 역문제에 대한 딥 균형 모델(DEQ)에서 측정 모델에 대한 확률적 근사값을 적용하는 메모리 효율적인 프레임워크인 온라인 딥 균형 정규화를 제안한다. 배치 처리 방식이 아닌 온라인으로 측정값을 처리함으로써, 희박한 시야 CT, MRI, 강도 회절 톰그래피 응용 분야에서 메모리 및 계산 복잡도를 크게 감소시키면서도 배치 DEQ와 유사한 재구성 품질을 달성한다.

ABSTRACT

Plug-and-Play Priors (PnP) and Regularization by Denoising (RED) are widely-used frameworks for solving imaging inverse problems by computing fixed-points of operators combining physical measurement models and learned image priors. While traditional PnP/RED formulations have focused on priors specified using image denoisers, there is a growing interest in learning PnP/RED priors that are end-to-end optimal. The recent Deep Equilibrium Models (DEQ) framework has enabled memory-efficient end-to-end learning of PnP/RED priors by implicitly differentiating through the fixed-point equations without storing intermediate activation values. However, the dependence of the computational/memory complexity of the measurement models in PnP/RED on the total number of measurements leaves DEQ impractical for many imaging applications. We propose ODER as a new strategy for improving the efficiency of DEQ through stochastic approximations of the measurement models. We theoretically analyze ODER giving insights into its convergence and ability to approximate the traditional DEQ approach. Our numerical results suggest the potential improvements in training/testing complexity due to ODER on three distinct imaging applications.

연구 동기 및 목표

• 영상 역문제에 대한 딥 균형 모델(DEQ)에서 배치 처리의 높은 메모리 및 계산 비용을 해결하기 위해.
• 각 반복에서 전체 측정 집합에 의존하는 것의 의존도를 줄이며, DEQ를 통한 이미지 사전(end-to-end) 학습을 가능하게 하기 위해.
• 수렴성과 성능을 유지하면서도 훈련 및 추론 효율성을 향상시키는, 확률적이고 온라인 형태의 DEQ 변종을 개발하기 위해.
• 대규모 영상 응용 분야에서 PnP/RED 프레임워크에 대한 온라인 처리의 실용적 타당성을 입증하기 위해.
• 온라인 접근법의 수렴성 및 근사 품질에 대한 이론적 분석을 제공하기 위해.

제안 방법

• ODER는 RED 기반 DEQ의 데이터 일관성 층에 대한 확률적 근사값을 도입하여, 전체 측정 집합이 아닌 각 반복에서 부분 집합만을 처리한다.
• 고정점 반복을 통해 암시적 네트워크 방정식을 해결하며, 미니배치 측정값을 사용하여 확률적 기울기를 계산한다.
• 고정점 방정식을 통한 암시적 미분을 활용하여, 중간 활성값을 저장하지 않고도 엔드 투 엔드 백프로파게이션을 가능하게 한다.
• CNN 사전는 측정 모델과 함께 함께 훈련되어 노이즈 및 확률적 처리로 인한 잔상 제거를 가능하게 한다.
• 각 응용 분야(예: CT, MRI)에 맞게 적응된 Adam 최적화 기법을 사용하며, 학습률 스케줄링 및 가중치 감쇠를 적용한다.
• 측정값은 겹치지 않는 청크로 나누어지며, 각 반복에서 각 청크에서 무작위 부분 집합을 선택하여 데이터 일관성 업데이트를 계산한다.

실험 결과

연구 질문

• RQ1 DEQ 기반 PnP/RED 프레임워크에서 측정값의 확률적 처리가 배치 처리 수준의 재구성 품질을 유지할 수 있는가?
• RQ2 암시적 네트워크의 고정점 반복에서 온라인 처리가 수렴성과 안정성에 어떤 영향을 미치는가?
• RQ3 대규모 영상 역문제에서 ODER는 훈련 및 추론 복잡도를 얼마나 줄일 수 있으며, 성능 유지는 어떻게 유지하는가?
• RQ4 같은 반복당 계산 예산에서 ODER는 표준 DEQ 및 기준 방법보다 SNR 및 SSIM 측면에서 우수한 성능을 보일 수 있는가?
• RQ5 측정값 부분 집합 크기 및 샘플링 전략의 선택이 ODER의 최종 재구성 품질에 어떤 영향을 미치는가?

주요 결과

• ODER는 전체 배치 DEQ와 유사한 재구성 품질을 달성하며, 희박한 시야 CT에서 반복당 50개의 시야만으로도 RED(DEQ 180/180) 수준의 SNR 및 SSIM 성능을 보였다.
• 희박한 시야 CT에서, 동일한 반복당 계산 예산 하에 ODER는 50/180 시야를 사용하는 배치 DEQ보다 약 3 dB 높은 SNR 성능을 기록했다.
• 10% 샘플링을 적용한 평행 MRI에서, ODER는 RED(노이즈 제거) 및 RED(언폴드)보다 뛰어난 시각적 재구성 품질과 낮은 오차 잔여치를 확보했다.
• 500개의 측정값을 사용하는 강도 회절 톰그래피에서, ODER는 SNR 및 SSIM 측면에서 전체 배치 RED(DEQ)의 성능을 따라잡으며 더 나은 잔상 억제 효과를 보였다.
• 측정값 부분 집합의 청크화 및 무작위 샘플링 전략은 ODER의 경험적 재구성 품질과 안정성을 향상시켰다.
• 이론적 분석을 통해 ODER가 배치 DEQ 고정점에 수렴하는 해를 근사함을 확인하였으며, 약한 가정 하에 수렴 보장이 있음을 입증했다.