DiffDIS : High-Precision Dichotomous Image Segmentation via Probing Diffusion Capacity

ICLR 2025

 

Author : Qian Yu, Peng-Tao Jiang, Hao Zhang, Jinwei Chen, Bo Li, Lihe Zhang, Huchuan Lu
Dalian University of Technology, vivo Mobile Communication Co., Ltd

paper : https://proceedings.iclr.cc/paper_files/paper/2025/file/fcefc6e6370d49b9bad4fa4881aa1303-Paper-Conference.pdf

Github : https://github.com/qianyu-dlut/diffdis

 

GitHub - qianyu-dlut/DiffDIS: High-Precision Dichotomous Image Segmentation via Probing Diffusion Capacity (ICLR2025)

 

1. Introduction

High-accuracy dichotomous image segmentation (DIS, 이분 이미지 분할)은 natural scenes 내에서 카테고리에 상관없이 전경 객체를 정확하게 식별하는 것을 목표로 한다. DIS는 객체의 특성에 관계없이 고해상도의 세밀한 객체 분할에 집중한다. 이러한 객체들을 위해서는 더욱 정교한 feature selection과 알고리즘이 필요하다. 현재의 CNN/Transformer 기반 방법들은 강력한 특징 추출 능력을 갖추고 있지만, 고해상도 이미지에서 receptive field (수용영역) 확장과 세부 정보 보존 사이의 균형을 맞추는데 어려움을 겪는다. 즉, 넓은 범위의 문맥을 보면서도 작은 세부 정보를 유지하는 것이 Trade-off다.

 

Diffusion probabilistic models (DPMs, 디퓨전 모델)은 이미지 전체에 걸쳐 노이즈 변수를 예측함으로써, 목표 분포를 더욱 정확하게 학습하면서 global receptive field를 유지하는데 유망하다. 더욱이, stable diffusion (SD)은 수십억개의 이미지로 구성된 방대한 데이터셋을 통해 학습되어, 뛰어난 일반화 능력을 보여주며 풍부하고 다재다능한 image representation을 제공하여 거시적인 맥락과 미시적인 정밀도를 모두 요구하는 작업에 이상적인 후보이다. 최근 연구들의 발전은 diffusion이 고해상도 이미지 분할의 정확성과 견고성을 향상시키는 강력한 도구가 될 수 있음을 시사한다.

 

하지만, diffusion model을 DIS에 활용하기에는 몇가지 챌린지가 있다. 

(1) : DPM의 고유한 고차원 연속 특징 공간 (high-dimensional continuous feature space)은 이진 분할의 이산적 특성과 충돌하여 예측 과정에서 불일치가 발생할 수 있음. (연속적인 결과 -> 이진적인 mask로의 변환 과정에서의 불일치)

(2) : 확산 모델은 추론 시간이 오래 걸림. 확산 모델의 반복적인 특성으로 인해 제대로된 결과에는 일반적으로 100 step 이상 소요되며, 이미 데이터가 방대해서 느린 HR 영상의 추론 속도를 더욱 악화시킴.

(3) : 확산모델의 확률적 속성과 이미지 예측 task의 결정론적 결과 사이의 충돌. (같은 입력에 대해 항상 결과가 동일해야 함)

 

저자들은 앞서 언급한 문제점과 작업별 복잡성을 모두 해결하는 DiffDIS를 제안한다. DiffDIS는 다음 전략들을 통해 처리 속도 향상, 세부 정보 인식 강화, 결정성 증대 (higher determinisim)에 초점을 맞춘다.

 

1. VAE를 latent space <-> binary mask의 변환기로 사용. (feature space 안맞는 문제 해결)

2. Stable diffusion 처럼 pixel space가 아닌 latent space에서 diffusion 수행. (연산량 문제 해결)

3. pretrained SD-Turbo를 이용한 one-step denoising. (긴 추론시간 & 출력 변동성 문제 해결)

4. Mask와 edge를 동시에 예측하는 전략. (edge constraint이 출력범위를 제한 -> 얇고 복잡한 경계 검출에 효과적)

5. RGB latent representation을 확산 모델의 condition으로 사용. (RGB 이미지에서 seg를 만드는거니까... 어쩌면 당연한 일.)

6. Scale-wise Conditional Injection 제안. (RGB condition을 network 여러 지점에 해상도를 고려하여 주입. -> 해상도에 따라 local detail + global structure 둘 다 활용됨)

 

쉽게 정리하면, 고해상도 RGB 이미지와 정답 binary mask를 VAE latent 공간으로 압축하고, SD-Turbo 기반 diffusion network가 RGB 조건을 이용해 mask latent를 한 번에 예측하도록 한다. 이때 mask뿐 아니라 edge도 함께 예측하고, 여러 해상도에서 RGB 조건을 주입하여 전체 형태와 미세 경계를 동시에 보존한다.

 

