GuidedMix-Net 리뷰

오늘 정리할 논문은  AAAI'22에 accept 된 guidedmix-net이다. 
semi supervised learning으로 segmentation 성능 개선 방법은 제안 한다.
 
title: GuidedMix-Net: Semi-Supervised Semantic segmentation by using Labeled Image as Reference

 

 official repo:https://github.com/yh-pengtu/GuidedMix-Net/tree/main

GitHub - yh-pengtu/GuidedMix-Net: GuidedMix-Net: Learning to Improve Pseudo Masks Using Reference Labeled Images

문제 제기:

기존 제안된 방식은 unlabeled data를 모델의 consistency만을 강화하는 방식으로 사용해 왔고 이는 unlabeled data로 부터 얻을 수 있는 많은 양의 prior knowledge 를 활용하지 못하는 방식이다.
 
1. 학습시 labeled data와 unlabeled data 를 분리해서 따로 사용하고 unlabeled data 는 다양한 manual perturbation 을 이용해 consistency 학습에만 사용
-> 이런 방식으로는 unlabeled data 가 가지는 정보를 모두 활용 하지 못하고 이는 모델이 sub optimal(local minima)에 도달하는 결과를 초래할 수 있다.
2. CCT와 같은 방식은 다양한 perturbation을 생성해 consistency를 학습하는 방식이기 때문에 학습 시간이 오래걸릴 수 있음
 
아래 Figure 1. (c)를 보면 CCT 로 생성한 pseudo mask 는 부정확 한데 본 논문의 저자들은 이것을 자신들의 주장을 뒷바침 하는 근거로 제시한다. 

 

Proposed method

GuidedMix-Net에서는 labeled data 의 정보를 unlabeled data 에 전달해 학습을 guide 하는 방법을 제안하고 이를  mutual information transfer 라 하고 전체 아키텍쳐는 아래 그림과 같다. 
(우선 이 그림엔 틀린 부분이 있는것으로 보이는데 official code 를 보면 parameter sharing 이 되는 부분은 Figure 2 의 붉은 색으로 표시된 encoder 부분만 이다. decoder 는 labeled data만 이용해 학습 하는 decoder 와 mix input을 이용해 학습하는 decoder가 서로 다르다.)
 

labeled data의 정보를 이용해 unlabeled data을 이용한 학습을 guide 하기 위해 이논문에서는 3step process 를 제안한다. 

1. labeled-unlabeled pair interpolation

2. Mutual information transfer

3. pseudo mask generation

 

Labeled-unlabeled pair interpolation

이 단계는 labeled image ($x^{i}_{l}$) 와 unlabeled image($x^{j}_{u}$)을 혼합하해 Figure 2 의 아래쪽 입력에 해당하는 혼합이미지 $x^{j}_{mix}$를 생성하는 단계이다. 

Function 1.

혼합 이미지를 생성하는 방법은 위와 같다. labeled image ($x^{i}_{l}$) 와 unlabeled image($x^{j}_{u}$)를 $\lambda$ 비율로 더하는 과정이 전부이다. $\lambda \in (0,1)$ 로 $Beta(\alpha, \alpha), \alpha=1$ 로 부터 생성하고 $\lambda=min(\lambda, 1-\lambda)$ 와 같이 설정한다. 

 

이때 랜덤으로 labeled-unlabeled image pair를 선택하는 것보다 유사 이미지를 선택하는 것이 더 효과적이라고 하는데 그 이유는 동일/유사한 object 를 포함한 두 이미지가 섞여야 labeled image 로 부터의 정보가 unlabeled image 로 전파 더 잘 전파 될 수 있기 때문이다. 유사 이미지를 선택하는 방법은 mini batch 내에서 encoder의 output 으로부터 계산한 euclidean distance 가 가장 작은 labeled-unlabaled image 쌍을 선택하는 것이다. 아래 식에서 $\Gamma$ 는 model의 encoder 를 의미한다.

 한가지 트릭으로 이 방식으로 유사한 이미지 쌍을 선택하기 위해 encoder에 fullly connected layer 를 추가해 classification network을 만들고 labeled image만을 이용해 네트워크를 미리 트래이닝 시켜 비슷 한 이미지들이 feature space에서 유사한 위치에 있도록 한다. 

 

이를 코드로 나타내면 아래와 같다. Official repo에서 발췌한 건데 mini batch 내에서 유사한 이미지를 선택하는 부분은 빠져 있고 mini batch random image pair를 만드는 방식이다. 또한 가지 Figure 2에서와 다른 부분은 $x_{mix}$를 생성할 때 사용 되는 것은 실제 입력 이미지가 아니라 encoder의 output 에 해당한 feature map 이라는 것이다. 

lam = np.random.beta(alpha, alpha)
lam = min(lam, 1-lam)
while lam >= self.lambda_factor:
    lam = np.random.beta(alpha, alpha)
    lam = min(lam, 1-lam)

