[Review] Medical Diffusion: Denoising Diffusion Probabilistic Models for 3D Medical Image Generation

유방 MRI만 관심이 있기에, 이를 초점으로 논문을 정리했다.

1. Introduction

이전 연구는 GAN을 사용 했지만, 이 기법에는 한계가 있다.

1. 모델 훈련의 어려움과 모드 붕괴가 일반적인 문제라 다양한 샘플 생성의 어려움.
2. 모드 붕괴가 발생하지 않더라도 이미지의 다양성이 제한적임.
3. GAN 및 유사 모델은 이미지 도메인에만 초점을 맞추기 때문에 텍스트에서 이미지를 생성하거나 그 반대로 생성하는 것이 간단하지 않지만, diffusion 모델은 다향한 이미지를 생성할 수 있다.

structural similarity index를 통해 이미지의 다양성을 정량적으로 조사하고, 생성된 데이터을 segmentation 모델에 사용할 수 있는지 조사한다.

2. Materials and Methods

2.1 Description of Dataset

모델의 일반화 성능을 측정하기 위해 아래 4개 데이터 셋을 각각 학습시켜 확인했다.

1. MRNet (무릎 mri)
2. ANDI (Alzheimer’s Disease Neuroimaging Initiative )
3. Duke dataset (Breast cancer MRI Dataset)
   • 각 환자의 non-fat saturated T1-weighted sequences를 사용한다.
4. LIDC, IDRI의 low-dose(저선량) 폐 CT사용

Duke datset에서 breast region에 해당하는 ground truth mask가 있는 200개(명)의 T1 weighted breast MRI 데이터 세트를 사용하여 self-supervised pretraining 방식에서 합성된 유방 이미지의 사용을 평가했다.

2.2 Data Preprocessing

유방암 데이터 세트의 이미지는 먼저 모든 이미지를 공통 복셀 간격 (0.66mm, 0.66mm, 3mm)으로 리샘플링한다. 그 뒤 윤곽이 그려진 유방의 해당 segementation mask를 사용하여 관심 영역을 잘라내는 방식으로 전처리했다. 그리고 이미지를 반으로 분할하여 왼쪽 유방과 오른쪽 유방이 별도의 이미지에 표시되도록 했다. 마지막으로 이미지의 크기를 256x256x32의 균일한 모양으로 조정했다.

2.3 Architecture

이 논문에서 제안하는 Medical Diffusion archtiecture에서는 2단계로 접근 방식을 구분한다.

1. image를 low-dimensional한 latent space로 인코딩한다.
2. 데이터의 latent representation에 대해 DDPM 모델을 훈련시킨다.

vector quantized autoencoder, 특히, VQ-GAN과 DDPM 대한 배경지식을 제공한다.

VQ-GAN

이미지를 의미있는 latent representation으로 인코딩하기 위해 vector quantized autoencoder는 VAE의 흐릿한 출력 문제를 완화할 수 있는 사용가능한 선택지로 입증되었다. autoeocnder의 bottleneck에 있는 latent feature vector를 quantized representation으로 매핑하는 방식으로 작동한다. VQ-GAN 아키텍처는 vector quantized autoencoder의 한 종류로 볼 수 있으며, 출력에 discriminator loss를 부과함으로써 image reconsturction 품질을 더욱 향상시킨다. 보다 정확하게는 이미지가 인코더에 공급되어 latent code \(Z_e\)를 구성한다.

• H는 높이, W는 너비, C는 채널 수, k는 latent feature map의 수, s는 compression factor를 나타낸다.

