한 줄 요약
CONFLUX는 3D 흉부 CT 이미지를 생성하는 조절 가능한 라티언트 디퓨전 모델로, 강화 학습을 통한 후학습으로 조건 부합도를 47% 개선한다.
핵심 기여도
- 3D 변분 자동인코더(VAE)와 라티언트 공간에서 생성하는 리니어 플로우 변형기(rectified-flow transformer)를 결합한 CONFLUX 모델 제안.
- 18개 이상 병변 진단, 성별, 연령, 재구성 커널 등 구조화된 메타데이터를 조건으로 사용.
- 3D 볼륨 생성 기준에서 FID 32.3 (MAISI 대비 74.6 개선) 달성.
- 강화 학습 후학습(GRPO)을 통해 생성 볼륨의 조건 부합도를 47% 향상.
핵심 아이디어
CONFLUX는 3D 흉부 CT 이미지를 생성하는 데 있어 높은 품질과 조절 가능성을 동시에 달성하기 위해 라티언트 공간에서 생성하는 디퓨전 모델을 채택했다. 기존 3D 생성 모델은 높은 해상도와 볼륨 구조를 유지하면서 조건에 따라 생성하는 것이 어려웠다. 이를 해결하기 위해 CONFLUX는 3D VAE를 사용해 볼륨을 압축하고, 라티언트 공간에서 리니어 플로우 변형기를 통해 생성한다. 생성 조건은 adaptive layer normalization을 통해 18개 이상 병변 진단, 성별, 연령, 재구성 커널 등 구조화된 메타데이터로 전달된다.
또한, 생성 볼륨이 요청된 조건을 얼마나 정확히 반영하는지를 강화 학습(GRPO)을 통해 후학습 단계에서 개선한다. 이는 기존의 확률 기반 손실 함수가 조건 부합도를 보장하지 못한다는 문제를 해결하기 위한 접근이다.
기술적 접근법
- **모델 구조**: 3단계 라티언트 디퓨전 모델 (3D VAE + 리니어 플로우 변형기 + 강화 학습 후학습).
- **데이터 전처리**: CT-RATE 데이터셋 사용, 볼륨 크기 216×176×200, 강도 정규화 [-1, 1].
- **조건 조절**: adaptive layer normalization을 통해 18개 이상 병변 진단, 성별, 연령, 재구성 커널 조건 전달.
- **후학습 알고리즘**: group-relative policy optimization (GRPO)를 사용한 강화 학습.
- **학습 데이터**: 40,800개의 품질 통과 볼륨 중 18,417개 사용, 3,039개는 검증용으로 분리.
주요 결과
- **FID 성능**: 32.3 (MAISI 대비 74.6 개선).
- **조건 부합도 개선**: 강화 학습 후학습으로 47%의 부족함을 감소.
- **데이터셋**: 약 200,000개의 흉부 CT 볼륨을 생성한 합성 데이터셋 공개.
의의 및 한계
CONFLUX는 3D 흉부 CT 생성에서 높은 품질과 조절 가능성을 동시에 달성한 첫 번째 라티언트 디퓨전 모델로, 병원 데이터 부족 문제를 해결하고 연구용 데이터셋을 확장하는 데 기여한다. 특히, 강화 학습 후학습(GRPO)을 3D 의료 이미지 생성에 처음 적용한 점이 학술적으로 의미가 크다.
그러나, 생성 볼륨의 조건 부합도는 NLP 추출된 라벨에 기반하기 때문에 전문가가 직접 검증한 라벨과의 차이가 있을 수 있다. 또한, 생성된 볼륨이 실제 병리학적 진단에 얼마나 정확하게 반영되는지는 추가 연구가 필요하다.
실용적 활용
CONFLUX는 의료 데이터 부족한 집단을 보완하거나, 프라이버시 제약 하에서 데이터를 공유할 때 유용하게 사용될 수 있다. 또한, 다양한 병변 조건을 반영한 대규모 합성 데이터셋은 의료 AI 모델 훈련 및 임상 연구 설계에 활용 가능하다.