Diffusion Model as a Generalist Segmentation Learner

"Diffusion Model as a Generalist Segmentation Learner"(DiGSeg) 논문에 대한 상세한 기술적 요약입니다.

1. 문제 제기

현재 분할 시스템은 매우 파편화되어 있습니다. 모델들은 일반적으로 특정 작업 (시맨틱, 인스턴스, 팬옵틱), 특정 어휘 (클로즈드셋 대 오픈보카불러리), 또는 특정 도메인 (자연, 의료, 원격 감지) 에 맞춰 설계됩니다. 이러한 설정 간 전환은 종종 아키텍처의 재설계나 처음부터의 재학습을 요구합니다.

최근 연구들은 어텐션 맵을 추출하여 사전 훈련된 텍스트 - 이미지 확산 모델을 분할에 재사용하려는 시도를 해왔으나, 이러한 접근법은 다음과 같은 중대한 한계를 겪고 있습니다:

• 노이즈와 불일치: 원시 어텐션 맵은 종종 저해상도이며 노이즈가 많고 레이어 간 일관성이 부족합니다.
• 후처리 의존성: 사용 가능한 마스크를 생성하기 위해 종종 무거운 휴리스틱 후처리가 필요합니다.
• 일반화 부족: 기존 방법들은 단일 작업에 집중하는 경향이 있으며, 아키텍처 커스터마이징 없이 다양한 도메인과 오픈보카불러리 쿼리에 걸쳐 견고한 일반화를 입증하지 못합니다.

이 논문이 다루는 핵심 질문은 다음과 같습니다: 단일 통합 확산 기반 모델이 도메인별 아키텍처 변경 없이 시맨틱, 오픈보카불러리, 그리고 크로스도메인 작업을 처리할 수 있는 범용 분할 학습기로 기능할 수 있는가?

2. 방법론: DiGSeg 프레임워크

저자들은 사전 훈련된 확산 모델 (특히 Stable Diffusion v2) 을 생성적 사전 지식에서 통합 미세 조정 프로토콜을 통해 명시적인 분할 학습기로 변환하는 DiGSeg프레임워크를 제안합니다.

핵심 아키텍처

이 프레임워크는 세 가지 주요 구성 요소로 이루어져 있습니다:

1. 시각 잠재 경로 (Visual Latent Pathway):

   • 이미지를 원시 픽셀로 취급하는 대신, 입력 RGB 이미지 ($x$) 와 정답 분할 마스크 ($y$) 를 Stable Diffusion 의 고정된 VAE 인코더를 사용하여 컴팩트한 잠재 공간으로 인코딩합니다.
   • 3 채널 VAE 와의 호환성을 유지하기 위해 단일 채널 마스크는 세 채널로 복제됩니다.
   • 이는 공간 구조를 보존하면서 잠재 공간에서 정교한 대응 관계의 효율적인 학습을 가능하게 합니다.

2. CLIP 정렬 텍스트 조건부 (CLIP-Aligned Text Conditioning):

   • 오픈보카불러리 기능을 활성화하기 위해 고정된 CLIP 텍스트 인코더가 클래스 이름이나 자연어 설명에 대한 임베딩 ($t_{clip}$) 을 생성합니다.
   • 이러한 임베딩은 여러 노이즈 제거 스케일에서 크로스 어텐션을 통해 확산 U-Net 에 주입됩니다. 이를 통해 모델은 텍스트 쿼리를 전역적 및 지역적으로 진화하는 시각적 표현과 정렬할 수 있습니다.

3. 분할 일관성 노이즈 제거 U-Net:

   • 학습 목표: 모델은 이미지 잠재 ($z_x$) 와 텍스트 임베딩 ($t_{clip}$) 에 조건부인 분할 잠재 ($z_y$) 에 추가된 노이즈를 예측하도록 훈련됩니다.
   • 생성적 공식화: 이전 작업들이 확산을 특징 추출기로 사용하는 것과 달리, DiGSeg 는 분할을 조건부 노이즈 제거 생성 작업 ($p_\theta(y|x)$) 으로 취급합니다. U-Net 은 무작위 노이즈에서 분할 일관성 잠재 표현으로 노이즈를 제거하는 법을 학습합니다.
   • 노이즈 전략: 새로운 어닐링 다중 스케일 노이즈 스케줄이 도입됩니다. 고주파수 교란이 초기 단계에서 지배적이어서 정교한 구조를 학습하게 하고, 저주파수 구성 요소는 나중에 영향력을 발휘하여 시맨틱 일관성을 보장합니다.

추론 과정