vector quantization step에서는 latent feature vector를 학습된 codebook z에 포함된 가장 가까운 해당 코드북 벡터 \(e_n\)으로 각각을 대체하여 양자화한다. 그런 다음 양자화된 특징 벡터를 decoder G에 공급하여 이미지를 reconstruction한다. 학습 목표는 reconstruction loss \(L_{rec}\), codebook loss \(L_{codebook}\), commitent loss\(L_{commit}\)의 최소화로 정의된다. 원저자가 정의한 대로, 우리는 preceptual loss를 reconstruction loss로 사용하고, 미분이 되지 않는 quanization step을 극복하기 위해 straight-through estimator를 사용한다. commitment loss는 양자화되지 않은 latent feature vecotr와 해당 codebook vector의 MSE로 정의된다. gradient는 양자화된 codebook vector에 더 근접해지기 위해 연속적인 latent feature vector에 대해서만 계산된다. learnable codebook vector는 매핑되는 모든 latent vector에 대해 exponential moving average를 유지함으로 최적화된다. 또한 더 나은 reconstruction 성능을 위해 출력에 patch-based discriminator가 사용된다. 이 아키텍처를 3D 입력을 지원하도록 확장하기 위해 2D convolution을 3D convolution으로 대체한다. 또한 원래 VQ-GAN 모델의 discriminator를 이미지 볼륨의 무작위 슬라이스를 입력으로 사용하는 슬라이스별 discriminator를 전체 reconsturction volume을 입력으로 사용하는 3D discriminator로 대체한다. 또한 GAN 훈련의 안정화를 위해 feature matching loss를 추가할 때에도 원 저자의 접근 방식을 따른다.

codebook은 VQ-VAE, VQ-GAN에서 설명이 나오는데 일련의 임베딩 벡터의 모음이다. 임베딩 벡터는 데이터의 고차원 표현을 저차원의 실수 벡터로 변환하는 방식을 의미하는데, 복잡한 구조가 관계를 가진 데이터를 저차원의 공간에 표현한다. 고차원의 내용을 저차원의 벡터로 압축하여 데이터의 복잡성과 계산의 효율성을 관리할 수 있다. 데이터 간의 유사성을 캡처하는 것도 가능하다. 이것들을 묶음으로 들고 있는 구조가 codebook이다. VQ-VAE의 encoder를 통해 latent representation으로 변환한 뒤 codebook의 벡터에 매핑하는 것을 양자화 과정이라 하는데, 이는 연속적인 latent representation을 이산적인 표현으로 변환하는데 있어서 중요한 역할을 한다. 이를 통해 데이터의 압축을 통해 효과적인 학습이 가능하다. 즉, 일종의 압축과정으로 이해하면 된다. (input -> encoder -> latent representation -> codebook의 벡터로 mapping(compression)) 그리고 codebook의 각 벡터는 특정 패턴이나 특성을 표현하게 되므로, codebook 자체는 모델이 학습한 데이터의 다양한 특성을 갭처하는 dictionary역할을 한다. 이렇게 학습된 코드북은 새로운 데이터에 대한 예측을 생성하거나, 학습 데이터의 특성을 분석하고 이해하는데 도움이 된다.

Diffusion Models

Diffusion model은 latent variable \( x_1 \cdots x_T \)를 통해 정의되는 생성 모델의 한 종류이다. 주요 아이디어는 이미지 \( x_0 \)에서 시작하여 여러 \(T\)의 timestep 동안 가우시안 노이즈를 이미지에 지속적으로 추가하여 분산을 증가시킨다. timestep \(t\)에서 이미지에 노이즈가 씌워진 버전과 timestep자체를 condition으로 하는 신경망을 노이즈 분포(이미지를 교란하는 데 사용된)를 배우도록 학습시켜 time \(t - 1\)에서 데이터 분포 \( p(x_{t-1} | x_t) \)를 inference할 수 있도록 한다.

T가 충분히 커지면 prior distribution \( \mathcal{N}(0,I) \)로 \( p(x_T) \)를 근사하고 이 분포에서 샘플링한 다음 마르코프 체인을 역방향으로 횡단하여 학습된 분포 \(p_θ(x_0) := ∫ p_θ(x_{0:T} )dx_{1:T}\) 로부터 새로운 이미지를 샘플링할 수 있다. 노이즈를 모델링하는데 사용된 신경망은 일반적으로 U-Net이 선택된다. 3D 데이터를 지원하기 위해 U-Net의 2D convolution을 3D convolution으로 대체한다. 또한 Vidieo diffusion의 접근 방식을 따라 고해상도 평면(3x3x1 크기의 커널)에서만 convolution을 사용한 다음 이 고해상도 평면(따라서 depth 차원을 배치 크기의 확장으로 취급)에 spatial attention block을 사용하여 계산 효율성을 높인다. spatial attention block에 이어 고해상도 이미지 평면 축이 batch 축으로 처리되는 depth attention block이 이어진다.

