범용 분할 학습기로서의 확산 모델

이미지 생성에 특화된 확산 모델이 가진 시각적 이해 능력을 극대화하여 별도의 복잡한 구조 없이도 정교한 이미지 분할이 가능함을 입증했다. 이는 생성 모델이 단순한 화가 역할을 넘어 의료, 위성, 농업 등 다양한 분야에서 범용적인 시각 이해 도구로 활용될 수 있는 새로운 방향을 제시한다.

기존의 확산 모델 기반 분할 방식은 모델 내부의 Attention Map을 추출하여 후처리하는 방식에 의존했으나, 이는 해상도가 낮고 노이즈가 많아 경계가 불분명한 한계가 있었다. DiGSeg는 확산 모델이 이미 이미지의 구조적 특징을 잘 이해하고 있다는 점에 착안하여, 모델이 직접 분할 마스크를 '생성'하도록 학습시킨다.

먼저 RGB 이미지와 정답 마스크를 동일한 잠재 공간(Latent Space)으로 투영한다. 이때 확산 모델은 이미지 생성 시 노이즈를 제거하며 구체적인 형상을 만들어가는 과정을 학습하는데, DiGSeg는 이 과정을 역이용하여 노이즈가 섞인 마스크 잠재값에서 원본 마스크를 복원하도록 유도한다.

결과적으로 모델은 텍스트 프롬프트와 입력 이미지의 특징을 결합하여 픽셀 단위의 정밀한 마스크를 직접 출력하게 된다. 이는 생성 모델의 강력한 시각적 사전 지식(Prior)을 유지하면서도 분할 작업에 필요한 정교한 경계 인식 능력을 명시적으로 학습시키는 결과를 낳는다.

DiGSeg는 사전 학습된 Stable Diffusion v2의 U-Net을 기반으로 세 가지 핵심 모듈을 구성한다. 첫째, Visual Latent Pathway는 VAE 인코더를 사용하여 이미지 x와 마스크 y를 잠재 표현 zx, zy로 변환한다. 이때 단일 채널인 마스크를 3채널로 복제하여 VAE와의 호환성을 유지한다.

둘째, CLIP-Aligned Text Conditioning 모듈은 텍스트 프롬프트를 CLIP 텍스트 인코더로 임베딩한 후, U-Net의 여러 스케일에 걸쳐 Cross-Attention을 통해 주입한다. [텍스트 임베딩 → Cross-Attention 연산 → 시각 특징과 결합] 과정을 통해 모델은 특정 단어가 이미지의 어느 영역에 해당하는지 학습한다.

셋째, Segmentation-Consistent Denoising 과정에서는 노이즈가 추가된 마스크 잠재값 zyt를 입력받아 원래의 zy를 예측한다. 이때 MSE Loss를 사용하여 [예측된 노이즈와 실제 추가된 노이즈의 차이]를 계산하고, 이 오차를 줄이는 방향으로 U-Net의 가중치를 업데이트하여 분할 성능을 최적화한다.

DiGSeg는 주요 세그멘테이션 벤치마크에서 기존 SOTA 모델들을 압도하는 성과를 거두었다. ADE20K 데이터셋에서 mIoU 58.6점을 기록하며 이전 최고 기록 대비 +1.3점 향상되었으며, COCO 데이터셋에서도 50.8점을 달성하여 +2.1점의 성능 개선을 보였다.

개방형 어휘(Open-Vocabulary) 설정에서도 A-847, PC-459 등 5개 벤치마크 모두에서 최고 성능을 기록했다. 특히 데이터 효율성 측면에서 전체 데이터의 50%만 사용하고도 전체 데이터를 사용한 기존 모델들과 대등한 성능을 유지하는 강력한 일반화 능력을 입증했다.

도메인 전이 실험에서는 추가적인 구조 변경 없이도 의료(REFUGE-2), 위성(DeepGlobe), 농업(Phenobench) 데이터셋에서 우수한 성능을 보여 범용 분할 학습기로서의 가치를 증명했다. 특히 DeepGlobe 도로 분할 작업에서는 IoU 65.78점을 기록하며 전문 모델들을 제치고 가장 높은 정확도를 보였다.

DiGSeg 아키텍처는 SDv2의 U-Net 백본을 활용하며 VAE 인코더/디코더는 고정(Frozen) 상태로 유지한다. 학습 시에는 U-Net의 Cross-Attention 블록과 일부 MLP 레이어만 미세 조정하여 연산 효율성을 높였다. 텍스트 임베딩은 CLIP ViT-L/14를 통해 추출되어 1/4, 1/8 해상도의 어텐션 블록에 주입된다.

수학적으로는 조건부 확산 과정 p(y|x)를 모델링하며, 추론 시에는 DDIM 샘플러를 사용하여 1~50단계 내에서 마스크를 생성한다. 특히 DDIM-trailing 스케줄을 도입하여 단 1회의 디노이징 단계만으로도 준수한 분할 결과를 얻을 수 있도록 최적화했다.

노이즈 전략에서는 Annealed Multi-scale Noise를 사용하여 학습 초기에는 미세한 경계(고주파)에 집중하고, 학습이 진행됨에 따라 전체적인 객체의 형태(저주파)를 정교화하도록 설계했다. 이는 확산 모델이 생성 작업에서 이해 작업으로 전이될 때 발생하는 불일치 문제를 효과적으로 해결한다.

의료 도메인(REFUGE-2 등)에서의 성능은 일반 자연 이미지에 비해 상대적으로 낮게 나타났다. 이는 백본으로 사용된 CLIP이 전문적인 의료 용어 및 시각적 특징에 대한 이해도가 부족하기 때문인 것으로 분석되며, 특수 도메인에 대한 추가적인 정렬 연구가 필요함을 시사한다.