CONFLUX: A Latent Diffusion Model for 3D Chest-CT Synthesis with RL Post-Training

기술 요약: CONFLUX

문제 정의

제어 가능한 3D 의료 영상 생성 모델은 코호트 증강, 개인정보 보호를 위한 데이터 공유, 반사실적 연구 설계(counterfactual study design)와 같은 작업을 위해 특정 임상 속성을 가진 볼륨을 합성하는 것을 목표로 합니다. 그러나 고충실도(high-fidelity)의 네이티브 3D 합성을 수행하면서 동시에 요청된 조건(conditioning)에 충실하게 만드는 것은 여전히 큰 과제입니다. 기존 모델들은 표준 학습 목적 함수(예: flow matching)가 개별 샘플의 속성 실현보다는 집합적 분포의 현실성을 최적화하기 때문에, 생성된 개별 샘플이 요청된 특정 임상 소견을 반드시 나타내도록 보장하는 데 어려움을 겪는 경우가 많습니다. 이러한 격차는 구조화된 방사선학적 메타데이터에 의해 설명되어야 하는 흉부 CT에서 특히 두드러지며, 이는 높은 해상도와 볼륨 일관성을 요구합니다.

방법론

CONFLUX는 변이형 오토인코더(VAE), 직교 흐름 트랜스포머(rectified-flow transformer), 그리고 강화 학습(RL) 사후 학습 단계를 통합한 3단계 잠재 확산 프레임워크입니다.

1. 잠재 오토인코더 (1단계)

직접적인 복셀 공간(voxel-space) 생성의 계산 비용을 완화하기 위해, 3D 컨볼루션 VAE를 사용하여 전처리된 CT 볼륨($216 \times 176 \times 200$)을 저해상도 잠재 공간($16 \times 27 \times 22 \times 25$)으로 압축합니다.

• 아키텍처: 인코더와 디코더는 최저 해상도에서 3D 셀프 어텐션을 사용하는 그룹 정규화 잔차 블록(group-normalized residual blocks)을 활용합니다.
• 학습: 모델은 $\ell_1$ 재구성 손실, 쿨백-라이블러(Kullback–Leibler) 정규화 항, 그리고 삼평면(tri-planar) LPIPS 지각 손실을 최소화합니다.
• 정규화: 학습 후, flow matching을 용이하게 하기 위해 잠재 공간을 근사 단위 분산(near-unit variance)으로 정규화합니다.

2. 조건부 직교 흐름 트랜스포머 (2단계)

단일 스트림 트랜스포머는 학습된 잠재 공간 내에서 가우시안 노이즈로부터 잠재 임베딩을 생성합니다.

• 흐름 매칭(Flow Matching): 모델은 직교 흐름 목적 함수 하에서 학습되며, 노이즈에서 데이터 잠재값으로 이어지는 직선 경로를 따라 일정한 속도 벡터장을 예측하는 법을 배웁니다.
• 조건화(Conditioning): 구조화된 방사선학적 메타데이터(18개의 이진 이상 소견, 성별, 연령대, 재구성 커널)는 42차원 벡터로 결합됩니다. 이 벡터는 적응형 레이어 정규화(adaLN-zero)를 통해 트랜스포머를 조절합니다.
• 샘플링: 생성은 시간 이동 오일러 그리드(time-shifted Euler grid)를 사용하여 $t=1$에서 $t=0$까지 확률 흐름 ODE를 적분하는 과정을 거치며, 이후 동결된 VAE 디코더로 디코딩됩니다.

3. 강화 학습 사후 학습 (3단계)

집합적 분포 매칭과 개별 샘플의 속성 충실도 사이의 괴리를 해결하기 위해, CONFLUX는 그룹 상대적 정책 최적화(GRPO)를 채택합니다.

• 동기: 표준 flow matching은 특정 요청된 소견이 단일 생성 볼륨에 실제로 나타나는지 여부를 명시적으로 보상하지 않습니다.
• 메커니로즘: 결정론적 ODE 샘플러를 확률적 샘플러(SDE)로 교체하여 정책 경사 업데이트를 위한 다루기 쉬운 로그 확률(log-probabilities)을 제공합니다.
• 보상 신호: 실제 잠재 표현(latent representations)으로 학습된 독립적인 분류기가 요청된 소견과 실현된 소견 사이의 일치도를 기준으로 생성된 볼륨을 평가합니다. 보상은 분류기의 예측값과 조건 벡터 사이의 음의 가중 크로스 엔트로피로 정의됩니다.
• 최적화: 모델은 그룹당 프롬프트별 롤아웃(rollout)의 평균 보상을 빼고 표준 편차로 나누는 방식으로 그룹 상대적 이점(group-relative advantages)을 최대화하도록 미세 조정됩니다. 업데이트에는 사전 RL 참조 정책에 대한 KL-다이버전스 페널티를 사용하는 클리핑된 PPO 대리 함수(clipped PPO surrogate)를 사용하여 드리프트(drift)를 방지합니다.

