3D MedDiffusion: A 3D Medical Latent Diffusion Model for Controllable and High-quality Medical Image Generation (IEEE TMI 2025)

이번 포스트에서는 IEEE TMI (2025년 7월!!) 에 accept 된 3D Medical Image 전용 diffusion model 인 MedDiffusion 모델을 리뷰해보겠습니다.

Preliminaries

기본지식을 위한 섹션입니다. 3D medical image generation 을 잘 아시는 분은 바로 methods 로 넘어가셔도 좋습니다.

3D vs 2D?

2D 사진은 224x224, 256x256 (65536) 등의 resolution 에서 작동하기 때문에 큰 모델도 GPU 에 부담없이 들어갈 수 있지만, 의료 영상은 256x256x256 (16,777,216 !!) 같이 엄청 큰 dimension이 필요하기 때문에 기존 모델을 그대로 3D 연산으로 바꿀 수 없습니다. 현재 이 문제를 해결하려면 2D 로 나누어서 순차적으로 generation 을 수행하거나, 더 작은 patch 로 나누어서 하는 방법밖에 없는데, 이 논문은 후자입니다.

VQ-VAE

고 해상도 diffusion model 은 latent space 에서 작동합니다.

MNIST 같이 작은 데이터는 직접 이미지에 noise 를 추가해서 해볼만한데, $3\times 512 \times 512$ 같은 초고해상도 이미지는 커버해야하는 픽셀 수가 너무 많습니다. 대신에, 이 이미지를 압축한 "latent space" 에서 denoising 을 수행하는데, stable diffusion 같은 경우는 $4\times 64 \times 64$ (무려 원본의 2%) 까지 압축을 해서 진행합니다. 하지만, 이렇게 압축을 했으면 결국 다시 원 해상도로 복원을 해야하는데, 이 역할은 diffusion model 이 아닌 variational autoencoder (VAE) 이 담당하게 됩니다.

다만, 원래 autoencoder 은 L1/L2 Reconstruction Loss 만 활용해서 결과물이 굉장히 blurry 한 문제가 발생합니다. 이 문제를 해결하기 위해 perceptual loss 와 adversarial loss (discriminator) 를 추가하여 reconstruction quality 를 개선합니다. 결국, latent diffusion model 은 이 VAE 의 quality 가 좌우한다고 봐도 무방합니다.

여기서 blurring 을 더욱 해결하려고 도입한게 Vector Quantised-Variational AutoEncoder (VQ-VAE) 인데, 기존의 VAE 는 continuous 한 latent space 를 가지고 있어서 어떤 두 영상의 latent representation 사이가 아래 사진처럼 "interpolation" 되는 것 처럼 나오는 문제점이 있었습니다.

중간에 0, 6, 9 같이 잘 보이는 사진도 있지만, 각 숫자 사이의 경계에 있는 형체를 알 수 없는 숫자들이 많이 껴 있는걸 보실 수 있습니다. 사실 생각해보면, 1 과 2 사이의 숫자는 뭐임??? 이라고 물어보면 대답하기가 어렵습니다. 1 아니면 2지, 두 숫자 사이는 뭔데? 라고 되묻겠죠.

따라서, latent space 를 discrete 하게 만든것이 VQ-VAE 입니다. 자세한 설명은 이 유튜브 영상이 기가막히게 하니, 여기로 찾아가시길 바랍니다.

Methods

MedDiffusion 은 3가지 component 로 나뉩니다.

1) Patch-volume autoencoder: 3D medical image 를 compact latent space 로 encoding 하는 역할
2) BiFlowNet: latent diffusion model
3) ControlNet: for downstream tasks such as conditional generation

Patch-volume autoencoder

MedDiffusion 에서는 VQ-VAE를 Patch-wise training과 Volume-wise training 으로 나누었습니다. 먼저 조그만한 (64x64x64) 패치를 복원하는 모델을 만들고, 그 모델을 활용해 전체 volume (256x256x256) 을 다루는 모델을 제작하겠다는거죠.

Patch-wise training

별 거 없습니다. 조그만한 N 개의 패치를 뽑아서 VQ-VAE 를 훈련합니다.

Volume-wise training

위 피규어에서 보이는 대로, volume 을 N개의 패치로 나눈 뒤 patch-wise training 에서 훈련한 patch encoder 을 활용해서 각각 quantized patch feature 를 추출합니다. 패치별로 추출한 feature 을 concatenate 해서 $\tilde{\mathbf{Z}} \in \R^{N\cdot M\times C}$ 를 추출합니다 (M: flattened spatial dimension, C: codebook feature dimension).

Joint Decoder 은 patch-wise training 에서 훈련한 patch decoder 을 fine-tuning 합니다. 이 디코더를 활용해서 최종 reconstruction $\tilde{\mathbf{X}} =D_J(\tilde{\mathbf{Z}})\in \R^{H\times W\times D}$ 를 한번에 수행해버립니다.

왜 joint decoder 이라고 하는지는 의문입니다. Patch-wise training 이랑 volume-wise training 을 한 에폭에 같이하는걸까요?

Training loss

3개의 loss term 으로 나뉩니다.