1. 입력 이미지가 잠재 표현으로 인코딩됩니다.
2. 분할 잠재는 가우시안 노이즈로 초기화됩니다.
3. 텍스트 조건부 하에서 U-Net 은 잠재 공간을 반복적으로 노이즈 제거합니다 (효율을 위해 'trailing' 타임스텝 전략을 사용하는 DDIM 샘플러 사용).
4. 최종 깨끗한 잠재는 고정된 VAE 디코더를 사용하여 픽셀 공간으로 디코딩되어 분할 마스크를 생성합니다.
5. 테스트 시간 앙상블: 공간 일관성을 향상시키고 아티팩트를 줄이기 위해 잠재 공간에서 서로 다른 노이즈 시드로 수행한 여러 추론 패스를 평균화합니다.

3. 주요 기여

• 통합 프레임워크: DiGSeg 는 클로즈드보카불러리, 오픈보카불러리, 그리고 크로스도메인 작업을 작업별 아키텍처 변경 없이 처리할 수 있는 범용 분할 학습기로 사전 훈련된 확산 백본을 성공적으로 재사용한 최초의 프레임워크입니다.
• 직접 잠재 학습: 노이즈가 많은 어텐션 맵에 의존하는 대신, 노이즈 제거 U-Net 을 미세 조정하여 분할 일관성 잠재를 직접 생성하도록 함으로써 무거운 후처리 없이 고품질 마스크를 생성합니다.
• 다중 스케일 텍스트 조건부: 여러 노이즈 제거 스케일에 걸쳐 언어 특징을 주입하는 메커니즘을 도입하여 임의의 텍스트 프롬프트와 시각적 특징 간의 견고한 정렬을 가능하게 합니다.
• 새로운 노이즈 전략: 정교한 경계 학습과 전역 시맨틱 일관성 사이의 균형을 맞추는 어닐링 다중 스케일 노이즈 스케줄을 제안합니다.

4. 실험 결과

저자들은 DiGSeg 를 표준 벤치마크와 특수 도메인에서 평가하여 최첨단 (SOTA) 성능을 입증했습니다.

• 오픈보카불러리 분할:
  • A-847, PC-459, A-150, PC-59, Cityscapes에서 13 개의 베이스라인 (ODISE, OVSeg, CAT-Seg 포함) 을 능가했습니다.
  • 두 번째로 좋은 방법 대비 상당한 향상 (예: PC-59 에서 +2.8 mIoU, PC-459 에서 +2.2) 을 달성했습니다.
• 표준 시맨틱 분할:
  • COCO(50.8 mIoU) 와 ADE20K(58.6 mIoU) 에서 SOTA 를 달성하여 Mask2Former 와 OneFormer 와 같은 강력한 트랜스포머 기반 모델을 능가했습니다.
• 크로스도메인 전이:
  • 아키텍처 수정 없이 의료(REFUGE-2), 농업(PhenoBench), 원격 감지(DeepGlobe) 데이터셋으로 강력한 제로샷 전이를 입증했습니다.
  • 참고: 의료 도메인의 성능은 자연 장면보다 낮았는데, 이는 CLIP 의 사전 훈련 데이터에 있는 도메인 격차에 기인하지만 여전히 경쟁력 있습니다.
• 효율성 및 데이터 사용:
  • 모델은 매우 데이터 효율적이며, 훈련 데이터의 **50%**만으로 거의 동일한 성능을 달성합니다.
  • DDIM-trailing 을 통해 추론 속도가 향상되어 (1 단계에서 포화 도달) 하지만 확산의 특성상 피드포워드 모델보다는 여전히 느립니다.

5. 의의

이 작업은 시각 생성과 시각 이해 사이의 격차를 근본적으로 연결합니다.

• 패러다임 전환: 확산 모델이 합성만을 위한 것이라는 개념에 도전하여, 적절히 조건부화될 때 강력한 범용 판별 학습기로 기능할 수 있음을 입증합니다.
• 일반화: 대규모 생성 모델에 인코딩된 풍부한 공간 및 시맨틱 사전 지식이 다양한 도메인 (자연, 의료, 농업) 의 밀집 예측 작업으로 이전 가능함을 보여줍니다.
• 간소화: 복잡한 다단계 파이프라인이나 특수 헤드를 피함으로써 DiGSeg 는 분할을 위한 간소화된 통합 인터페이스를 제공하며, 단일 생성 백본이 컴퓨터 비전 작업의 광범위한 배열을 처리할 수 있는 미래를 시사합니다.

결론적으로, DiGSeg 는 잠재 조건부 프로토콜과 텍스트 정렬로 미세 조정된 현대적인 확산 백본이 오픈보카불러리 및 크로스도메인 시나리오에서 우수한 유연성을 제공하면서도 전문 분할 모델을 능가할 수 있음을 입증합니다.