확산 기반 프로토타입 강화로 한 번 학습 의료 영상 분할 혁신

초록

DPL은 단일 지원 이미지에서 추출한 프로토타입을 확산 모델로 확장·다양화하고, 공간 조건을 주입해 해부학적 일관성을 유지한다. 전·후 확산 과정을 통해 생성된 다변량 프로토타입을 보수적 융합하여 원본 프로토타입의 신뢰성을 보존하면서 다양성을 높인다. MRI·CT 복부 데이터셋에서 기존 방법들을 크게 앞서며, 1‑shot 의료 영상 분할의 새로운 기준을 제시한다.

상세 분석

본 논문은 의료 영상의 고유한 해부학적 변이성과 라벨 부족 문제를 해결하기 위해 프로토타입을 확률 분포로 모델링하는 Diffusion Prototype Learning(DPL)을 제안한다. 기존 프로토타입 기반 1‑shot 방법은 지원 이미지의 특징을 단순 평균해 하나의 고정점으로 만든다. 이는 환자마다 다른 장기 형태·크기·조직 대비를 반영하지 못해 일반화에 한계를 보인다. DPL은 이러한 한계를 세 단계 혁신으로 극복한다.

첫째, 확산 기반 프로토타입 강화 모듈은 초기 프로토타입 p₀를 정규 가우시안 노이즈와 결합해 전방 확산(pₜ = √ᾱₜ·p₀ + √(1‑ᾱₜ)·ε)으로 다양한 변형을 만든다. 여기서 T=20, 코사인 β 스케줄을 사용해 노이즈 강도를 조절한다. 둘째, 공간‑조건 인코딩은 프로토타입의 통계적 특성(중심 좌표, 분산, compactness, elongation)을 4‑차원 기하학적 벡터 g_geom으로 변환하고, 이를 f_enc를 통해 c_spatial에 매핑한다. 역방향 확산 단계에서 ε̂ₜ = ε_θ(pₜ, t) + αₜ·c_spatial 형태로 주입함으로써 초기 단계에서 강한 공간 가이드를 제공하고, 후반부에서는 점차 감소시켜 세밀한 디테일 복원을 돕는다.

셋째, 보수적 융합 전략은 원본 프로토타입 p₀와 확산으로 얻은 p_diff를 가중합(w_fidelity·p₀ + w_diversity·p_diff) / (w_fidelity + w_diversity) 형태로 결합한다. 여기서 w_fidelity와 w_diversity는 시그모이드 활성화된 학습 파라미터이며, 모델이 상황에 따라 신뢰성(원본 보존)과 다양성(변형 활용) 사이의 최적 균형을 스스로 찾게 만든다.

학습 목표는 세 가지 손실의 가중합 L_total = L_seg + L_align + β·L_diffusion이다. L_seg는 강화된 프로토타입으로 쿼리 이미지와 매칭해 얻은 픽셀‑레벨 교차 엔트로피, L_align은 지원·쿼리 양쪽에서 생성된 강화 프로토타입 간 L2 정렬, L_diffusion은 전통적인 DDPM 손실로 노이즈 예측 정확도를 촉진한다. β=0.02로 설정해 확산 정규화가 과도하게 세분화 손실을 압도하지 않도록 조정하였다.

실험에서는 복부 MRI(20건)와 CT(30건) 데이터셋에 1‑shot 설정으로 3 epoch 훈련 후 Dice 평균을 보고하였다. DPL은 기존 프로토타입 기반(PANet, ALPNet), 쿼리‑가이드(Q‑Net) 등과 비교해 평균 Dice를 4~7%p 상승시켰으며, 특히 신장·간·비장 등 형태가 크게 변하는 장기에서 큰 개선을 보였다. Ablation 연구에서는(1) 공간‑조건 없이 확산만 적용했을 때 성능 저하, (2) 보수적 융합을 제거하고 단순 평균을 사용했을 때 원본 프로토타입의 신뢰성 손실, (3) β 값을 크게 늘렸을 때 과도한 노이즈가 세그멘테이션 정확도를 감소시키는 현상을 확인했다.

이러한 설계는 학습‑추론 일관성을 유지한다는 점에서도 의의가 있다. 동일한 확산‑강화‑융합 파이프라인을 훈련과 추론 모두에 적용함으로써, 테스트 시 프로토타입 분포가 훈련 시와 동일한 방식으로 생성·융합된다. 결과적으로 모델은 라벨이 전혀 없는 새로운 환자에서도 안정적인 프로토타입을 생성해 높은 분할 정확도를 달성한다.

전반적으로 DPL은 확산 모델을 데이터 생성이 아닌 특징 공간 탐색 도구로 전환함으로써, 의료 영상 1‑shot 세그멘테이션에서 필수적인 해부학적 다양성을 효과적으로 포착한다. 향후 다중‑shot, 다기관 데이터, 그리고 3‑D 전역 컨텍스트 통합 등으로 확장 가능성이 크다.