[논문리뷰] (2022) SegDiff: Image Segmentation with Diffusion Probabilistic Models

Paper : (2022) SegDiff: Image Segmentation with Diffusion Probabilistic Models

Researchers : Tomer Amit, Tal Shaharbany, Eliya Nachmani, and Lior Wolf

Organizations : Tel-Aviv University, Facebook AI Research

Contents 

1. Introduction
2. Architecture & Training Algorithm 
3. Employing multiple generations 
4. Datasets
5. Training Details 
6. Result

1.  Introduction

 Diffusion 방식을 사용해서 Segmentation Task 를 수행한 논문이다. Conditional Generation을 활용해서 Segmentation 을 풀어냈다. 

 

 

   2.  Architecture & Training Algorithm

 

< Architecture >

 

Conditional Diffusion 을 응용한 아키텍처로 보인다. 아키텍처 상의 x0 를 reconstruction 할 때, 이미지 I 를 conditional 한 label y 로 취급한 형태이다. 논문에 나온 모델 아키텍처를 간단히 정리하면은 아래와 같다.

• G Encoder
  • RRDBs(Residual Dense Blocks)
    • w/o batchNorm ( groupNorm maybe .. )
    • 2d-conv
    • residual connection
    • leaky RELU
• F Encoder
  • 2d conv
  • single-channel input
  • C channels output
• D & E
  • U-Net ( Encoder , Decoder )
    • residual block
      • two conv block
        • groupNorm
        • Silu activation
        • 2d-conv
        • get time embedding through a linear layer
    • resolution 16 x 16 , 8 x 8
    • attention layer
    • bottleneck
      • two residual blocks with attention in between
      • multiple attention head
    • Encoder side (E) ⇒ Decoder side (D)

아키텍처 자체를 improved DDPM 에서 가져왔기 때문에 크게 다른점은 보이지 않는다.

 

 

 

< Training Algorithm >

아키텍처에서 보이는 image 와 segmentation mask 를 데이터셋 D = { I, M } 으로 둔다.

• I, M 을 샘플링한다
• t 스텝은 uniform{ 1 ~ T } 에서 결정한다
• βt 는 linear scheduler 를 사용하였다. (improved ddpm 에서 나온 cosine scheduler 를 왜 사용하지 않았을까는 의문이다)
• αt 는 기존에 ddpm 에서 정의 그대로를 따른다. α¯t 도 마찬가지.
• noise 를 target 분포로 가져가는 것 또한 동일하다.

 

 

< Objective function >

improved DDPM 을 기본으로 하였고, 역시 noise 를 모델링하였다.

 

 

 

 

 3. Employing multiple generations 

 

< Multiple generations >

<paper> Obtaining multiple segmentation results for Cityscapes, Vaihingen, and MoNuSeg

inference 를 해내는 작업은 standard distribution 인 N(xt, t)z 에서 샘플링을 해내는 과정에서 randomness가 발생하게 된다. 따라서 이러한 현상을 커버하기 위해서 같은 input으로 여러번 추론해낸 다음에 추론 결과에 average 를 취하는 방식을 선택하였다.

 

< Generation instances testing >

4. Datasets 

1. The Cityscapes
   1. 5,000 total
   2. 2,975 training, 500 for validation, 1,525 testing
2. The Vaihingen
   1. 168 aerial images (Vaihingen, in Germany)
   2. 100 train and 67 test
3. MoNuSeg
   1. 30 microscopic images
      1. seven organs
      2. 21,623 individual nuclei
      3. 14 test imgs 

 

5. Training Details

기존의 diffusion 에서 사용했던 step들이 1000 ~ 4000 쯤 되었는데 해당 논문 에서는 inference time 을 많이 의식해서 100 diffusion step 을 사용하였다. 결국은 step 수에 대한 scaling 테스트를 진행해서 뒤에 표로 첨부하였다.

Optimizer 는 AdamW 를 사용하였고, RRDB block 들을 적극적으로 활용하였다. Unet 을 256 x 256 사이즈에 맞춰서 설계하였다.

 

⇒ Vaihingen Dataset

• Unet
  • 256 x 256 with additional layer, RRDB blocks with 7 depth
  • channel size [C, C, 2C, 2C, 4C, 4C] C =128
• Batch size
  • 8
• Optimizer
  • AdamW
• Augmentations
  • random scaling [0.75, 1.25]
  • rotation sampled uniformly 0~360 degree
  • independent horizontal and vertical flips ( pb 0.5 )
  • random color jitter ( max 0.6 brightness, 0.5 contrast, 0.4 saturation, 0.025 hue )

 

6. Result