저자들의 Contribution은 다음과 같다.

(1) : 확산 모델의 강력한 사전 지식을 활용하여 DIS 작업을 수행하는 DiffDIS 제안. 이 모델은 수용 영역 확장 - 세부 정보 보존 사이의 균형 문제를 해결함.

(2) : 확산 모델의 반복적인 특성을 간단한 one-step denoising을 구현하여 end-to-end framework로 변환함으로써 추론 속도를 크게 향상시킴.

(3) : 보조 edge generation task 도입. 이는 생성된 마스크의 결정성을 강화하면서 세부 표현의 미묘한 균형을 달성함.

(4) : DIS 벤치마크 데이터셋에서 거의 모든 지표를 능가하여 새로운 SoTA를 확립.

 

2. Method

2.1 Overall architecture

2.2 Edge-assisted training strategy

고해상도 작업, 특히 확산 기반 모델의 경우 잠재공간이 입력 크기보다 작게 축소되기 때문에 디테일한 특징을 정확히 표현하는것은 어렵다. 몇몇 연구들은 추가적인 edge constraints을 도입하여 mask의 boundary segmentation 능력을 향상시켰다. 저자들은 mask와 edge 둘 다를 위한 noise를 예측하게 함으로써, 보조적인 edge 정보로 network에 제약을 준다. 이 두 개의 예측 스트림은 동일한 네트워크 구조 내에서 작동하며 파라미터를 공유하고, 이들을 효과적으로 구분하기 위해 batch discriminative embedding (BDE) 이 적용된다. U-net 아키텍쳐, 특히 semantic information이 가장 집중된 mid-block에서, 저자들은 기존의 attention 모듈을 Detail-Balancing Interactive Attention으로 업그레이드하여 작업별 성능 향상을 구현하였다. 이 매커니즘은 mask stream과 edge stream의 attention 영역을 정렬하여, 두 스트림 간의 보다 효율적인 상호 작용과 상호 보완성을 촉진한다.

 

 

Batch-Discriminative Embedding (BDE)

저자들은 batch 내에 추가적인 discriminative labels d_lab을 통합하여, single stable diffusion model이 여러 타입의 출력을 동시에 생성할 수 있게하였다. 이는 상호간섭없이 원활한 도메인 처리를 보장한다. 구체적으로, d_lab은 먼저 binary form으로 표현된 다음, positional encoding을 통해 인코딩 된다. 이후 learnable projection head를 통과하고, time embedding과 element-wise로 결합된다. 결합된 embedding은 ResBlocks에 입력되어 모델의 배치별 출력 생성 능력을 향상시킨다.

이는 기존에 있던 기법을 가져와 적절하게 활용한 부분.

• Wonder3D: RGB와 normal 중 무엇을 생성할지 알려주는 domain switcher
• GeoWizard: depth와 normal 중 무엇을 생성할지 알려주는 geometry switcher
• DiffDIS: mask와 edge 중 무엇을 예측할지 알려주는 batch-discriminative embedding

 

쉽게 말하자면, 하나의 공유 U-Net이 mask와 edge를 각각 별도의 batch sample로 동시에 처리해야하니까, 각 sample이 mask인지 edge인지 구분할 수 있도록 label embedding을 제공하는 것. 기존에 batch = 2라면 Image 2장에 mask 2장, edge 2장이 있는 형태인데 mask인지 edge인지 레이블을 달아서 batch = 4로 만드는 셈. (하나의 batch에 mask1, mask2, edge1, edge2)

중간에 Time embedding이 들어가는 이유는, 원래 diffusion condition에 time이 들어가기 때문.

Learnable projection head는 보통 작은 신경망인데, label encoding을 U-Net이 사용할 수 있는 크기의 embedding으로 변환하는 역할을 함. (여기서는 아마 time-embedding과 같은 차원으로 변환될 것)

 

Detail-Balancing Interactive Attention (DBIA)

Mask와 edge의 연속성과 정렬을 보장하고, 이들의 의미적 단서와 attention이 집중되는 영역을 조화시키기 위해, 저자들은 DBIA를 도입하였다. DBIA는 두 도메인 사이에서 단순하지만 효과적인 정보 교환을 촉진하도록 설계되었으며, 궁극적으로 edge의 세부 정보와 의미적 내용 측면에서 모두 잘 정렬된 출력을 생성하는 것을 목표로 한다. 쉽게말하면, DiffDIS는 하나의 U-Net으로 두 출력 (edge, mask)을 동시에 예측하는데, 둘은 서로 관련되지만 강조하는 정보가 다르다. (mask - 객체 전체는 찾지만 경계가 뭉개짐 / edge -  세밀한 선은 찾지만 어떤 선이 실제 객체 경계인지 의미적으로 모호함) 그래서 두 정보가 상호 보완하여 서로 좋아지게 하고자 함.

 

