[논문 리뷰] Fourier Diffusion Models: A Method to Control MTF and NPS in Score-Based Stochastic Image Generation

Matthew Tivnan, Jacopo Teneggi외 7명|arXiv (Cornell University)|2023. 03. 23.

Generative Adversarial Networks and Image Synthesis인용 수 4

한 줄 요약

이 논문은 스칼라 확산 프로세스를 이동 불변 컨볼루션과 정적 노이즈로 대체함으로써 각 시간 단계에서 MTF와 NPS를 직접 제어할 수 있는 새로운 스코어 기반 생성 모델인 푸리에 확산 모델(Fourier Diffusion Models)을 제안한다. 이는 진짜 이미지에서 측정된 흐린 이미지와 노이즈가 있는 이미지로의 연속적인 확률 흐름을 알려진 시스템 특성과 함께 모델링함으로써, 훨씬 적은 시간 단계로도 뛰어난 이미지 생성 성능을 달성한다. 스칼라 모델이 128 또는 1024단계를 사용할 때와 동일한 성능을 달성하기 위해 단 8단계만으로도 가능하다.

ABSTRACT

Score-based stochastic denoising models have recently been demonstrated as powerful machine learning tools for conditional and unconditional image generation. The existing methods are based on a forward stochastic process wherein the training images are scaled to zero over time and white noise is gradually added such that the final time step is approximately zero-mean identity-covariance Gaussian noise. A neural network is then trained to approximate the time-dependent score function, or the gradient of the logarithm of the probability density, for that time step. Using this score estimator, it is possible to run an approximation of the time-reversed stochastic process to sample new images from the training data distribution. These score-based generative models have been shown to out-perform generative adversarial neural networks using standard benchmarks and metrics. However, one issue with this approach is that it requires a large number of forward passes of the neural network. Additionally, the images at intermediate time steps are not useful, since the signal-to-noise ratio is low. In this work we present a new method called Fourier Diffusion Models which replaces the scalar operations of the forward process with shift-invariant convolutions and the additive white noise with additive stationary noise. This allows for control of MTF and NPS at intermediate time steps. Additionally, the forward process can be crafted to converge to the same MTF and NPS as the measured images. This way, we can model continuous probability flow from true images to measurements. In this way, the sample time can be used to control the tradeoffs between measurement uncertainty and generative uncertainty of posterior estimates. We compare Fourier diffusion models to existing scalar diffusion models and show that they achieve a higher level of performance and allow for a smaller number of time steps.

연구 동기 및 목표

• 순수한 흰색 노이즈로의 열화로 인해 많은 시간 단계가 필요한 전통적인 스코어 기반 확산 모델의 비효율성을 해결하기 위해.
• 전진 및 역방향 확산 과정에서 조절 가능한 조절 전송 함수(MTF)와 노이즈 파wer 스펙트럼(NPS)를 직접 제어하기 위해.
• 실제 영상 시스템 특성(예: 공간 블러 및 무늬 노이즈)과 일치하는 전진 과정을 모델링하여 이미지 생성 성능을 향상시키기 위해.
• 역방향 샘플링 단계 수를 줄여 계산 비용을 절감하기 위해.
• 임상적으로 관련성이 높은 응용 분야인 저선량 CT 영상 복원에서 이 방법의 효과성을 입증하기 위해.

제안 방법

• 스칼라 전진 확산을 원형 행렬을 사용한 이동 불변 컨볼루션 연산으로 대체하여 푸리에 도메인에서 시스템 블러(MTF)를 모델링한다.
• 지정된 노이즈 파워 스펙트럼(NPS)을 가진 정적 가우시안 노이즈로 추가 노이즈를 모델링함으로써 주파수 도메인에서 노이즈 특성 제어를 가능하게 한다.
• LSI 시스템 응답과 ASGN을 갖는 시퀀스로 전진 과정을 정의함으로써 최종 상태가 순수한 노이즈가 아닌 측정된 영상 통계와 일치하도록 한다.
• 시간에 따라 변화하는 스코어 함수를 푸리에 도메인에서 추정하기 위해 스코어 매칭 신경망을 훈련시킨다. 이는 반복적 노이즈 제거 및 선명함 향상에 기여한다.
• 학습된 스코어 함수를 사용하여 역방향 샘플링을 수행함으로써 측정된 데이터에서 고품질 이미지를 생성한다. 시간 단계는 노이즈가 많고 뿌연 측정값에서부터 깨끗한 이미지로의 진행을 나타낸다.
• 모든 선형 시스템과 추가 가우시안 노이즈를 갖는 가우시안 프로세스에 적용 가능한 스코어 매칭 손실의 일반화 및 역방향 SDE 설정을 제안하며, 이는 이동 불변 시스템에 국한되지 않는다.

Figure 1: A probabilistic graphical model for the stochastic process, consisting of linear shift invariant systems and additive stationary Gaussian noise. The measurement-conditioned diffusion model, in light purple, is trained to approximate the reverse process, in dark purple.

실험 결과

연구 질문

• RQ1 전진 확산 과정에서 MTF와 NPS를 제어함으로써 스코어 기반 이미지 생성의 효율성과 품질을 향상시킬 수 있는가?
• RQ2 측정된 영상 통계로 수렴하도록 전진 과정을 모델링함으로써 순수한 노이즈가 아닌 상태로의 수렴이 고품질 생성을 위한 역방향 샘플링 단계 수를 줄일 수 있는가?
• RQ3 초기 측정된 영상과 전진 과정의 최종 상태 간의 유사성이 역방향 샘플링 과정의 성능에 어떤 영향을 미치는가?
• RQ4 푸리에 확산 모델이 더 적은 시간 단계와 낮은 계산 비용으로 스칼라 확산 모델보다 이미지 복원 작업에서 뛰어난 성능을 보일 수 있는가?
• RQ5 향상된 성능은 역방향 과정 시작 시점과 진짜 데이터 분포 간의 분포 이격 감소 때문인가, 아니면 컨볼루션 네트워크와의 아키텍처 호환성 때문인가?

주요 결과

• 푸리에 확산 모델은 스칼라 모델이 128 또는 1024단계를 사용할 때와 동일한 성능을 달성하기 위해 단 8단계만으로도 동등하거나 더 뛰어난 이미지 품질을 달성한다.
• 모든 시간 단계 수에서 푸리에 확산 모델의 평균 제곱 오차(MSE)는 스칼라 모델보다 유의미하게 낮으며, 특히 낮은 시간 단계에서 가장 큰 향상이 관찰된다.
• 평균 제곱 편향 및 평균 분산 지표 역시 푸리에 확산 모델에서 일관되게 향상되어 더 높은 정밀도와 낮은 추정 오차를 나타낸다.
• 시각적 비교 결과, 푸리에 확산 모델은 특히 낮은 시간 단계에서 더 높은 품질의 재구성을 생성하며, 잡음과 구조 손실을 줄였다.
• 푸리에 확산 모델의 역방향 과정은 스칼라 모델 대비 시간이 지남에 따라 더 적은 열화를 보이며, 더 안정적이고 정확한 스코어 추정이 가능함을 시사한다.
• 결과적으로 향상된 성능는 초기 측정된 영상과 전진 과정의 최종 상태 간의 분포 유사도가 높아져 신경망이 인위적인 열화 과정을 뒤집는 데 부담을 덜 받기 때문임을 시사한다.

Figure 2: Diagram of the score-matching neural network. The inputs are the forward process sample image, low-dose CT measured image, and sample time. The output is the estimated score function.