indices = torch.randperm(img.size(0), device=img.device, dtype=torch.long)
feats_mixed = lam * feats_aspp_labeled + (1-lam) * feats_aspp_unlabeled[indices]

 

Mutual information transfer

단순히 labeled-unlabeled image 를 interpolation 하는 것 만으로는 labeled image의 신뢰할 수 있는 정보(labeled 이 있기때문에 이렇게 표현한 것 같다)가 unlabeled image로 흐르는 것이 제한 적이기 때문에 이를 보완 할 방법으로 non-local block을 제안한다. 

 

 왜 단순 pairing 만으로 신뢰할 수 잇는 정보의 전파가 안되는 것일까? 이건 convolution filter의 특징때문으로 볼 수 있다. convolution filter 는 지역성을 가진다. 따라서 pairing된 이미지에서 같은 class에 속한 object 가 서로 유사한 위치에 있을 때(short-range)는 convolution만으로도 labeled image의 정보가 unlabeled image 로 흐를 수 있지만  mixed 된 이미지에서 정보를 전달할 object들이 서로 원거리에 있다면(long-range) 단순 convolution연산 만으로는 이 논문에서 원하는 효과가 제한적이다. 

 

그래서 이 논문에서는 non-local block을 사용 했다. Figure 3 의 블러도만 봐도 단적으로 안수 있듯이 self-attention과 같은 역할을 하는 non-local block을 이용해 long-range 정보를 전달하겠다는 전략이다. (self-attention 은 임의의 두 포인트의 interaction을 계산하기 때문에 long range dependency를 직접적으로 계산할 수 있다.)

논문에서는 이 모듈을 MITrans 라 명명 하는데 Mutual Information Transfer 이라는 의미이다.

왜 이 블럭이 labeled image 의 정보를 이용해 unlabled-image의 학습을 가이드 할 수 있다고 주장하는 것인지 의문이 들수 있는데 간단히 생각해 보자. labeled image를 이용해 supervised learning 을 진행 하므로 labeled image 가 encoder 를 거쳐 생성 되는 feature map에 있는 정보는 꽤나 신뢰 할 만 한 각 class 에 속하는 pixel들의 정보라 할 수 있다.  이 feature 과의 convolution 연산 이나 self-attention 연산을 거쳤을 때 그 response 가 크면 smoothness assumption 에 의해 유사한 feature 라고 볼 있다. 작으면 유사하지 않은 feature 라고 볼 수 있다. unlabeled image 로 부터 추출한 feature 가 labeled image로 부터 추출한 feature를 통해 correction 되는 효과를 기대할 수 있다는 것이다. 이 때문에 mutual information transfer라고 이 개념을 설명 한다. 

 

한가지 정말 맘에 드는 것은 MITrans block은 모델을 학습 할때만 사용 된다. 단순 inference time에는 이 블럭이 사용되지 않는다. 이게 왜 중요한 부분인지 edge device에서 서비스를 해본 사람은 이해 할 수 있을 것이다. (23년 기준..). Edge device의 경우 chip vender에 따라 self attention에 필요 한 연산을 지원하는 sdk를 아직 배포 하지 않는 곳도 있고 오래된 sdk version을 사용 하는 경우 self attention 을 inference time에 사용하는게 불가능 한 경우가 있기 때문이다. 

 

중요한 블럭인것은 맞지만 Figure 3 만 봐도 의도를 직관적으로 파악 할 수 있으므로 수식을 생략한다.

 

pseudo mask generation

마지막으로 labeled-unlabeled image pair 를 이용해 생성한 $x_{mix}$ 를  MITrans 에 통과 시킨후 decoder 를 통해 생성한 segmentation mask 는 labeled image에 존재하는 object에 해당하는 부분이 존재한다.

이 부분은 unlabeled image의 pseudo gt 로 사용 하기에 필요 없을 뿐더러 해당 부분에 대한 정확한 정보는 labeled image 의 GT를 이용해 supervised learning process(Figure 1의 윗 부분)에서 학습 되므로 필요가 없다. 

 

 이러한 이유로 이 논문에서는 mixed image를 이용해 생성한 pseudo mask 에서 labeled image 에 해당 하는 부분와 unlabeled image 에 해당 하는 부분을 분리 하는 방법(decoupling method)을 생각해 냈다.  방법 자체는 아주 단순하고 직관적이다.  mixed image를 네트워크에 통과시켜 얻은 prediciton 에서 labeled image를 네트워크에 통과시켜 얻은 prediction 을 빼는 것이다. 

 

 왜 이렇게 한 것일까? mixed image는 labeled/unlabeled image의 가중 합으로 생성된 것이므로 mixed image를 네트워크에 통과시켜 나온 prediction은 labeled image와 unlabeled image 를 네트워크에 각각 통과시켜 나온 prediction $M_{l}, M_{u}$ 을 더한 것과 유사할 것이라는 가정에서 시작 된다.  

