MICCAI 2025에 accept된 AVDM: Controllable Adversarial Diffusion Model for Vessel-to-Volume Synthesis라는 논문이다. diffusion training framework에 adversarial supervision을 통합시키는 연구인데 흥미로워서 리뷰해보려 한다.
Introduction
혈관 segmentation은 medical image analysis에서 중요한 task이다. 특히 뇌졸증, 뇌동맥류, 관상동맥질환 같은 장애 진단과 치료에 중요한 역할을 한다. medical image analysis에 진보에도 불구하고, 정확하고 robust한 segmentation은 여전히 어려운 문제로 남아있다. 이것은 주로 정교하고 미세한 혈관 구조의 복잡성과, 영상 촬영 장비나 프로토콜에 따라 달라지는 SNR, 혈관 패턴, 배경 조직의 차이로 발생하는 큰 domain gap 때문이다. 이러한 변이는 지도학습 기반 방법들이 보지 못한 3차원 혈관 domain에 대해 일반화 능력을 심각하게 제한한다. 따라서 연구자와 임상의들은 혈관 영상을 분석하기 위해, 처음부터 픽셀 단위로 일관된 label을 수작업으로 annotation하는 과정에 의존할 수 밖에 없다고 한다.
최근에, 실제 의료 이미지들의 이용 제한으로 인해 GAN이나 DDPM과 같은 생성형 모델이 의료 영상 합성에 널리 적용되고있다. 그러나 존재하는 표준 모델들은 직접적인 3D 혈관 마스크들을 생생해내는데 실패한다. 강력한 사전학습된 Latent diffusion models (LDMs)을 vessel-to-volume 합성에 적용(fine-tuning)시킬때 real-word noise로부터의 domain shift가 나타난다. 이는 모델이 혈관의 고유한 특성과 이질적인 기하학적 구조를 고려하지 않아 mask controllability를 읽기 때문이라고 한다. 그 결과, 모델은 기존 훈련 유사하거나 중복되는 샘플만 생성하게 되며, 이는 다양한 데이터가 필수적인 downstream task에서의 유용성을 떨어뜨린다. annotated된 데이터는 오직 real-word 환경을 부부적으로만 반영하기 때문에, 추가적인 다양성을 제공하여 실제 데이터를 보완할 합성된 샘플이 필요하다.
Adversarial 예시들은 조건 입력과의 alignment 불량으로 인해 발생하는 노이즈 예측 오류를 완화한다. 이러한 맥락에서, 저자들은 diffusion training process에서 제한된 adversarial supervision을 도입하는 Controllable Adversarial Diffusion Modeling (AVDM)을 제안한다. AVDM은 image들을 저차원 latent 공간으로 매핑하고 gradient-optimized directions으로 이동시켜 높은 structural fidelity와 domain adaptability을 가진 adversarial examples를 생성한다. 생성된 adversarial examples를 다시 원래 이미지 공간으로 되돌릴 때, texture와 구조적 일관성을 유지하기 위해 semantic segmentation 기반 discriminator를 사용한다. 이 discriminator는 조건 정보를 활용하여 diffusion model에 픽셀 단위의 피드백을 제공함으로써, 원래 혈관 이미지의 특징이 그대로 보존되도록 돕는다. 이 통합은 adversarial examples들이 실제처럼 정렬되어지고 정확하게 original mask label에 해당되도록 보장한다.
저자들은 혈관 마스크를 조건으로 주어 high-fidelity volumetric image들은 생성해내는 diffusion-base framework인 AVDM을 제안한다. 이 모델은 mask-conditional 의료 영상 생성 분야에서 상당한 발전을 나타내고, 사실성을 유지하면서 해부학적 정확도를 달성한다고 한다. 저자들은 다양한 데이터셋에 걸쳐 AVDM의 성능을 평가했고, 그 결과 AVDM이 해부학적 마스크에 대해 fidelity 면에서 최신 생성 모델들을 능가했다고 한다.
Method
Projecting Volumetric to Diffusion Latent
다양한 도메인들에 대해 domain shift를 처리할 수 있는 adversarial examples를 생성하기 위해, 저자들은 volumetric image들은 저차원 공간으로 project하기 위해 Stable Diffusion과 같은 생성 모델을 활용한다.
이 저차원 공간을 최적화함으로써, 일관성있고 다양한 adversarial samples을 합성하기 위해 이미지 공간으로 다시 투영되는 adversarial representation을 효율적으로 찾아낸다.
주어진 input volumetric image $x_0$에 대해 diffusion inversion을 사용해 저차원 공간으로 매핑한다. inversion은 schedule $\{\beta_1,...,\beta_T\} ∈ (0,1)$을 활용한다, $\alpha_t = \prod_{i=1}^{t}(1 - \beta_i)$이고 t는 timestep이다. AVDM은 foward diffusion procedure을 따른다.
AVDM의 framework는 denoising process의 역 방향으로 동작한다. $x_{t} → x_0$가 아닌 $x_0 → x_{t}$로. 점진적으로 random noise에서 clean image로 denoising하는 것 대신에, diffusion inversion을 통해 input image $x_0$를 특정 time point에 $x_T$ 잠재 공간으로 투영시킨다.
Classifier-free guidance method는 unconditional 예측과 원활하게 특정 condition이 주어진 상태의 예측 모두 생성한다. 그리고 이 두 예측을 guidance scale factor w와 null text embedding ∅를 주었을 때, timestamp t에서 classifier-free guided 노이즈 예측은 다음과 같이 계산된다:
Stable Diffusion의 표준 세팅 값으로 $\omega=7.5$가 사용된다. DDIM sampling의 reverse process에서 classifier-free guidance는 학습된 노이즈 분포로부터의 편차를 일으켜 시각적 이상 현상을 만들어내고, 이는 결과물의 사실감을 저하시킬 수 있다. 이 문제를 완화하기 위해서, 저자들은 두 분포 일관성과 해부학적 충실도를 보존하기 위해 time-dependent null embedding ∅을 학습시킨다. 처음에는, $\omega=1$로 DDIM inverse sampling process를 실행하여 {x₀*, …, xₜ*}라는 연속적인 잠재 표현들을 얻는데, 이는 x₀* = x₀에서 시작한다. 그 후, t = {T, …, 1}에 대해 최적화 과정을 수행하는데, 이때 ω = 7.5를 사용하고 x̄_T = x*_T로 설정한다:
쉽게 이해하기 위해, $x_{t−1}(\bar{x_t},t,C,ϕ_t)$는 DDIM sampling 단계를 나타내고 $\bar{x_t}$는 input latent, $ϕ_t$는 null text embedding, 그리고 C는 text embedding을 나타낸다. 각 스텝이 완료되면, $\bar{x}_{t-1}$은 다음과 같이 업데이트 된다.
결국, 최적화된 null text embedding $ϕ_t$를 통해 $\bar{x_t} = x^*_T$라는 latent 표현에 도달할 수 있다. 이 latent는 diffusion 모델에 의해 생성된 것이며, low dimensional manifold에서 적대적 이미지를 생성하는 데 활용된다.
Discriminator Supervision with Segmentation Alignment Process
latent 표현을 기반으로, 저자들은 AVDM의 diffusion model의 denoising 과정을 다음과 같이 공식화한다. U-Net denoise $\epsilon_\theta$는 MSE 목적함수를 통해 추가된 노이즈를 추정하도록 학습된다.
노이즈 이미지 $x_t$, time step t 외에도 U-Net은 혈관 입력 y를 추가 입력으로 사용한다. y는 원본 이미지 $x_0$로부터 추출된 혈관 정보를 포함하고 있기 때문에, 노이즈 추정을 단순화하고 암묵적으로 디노이징 단계에서 이미지 합성을 guide한다. $x_t$와 $\epsilon_\theta$로부터 denoised된 clean image $\hat{x}^{(t)}_0$를 생성해낼 수 있다.
그러나, 목적함수 $L_{noise}$에는 layout fidelity에 대한 명시적인 감독(supervision)이 없기 때문에, 생성된 이미지 $\hat{x}^{(t)}_0$와 구조적 조건 y 사이의 불일치가 발생하는 경우가 많다. 그러므로 저자들은 $\hat{x}^{(t)}_0$에 직접적인 감독을 가하여 layout alignment를 보장하려한다. 또한 단순한 정합성 뿐 만 아니라 다양성을 확보하기 위해, segmenter를 U-Net과 함께 훈련 가능한 상태로 둔다. 이에 착안하여, 저자들은 U-Net과 segmenter 간의 adversarial game을 공식화 한다. 구체적으로, segmenter는 실제 이미지의 각 픽셀에 대한 클래스 레이블을 분류하도록 훈련된 discriminator의 역할을 수행한다.
쌍으로 주어진 ground-truth label map과 함께, 판별기는 U-Net이 생성한 가짜 이미지를 추가적인 fake class로 분류한다. 판별기의 주요 임무는 semeantic segmentation을 수행하는 것이므로, 목적함수는 cross-entropy loss를 사용한다.
Image Generation with Controllable Adversarial
adversarial latent 표현을 도출한 후, reverse diffusion process는 최종 advesarial 예시를 생성해내는데 적용될 수 있다. 저자들은 reverse 과정에서 ControlNet을 통합함으로써, 향상된 공간적 제어(spatial shaping) 능력을 얻는다.
ControlNet은 Stable Diffusion 모델의 denoising U-Net 아키텍처 내에서 task 특화 조건부 정보의 정밀도를 개선한다. Stable Diffusion 모델은 U-Net 구조를 기반으로 하며, 인코더(encoder), 미들 블록(middle block), 디코더(decoder)로 구성되어 있고, 각 부분은 12개의 블록으로 이루어진다.
ControlNet은 이 구조를 확장하기 위해, 원래 모델의 12개 인코더 블록과 미들 블록을 학습 가능한 형태로 복제한다. 이 블록들은 네 가지 해상도 수준(resolution level)에 걸쳐 분포하며, 각 수준에는 세 개의 블록이 포함된다.
이렇게 학습 가능한 블록들의 출력은, diffusion U-Net의 12개 skip connection과 미들 블록에 매끄럽게 통합된다. 이 통합은 모델이 이미지의 특성을 더욱 정밀하게 조정하고 맞춤화할 수 있는 능력을 크게 향상시킨다.
Experiments
저자들은 공개적으로 이용할 수 있는 3D 혈관 datasets을 사용하여 데이터를 합성하는 실험을 진행했다. 데이터셋은 127 TOF-MRA volume으로 이루어져있고, ADAM, CoW, IXI-HH 데이터 센터에서 수집되어졌다. 이 데이터셋들은 취득 장치, 스캐닝 프로토콜, 자기장 세기, 그리고 공간 해상도 측면에서 상당한 다양성을 보인다. 추가적으로, 저자들은 downstream task를 위해 zero-shot, one-shot, few-shot segmentation 작업을 평가하였으며, 이 과정에서 세 개의 unseen 3D 혈관 데이터셋인 IXI-Guys, ICBM, LocH를 활용하였다.
이 맥락에서, 전체 뇌 스캔으로부터 크기 $128^3$의 3D 패치를 추출하였다. 우리는 혈관 밀도 기반(vessel-density-driven) 접근법을 사용하였는데, 이는 ground-truth 혈관 마스크에서 식별된 혈관 구조의 존재 여부에 따라 패치를 선택하는 방식이다. 슬라이딩 윈도우 알고리즘이 3D 볼륨을 스캔하면서, 혈관 밀도가 높은 영역(패치 내 혈관 voxel의 비율로 계산됨)을 우선적으로 선택하였다. 각 환자마다, 혈관 밀도 순으로 상위 50개의 패치를 선택하였고, 각 패치가 최소 20% 이상의 혈관 커버리지를 가지도록 하여 표현력과 다양성의 균형을 맞추었다.
adversarial example 생성을 위한 과정에서는, 우리는 DDIM(Denoising Diffusion Implicit Models)을 T=100 단계로 설정하여 사용하였다.