Putting It All Together

첫 번째 단계에서는 데이터의 의미있는 low-dimensional latent representation을 학습하기 위해 전체 데이터 셋으로 VQ-GAN 모델을 훈련한다. diffusion model에 주어지는 입력은 -1과 1사이로 정규화되어야 하므로, 이미지의 latent representation도 이 범위에 속하도록 만들어야 한다. VQ-GAN 모델의 vector quantization 단계에서 학습된 codebook의 vector가 quantization이전의 latent feature vecotr에 가까워지도록 강제한다고 가정하면, quantization되지 않은 latent feature representation의 최대값을 학습된 codebook의 최대값으로 근사화한다.

마찬가지로, 학습된 codebook의 최소값으로 양자화되지 않은 latent representation의 최소값을 근사화한다. 따라서 정량화되지 않은 latent vector에 대해 간단한 min-max 정규화를 수행하여 -1과 1 범위에 가까운 값을 가진 latent representation을 얻는다. 그런 다음 3D diffusion model을 훈련하는 데 사용할 수 있다. 그 다음 standard gaussian에서 노이즈 샘플링을 시작하여 reverse direction diffusion process를 실행하여 새 이미지를 생성할 수 있다.

이 프로세스의 출력은 VQ-GAN의 학습된 codebook을 사용하여 양자화되고 해당 이미지를 생성하기 위해 디코더에 순차적으로 공급된다. 24GB GPU RAM이 장착된 RTX6000에서 훈련되었고, 각 모델당 약 7일이 소요되었다. 각 모델에 대한 내용은 Table1을 참조.

3. Result

3.1 Medical Diffusion Models can be Robustly Trained

적은 데이터 세트에서도 성능을 보여주기 위해 breast MRI의 데이터는 1,844개의 이미지를 사용했다고 하는데, Duke데이터는 990개 정도 아니었나..? 왼쪽의 brain MRI의 998개와 숫자를 착각해서 잘못작성했나 싶긴하다.

데이터 세트가 비교적 작았음에도 불구하고, 각 모델은 수렴에 성공했고 어느 하이퍼파라미터의 fine-tuning없이도 진짜같은 합성 이미지를 생성했다(그림 1 참조). 특히 어떤 훈련 세션에서도 mode collpase는 발견되지 않았고, 모델 아키텍처는 넓은 범위의 해상도를 수용 가능했으며 64x64x64 복셀 Brain MRI, 128x128x128 복셀의 흉부 CT, 256x256x32 복셀을 가진 유방 및 무릎 MRI의 anisotropic 해상도를 모두 커버하고 적응할 수 있었다. 네 가지 데이터 세트 각각에서 사실적인 3차원 데이터를 생성할 수 있었다.

즉, 내가 얻어갈점은 256x256x32 해상도로 Duke 데이터 셋을 사용하여 학습했다는 것.

3.2 Medical Diffusion Models can Generate High Quality Medical 3D Data

합성 이미지가 사실적인지 외관을 전문의들이 평가했다.

1. 전체 이미지 외관의 품질
2. 슬라이스 간의 일관성
3. 해부학적 정확성

Reader A는 9년차, B는 5년차 전문의다. Duke 데이터셋에 대해서만 살펴본다.

• 외관이 사실적인지 판단하는 경우 조금 비현실적인 부분이 있다는 평이 90%이상이므로, 애매한 부분이 있는 것 같다. (저자들은 괜찮다고 얘기했지만, 합성 이미지인지 실제 이미지인지 구분 못하는 비율이 더 높아야 한다고 생각한다.)

• 슬라이스간의 일관성 부분이며, 50% 가깝게 일관성이 챙겨진다고 판독했다.

• 해부학적으로 정확하다는 의견은 40% 초중반으로 판독되었다.