• Vector quantization loss

원본 이미지를 복구하는 reconstruction (l2) loss 와 코드북 임베딩이 discrete 하도록 만드는 핵심 loss 입니다.

• Adversarial loss

이것도 VQ-VAE 에 나왔던 discriminator loss 입니다.

• Tri-plane loss

이건 그냥 perceptual loss 인데, 2D pretrained 모델을 그대로 사용할 수 없어서 volume 에서 3 orthogonal slices (즉, axial, sagittal, coronal 모두) 를 무작위로 선정해 input 과 reconstruction 에서 VGG 16 을 활용해 perceptual loss 를 구했습니다. 3장인 이유는 아마 VGG가 3채널을 입력받기 때문이라 생각합니다. 저는 이녀석이 patch artifact 를 줄이는 데에 핵심 역할을 하지 않았나 생각합니다 (이거 진짜 뇌피셜입니다 ㅋㅋ 추후에 다룸).

3D Medical Diffusion Model

VQ-VAE 를 잘 훈련했다면 그 다음은 latent diffusion model 입니다.

BiFlowNet Architecture

Diffusion 을 잘 수행하기 위해 모델을 두 파트로 나눕니다. 하나는 패치를 복원하는데에 집중하는 intra-patch flow, 그리고 다른 하나는 전체 volume 을 (그래서 화살표에 Merge 써있음) 조화롭게 복원하는데 에 집중하는 inter-patch flow 입니다.

intra-patch flow 에서는 작은 patch 의 디테일을 잘 살려야 하기 때문에 조금 더 capacity 가 큰 DiT 모델을 활용합니다. 반면, inter-patch flow 는 전체 볼륨의 일관성을 살려야 하고 처리해야하는 데이터가 더 크기 때문에 UNet 을 활용합니다. 두 flow 모두 한번에 훈련하기 위해서 중간중간에 element-wise addition 을 수행합니다.

Implementation details

간략한 요약입니다.

• Patch size: $64\times 64\times 64$
• Codebook: 8192 개의 codes, 8 dimensions
• Discriminator trained after 20000 iterations
• 1 A100 (80GB) for VQ-VAE training, 8 A100 for diffusion

Results

먼저 quantitative results 를 보겠습니다. FID 가 main metric 인데, 3D ResNet (pretrained on medical dataset) 을 활용해서 구했다고 합니다.

Competing baseline 인 MAISI 와 MedSyn 의 절반보다 낮은 FID 를 보이는 것을 알 수 있습니다.

다음은 qualitative 평가를 보겠습니다.

Chest CT 도 있는데 brain MRI 에서 차이점이 더 드러난다고 생각해서 가져왔습니다. Competing baseline 인 MAISI 와 MedSyn 에 대비해서 조금 더 선명한 cortical contrast 가 보이는 것을 확인할 수 있습니다. 다만, fidelity 가 높은지는 결론내릴 수 없습니다 (이는 paired data 로 확인해야함).

Ablation Study

먼저 generated images 를 t-SNE 분석해본 결과입니다. 흥미롭게도, MedDiffusion 이 Real embeddings 와 제일 가깝게 clustering 을 이루는 것을 볼 수 있습니다. WDM 또한 real 과 교집합을 어느정도 이루고 있으나, real 의 전체 space 와 겹치지는 않네요. 이는 WDM 이 MedDiffusion 에 비해 diversity 가 떨어진다라고 해석할 수 있습니다. 반면에, MedSyn 은 아예 real 데이터와 겹치는 부분이 적은데, 이는 fidelity 가 부족하다고 볼 수 있습니다.

다음은 joint decoder 의 효과입니다. 눈으로 보았을 때 patch artifact 가 유의미하게 사라진 것을 볼 수 있으나, 우측 하단 척추뼈에 희미하게 반짝이는 artifact 가 아직 남아있는 것을 확인할 수 있습니다. Joint decoder 을 활용했을 때 PSNR 이 34.2 에서 35.8 (+1.6) 올랐다고 합니다.

VQ-VAE 의 ablation study 입니다. Adversarial loss 를 추가했을 때 PSNR 이 0.37 올랐고, perceptual loss 를 썼을 때 2.52 (!!) 올랐습니다.

BiFlowNet 의 ablation study 입니다. 주목할 것은 intra-patch 의 성능 기여도인데, 전체 volume 만 신경쓰는게 아니라 각 패치의 denoising 도 신경써야 전반적인 resolution 을 상승시킬 수 있다는 결과입니다.

Conclusion

이전에도 온전한 3D framework 로 medical image generation 을 한 논문은 많았지만, 이 논문은 resolution을 더 크게 끌어올린 engineering 이 돋보입니다.

또한, 패치 기반의 generation 의 가장 큰 단점인 boundary 의 patch artifact 를 joint decoder 을 활용해서 상당히 mitigate 한 점도 높게 평가할 수 있겠습니다.

전반적으로 이 논문은 3D medical image generation 의 SOTA 반열 (25년 8월 기준)에 올랐다고 생각하며, 비슷한 주제로 논문을 쓸 때 반드시 baseline 으로 포함할 필요가 있겠습니다.