기존의 SD에는 self-attention, cross-attention이 존재하는데, 저자들은 여기에 mutual cross-domain attention을 추가한다. 

(mask는 edge에서, edge는 mask에서 각각 참조하고 있음) 이 방법은 두 영역간의 직접적이고 효율적인 정보 교환을 보장하여, 전반적인 특징 표현을 향상시키는 상호작용을 촉진한다.

 

2.3 Scale-wise conditional injection

Dense prediction task를 위한 기존 확산 모델에서는 일반적으로 input 단계에서 condition을 준다. 이러한 방식은 후속 단계에서 정보 손실을 초래할 수 있다. 장기적이고 심층적인 conditional guidance를 구축하기 위해, 저자들은 multi-granular perception (다양한 세밀도 수준의 정보 인식)을 향상시키고 심층적인 시각적 상호작용을 촉진하는 SWCI 기법을 도입한다. 구체적으로, U-Net encoder의 해당 layer에 multi-scale condition을 통합한다. 그리고 3개의 injector heads를 사용하여 중간중간의 feature 크기 및 채널 수를 맞춰준다.

 

3개의 Injection heads는 1개의 simple convolution laeyr와 2개의 zero convolution layers로 구성된다. 각각의 헤드는 입력으로 해당 스케일에 맞게 크기가 조정된 conditional latent code를 받는다. 여기서 zero convolution은 모든 weight, bias가 0으로 초기화된 레이어로, ControlNet에서 사용된 방식이다. 초기상태에서는 값이 0이 나오므로 feature를 처음부터 건드리지 않지만, 학습이 진행되면서 weight가 업데이트 되어 RGB conditioning 경로를 안정적으로 학습하게 된다.

 

2.4 One-step mask sampling (Inference)

 

Inference 과정에서, 먼저 RGB image는 VAE encoder를 통해 latent space로 인코딩된다. 다음으로, standard gaussian noise에서 시작 변수를 샘플링하여, mask와 edge의 초기화로 사용한다. (아마 복제해서 사용하는듯?) Training 단계와 유사하게, 이 두 구성요소는 batch로 연결되며, batch-discriminative embedding (BDE)가 적용되어 효과적으로 구분된다. (쉽게 말해 그냥 레이블이 달려서 edge인지 mask인지 구분이 된다.) RGB latent representation에 따라 조건화된 concatnated components들은 노이즈를 예측하기 위해 U-net에 입력된다. 이후 알고리즘 2에 따라 기존 DDPM 방식처럼 샘플링 프로세스가 진행된다. 마지막으로, mask, edge map은 VAE 디코더를 사용하여 latent code에서 디코딩되고, channel 평균화를 통해 후처리된다. 

 

3. Experiments

3.1. Dataset, Metrics, and Implementation Details

 

• 데이터셋: DIS5K 사용
  • 학습: 3,000장, 225개 카테고리
  • 검증: DIS-VD
  • 테스트: DIS-TE1~4, 각 500장
• 학습 환경: PyTorch, NVIDIA H800 GPU 1장
• 입력 크기: 모든 이미지를 1024×1024로 resize
• 모델 초기화: SD V2.1을 backbone으로 사용하고, SD-Turbo pretrained parameter로 초기화
• 데이터 증강: horizontal flip, crop, rotation, CutMix
• 학습 설정: Adam, learning rate 3×10^−5, batch size 4, 최대 90 epoch
• 평가 시 후처리: 예측 맵을 이진화하여 작은 noise를 제거한 뒤 성능 계산
• 평가지표

 

3.2 Comparision

3.3 Ablation

각 컴포넌트의 effectiveness

적용할수록 점진적으로 결과가 좋아지고, 다 적용하니까 결과가 좋더라.

 

 

Pre-trained parameters와 denoising steps의 다양성

1, 2, 5행을 봤을 때, SD-Turbo 파라미터 가져다 쓰는게 제일 좋았음.

2, 3, 4행은 서로 다른 few-step denoising paradigms을 보여줌. 4행은 정확도가 조금 더 높지만, 시간이 두배로 증가함. 즉. step을 늘린다고 반드시 성능이 좋아지지는 않음.

 

 

DBIA에서 다양한 상호작용 방법 비교

DBIA의 추가적인 cross-attention 매커니즘에 있는 상호작용 방법을 Fu et al과 유사한 fusion-oriented approach로 대체했는데, 저자들께 더 좋다고 함.

 

---

dichotomous : 이분법적인

encompass : 포함하다. 둘러싸다.

versatile : 다재다능한, 다용도의, 다목적

streamline : 유선형으로 하다. 간소화 하다.

tough to + v : ~ 하기 어렵다.

seamless : 원활할, 아주 매끄러운