3.3 The Dimension of the Latent Space is Important for High Quality Image Generation

이미지 생성 품질에 관해 latent 차원이 미치는 영향에 대해 분석하기 위해, 두 가지 다른 압축 factor로 VQ-GAN autoencoder를 학습시켰다. 긱 spatial dimension을 8배로 압축시키면(256x256x32를 32x32x4의 latent 차원으로) 해부학적 특징들이 손실되는 것을 발견했다(Fig 3).

• 유방 MRI에서 섬유선 조직의 연속성이 왜곡된다.

더 작은 압축 계수인 4를 사용하여(즉, 이미지 256x256x32크기를 64x64x8의 latent 차원으로) VQ-GAN autoencoder를 학습시키면 해부학적 특징이 좀 더 정확하게 reconstruction된다. 이 논문에서 사용한 4개의 데이터 셋 모두에서 최대 압축 계수 4의 경우 정확한 해부학적 디테일이 포함되어있다고 나타났다. (즉 latent space가 이미지의 원래 차수에 비해 4배 작음.)

3.4 Medical Diffusion Models outperform GANs in Terms of Image Diversity

해당 논문의 diffusion model 성능을 비교하기 위해 Wassersten GAN(WGAN)을 baseline으로 선택했다. 고해상도 이미지로 WGAN-GP를 훈련할 때 서로 다른 동작을 발견했기 때문에, 비교 대상을 64x64x64 크기의 생성된 뇌 MRI 이미지로 제한했다. 동일한 데이터 세트의 1000개 합성 샘플 쌍에 대한 결과를 평균화하여 multi-scale structural similarity metric(MS-SSIM)측면에서 비교했다.

MS-SSIM 점수가 낮을수록 이미지의 다양성이 높다는 것을 나타내는데 GAN 모델은 0.9996이었고, diffusion model은 0.8577이었다. 이 수치는 diffusion 모델이 더 다양한 샘플을 생성할 수 있다는 것을 나타낸다고 저자는 말한다.

3.5 Synthetic Data can be Used to Train Neural Networks

Duke 데이터 세트에서 훈련된 diffusion model을 사용하여 2000개의 합성 이미지를 생성한 뒤, 합성 데이터에 대한 self-supervised setting에서 Swin UNETR을 pre-trained 했다. 그런 다음 B 기관의 사용 가능한 segmentation data로 pre-trained된 신경망을 파인 튜닝하여 MRI scan에서 유방 영역을 segmentation했다. 제한된 데이터 환경에서 성능 향상을 보여주기 위해 B기관에서 사용 가능한 데이터 양을 증가(5%, 10%, 20%, 40%, 80%, 100%)시켜 여러 차례 학습을 실행했다. 비교를 위해 합성 데이터로 pre-train을 수행하지 않은 상태에서 동일한 작업을 수행하는 똑같은 네트워크를 학습 시켰다.

다른 기관의 합성 데이터로 pre-training하면 라벨링된 학습 데이터가 적은 환경에서 Dice score를 기준으로 세그멘테이션 성능을 크게 향상시킬 수 있었다(pre-training없이는 0.91, 5%의 데이터를 가지고 pre-training했을 때는 0.95).

• 훈련에 사용할 수 있는 데이터가 거의 없는 경우 합성 데이터로 pre-training을 하면 dice score를 크게 향상시킬 수 있다.

4. Discussion

diffusion model을 사용하여 3D 의료 영상을 합성할 때, 구조적인 일관성 및 해부학적 특징을 정확하게 반영할 수 있음을 나타냄. 약 1,000개의 샘플로 구성된 작은 데이터 셋으로 훈련하더라도 모델이 수렴하는 것을 보여줌.

연구의 한계

1. 약 1,000건의 비교적 작은 데이터 세트를 대상으로 모델을 평가했다.
2. 생성된 3D 볼륨은 완전한 diagnostic resolution을 갖지 못한다. 이는 사용 가능한 공개 데이터의 해상도가 제한되어 있고, 모든 경우에 최신 이미지 해상도를 반영하지 못하기 때문이다.