주요 기여

1. 제어 가능한 3D 잠재 직교 흐름 모델: adaLN을 통해 구조화된 방사선학적 메타데이터에 직접 조건화되는 흉부 CT를 위한 네이티브 3D 아키텍처로, 강력한 3D 확산 베이스라인들과 경쟁할 만한 합성 품질을 달달성합니다.
2. 충실도를 위한 RL 사후 학습: 3D 의료 흐름 모델에 GRPO를 적용한 첫 사례입니다. 이 단계는 독립적인 판사 모델을 통해 검증된 바와 같이, 생성된 볼륨 내에서 요청된 소견의 신뢰성을 명시적으로 최적화합니다.
3. 오픈 릴리스: 학습된 모델과 약 200,000 볼륨 규모의 합성 흉부 CT 데이터셋을 공개합니다. 이 데이터셋은 실제 메타데이터 비율에 따라 조건화되었으며, 다양한 임상 소견을 포괄하여 코호트 증강 및 제어된 연구 설계에 활용될 수 있도록 설계되었습니다.

결과

합성 품질

홀드아웃 검증 세트($N=501$)에서 삼평면 FID(Fréchet Inception Distance)로 평가했을 때, CONFLUX는 기존 베이스라인을 크게 상회했습니다.

• 삼평면 FID: CONFLUX 32.3 vs MAISI 74.6 vs GenerateCT 145.4.
• 커버리지 및 밀도: CONFLUX 샘플은 MAISI에 비해 3.5배 더 밀도가 높고, 실제 매니폴드를 6배 더 많이 커버합니다.
• 다양성: 모델은 높은 다양성(MS-SSIM 0.507)을 유지하며 실제 데이터(0.41)와 유사한 수준을 보여 모드 붕괴(mode collapse)를 피했습니다.

조건화 충실도

GRPO를 통한 사후 학습은 독립적인 이미지 공간 분류기로 측정했을 때 요청된 속성과 실현된 속성 간의 일치도를 향상시켰습니다.

• 성능 향상: 평균 정밀도(AP)는 0.330에서 0.344로, AUROC는 0.684에서 0.699로 증가했습니다.
• 격차 회복: 절대적인 이득은 미미해 보일 수 있으나, 이는 베이스 모델과 실제 데이터 천장(AP 0.360) 사이의 성능 격차 중 47%를 회복한 것입니다.
• 특이성: 개선 사항은 최적화된 소견 그룹에 국한되었습니다. 가중치가 부여되지 않은 그룹(성별, 연령, 커널)은 노이즈 범위 내에서 변화가 없었습니다.

보상 분류기 성능

RL 보상 신호로 사용된 분류기는 실제 검증 분할에서 매크로 AUROC 0.793을 달성하였으며, 이는 CT-RATE 벤치마크에서 CT-CLIP 변형 및 CT-Net을 능가하는 성능으로, 압축된 잠재 공간에서도 유효한 보상 신호로서의 효용성을 입증했습니다.

의의 및 주장

본 논문은 CONFL서가 고충실도 생성과 속성 제어를 분리함으로써 제어 가능한 3D 의료 영상 합성에 있어 중요한 진전을 이루었다고 주장합니다. RL 사후 학습을 통합함으로써, 저자들은 표준 flow-matching 목적 함수로는 보장되지 않는 특성인 "개별 샘플 내 임상 속성의 충실도"를 직접 최적화할 수 있음을 입증했습니다. 대규모의 조건부 다양성을 갖춘 합성 데이터셋의 공개는 실제 데이터가 부족하거나 개인정보 보호가 제한적인 시나리오에서의 연구를 촉진하는 것을 목표로 합니다. 이 연구는 3D 의료 흐름 모델에 GRPO를 적용한 첫 사례로서, 더 신뢰할 수 있고 제어 가능한 임상 연구 설계 도구로 가는 경로를 제시합니다.