수식으로 표현 하자면 아래와 같다. 

여기서 $\Gamma$는 segmentation network을 나타내고 $M_{l},M_{u}$는 각각  labeled image , unlabeled image의 prediction 에 대응 된다. 

 

저자는 hard decoupling과 soft decoupling  두 가지를 테스트 해봤는데 soft-decoupling 이 결과는 더 좋았다. 

 

hard decoupling  

Hard decoupling 은 정말 mixed image와 labeled image의 prediction을 단순 뺄샘한다.  하지만 mixed image를 생성 할때 

$x_{mix} = \lambda x_{l} + (1-\lambda)x_{u}$ 와 같이 두 이미지를 weighted sum  하기 때문에 이는 적절하지 않아 보인다. 

 

Soft decoupling

soft decoupling은 입력인 mixed image를 생성할 때 사용했던 weight $\lambda$를 고려해 위와 같이 pseudo mask 를 생성한다. soft decoupling을 사용 할때 얻을 수 있는 장점은 같은 class에 속한 object 가 labeled, unlabeled image 의 동일 한 위치에 있을때 labeled image prediction 을 mixed image prediction 에서 빼도 unlabeled image object에 대한 추론 정보가 남아있다는 것이다. 

 

위와 같이 생성된 psedu mask $M_{u-dec}$는 unlabeled image의 pseudo GT 로 활용하며 unlabeled image의 loss는 아래와 같이 mean squred error로 계산 된다. 

$M_{u-dec}$ 는 soft-decoupling 또는 hard decoupling 을 통해 생성한 pseudo mask 이고 $M_{u}$ 는 unlabeled image의 prediction 결과 이다. 

 

 

Loss

GuidedMix-Net 의 loss function 은 아래와 같다. 

$L_{usup} 은 바로 이전에 unlabeled image의 loss 이고 $\omega$ 는 unsupervised loss weight 으로 하이퍼 파라미터이다. 

$L_{sup}$는 labeled image의 superised loss 이다. 직관 적으로 알 수 있듯이 $L_{ce}$ 는 labeled image의 cross entropy이고 $L_{cla}$는 image level의 classification loss 이다. 이전 섹션인 proposed method의 similar image selection 을 위해 encoder에 fully connected layer 를 추가해 classification network 을 학습 시킨다고 했는데 이 부분에 해당하는 loss 이다. 

 

$L_{dec}$ 에 대한 아래 부분은 설명이 틀릴 수 있음을 밝힌다. 코드 상에서 해당 하는 부분을 정확히 파악 할 수 없어 정확한 의미 파악이 힘들다. 틀린 부분은 지적바란다. 

 

$L_{dec}는 labeled image의 consistency loss term 으로 볼 수 있다. labeled-unlabeled image pair 를 이용해 mixed image를 생성했던 것을 그대로 labeled-labeled image pair 에 적용해 mixed image 를 생성한다. (단 여기서 labeled image pair 에 속한 두 장의 이미지는 같은 이미지인것으로 보인다.)

$\hat{x_{l}}, x_{l}$가 모두 labeled image 일때  $x^{l}_{mix}= \lambda x_{l} + (1-\lambda)\hat{x_{l}}$과 같이 생성한다. 이를 network에 통과시켜 prediction $M_{mix}$를 얻는다. 그 후 labeled-unlabeled image pair의 prediction decoupling 방법을 여기에 적용한다. $M_{dec}=M_{mix} -M_{l}$  여기서 M_{l}은 labeld image $x_{l}$의 prediction 이다. 

 

이렇게 구한 $M_{dec}$ 를 이용해 아래와 같이 구한다. 

Result 

Table 1은 $\lambda$ 값의 변화가 성능에 미치는 영향을 보여준다. $\lambda$ 의 제한 범위에 따라 성능이 굉장히 상이 하니 참조 하자. 

 

Table 2는 similar labeled-unlabeled image pair selection, MITrans(mutual information transfer) block, hard/soft decoupling 의 적용에 따른 결과를 보여 준다. 

labeled-unlabeled image pair selection 의 경우 random selection 보다 eclidean distance 를 이용한 silimar image pair selection 을 할 경우 성능이 더 좋았다. MUTrans block 도 사용하는 경우 성능이 더 좋았다. Hard decoupling 보다는 soft decoupling 을 사용하는게 더 성능이 좋았다. 

다음으로 다른 SOTA 들과의 비교이다. 다른 semi supervised learning 기반 방식들 보다도 동일한 비율의 labeled 데이터를 사용 했을 때 더 좋은 성능을 보였고 개인적으로 인상깊은 것은 이전에 읽은 논문인 CCT와의 비교 였다. 

CCT대비 큰 성능 이득을 보여 주는게 인상적이었다. 내 데이터 셋에 적용해 비교해 봐야겠다. 

 

-끝-